change the benchmark executable to slurp keys from the database instead of stdin

1 files changed, 40 insertions, 7 deletions

diff --git a/src/main.cpp b/src/main.cpp
index 731323a..81c56b6 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -1,33 +1,66 @@
 #include <chrono>
-#include <iostream>
+#include <cinttypes>
 #include <print>
+#include <random>
 #include <string>
+#include <string_view>
 #include <vector>
 
 #include <boost/sqlite.hpp>
 
+void parse_int(std::string_view arg, std::uint64_t &i) {
+  auto [_, ec] = std::from_chars(arg.data(), arg.data() + arg.size(), i);
+
+  if (ec == std::errc{}) {
+    return;
+  } else {
+    throw std::runtime_error(std::format("failed to parse {} as u64", arg));
+  }
+}
+
 int main(int argc, char **argv) {
-  if (argc < 2) {
-    std::println(stderr, "usage: bench <database>");
+  if (argc < 3) {
+    std::println(stderr, "usage: bench <database> <limit> [seed]");
     return 1;
   }
 
+  auto *path = argv[1];
+  auto *limit = argv[2];
+
+  std::uint64_t seed;
+  if (argc > 3) {
+    try {
+      parse_int(argv[3], seed);
+    } catch (std::exception &e) {
+      std::println(stderr, "{}", e.what());
+      return 1;
+    }
+  } else {
+    seed = std::chrono::system_clock::now().time_since_epoch().count();
+  }
+
+  std::default_random_engine rng{seed};
+
   std::println(stderr, "starting benchmark");
 
   boost::sqlite::connection connection{argv[1]};
 
   std::vector<std::string> keys;
 
-  std::string line;
-  while (std::getline(std::cin, line)) {
-    keys.push_back(line);
+  auto st = connection.prepare("select kv.key from kv limit ?");
+
+  for (const auto row : st.execute({limit})) {
+    auto key = row.at(0).get_text();
+    keys.push_back(key);
   }
 
+  std::shuffle(keys.begin(), keys.end(), rng);
+
   std::println(stderr, "slurped {} keys into memory", keys.size());
 
   auto start = std::chrono::system_clock::now();
 
-  auto st = connection.prepare("select kv.value from kv where kv.key=?");
+  st = connection.prepare("select kv.value from kv where kv.key=?");
 
   for (const auto &key : keys) {
     auto row = st.execute({key});