Upgrade to thread-local-storage>=0.2

polytypic · polytypic · commit b910b92cbd1f · 2024-08-24T22:22:31.000+03:00
diff --git a/bench/bench_tls.ml b/bench/bench_tls.ml
@@ -1,6 +1,6 @@
 open Multicore_bench
 
-let key = Picos_thread.TLS.new_key (fun () -> -1)
+let key = Picos_thread.TLS.create ()
 
 let run_one ~budgetf ~n_domains ~op () =
   let n_ops =
@@ -9,7 +9,10 @@ let run_one ~budgetf ~n_domains ~op () =
 
   let n_ops_todo = Countdown.create ~n_domains () in
 
-  let init _ = Countdown.non_atomic_set n_ops_todo n_ops in
+  let init _ =
+    Picos_thread.TLS.set key (-1);
+    Countdown.non_atomic_set n_ops_todo n_ops
+  in
   let work domain_index () =
     match op with
     | `Get ->
@@ -18,7 +21,7 @@ let run_one ~budgetf ~n_domains ~op () =
           if n <> 0 then
             let rec loop n =
               if 0 < n then begin
-                let d = Picos_thread.TLS.get key in
+                let d = Picos_thread.TLS.get_exn key in
                 loop (n + d)
               end
               else work ()
diff --git a/dune-project b/dune-project
@@ -29,7 +29,7 @@
    (>= 0.1.0))
   ;; For OCaml 4 compatible defaults that are not used on OCaml 5 by default
   (thread-local-storage
-   (>= 0.1))
+   (>= 0.2))
   (mtime
    (>= 2.0.0))
   (psq
diff --git a/lib/picos/ocaml4/picos_ocaml.ml b/lib/picos/ocaml4/picos_ocaml.ml
@@ -14,11 +14,14 @@ module Handler = struct
     and await _ = error in
     E { context = (); handler = { current; spawn; yield; cancel_after; await } }
 
-  let key = Picos_thread.TLS.new_key @@ fun () -> default
-  let get () = Picos_thread.TLS.get key
+  let key = Picos_thread.TLS.create ()
+  let get () = Picos_thread.TLS.get_exn key
 
   let using handler context main =
-    let old = Picos_thread.TLS.get key in
+    let old =
+      try Picos_thread.TLS.get_exn key
+      with Picos_thread.TLS.Not_set -> default
+    in
     Picos_thread.TLS.set key (E { context; handler });
     match main (handler.current context) with
     | value ->
diff --git a/lib/picos_randos/picos_randos.ml b/lib/picos_randos/picos_randos.ml
@@ -53,7 +53,16 @@ type t = {
   mutable run : bool;
 }
 
-let fiber_key = Picos_thread.TLS.new_key @@ fun () -> ref Fiber.Maybe.nothing
+let fiber_key : Fiber.Maybe.t ref Picos_thread.TLS.t =
+  Picos_thread.TLS.create ()
+
+let get () =
+  match Picos_thread.TLS.get_exn fiber_key with
+  | p -> p
+  | exception Picos_thread.TLS.Not_set ->
+      let p = ref Fiber.Maybe.nothing in
+      Picos_thread.TLS.set fiber_key p;
+      p
 
 let rec next p t =
   match Collection.pop_exn t.ready with
@@ -158,21 +167,21 @@ let context ?fatal_exn_handler () =
   and current =
     Some
       (fun k ->
-        let p = Picos_thread.TLS.get fiber_key in
+        let p = Picos_thread.TLS.get_exn fiber_key in
         let fiber = Fiber.Maybe.to_fiber !p in
         Collection.push t.ready (Current (fiber, k));
         next p t)
   and yield =
     Some
       (fun k ->
-        let p = Picos_thread.TLS.get fiber_key in
+        let p = Picos_thread.TLS.get_exn fiber_key in
         let fiber = Fiber.Maybe.to_fiber !p in
         Collection.push t.ready (Continue (fiber, k));
         next p t)
   and return =
     Some
       (fun k ->
-        let p = Picos_thread.TLS.get fiber_key in
+        let p = Picos_thread.TLS.get_exn fiber_key in
         let fiber = Fiber.Maybe.to_fiber !p in
         Collection.push t.ready (Return (fiber, k));
         next p t)
@@ -182,7 +191,7 @@ let context ?fatal_exn_handler () =
     function
     | Fiber.Current -> t.current
     | Fiber.Spawn r ->
-        let p = Picos_thread.TLS.get fiber_key in
+        let p = Picos_thread.TLS.get_exn fiber_key in
         let fiber = Fiber.Maybe.to_fiber !p in
         if Fiber.is_canceled fiber then t.yield
         else begin
@@ -193,7 +202,7 @@ let context ?fatal_exn_handler () =
         end
     | Fiber.Yield -> t.yield
     | Computation.Cancel_after r -> begin
-        let p = Picos_thread.TLS.get fiber_key in
+        let p = Picos_thread.TLS.get_exn fiber_key in
         let fiber = Fiber.Maybe.to_fiber !p in
         if Fiber.is_canceled fiber then t.yield
         else
@@ -211,7 +220,7 @@ let context ?fatal_exn_handler () =
     | Trigger.Await trigger ->
         Some
           (fun k ->
-            let p = Picos_thread.TLS.get fiber_key in
+            let p = Picos_thread.TLS.get_exn fiber_key in
             let fiber = Fiber.Maybe.to_fiber !p in
             if Fiber.try_suspend fiber trigger fiber k t.resume then next p t
             else begin
@@ -221,14 +230,14 @@ let context ?fatal_exn_handler () =
     | _ -> None
   and retc () =
     Atomic.decr t.num_alive_fibers;
-    let p = Picos_thread.TLS.get fiber_key in
+    let p = Picos_thread.TLS.get_exn fiber_key in
     next p t
   in
   t
 
 let runner_on_this_thread t =
   Select.check_configured ();
-  next (Picos_thread.TLS.get fiber_key) t
+  next (get ()) t
 
 let rec await t =
   if !(t.num_waiters_non_zero) then begin
@@ -257,7 +266,7 @@ let run_fiber ?context:t_opt fiber main =
     t.run <- true;
     Mutex.unlock t.mutex;
     Collection.push t.ready (Spawn (fiber, main));
-    next (Picos_thread.TLS.get fiber_key) t;
+    next (get ()) t;
     Mutex.lock t.mutex;
     await t
   end
diff --git a/lib/picos_select/picos_select.ml b/lib/picos_select/picos_select.ml
@@ -101,9 +101,7 @@ let cleared =
   Computation.return computation Cleared;
   computation
 
-let intr_key =
-  Picos_thread.TLS.new_key @@ fun () : [ `Req ] tdt ->
-  invalid_arg "has not been configured"
+let intr_key : [ `Req ] tdt Picos_thread.TLS.t = Picos_thread.TLS.create ()
 
 let key =
   Picos_domain.DLS.new_key @@ fun () ->
@@ -289,7 +287,7 @@ let handle_signal signal =
        Computation.return r.computation r.value
   end
   else if signal = config.intr_sig then
-    let (Req r) = Picos_thread.TLS.get intr_key in
+    let (Req r) = Picos_thread.TLS.get_exn intr_key in
     Computation.return r.computation Signaled
 
 let reconfigure_signal_handlers () =
diff --git a/lib/picos_thread/picos_thread.mli b/lib/picos_thread/picos_thread.mli
@@ -12,21 +12,28 @@ module TLS : sig
       [threads.posix] library, is not available, this will use
       {!Picos_domain.DLS}. *)
 
-  type 'a key
-  (** Represents a key for storing values of type ['a] in storage associated
-      with threads. *)
-
-  val new_key : (unit -> 'a) -> 'a key
-  (** [new_key compute] allocates a new key for associating values in storage
-      associated with threads.  The initial value for each thread is [compute]d
-      by calling the given function if the [key] is {{!get}read} before it has
-      been {{!set}written}. *)
-
-  val get : 'a key -> 'a
-  (** [get key] returns the value associated with the [key] in the storage
-      associated with the current thread. *)
-
-  val set : 'a key -> 'a -> unit
-  (** [set key value] sets the [value] associated with the [key] in the storage
-      associated with the current thread. *)
+  type 'a t
+  (** Represents a key for associating values with threads. *)
+
+  val create : unit -> 'a t
+  (** [create ()] allocates a new key for associating values with threads.
+
+      ⚠️ Keys should not be created dynamically as each key will potentially
+      increase the space taken by every thread. *)
+
+  exception Not_set
+  (** Exception raised by {!get_exn} when no value is associated with the
+      specified key for the current thread. *)
+
+  val get_exn : 'a t -> 'a
+  (** [get_exn key] returns the value associated with the specified key for the
+      current thread or raises {!Not_set} in case no value has been {!set} for
+      the key.
+
+      ⚠️ The {!Not_set} exception is raised with no backtrace.  Always catch the
+      exception unless it is known that a value has been set. *)
+
+  val set : 'a t -> 'a -> unit
+  (** [set key value] associates the value with the specified key for the
+      current thread. *)
 end
diff --git a/lib/picos_thread/picos_thread.none.ml b/lib/picos_thread/picos_thread.none.ml
@@ -1,3 +1,13 @@
-let is_main_thread = Picos_domain.is_main_domain
+open Picos_domain
 
-module TLS = Picos_domain.DLS
+let is_main_thread = is_main_domain
+
+module TLS = struct
+  type 'a t = 'a DLS.key
+
+  exception Not_set
+
+  let create () = DLS.new_key @@ fun () -> raise_notrace Not_set
+  let get_exn = DLS.get
+  let set = DLS.set
+end
diff --git a/picos.opam b/picos.opam
@@ -11,7 +11,7 @@ bug-reports: "https://github.com/ocaml-multicore/picos/issues"
 depends: [
   "dune" {>= "3.14"}
   "backoff" {>= "0.1.0"}
-  "thread-local-storage" {>= "0.1"}
+  "thread-local-storage" {>= "0.2"}
   "mtime" {>= "2.0.0"}
   "psq" {>= "0.2.1"}
   "multicore-magic" {>= "2.3.0"}
