Further tweaks to reduce stack usage

Author: polytypic

lib/picos/bootstrap/picos_bootstrap.ml

@@ -489,6 +489,10 @@ module Fiber = struct
         Sys.opaque_identity (Obj.magic (Array.unsafe_get r.fls key) : a)
       else default
 
+    let reserve (type a) (Fiber r : fiber) (key : a t) =
+      let fls = r.fls in
+      if Array.length fls <= key then r.fls <- grow fls key
+
     let set (type a) (Fiber r : fiber) (key : a t) (value : a) =
       let fls = r.fls in
       if key < Array.length fls then

lib/picos/picos.mli

@@ -1001,6 +1001,17 @@ module Fiber : sig
 
         ⚠️ It is only safe to call [remove] from the fiber itself or when the
         fiber is known not to be running. *)
+
+    val reserve : fiber -> 'a t -> unit
+    (** [reserve fiber key] ensures that sufficient space has been allocated to
+        associate a value with the specified [key] such that a subsequent {!set}
+        with the [key] will not allocate.
+
+        ℹ️ This can be used to optimize the population of the FLS and to avoid
+        performing memory allocations in critical sections.
+
+        ⚠️ It is only safe to call [reserve] from the fiber itself or when the
+        fiber is known not to be running. *)
   end
 
   (** {2 Interface for spawning} *)

lib/picos_std.finally/picos_std_finally.ml

@@ -51,12 +51,15 @@ let forbidden release x =
          We also do not reraise the exception! *)
       release x
 
-(** This helps to reduce CPU stack usage with the native compiler. *)
-let[@inline never] release_and_reraise exn release x =
+let[@inline never] release_and_reraise exn x release =
   let bt = Printexc.get_raw_backtrace () in
   forbidden release x;
   Printexc.raise_with_backtrace exn bt
 
+let[@inline never] release_and_return value x release =
+  forbidden release x;
+  value
+
 (* *)
 
 let rec drop instance =
@@ -70,16 +73,19 @@ let rec drop instance =
       end
       else drop instance
 
-(** This helps to reduce CPU stack usage with the native compiler. *)
-let[@inline never] drop_and_reraise bt instance exn =
+let[@inline never] drop_and_reraise_as bt instance exn =
   drop instance;
   Printexc.raise_with_backtrace exn bt
 
+let[@inline never] drop_and_reraise exn instance =
+  let bt = Printexc.get_raw_backtrace () in
+  drop_and_reraise_as bt instance exn
+
 (* *)
 
-let await_transferred_or_dropped instance =
+let await_transferred_or_dropped instance result =
   match Atomic.get instance with
-  | Transferred | Dropped -> ()
+  | Transferred | Dropped -> result
   | Borrowed as case ->
       (* This should be impossible as [let@ _ = borrow _ in _] should have
          restored the state. *)
@@ -88,12 +94,18 @@ let await_transferred_or_dropped instance =
       match Trigger.await r.transferred_or_dropped with
       | None ->
           (* We release in case we could not wait. *)
-          drop instance
+          drop instance;
+          result
       | Some (exn, bt) ->
           (* We have been canceled, so we release. *)
-          drop_and_reraise bt instance exn
+          drop_and_reraise_as bt instance exn
     end
 
+let[@inline never] instantiate instance scope =
+  match scope instance with
+  | result -> await_transferred_or_dropped instance result
+  | exception exn -> drop_and_reraise exn instance
+
 let[@inline never] instantiate release acquire scope =
   let instance =
     Sys.opaque_identity
