From bfd70dc5c25fc156b301931970d518f19361ff37 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Tue, 27 Feb 2024 22:11:50 -0500
Subject: [PATCH] fix fiber: spawn sub-fiber with a copy of current local
 storage

---
 src/core/task_local_storage.ml  |  1 +
 src/core/task_local_storage.mli |  1 +
 src/fib/fiber.ml                | 13 ++++++++-----
 src/fib/fiber.mli               |  8 ++++++++
 4 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/src/core/task_local_storage.ml b/src/core/task_local_storage.ml
index 87ca1424..c6920261 100644
--- a/src/core/task_local_storage.ml
+++ b/src/core/task_local_storage.ml
@@ -65,6 +65,7 @@ module Private_ = struct
     let k_storage = k_ls_values
     let[@inline] create () = [||]
     let copy = Array.copy
+    let[@inline] copy_of_current () = copy @@ !(get_cur_ ())
     let dummy = [||]
   end
 end
diff --git a/src/core/task_local_storage.mli b/src/core/task_local_storage.mli
index c2ce778a..a1f96fab 100644
--- a/src/core/task_local_storage.mli
+++ b/src/core/task_local_storage.mli
@@ -54,6 +54,7 @@ module Private_ : sig
     val k_storage : t ref option Thread_local_storage_.key
     val create : unit -> t
     val copy : t -> t
+    val copy_of_current : unit -> t
     val dummy : t
   end
 end
diff --git a/src/fib/fiber.ml b/src/fib/fiber.ml
index aa724f95..6d888c8b 100644
--- a/src/fib/fiber.ml
+++ b/src/fib/fiber.ml
@@ -154,7 +154,7 @@ let add_child_ ~protect (self : _ t) (child : _ t) =
     ()
   done
 
-let spawn_ ~on (f : _ -> 'a) : 'a t =
+let spawn_ ~ls ~on (f : _ -> 'a) : 'a t =
   let id = Handle.generate_fresh () in
   let res, _promise = Fut.make () in
   let fib =
@@ -169,7 +169,6 @@ let spawn_ ~on (f : _ -> 'a) : 'a t =
   let run () =
     (* make sure the fiber is accessible from inside itself *)
     Task_local_storage.set k_current_fiber (Some (Any fib));
-    assert (Task_local_storage.get k_current_fiber |> Option.is_some);
     try
       let res = f () in
       resolve_ok_ fib res
@@ -179,11 +178,11 @@ let spawn_ ~on (f : _ -> 'a) : 'a t =
       resolve_as_failed_ fib ebt
   in
 
-  Runner.run_async on run;
+  Runner.run_async ?ls on run;
 
   fib
 
-let[@inline] spawn_top ~on f : _ t = spawn_ ~on f
+let[@inline] spawn_top ~on f : _ t = spawn_ ~ls:None ~on f
 
 let[@inline] self () : any =
   match Task_local_storage.get k_current_fiber with
@@ -194,7 +193,9 @@ let spawn_link ~protect f : _ t =
   match Task_local_storage.get k_current_fiber with
   | None -> failwith "Fiber.spawn_link: must be run from inside a fiber."
   | Some (Any parent) ->
-    let child = spawn_ ~on:parent.runner f in
+    (* spawn [f()] with a copy of our local storage *)
+    let ls = Task_local_storage.Private_.Storage.copy_of_current () in
+    let child = spawn_ ~ls:(Some ls) ~on:parent.runner f in
     add_child_ ~protect parent child;
     child
 
@@ -235,6 +236,8 @@ let with_self_cancel_callback cb (k : unit -> 'a) : 'a =
   Fun.protect k ~finally:(fun () -> remove_top_cancel_cb_ self)
 
 let[@inline] await self = Fut.await self.res
+let[@inline] wait_block self = Fut.wait_block self.res
+let[@inline] wait_block_exn self = Fut.wait_block_exn self.res
 
 module Suspend_ = Moonpool.Private.Suspend_
 
diff --git a/src/fib/fiber.mli b/src/fib/fiber.mli
index 80863bc7..ebbe5e96 100644
--- a/src/fib/fiber.mli
+++ b/src/fib/fiber.mli
@@ -64,6 +64,14 @@ val is_success : _ t -> bool
 val await : 'a t -> 'a
 (** [await fib] is like [Fut.await (res fib)] *)
 
+val wait_block_exn : 'a t -> 'a
+(** [wait_block_exn fib] is [Fut.wait_block_exn (res fib)].
+    {b NOTE}: See {!Fut.wait_block} for warnings about deadlocks. *)
+
+val wait_block : 'a t -> 'a Fut.or_error
+(** [wait_block fib] is [Fut.wait_block (res fib)].
+    {b NOTE}: See {!Fut.wait_block} for warnings about deadlocks. *)
+
 val check_if_cancelled : unit -> unit
 (** Check if the current fiber is cancelled, in which case this raises.
     Must be run from inside a fiber.