refactor: Refactor pending target cleanup into Sandbox (#14483)

The build system used to special-case non-sandboxed actions by manually
adding and removing their targets from Pending_targets while real
sandboxes went through Sandbox.create and Sandbox.destroy. That kept the
cleanup lifecycle split between Build_system and Sandbox.

Move pending target tracking into Sandbox and make Sandbox.with_ take an
optional sandbox mode. A missing mode now creates a fake sandbox
context: paths map to themselves, target extraction is a no-op, and
destroying the context only removes the action from the pending-target
set. Real sandbox modes still create and destroy actual sandbox
directories.

Keep the fake non-sandbox context outside the live-sandbox throttle.
Only real sandboxes consume one of the 100 global sandbox slots.

Signed-off-by: Rudi Grinberg <me@rgrinberg.com>

Author: rgrinberg

src/dune_engine/build_system.ml

@@ -120,25 +120,6 @@ let rec with_locks ~f = function
       ~f:(fun () -> with_locks ~f mutexes)
 ;;
 
-module Pending_targets = struct
-  (* All file and directory targets of non-sandboxed actions that are currently
-     being executed. On exit, we need to delete them as they might contain
-     garbage. *)
-
-  let t = ref Targets.empty
-  let remove targets = t := Targets.diff !t (Targets.Validated.unvalidate targets)
-  let add targets = t := Targets.combine !t (Targets.Validated.unvalidate targets)
-
-  let cleanup () =
-    let targets = !t in
-    t := Targets.empty;
-    Targets.iter
-      targets
-      ~file:(fun p -> p |> Path.Build.to_string |> Fpath.unlink_no_err)
-      ~dir:(fun p -> Path.rm_rf (Path.build p))
-  ;;
-end
-
 type rule_execution_result =
   { facts : Dep.Fact.t Dep.Map.t
   ; targets : Digest.t Targets.Produced.t
@@ -403,18 +384,9 @@ module Internal = struct
       =
       action
     in
-    let deps =
-      Dep.Facts.paths
-        ~expand_aliases:
-          (Execution_parameters.expand_aliases_in_sandbox execution_parameters)
-        facts
-    in
     let execute_action sandbox =
-      let action =
-        match sandbox with
-        | None -> action
-        | Some sandbox -> Action.sandbox action sandbox
-      in
+      let is_sandboxed = Sandbox.is_sandboxed sandbox in
+      let action = if is_sandboxed then Action.sandbox action sandbox else action in
       let action =
         (* We must add the creation of the stamp file after sandboxing it, as
            otherwise the stamp file would end up inside the sandbox. This is
@@ -438,12 +410,7 @@ module Internal = struct
         | None -> Path.Build.root
         | Some context -> context.build_dir
       in
-      let root =
-        Path.build
-          (match sandbox with
-           | None -> root
-           | Some sandbox -> Sandbox.map_path sandbox root)
-      in
+      let root = Path.build (Sandbox.map_path sandbox root) in
       let action_trace = Action_trace.create rule_digest in
       with_locks locks ~f:(fun () ->
         let* action_exec_result =
@@ -464,17 +431,17 @@ module Internal = struct
         let* action_exec_result = Action_exec.Exec_result.ok_exn action_exec_result in
         let* () = Action_trace.collect action_trace in
         let* () =
-          match sandbox with
-          | None -> Fiber.return ()
-          | Some sandbox ->
+          if not is_sandboxed
+          then Fiber.return ()
+          else (
             (* The stamp file for anonymous actions is always created outside
                the sandbox, so we can't move it. *)
             let should_be_skipped =
               match rule_kind with
               | Normal_rule -> fun (_ : Path.Build.t) -> false
               | Anonymous_action { stamp_file; _ } -> Path.Build.equal stamp_file
             in
-            Sandbox.move_targets_to_build_dir sandbox ~should_be_skipped ~targets
+            Sandbox.move_targets_to_build_dir sandbox ~should_be_skipped ~targets)
         in
         let+ produced_targets =
           maybe_async_rule_file_op (fun () -> Targets.Produced.of_validated targets)
@@ -483,35 +450,26 @@ module Internal = struct
         | Ok produced_targets -> { Exec_result.produced_targets; action_exec_result }
         | Error error -> User_error.raise ~loc (Targets.Produced.Error.message error))
     in
-    match sandbox_mode with
-    | Some mode ->
-      Sandbox.with_live_sandbox_slot ~f:(fun () ->
-        let* sandbox =
-          Sandbox.create
-            ~mode
-            (Option.value ~default:Ignore corrections)
-            ~dirs:(Dep.Facts.necessary_dirs_for_sandboxing facts)
-            ~deps
-            ~rule_dir:targets.root
-            ~rule_loc:loc
-            ~rule_digest
-        in
-        (* CR-someday rgrinberg: Dynamic actions may discover dependencies
-           while this slot is held. If all sandbox slots are held by such
-           actions, sandboxed rules for the discovered dependencies cannot
-           start. *)
-        Fiber.finalize
-          ~finally:(fun () -> Sandbox.destroy sandbox)
-          (fun () -> execute_action (Some sandbox)))
-    | None ->
-      (* If the action is not sandboxed, we use [pending_file_targets] to
-         clean up the build directory if the action is interrupted. *)
-      Pending_targets.add targets;
-      Fiber.finalize
-        ~finally:(fun () ->
-          Pending_targets.remove targets;
-          Fiber.return ())
-        (fun () -> execute_action None)
+    let deps, sandbox_dirs =
+      match sandbox_mode with
+      | None -> Path.Set.empty, Path.Build.Set.empty
+      | Some _ ->
+        ( Dep.Facts.paths
+            ~expand_aliases:
+              (Execution_parameters.expand_aliases_in_sandbox execution_parameters)
+            facts
+        , Dep.Facts.necessary_dirs_for_sandboxing facts )
+    in
+    Sandbox.with_
+      ~mode:sandbox_mode
+      (Option.value ~default:Ignore corrections)
+      ~rule_loc:loc
+      ~dirs:sandbox_dirs
+      ~deps
+      ~rule_dir:targets.root
+      ~rule_digest
+      ~targets
+      ~f:execute_action
 
   and promote_targets ~rule_mode ~targets ~promote_source =
     match rule_mode, !Clflags.promote with
@@ -1176,7 +1134,7 @@ let run f =
         Memo.run_with_error_handler f ~handle_error_no_raise:report_early_exn)
     in
     Dtemp.clear ();
-    Pending_targets.cleanup ();
+    Sandbox.cleanup_pending_targets ();
     Target_promotion.save ();
     let* outcome =
       match outcome with

src/dune_engine/sandbox.ml

@@ -5,6 +5,27 @@ let max_live_sandboxes = 250
 let live_sandbox_throttle = lazy (Fiber.Throttle.create max_live_sandboxes)
 let with_live_sandbox_slot ~f = Fiber.Throttle.run (Lazy.force live_sandbox_throttle) ~f
 
+module Pending_targets = struct
+  (* All file and directory targets of non-sandboxed actions that are currently
+     being executed. On exit, we need to delete them as they might contain
+     garbage. *)
+
+  let t = ref Targets.empty
+  let remove targets = t := Targets.diff !t (Targets.Validated.unvalidate targets)
+  let add targets = t := Targets.combine !t (Targets.Validated.unvalidate targets)
+
+  let cleanup () =
+    let targets = !t in
+    t := Targets.empty;
+    Targets.iter
+      targets
+      ~file:(fun p -> p |> Path.Build.to_string |> Fpath.unlink_no_err)
+      ~dir:(fun p -> Path.rm_rf (Path.build p))
+  ;;
+end
+
+let cleanup_pending_targets = Pending_targets.cleanup
+
 let maybe_async f =
   (* It would be nice to do this check only once and return a function, but the
      type of this function would need to be polymorphic which is forbidden by the
@@ -47,16 +68,30 @@ let init =
 
 type snapshot = [ `Dir | `File of Stat.t ] Path.Map.t
 
-type t =
+type real =
   { dir : Path.Build.t
   ; snapshot : snapshot option
   ; corrections : Corrections.t
   ; deps : Path.Set.t option
   ; loc : Loc.t
   }
 
-let dir t = t.dir
-let map_path t p = Path.Build.append t.dir p
+type t =
+  | Sandboxed of real
+  | No_sandbox of { targets : Targets.Validated.t }
+
+let is_sandboxed = function
+  | Sandboxed _ -> true
+  | No_sandbox _ -> false
+;;
+
+let map_real_path t p = Path.Build.append t.dir p
+
+let map_path t p =
+  match t with
+  | Sandboxed t -> map_real_path t p
+  | No_sandbox _ -> p
+;;
 
 let copy_recursively =
   let chmod_file = Permissions.add Permissions.write in
@@ -105,7 +140,7 @@ let copy_recursively =
       ()
 ;;
 
-let create_dir t dir = Path.mkdir_p (Path.build (map_path t dir))
+let create_dir t dir = Path.mkdir_p (Path.build (map_real_path t dir))
 
 let create_dirs t ~dirs ~rule_dir =
   create_dir t rule_dir;
@@ -150,7 +185,7 @@ let link_deps t ~mode ~deps =
           "Action depends on source tree. All actions should depend on the copies in the \
            build directory instead."
           [ "path", Path.to_dyn path ]
-    | Some p -> link path (Path.build (map_path t p)))
+    | Some p -> link path (Path.build (map_real_path t p)))
 ;;
 
 let snapshot t =
@@ -177,7 +212,7 @@ let snapshot t =
   snapshot
 ;;
 
-let find_corrected_files (t : t) ~deps =
+let find_corrected_files (t : real) ~deps =
   (* CR-someday rgrinberg: fuse this step with deletion *)
   let start = Time.now () in
   let corrected =
@@ -244,7 +279,7 @@ let register_corrected_file_promotions t ~deps =
       })
 ;;
 
