1 files changed, 26 insertions, 4 deletions

diff --git a/doc/bench/db_bench_tree_db.cc b/doc/bench/db_bench_tree_db.cc
index 214d9b7..ed86f03 100644
--- a/doc/bench/db_bench_tree_db.cc
+++ b/doc/bench/db_bench_tree_db.cc
@@ -68,6 +68,9 @@ static bool FLAGS_use_existing_db = false;
 // is off.
 static bool FLAGS_compression = true;
 
+// Use the db with the following name.
+static const char* FLAGS_db = NULL;
+
 inline
 static void DBSynchronize(kyotocabinet::TreeDB* db_)
 {
@@ -292,11 +295,16 @@ class Benchmark {
     bytes_(0),
     rand_(301) {
     std::vector<std::string> files;
-    Env::Default()->GetChildren("/tmp", &files);
+    std::string test_dir;
+    Env::Default()->GetTestDirectory(&test_dir);
+    Env::Default()->GetChildren(test_dir.c_str(), &files);
     if (!FLAGS_use_existing_db) {
       for (int i = 0; i < files.size(); i++) {
         if (Slice(files[i]).starts_with("dbbench_polyDB")) {
-          Env::Default()->DeleteFile("/tmp/" + files[i]);
+          std::string file_name(test_dir);
+          file_name += "/";
+          file_name += files[i];
+          Env::Default()->DeleteFile(file_name.c_str());
         }
       }
     }
@@ -385,8 +393,12 @@ class Benchmark {
     db_ = new kyotocabinet::TreeDB();
     char file_name[100];
     db_num_++;
-    snprintf(file_name, sizeof(file_name), "/tmp/dbbench_polyDB-%d.kct",
-                               db_num_);
+    std::string test_dir;
+    Env::Default()->GetTestDirectory(&test_dir);
+    snprintf(file_name, sizeof(file_name),
+             "%s/dbbench_polyDB-%d.kct",
+             test_dir.c_str(),
+             db_num_);
 
     // Create tuning options and open the database
     int open_options = kyotocabinet::PolyDB::OWRITER |
@@ -470,6 +482,7 @@ class Benchmark {
 }  // namespace leveldb
 
 int main(int argc, char** argv) {
+  std::string default_db_path;
   for (int i = 1; i < argc; i++) {
     double d;
     int n;
@@ -494,12 +507,21 @@ int main(int argc, char** argv) {
     } else if (sscanf(argv[i], "--compression=%d%c", &n, &junk) == 1 &&
                (n == 0 || n == 1)) {
       FLAGS_compression = (n == 1) ? true : false;
+    } else if (strncmp(argv[i], "--db=", 5) == 0) {
+      FLAGS_db = argv[i] + 5;
     } else {
       fprintf(stderr, "Invalid flag '%s'\n", argv[i]);
       exit(1);
     }
   }
 
+  // Choose a location for the test database if none given with --db=<path>
+  if (FLAGS_db == NULL) {
+      leveldb::Env::Default()->GetTestDirectory(&default_db_path);
+      default_db_path += "/dbbench";
+      FLAGS_db = default_db_path.c_str();
+  }
+
   leveldb::Benchmark benchmark;
   benchmark.Run();
   return 0;