@@ -108,13 +120,7 @@ let[@inline never] instantiate release acquire scope =
   (* After this point there must be no allocations before [acquire ()]. *)
   let (Resource r : (_, [ `Resource ]) tdt) = Obj.magic (Atomic.get instance) in
   r.resource <- acquire ();
-  match scope instance with
-  | result ->
-      await_transferred_or_dropped instance;
-      result
-  | exception exn ->
-      let bt = Printexc.get_raw_backtrace () in
-      drop_and_reraise bt instance exn
+  instantiate instance scope
 
 (* *)
 
@@ -130,12 +136,8 @@ let[@inline never] rec transfer from scope =
           Trigger.signal r.transferred_or_dropped;
           scope into
         with
-        | result ->
-            await_transferred_or_dropped into;
-            result
-        | exception exn ->
-            let bt = Printexc.get_raw_backtrace () in
-            drop_and_reraise bt into exn
+        | result -> await_transferred_or_dropped into result
+        | exception exn -> drop_and_reraise exn into
       end
       else transfer from scope
 
@@ -172,25 +174,24 @@ let[@inline never] rec move from scope =
         | result ->
             forbidden r.release r.resource;
             result
-        | exception exn -> release_and_reraise exn r.release r.resource
+        | exception exn -> release_and_reraise exn r.resource r.release
       end
       else move from scope
 
 (* *)
 
+let[@inline never] finally x scope release =
+  match scope x with
+  | y -> release_and_return y x release
+  | exception exn -> release_and_reraise exn x release
+
 let[@inline never] finally release acquire scope =
   let x = acquire () in
-  match scope x with
-  | y ->
-      forbidden release x;
-      y
-  | exception exn -> release_and_reraise exn release x
+  finally x scope release
 
 let[@inline never] lastly action scope =
   match scope () with
-  | value ->
-      forbidden action ();
-      value
-  | exception exn -> release_and_reraise exn action ()
+  | value -> release_and_return value () action
+  | exception exn -> release_and_reraise exn () action
 
 external ( let@ ) : ('a -> 'b) -> 'a -> 'b = "%apply"

lib/picos_std.structured/bundle.ml

@@ -53,21 +53,43 @@ let get_flock fiber =
   | Bundle _ as t -> t
   | Nothing -> no_flock ()
 
-let await (type a) (Bundle r as t : t) fiber packed canceler outer
-    (pass : a pass) =
+let await (Bundle r as t : t) fiber packed canceler outer =
   decr t;
   Fiber.set_computation fiber packed;
   let forbid = Fiber.exchange fiber ~forbid:true in
   Trigger.await r.finished |> ignore;
   Fiber.set fiber ~forbid;
-  begin
-    match pass with FLS -> Fiber.FLS.set fiber flock_key outer | Arg -> ()
-  end;
+  if Fiber.FLS.get fiber flock_key ~default:Nothing != outer then
+    Fiber.FLS.set fiber flock_key outer;
   let (Packed parent) = packed in
   Computation.detach parent canceler;
   Control.Errors.check r.errors;
   Fiber.check fiber
 
+let[@inline never] raised exn t fiber packed canceler outer =
+  let bt = Printexc.get_raw_backtrace () in
+  error t exn bt;
+  await t fiber packed canceler outer;
+  Printexc.raise_with_backtrace exn bt
+
+let[@inline never] returned value (Bundle r as t : t) fiber packed canceler
+    outer =
+  let config = Atomic.get r.config in
+  if config land config_terminated_bit <> 0 then begin
+    let callstack =
+      let n = (config land config_callstack_mask) lsr config_callstack_shift in
+      if n = 0 then None else Some n
+    in
+    terminate ?callstack t
+  end;
+  await t fiber packed canceler outer;
+  value
+
+let join_after_realloc x fn t fiber packed canceler outer =
+  match fn x with
+  | value -> returned value t fiber packed canceler outer
+  | exception exn -> raised exn t fiber packed canceler outer
+
 let join_after_pass (type a) ?callstack ?on_return (fn : a -> _) (pass : a pass)
     =
   (* The sequence of operations below ensures that nothing is leaked. *)
@@ -93,38 +115,20 @@ let join_after_pass (type a) ?callstack ?on_return (fn : a -> _) (pass : a pass)
     Bundle { config; bundle; errors; finished }
   in
   let fiber = Fiber.current () in
-  let outer =
-    match pass with
-    | Arg -> Nothing
-    | FLS -> Fiber.FLS.get fiber flock_key ~default:Nothing
-  in
+  let outer = Fiber.FLS.get fiber flock_key ~default:Nothing in
+  begin
+    match pass with FLS -> Fiber.FLS.reserve fiber flock_key | Arg -> ()
+  end;
   let (Packed parent as packed) = Fiber.get_computation fiber in
   let (Packed bundle) = r.bundle in
   let canceler = Computation.attach_canceler ~from:parent ~into:bundle in
   (* Ideally there should be no poll point betweem [attach_canceler] and the
-     [match ... with] below. *)
-  match
-    Fiber.set_computation fiber r.bundle;
-    fn (match pass with FLS -> Fiber.FLS.set fiber flock_key t | Arg -> t)
-  with
-  | value ->
-      let config = Atomic.get r.config in
-      if config land config_terminated_bit <> 0 then begin
-        let callstack =
-          let n =
-            (config land config_callstack_mask) lsr config_callstack_shift
-          in
-          if n = 0 then None else Some n
-        in
-        terminate ?callstack t
-      end;
-      await t fiber packed canceler outer pass;
-      value
-  | exception exn ->
-      let bt = Printexc.get_raw_backtrace () in
-      error t exn bt;
-      await t fiber packed canceler outer pass;
-      Printexc.raise_with_backtrace exn bt
+     [match ... with] in [join_after_realloc]. *)
+  Fiber.set_computation fiber r.bundle;
+  let x : a =
+    match pass with FLS -> Fiber.FLS.set fiber flock_key t | Arg -> t
+  in
+  join_after_realloc x fn t fiber packed canceler outer
 
 let rec incr (Bundle r as t : t) backoff =
   let before = Atomic.get r.config in
@@ -136,14 +140,12 @@ let finish (Bundle { bundle = Packed bundle; _ } as t : t) canceler =
   Computation.detach bundle canceler;
   decr t
 
-(** This helps to reduce CPU stack usage with the native compiler. *)
 let[@inline never] raised exn child t canceler =
   let bt = Printexc.get_raw_backtrace () in
   Computation.cancel child exn bt;
   error t exn bt;
   finish t canceler
 
-(** This helps to reduce CPU stack usage with the native compiler. *)
 let[@inline never] returned value child t canceler =
   Computation.return child value;
   finish t canceler
@@ -181,14 +183,12 @@ let fork_as_promise_pass (type a) (Bundle r as t : t) thunk (pass : a pass) =
     decr t;
     raise canceled_exn
 
-(** This helps to reduce CPU stack usage with the native compiler. *)
 let[@inline never] raised_flock exn fiber =
   let t = get_flock fiber in
   let bt = Printexc.get_raw_backtrace () in
   error t exn bt;
   decr t
 
-(** This helps to reduce CPU stack usage with the native compiler. *)
 let[@inline never] raised_bundle exn t =
   error t exn (Printexc.get_raw_backtrace ());
   decr t