db_benchmark: properly clean up the generated file after each benchmark run

Author: glyh

src/test/db_benchmark/common.ml

@@ -101,17 +101,9 @@ module Ops = struct
     read_key (module Db) db key
 end
 
-(* Temporary directory management *)
-let make_temp_dir prefix =
-  let pid = Unix.getpid () |> Pid.to_int in
-  let dir_name = Printf.sprintf "%s_%d" prefix pid in
-  Unix.mkdir_p dir_name ; dir_name
-
-let cleanup_temp_dir dir =
-  match Sys.file_exists dir with
-  | `Yes ->
-      ignore
-        ( Core_unix.system (Printf.sprintf "rm -rf %s" (Filename.quote dir))
-          : Core_unix.Exit_or_signal.t )
-  | _ ->
-      ()
+let with_temp_dir prefix ~f =
+  let dir_name = Unix.mkdtemp prefix in
+  try f dir_name
+  with e ->
+    Mina_stdlib_unix.File_system.rmrf dir_name ;
+    raise e

src/test/db_benchmark/db_benchmark.ml

@@ -10,9 +10,8 @@ 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) ~dir (module Db : Common.Database with type t = db) name =
   (* Initialization: create DB and warmup *)
-  let dir = Common.make_temp_dir (Printf.sprintf "db_bench_%s" name) in
   let db = Db.create dir in
   Common.Ops.warmup (module Db) db ;
   eprintf "Warmup complete for %s\n" name ;
@@ -37,8 +36,9 @@ let make_read_bench (type db) (module Db : Common.Database with type t = db)
 
 let test ~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) )
+  Common.with_temp_dir name ~f:(fun dir ->
+      Bench.Test.create_with_initialization ~name (fun `init () ->
+          init_db ~dir (module Db) name |> f (module Db) ) )
 
 (* Create all benchmarks *)
 let all_benchmarks () =