From 5ad9e56281bdb1242657d2885928176246513ec5 Mon Sep 17 00:00:00 2001
From: glyh <lyhokia@gmail.com>
Date: Mon, 24 Nov 2025 13:19:53 +0800
Subject: [PATCH] db_benchmark: properly clean up the generated file after each
 benchmark run

---
 src/test/db_benchmark/common.ml       | 19 +++++--------------
 src/test/db_benchmark/db_benchmark.ml |  8 ++++----
 2 files changed, 9 insertions(+), 18 deletions(-)

diff --git a/src/test/db_benchmark/common.ml b/src/test/db_benchmark/common.ml
index 9c04519d54c3..8a5d99e24489 100644
--- a/src/test/db_benchmark/common.ml
+++ b/src/test/db_benchmark/common.ml
@@ -101,17 +101,8 @@ 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
+  Fun.protect
+    ~finally:(fun () -> Mina_stdlib_unix.File_system.rmrf dir_name)
+    (fun () -> f dir_name)
diff --git a/src/test/db_benchmark/db_benchmark.ml b/src/test/db_benchmark/db_benchmark.ml
index bf299c64592e..1776c5614c95 100644
--- a/src/test/db_benchmark/db_benchmark.ml
+++ b/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 () =