feat: optional dep on hmap, for inheritable FLS data

.github/workflows/main.yml

@@ -37,7 +37,7 @@ jobs:
       - run: opam exec -- dune build @install
 
       # install some depopts
-      - run: opam install thread-local-storage trace
+      - run: opam install thread-local-storage trace hmap
         if: matrix.ocaml-compiler == '5.2'
 
       - run: opam exec -- dune build --profile=release --force @install @runtest

dune-project

@@ -29,6 +29,7 @@
     (>= 1.9.0)
     :with-test)))
  (depopts
+   hmap
    (trace (>= 0.6)))
  (tags
   (thread pool domain futures fork-join)))

moonpool.opam

@@ -20,6 +20,7 @@ depends: [
   "mdx" {>= "1.9.0" & with-test}
 ]
 depopts: [
+  "hmap"
   "trace" {>= "0.6"}
 ]
 build: [

src/fib/dune

@@ -2,7 +2,11 @@
  (name moonpool_fib)
  (public_name moonpool.fib)
  (synopsis "Fibers and structured concurrency for Moonpool")
- (libraries moonpool picos)
+ (libraries moonpool picos
+            (select hmap_fls.ml from
+                    (hmap -> hmap_fls.real.ml)
+                    (-> hmap_fls.dummy.ml)
+                    ))
  (enabled_if
   (>= %{ocaml_version} 5.0))
  (flags :standard -open Moonpool_private -open Moonpool)

src/fib/fiber.ml

@@ -257,8 +257,10 @@ let spawn_ ~parent ~runner (f : unit -> 'a) : 'a t =
   let comp = Picos.Computation.create () in
   let pfiber = PF.create ~forbid:false comp in
 
-  (* inherit FLS from parent, if present *)
-  Option.iter (fun (p : _ t) -> PF.copy_fls p.pfiber pfiber) parent;
+  (* copy local hmap from parent, if present *)
+  Option.iter
+    (fun (p : _ t) -> Hmap_fls.Private_hmap_fls_.copy_fls p.pfiber pfiber)
+    parent;
 
   (match parent with
   | Some p when is_closed p -> failwith "spawn: nursery is closed"
@@ -328,3 +330,5 @@ let yield () : unit =
     check_if_cancelled_ self;
     PF.yield ();
     check_if_cancelled_ self
+
+include Hmap_fls

src/fib/fiber.mli

@@ -147,3 +147,9 @@ val spawn_ignore : ?protect:bool -> (unit -> _) -> unit
 (** [spawn_ignore f] is [ignore (spawn f)].
   The fiber will still affect termination of the parent, ie. the
   parent will exit only after this new fiber exits. *)
+
+(** {2 Local [Hmap.t]} *)
+
+include module type of struct
+  include Hmap_fls
+end

src/fib/hmap_fls.dummy.ml

@@ -0,0 +1,7 @@
+(**/**)
+
+module Private_hmap_fls_ = struct
+  let copy_fls _ _ = ()
+end
+
+(**/**)

src/fib/hmap_fls.real.ml

@@ -0,0 +1,47 @@
+open Moonpool.Private.Types_
+
+open struct
+  module FLS = Picos.Fiber.FLS
+end
+
+(** A local hmap, inherited in children fibers *)
+let k_local_hmap : Hmap.t FLS.t = FLS.create ()
+
+(** Access the local [hmap], or an empty one if not set *)
+let[@inline] get_local_hmap () : Hmap.t =
+  let fiber = get_current_fiber_exn () in
+  FLS.get fiber ~default:Hmap.empty k_local_hmap
+
+let[@inline] set_local_hmap (h : Hmap.t) : unit =
+  let fiber = get_current_fiber_exn () in
+  FLS.set fiber k_local_hmap h
+
+let[@inline] update_local_hmap (f : Hmap.t -> Hmap.t) : unit =
+  let fiber = get_current_fiber_exn () in
+  let h = FLS.get fiber ~default:Hmap.empty k_local_hmap in
+  let h = f h in
+  FLS.set fiber k_local_hmap h
+
+(** @raise Invalid_argument if not present *)
+let get_exn (k : 'a Hmap.key) : 'a =
+  let h = get_local_hmap () in
+  Hmap.get k h
+
+let get_opt (k : 'a Hmap.key) : 'a option =
+  let h = get_local_hmap () in
+  Hmap.find k h
+
+let[@inline] set (k : 'a Hmap.key) (v : 'a) : unit =
+  update_local_hmap (Hmap.add k v)
+
+(**/**)
+
+module Private_hmap_fls_ = struct
+  (** Copy the hmap from f1.fls to f2.fls *)
+  let copy_fls (f1 : Picos.Fiber.t) (f2 : Picos.Fiber.t) : unit =
+    match FLS.get_exn f1 k_local_hmap with
+    | exception FLS.Not_set -> ()
+    | hmap -> FLS.set f2 k_local_hmap hmap
+end
+
+(**/**)