-let create
+let create_real
       ~mode
       (corrections : Corrections.t)
       ~rule_loc
@@ -311,7 +346,7 @@ let register_snapshot_promotion t (targets : Targets.Validated.t) ~old_snapshot
   let add_copy_file p = diffs := p :: !diffs in
   let deletes = ref [] in
   let add_delete what file = deletes := (what, in_source_tree file) :: !deletes in
-  let target_root_in_sandbox = map_path t targets.root in
+  let target_root_in_sandbox = map_real_path t targets.root in
   let () =
     Path.Map.iter2 old_snapshot new_snapshot ~f:(fun p before after ->
       if
@@ -379,7 +414,7 @@ let hint_delete_dir =
   ]
 ;;
 
-let move_targets_to_build_dir t ~should_be_skipped ~(targets : Targets.Validated.t)
+let move_real_targets_to_build_dir t ~should_be_skipped ~(targets : Targets.Validated.t)
   : unit Fiber.t
   =
   let open Fiber.O in
@@ -401,9 +436,9 @@ let move_targets_to_build_dir t ~should_be_skipped ~(targets : Targets.Validated
       targets
       ~file:(fun target ->
         if not (should_be_skipped target)
-        then rename_optional_file ~src:(map_path t target) ~dst:target)
+        then rename_optional_file ~src:(map_real_path t target) ~dst:target)
       ~dir:(fun target ->
-        let src_dir = map_path t target in
+        let src_dir = map_real_path t target in
         (match Path.Untracked.stat (Path.build target) with
          | Error (Unix.ENOENT, _, _) -> ()
          | Error e ->
@@ -432,24 +467,54 @@ let move_targets_to_build_dir t ~should_be_skipped ~(targets : Targets.Validated
     Dune_trace.Event.sandbox `Extract ~start ~stop ~queued:None t.loc ~dir:t.dir)
 ;;
 
+let move_targets_to_build_dir t ~should_be_skipped ~(targets : Targets.Validated.t) =
+  match t with
+  | No_sandbox _ -> Fiber.return ()
+  | Sandboxed t -> move_real_targets_to_build_dir t ~should_be_skipped ~targets
+;;
+
 let failed_to_delete_sandbox dir reason =
   User_error.raise
     [ Pp.textf "failed to delete sandbox in %s" (Path.Build.to_string_maybe_quoted dir)
     ; User_error.reason reason
     ]
 ;;
 
-let destroy t =
-  let open Fiber.O in
-  let+ start, stop, queued =
-    maybe_async (fun () ->
-      try Path.rm_rf ~chmod:true (Path.build t.dir) with
-      | Sys_error e -> failed_to_delete_sandbox t.dir (Pp.verbatim e)
-      | Unix.Unix_error (error, syscall, arg) ->
-        failed_to_delete_sandbox
-          t.dir
-          (Unix_error.Detailed.pp (Unix_error.Detailed.create error ~syscall ~arg)))
-  in
-  Dune_trace.emit ~buffered:true Sandbox (fun () ->
-    Dune_trace.Event.sandbox `Destroy ~start ~stop ~queued t.loc ~dir:t.dir)
+let destroy = function
+  | No_sandbox { targets } ->
+    Pending_targets.remove targets;
+    Fiber.return ()
+  | Sandboxed t ->
+    let open Fiber.O in
+    let+ start, stop, queued =
+      maybe_async (fun () ->
+        try Path.rm_rf ~chmod:true (Path.build t.dir) with
+        | Sys_error e -> failed_to_delete_sandbox t.dir (Pp.verbatim e)
+        | Unix.Unix_error (error, syscall, arg) ->
+          failed_to_delete_sandbox
+            t.dir
+            (Unix_error.Detailed.pp (Unix_error.Detailed.create error ~syscall ~arg)))
+    in
+    Dune_trace.emit ~buffered:true Sandbox (fun () ->
+      Dune_trace.Event.sandbox `Destroy ~start ~stop ~queued t.loc ~dir:t.dir)
+;;
+
+let with_ ~mode corrections ~rule_loc ~dirs ~deps ~rule_dir ~rule_digest ~targets ~f =
+  match mode with
+  | None ->
+    Pending_targets.add targets;
+    let sandbox = No_sandbox { targets } in
+    Fiber.finalize ~finally:(fun () -> destroy sandbox) (fun () -> f sandbox)
+  | Some mode ->
+    with_live_sandbox_slot ~f:(fun () ->
+      let open Fiber.O in
+      let* sandbox =
+        create_real ~mode corrections ~rule_loc ~dirs ~deps ~rule_dir ~rule_digest
+      in
+      let sandbox = Sandboxed sandbox in
+      (* CR-someday rgrinberg: Dynamic actions may discover dependencies
+         while this slot is held. If all sandbox slots are held by such
+         actions, sandboxed rules for the discovered dependencies cannot start.
+      *)
+      Fiber.finalize ~finally:(fun () -> destroy sandbox) (fun () -> f sandbox))
 ;;

src/dune_engine/sandbox.mli

@@ -4,29 +4,31 @@ open Import
 
 type t
 
-val dir : t -> Path.Build.t
+(** [is_sandboxed t] is [true] when [t] represents a real sandbox. *)
+val is_sandboxed : t -> bool
 
 (** [map_path t p] returns the path corresponding to [p] inside the sandbox. *)
 val map_path : t -> Path.Build.t -> Path.Build.t
 
-(** Run [f] while holding one global live-sandbox slot.
+(** Delete targets left behind by interrupted non-sandboxed actions. *)
+val cleanup_pending_targets : unit -> unit
 
-    The slot should be acquired before creating a sandbox and released after it
-    has been destroyed, so rules cannot pre-create unbounded sandboxes while
-    waiting for their actions to run. *)
-val with_live_sandbox_slot : f:(unit -> 'a Fiber.t) -> 'a Fiber.t
+(** Run [f] with a sandboxing context.
 
-(** Create a new sandbox containing [dirs] and copy or link dependencies [deps]
-    inside it. *)
-val create
-  :  mode:Sandbox_mode.some
+    [mode = None] is represented as a sandboxing context that maps paths to
+    themselves and only tracks pending targets for cleanup. It is not affected
+    by the live-sandbox throttle. *)
+val with_
+  :  mode:Sandbox_mode.some option
   -> Corrections.t
   -> rule_loc:Loc.t
   -> dirs:Path.Build.Set.t
   -> deps:Path.Set.t
   -> rule_dir:Path.Build.t
   -> rule_digest:Digest.t
-  -> t Fiber.t
+  -> targets:Targets.Validated.t
+  -> f:(t -> 'a Fiber.t)
+  -> 'a Fiber.t
 
 (** Move all targets created by the action from the sandbox to the build
     directory, skipping the files for which [should_be_skipped] returns [true].
@@ -37,5 +39,3 @@ val move_targets_to_build_dir
   -> should_be_skipped:(Path.Build.t -> bool)
   -> targets:Targets.Validated.t
   -> unit Fiber.t
-
-val destroy : t -> unit Fiber.t