Disk Cache: design an API to disable GC collection

Author: glyh

src/lib/disk_cache/filesystem/disk_cache.ml

@@ -10,36 +10,39 @@ module Make (B : sig
   include Binable.S
 end) =
 struct
-  type t = string * int ref
+  type t = { root : string; next_idx : int ref; gc_freezed : bool ref }
 
   type id = { idx : int }
 
   let initialize path ~logger =
     Async.Deferred.Result.map (Disk_cache_utils.initialize_dir path ~logger)
-      ~f:(fun path -> (path, ref 0))
+      ~f:(fun root -> { root; next_idx = ref 0; gc_freezed = ref false })
 
   let path root i = root ^ Filename.dir_sep ^ Int.to_string i
 
-  let get ((root, _) : t) (id : id) : B.t =
+  let get ({ root; _ } : t) (id : id) : B.t =
     (* Read from the file. *)
     In_channel.with_file ~binary:true (path root id.idx) ~f:(fun chan ->
         let str = In_channel.input_all chan in
         Binable.of_string (module B) str )
 
-  let put ((root, next_idx) : t) x : id =
+  let freeze_gc { gc_freezed; _ } = gc_freezed := true
+
+  let put ({ root; next_idx; gc_freezed } : t) x : id =
     let idx = !next_idx in
     incr next_idx ;
     let res = { idx } in
     (* When this reference is GC'd, delete the file. *)
     Core.Gc.Expert.add_finalizer_last_exn res (fun () ->
-        (* Ignore errors: if a directory is deleted, it's ok. *)
-        try Core.Unix.unlink (path root idx) with _ -> () ) ;
+        if not !gc_freezed then
+          (* Ignore errors: if a directory is deleted, it's ok. *)
+          try Core.Unix.unlink (path root idx) with _ -> () ) ;
     (* Write the proof to the file. *)
     Out_channel.with_file ~binary:true (path root idx) ~f:(fun chan ->
         Out_channel.output_string chan @@ Binable.to_string (module B) x ) ;
     res
 
-  let count ((path, _) : t) = Sys.ls_dir path |> List.length
+  let count ({ root; _ } : t) = Sys.ls_dir root |> List.length
 end
 
 let%test_module "disk_cache filesystem" =

src/lib/disk_cache/identity/disk_cache.ml

@@ -14,4 +14,6 @@ struct
   let get () = ident
 
   let put () = ident
+
+  let freeze_gc () = ()
 end

src/lib/disk_cache/intf/disk_cache_intf.ml

@@ -19,6 +19,8 @@ module type S = sig
 
   (** Read from the cache, crashing if the value cannot be found. *)
   val get : t -> id -> Data.t
+
+  val freeze_gc : t -> unit
 end
 
 module type S_with_count = sig

src/lib/disk_cache/lmdb/disk_cache.ml

@@ -24,6 +24,7 @@ module Make (Data : Binable.S) = struct
           (** A list of ids that are no longer reachable from OCaml runtime, but
               haven't been cleared inside the LMDB disk cache *)
     ; queue_guard : Error_checking_mutex.t
+    ; gc_freezed : bool ref
     }
 
   (** How big can the queue [reusable_keys] be before we do a cleanup *)
@@ -38,15 +39,16 @@ module Make (Data : Binable.S) = struct
         ; counter = ref 0
         ; reusable_keys = Queue.create ()
         ; queue_guard = Error_checking_mutex.create ()
+        ; gc_freezed = ref false
         } )
 
   type id = { idx : int }
 
   let get ({ env; db; _ } : t) ({ idx } : id) : Data.t =
     Rw.get ~env db idx |> Option.value_exn
 
-  let put ({ env; db; counter; reusable_keys; queue_guard } : t) (x : Data.t) :
-      id =
+  let put ({ env; db; counter; reusable_keys; queue_guard; gc_freezed } : t)
+      (x : Data.t) : id =
     let idx =
       match
         Error_checking_mutex.critical_section queue_guard ~f:(fun () ->
@@ -63,12 +65,13 @@ module Make (Data : Binable.S) = struct
     let res = { idx } in
     (* When this reference is GC'd, delete the file. *)
     Gc.Expert.add_finalizer_last_exn res (fun () ->
-        (* The actual deletion is delayed, as GC maybe triggered in LMDB's
-           critical section. LMDB critical section then will be re-entered if
-           it's invoked directly in a GC hook.
-           This causes mutex double-acquiring and node freezes. *)
-        Error_checking_mutex.critical_section queue_guard ~f:(fun () ->
-            Queue.enqueue reusable_keys idx ) ) ;
+        if not !gc_freezed then
+          (* The actual deletion is delayed, as GC maybe triggered in LMDB's
+             critical section. LMDB critical section then will be re-entered if
+             it's invoked directly in a GC hook.
+             This causes mutex double-acquiring and node freezes. *)
+          Error_checking_mutex.critical_section queue_guard ~f:(fun () ->
+              Queue.enqueue reusable_keys idx ) ) ;
 
     Error_checking_mutex.critical_section queue_guard ~f:(fun () ->
         if Queue.length reusable_keys >= reuse_size_limit then (
@@ -77,6 +80,8 @@ module Make (Data : Binable.S) = struct
     Rw.set ~env db idx x ;
     res
 
+  let freeze_gc ({ gc_freezed; _ } : t) = gc_freezed := true
+
   let iteri ({ env; db; _ } : t) ~f = Rw.iter ~env db ~f
 
   let count ({ env; db; _ } : t) =