Properly remove tempdir after db bench

Author: glyh

src/test/db_benchmark/db_benchmark.ml

@@ -10,9 +10,11 @@ module Single_file_db = Single_file_impl.Make ()
 
 module Multi_file_db = Multi_file_impl.Make ()
 
-let init_db (type db) (module Db : Common.Database with type t = db) name =
+let init_db (type db) (module Db : Common.Database with type t = db) ~root name
+    =
   (* Initialization: create DB and warmup *)
-  let dir = Common.make_temp_dir (Printf.sprintf "db_bench_%s" name) in
+  let dir = root ^/ name in
+  Unix.mkdir_p dir ;
   let db = Db.create dir in
   Common.Ops.warmup (module Db) db ;
   eprintf "Warmup complete for %s\n" name ;
@@ -35,21 +37,23 @@ let make_read_bench (type db) (module Db : Common.Database with type t = db)
   let max_key = (Common.warmup_blocks * Common.keys_per_block) - 1 in
   fun () -> Common.Ops.random_read (module Db) db ~min_key ~max_key
 
-let test ~name (type db) (module Db : Common.Database with type t = db)
+let test ~root ~name (type db) (module Db : Common.Database with type t = db)
     (f : (module Common.Database with type t = db) -> db -> unit -> unit) =
   Bench.Test.create_with_initialization ~name (fun `init ->
-      init_db (module Db) name |> f (module Db) )
+      init_db ~root (module Db) name |> f (module Db) )
 
 (* Create all benchmarks *)
-let all_benchmarks () =
-  [ test ~name:"rocksdb_write" (module Rocksdb_db) make_write_bench
-  ; test ~name:"rocksdb_read" (module Rocksdb_db) make_read_bench
-  ; test ~name:"lmdb_write" (module Lmdb_db) make_write_bench
-  ; test ~name:"lmdb_read" (module Lmdb_db) make_read_bench
-  ; test ~name:"single_file_write" (module Single_file_db) make_write_bench
-  ; test ~name:"single_file_read" (module Single_file_db) make_read_bench
-  ; test ~name:"multi_file_write" (module Multi_file_db) make_write_bench
-  ; test ~name:"multi_file_read" (module Multi_file_db) make_read_bench
+let all_benchmarks ~root () =
+  [ test ~root ~name:"rocksdb_write" (module Rocksdb_db) make_write_bench
+  ; test ~root ~name:"rocksdb_read" (module Rocksdb_db) make_read_bench
+  ; test ~root ~name:"lmdb_write" (module Lmdb_db) make_write_bench
+  ; test ~root ~name:"lmdb_read" (module Lmdb_db) make_read_bench
+  ; test ~root ~name:"single_file_write"
+      (module Single_file_db)
+      make_write_bench
+  ; test ~root ~name:"single_file_read" (module Single_file_db) make_read_bench
+  ; test ~root ~name:"multi_file_write" (module Multi_file_db) make_write_bench
+  ; test ~root ~name:"multi_file_read" (module Multi_file_db) make_read_bench
   ]
 
 (* Filter benchmarks based on BENCHMARKS environment variable *)
@@ -83,6 +87,10 @@ let () =
     (Common.warmup_blocks * Common.keys_per_block) ;
   Printf.printf "\n" ;
 
-  (* Run benchmarks *)
-  let benchmarks = all_benchmarks () |> filter_benchmarks in
-  Command.run (Bench.make_command benchmarks)
+  let root = Common.make_temp_dir "db_bench" in
+  try
+    let benchmarks = all_benchmarks ~root () |> filter_benchmarks in
+    Command.run (Bench.make_command benchmarks)
+  with e ->
+    Common.cleanup_temp_dir root ;
+    raise e