Repository: ocaml-multicore/picos
Pull request #56 (2024-03-17 10:18:41)

Use `TLS` to store handlers on OCaml 4
Discussion with @c-cube after merging #54 quickly revealed that we actually want to allow multiple schedulers even on OCaml 4.  So, this PR changes the design for OCaml 4 support such that a "handler" for the effects is stored in `TLS`.  This allows different handlers to be used in OCaml 4 in different threads.  The mechanism is also implemented on OCaml 5 using effect handlers, but it is generally better to implement effect handlers directly in OCaml 5 &mdash; the API is really optimized for OCaml 4.


Patch:
diff --git a/bench/scheduler.ocaml4.ml b/bench/scheduler.ocaml4.ml
index 1c54aaf..03a9af7 100644
--- a/bench/scheduler.ocaml4.ml
+++ b/bench/scheduler.ocaml4.ml
@@ -1,2 +1,6 @@
-let () = Picos.set_picos_implementation Picos_threaded.implementation
-let run action = action ()
+open Picos
+
+let run action =
+  Handler.using Picos_threaded.handler
+    (Picos_threaded.create ~forbid:false (Computation.create ()))
+    action
diff --git a/lib/picos/bootstrap/picos_bootstrap.ml b/lib/picos/bootstrap/picos_bootstrap.ml
index 7edbd85..047b171 100644
--- a/lib/picos/bootstrap/picos_bootstrap.ml
+++ b/lib/picos/bootstrap/picos_bootstrap.ml
@@ -308,18 +308,14 @@ module Fiber = struct
   end
 end
 
-module type Implementation = sig
-  module Fiber : sig
-    val current : unit -> Fiber.t
-    val spawn : forbid:bool -> 'a Computation.t -> (unit -> unit) list -> unit
-    val yield : unit -> unit
-  end
-
-  module Computation : sig
-    val cancel_after : 'a Computation.t -> seconds:float -> Exn_bt.t -> unit
-  end
-
-  module Trigger : sig
-    val await : Trigger.t -> Exn_bt.t option
-  end
+module Handler = struct
+  type 'c t = {
+    current : 'c -> Fiber.t;
+    spawn :
+      'a. 'c -> forbid:bool -> 'a Computation.t -> (unit -> unit) list -> unit;
+    yield : 'c -> unit;
+    cancel_after :
+      'a. 'c -> 'a Computation.t -> seconds:float -> Exn_bt.t -> unit;
+    await : 'c -> Trigger.t -> Exn_bt.t option;
+  }
 end
diff --git a/lib/picos/ocaml4/dune b/lib/picos/ocaml4/dune
index adc2bc8..3ed27da 100644
--- a/lib/picos/ocaml4/dune
+++ b/lib/picos/ocaml4/dune
@@ -6,4 +6,5 @@
   (< %{ocaml_version} 5.0.0))
  (libraries
   (re_export picos.exn_bt)
+  picos.tls
   picos_bootstrap))
diff --git a/lib/picos/ocaml4/picos_ocaml.ml b/lib/picos/ocaml4/picos_ocaml.ml
index 4cee3fe..a5bb4dd 100644
--- a/lib/picos/ocaml4/picos_ocaml.ml
+++ b/lib/picos/ocaml4/picos_ocaml.ml
@@ -1,62 +1,59 @@
 open Picos_bootstrap
 
 let error () =
-  raise (Sys_error "Picos.set_picos_implementation must be called exactly once")
-
-type implementation = {
-  current : unit -> Fiber.t;
-  spawn : 'a. forbid:bool -> 'a Computation.t -> (unit -> unit) list -> unit;
-  yield : unit -> unit;
-  cancel_after : 'a. 'a Computation.t -> seconds:float -> Exn_bt.t -> unit;
-  await : Trigger.t -> Exn_bt.t option;
-}
-
-let implementation =
-  let current = error
-  and spawn ~forbid:_ _ _ = error ()
-  and yield = error
-  and cancel_after _ ~seconds:_ _ = error ()
-  and await _ = error () in
-  ref { current; spawn; yield; cancel_after; await }
+  raise (Sys_error "Picos.Handler.using not called for current thread")
+
+module Handler = struct
+  type entry = E : { context : 'a; handler : 'a Handler.t } -> entry
+
+  let default =
+    let current _ = error ()
+    and spawn _ ~forbid:_ _ _ = error ()
+    and yield _ = error ()
+    and cancel_after _ _ ~seconds:_ _ = error ()
+    and await _ _ = error () in
+    E { context = (); handler = { current; spawn; yield; cancel_after; await } }
+
+  let key = Picos_tls.new_key @@ fun () -> default
+  let get () = Picos_tls.get key
+
+  let using handler context thunk =
+    let old = Picos_tls.get key in
+    Picos_tls.set key (E { context; handler });
+    match thunk () with
+    | value ->
+        Picos_tls.set key old;
+        value
+    | exception exn ->
+        Picos_tls.set key old;
+        raise exn
+end
 
 module Trigger = struct
-  let await t = if Trigger.is_initial t then !implementation.await t else None
+  let await t =
+    if Trigger.is_initial t then
+      let (E r) = Handler.get () in
+      r.handler.await r.context t
+    else None
 end
 
 module Fiber = struct
-  let current () = !implementation.current ()
+  let current () =
+    let (E r) = Handler.get () in
+    r.handler.current r.context
 
   let spawn ~forbid computation mains =
-    let _ = current () in
-    !implementation.spawn ~forbid computation mains
+    let (E r) = Handler.get () in
+    r.handler.spawn r.context ~forbid computation mains
 
   let yield () =
-    let _ = current () in
-    !implementation.yield ()
+    let (E r) = Handler.get () in
+    r.handler.yield r.context
 end
 
 module Computation = struct
   let cancel_after computation ~seconds exn_bt =
     Computation.check_non_negative seconds;
-    let _ = Fiber.current () in
-    !implementation.cancel_after computation ~seconds exn_bt
+    let (E r) = Handler.get () in
+    r.handler.cancel_after r.context computation ~seconds exn_bt
 end
-
-let[@poll error] [@inline never] set_picos_implementation i =
-  !implementation.current == error
-  && begin
-       implementation := i;
-       true
-     end
-
-let set_picos_implementation (module I : Implementation) =
-  let i =
-    {
-      current = I.Fiber.current;
-      spawn = I.Fiber.spawn;
-      yield = I.Fiber.yield;
-      cancel_after = I.Computation.cancel_after;
-      await = I.Trigger.await;
-    }
-  in
-  if not (set_picos_implementation i) then error ()
diff --git a/lib/picos/ocaml5/picos_ocaml.ml b/lib/picos/ocaml5/picos_ocaml.ml
index 36e2d37..65b4c9a 100644
--- a/lib/picos/ocaml5/picos_ocaml.ml
+++ b/lib/picos/ocaml5/picos_ocaml.ml
@@ -51,6 +51,46 @@ module Computation = struct
     Effect.perform (Cancel_after { seconds; exn_bt; computation })
 end
 
-let set_picos_implementation _ =
-  raise
-    (Sys_error "Picos.set_picos_implementation must be called only on OCaml 4")
+module Handler = struct
+  let discontinue k exn =
+    Effect.Deep.discontinue_with_backtrace k exn (Printexc.get_raw_backtrace ())
+
+  let using (h : _ Handler.t) c =
+    let current =
+      Some
+        (fun k ->
+          match h.current c with
+          | fiber -> Effect.Deep.continue k fiber
+          | exception exn -> discontinue k exn)
+    and yield =
+      Some
+        (fun k ->
+          match h.yield c with
+          | unit -> Effect.Deep.continue k unit
+          | exception exn -> discontinue k exn)
+    in
+    let effc (type a) :
+        a Effect.t -> ((a, _) Effect.Deep.continuation -> _) option = function
+      | Fiber.Current -> current
+      | Fiber.Spawn r ->
+          Some
+            (fun k ->
+              match h.spawn c ~forbid:r.forbid r.computation r.mains with
+              | unit -> Effect.Deep.continue k unit
+              | exception exn -> discontinue k exn)
+      | Fiber.Yield -> yield
+      | Trigger.Await trigger ->
+          Some (fun k -> Effect.Deep.continue k (h.await c trigger))
+      | Computation.Cancel_after r ->
+          Some
+            (fun k ->
+              match
+                h.cancel_after c r.computation ~seconds:r.seconds r.exn_bt
+              with
+              | unit -> Effect.Deep.continue k unit
+              | exception exn -> discontinue k exn)
+      | _ -> None
+    in
+    let handler = Effect.Deep.{ retc = Fun.id; exnc = raise; effc } in
+    fun main -> Effect.Deep.match_with main () handler
+end
diff --git a/lib/picos/picos.common.ml b/lib/picos/picos.common.ml
index 57b99e7..9c5777a 100644
--- a/lib/picos/picos.common.ml
+++ b/lib/picos/picos.common.ml
@@ -29,6 +29,7 @@ module Fiber = struct
   include Picos_ocaml.Fiber
 end
 
-module type Implementation = Picos_bootstrap.Implementation
-
-let set_picos_implementation = Picos_ocaml.set_picos_implementation
+module Handler = struct
+  include Picos_bootstrap.Handler
+  include Picos_ocaml.Handler
+end
diff --git a/lib/picos/picos.mli b/lib/picos/picos.mli
index d424974..357fb5a 100644
--- a/lib/picos/picos.mli
+++ b/lib/picos/picos.mli
@@ -223,15 +223,19 @@
 
     {[open Picos]}
 
-    and set the Picos {{!implementations} implementation} to use on OCaml 4
+    and define a simple scheduler
 
     {[
-      if String.starts_with ~prefix:"4." Sys.ocaml_version then
-        Picos.set_picos_implementation Picos_threaded.implementation
+      let run main =
+        Handler.using
+          Picos_threaded.handler
+          (Picos_threaded.create ~forbid:false (Computation.create ()))
+          main
     ]}
 
-    to {{!Picos_threaded.implementation} the basic thread based implementation}
-    for the example code snippets in this documentation.
+    using {{!Picos_threaded.handler} the basic thread based implementation} that
+    comes with Picos for running the example code snippets in this
+    documentation.
 
     {3 Auxiliary modules} *)
 
@@ -254,25 +258,27 @@ module Trigger : sig
 
       Here is a simple example:
 
-      {@ocaml version<5.0.0[
-        let trigger = Trigger.create () in
-
-        let signaler =
-          Domain.spawn @@ fun () ->
-            Trigger.signal trigger
-        in
-        let finally () =
-          Domain.join signaler
-        in
-        Fun.protect ~finally @@ fun () ->
-
-        match Trigger.await trigger with
-        | None ->
-          (* We were resumed normally. *)
-          ()
-        | Some exn_bt ->
-          (* We were canceled. *)
-          Exn_bt.raise exn_bt
+      {[
+        run begin fun () ->
+          let trigger = Trigger.create () in
+
+          let signaler =
+            Domain.spawn @@ fun () ->
+              Trigger.signal trigger
+          in
+          let finally () =
+            Domain.join signaler
+          in
+          Fun.protect ~finally @@ fun () ->
+
+          match Trigger.await trigger with
+          | None ->
+            (* We were resumed normally. *)
+            ()
+          | Some exn_bt ->
+            (* We were canceled. *)
+            Exn_bt.raise exn_bt
+        end
       ]}
 
       ⚠️ Typically we need to cleanup after {!await}, but in the above example we
@@ -335,8 +341,8 @@ module Trigger : sig
       ℹ️ The behavior is that, {i unless [await] can return immediately},
 
       - on OCaml 5, [await] will perform the {!Await} effect, and
-      - on OCaml 4, [await] will call the {{!Implementation.Trigger.await}
-        [await] operation} of the {{!set_picos_implementation} implementation}.
+      - on OCaml 4, [await] will call the [await] operation of the {{!Handler}
+        current handler}.
 
       @raise Invalid_argument if the trigger was in the awaiting state, which
         means that multiple concurrent calls of [await] are being made. *)
@@ -530,39 +536,41 @@ module Computation : sig
 
       Here is an example:
 
-      {@ocaml version<5.0.0[
-        let computation =
-          Computation.create ()
-        in
-        let computer =
-          Domain.spawn @@ fun () ->
-            let rec fib i =
-              Computation.check computation;
-              if i <= 1 then
-                i
-              else
-                fib (i - 1) + fib (i - 2)
-            in
-            Computation.capture computation
-              fib 10
-        in
-        let finally () =
-          Domain.join computer
-        in
-        Fun.protect ~finally @@ fun () ->
-
-        let canceler =
-          Domain.spawn @@ fun () ->
-            Unix.sleepf 0.1;
-            Computation.cancel computation
-            @@ Exn_bt.get_callstack 2 Exit
-        in
-        let finally () =
-          Domain.join canceler
-        in
-        Fun.protect ~finally @@ fun () ->
-
-        Computation.await computation
+      {[
+        run begin fun () ->
+          let computation =
+            Computation.create ()
+          in
+          let computer =
+            Domain.spawn @@ fun () ->
+              let rec fib i =
+                Computation.check computation;
+                if i <= 1 then
+                  i
+                else
+                  fib (i - 1) + fib (i - 2)
+              in
+              Computation.capture computation
+                fib 10
+          in
+          let finally () =
+            Domain.join computer
+          in
+          Fun.protect ~finally @@ fun () ->
+
+          let canceler =
+            Domain.spawn @@ fun () ->
+              Unix.sleepf 0.1;
+              Computation.cancel computation
+              @@ Exn_bt.get_callstack 2 Exit
+          in
+          let finally () =
+            Domain.join canceler
+          in
+          Fun.protect ~finally @@ fun () ->
+
+          Computation.await computation
+        end
       ]}
 
       In this framework, a fiber is always associated with {{!Fiber.computation}
@@ -668,12 +676,9 @@ module Computation : sig
       ℹ️ The behavior is that [cancel_after] first checks that [seconds] is not
       negative, and then
 
-      - on OCaml 5, [cancel_after] will perform the {!Cancel_after}
-        effect, and
-      - on OCaml 4, [cancel_after] will next call {!Fiber.current} and will
-        finally call the {{!Implementation.Computation.cancel_after}
-        [cancel_after] operation} of the {{!set_picos_implementation}
-        implementation}.
+      - on OCaml 5, [cancel_after] will perform the {!Cancel_after} effect, and
+      - on OCaml 4, [cancel_after] will call the [cancel_after] operation of the
+        {{!Handler} current handler}.
 
       @raise Invalid_argument if [seconds] is negative or too large as
         determined by the scheduler. *)
@@ -805,9 +810,8 @@ module Fiber : sig
       ℹ️ The behavior is that
 
       - on OCaml 5, [yield] perform the {!Yield} effect, and
-      - on OCaml 4, [yield] will next call {!Fiber.current} and will finally
-        call the {{!Implementation.Fiber.yield} [yield] operation} of the
-        {{!set_picos_implementation} implementation}. *)
+      - on OCaml 4, [yield] will call the [yield] operation of the {{!Handler}
+        current handler}. *)
 
   (** {2 Interface for spawning} *)
 
@@ -830,9 +834,8 @@ module Fiber : sig
       ℹ️ The behavior is that
 
       - on OCaml 5, [spawn] performs the {!Spawn} effect, and
-      - on OCaml 4, [spawn] will next call {!Fiber.current} and will finally
-        call the {{!Implementation.Fiber.spawn} [spawn] operation} of the
-        {{!set_picos_implementation} implementation}. *)
+      - on OCaml 4, [spawn] will call the [spawn] operation of the {{!Handler}
+        current handler}. *)
 
   (** {2 Interface for current fiber} *)
 
@@ -854,9 +857,8 @@ module Fiber : sig
       ℹ️ The behavior is that
 
       - on OCaml 5, [current] performs the {!Current} effect, and
-      - on OCaml 4, [current] will call the {{!Implementation.Fiber.current}
-        [current] operation} of the {{!set_picos_implementation}
-        implementation}. *)
+      - on OCaml 4, [current] will call the [current] operation of the
+        {{!Handler} current handler}. *)
 
   val has_forbidden : t -> bool
   (** [has_forbidden fiber] determines whether the fiber {!forbid}s or
@@ -1043,47 +1045,38 @@ module Fiber : sig
       next. *)
 end
 
-(** {2 Implementations}
-
-    OCaml 4 does not have effect handlers.  It is, however, possible to
-    implement the effects based operations of Picos in terms of threads. *)
+module Handler : sig
+  (** Handler for the effects based operations of Picos for OCaml 4. *)
+
+  type 'c t = {
+    current : 'c -> Fiber.t;  (** See {!Picos.Fiber.current}. *)
+    spawn :
+      'a. 'c -> forbid:bool -> 'a Computation.t -> (unit -> unit) list -> unit;
+        (** See {!Picos.Fiber.spawn}. *)
+    yield : 'c -> unit;  (** See {!Picos.Fiber.yield}. *)
+    cancel_after :
+      'a. 'c -> 'a Computation.t -> seconds:float -> Exn_bt.t -> unit;
+        (** See {!Picos.Computation.cancel_after}. *)
+    await : 'c -> Trigger.t -> Exn_bt.t option;
+        (** See {!Picos.Trigger.await}. *)
+  }
+  (** A record of implementations of the primitive effects based operations of
+      Picos.  The operations take a context of type ['c] as an argument. *)
+
+  val using : 'c t -> 'c -> (unit -> 'a) -> 'a
+  (** [using handler context thunk] sets the [handler] and the [context] for the
+      handler of the primitive effects based operations of Picos while running
+      [thunk].
 
-module type Implementation = sig
-  (** Signature for an implementation of the primitive effects based operations of
-      Picos.
-
-      ℹ️ On OCaml 4 most of the effects based operations will call
-      {!Fiber.current} implicitly to check for cancelation before finally
-      calling the set {{!set_picos_implementation} implementation}. *)
-
-  (** Implements primitive effects based operations of {!Picos.Fiber}. *)
-  module Fiber : sig
-    val current : unit -> Fiber.t
-    (** See {!Picos.Fiber.current}. *)
-
-    val spawn : forbid:bool -> 'a Computation.t -> (unit -> unit) list -> unit
-    (** See {!Picos.Fiber.spawn}. *)
-
-    val yield : unit -> unit
-    (** See {!Picos.Fiber.yield}. *)
-  end
+      ℹ️ The behavior is that
 
-  (** Implements primitive effects based operations of {!Picos.Computation}. *)
-  module Computation : sig
-    val cancel_after : 'a Computation.t -> seconds:float -> Exn_bt.t -> unit
-    (** See {!Picos.Computation.cancel_after}. *)
-  end
+      - on OCaml 4, [using] stores the [handler] in {{!Picos_tls} [TLS]}, which
+        allows the operations to be accessed during the execution of the
+        [thunk], and
+      - on OCaml 5, [using] runs [thunk] with a deep effect handler that
+        delegates to the operations of the [handler].
 
-  (** Implements primitive effects based operations of {!Picos.Trigger}. *)
-  module Trigger : sig
-    val await : Trigger.t -> Exn_bt.t option
-    (** See {!Picos.Trigger.await}. *)
-  end
+      ⚠️ While this works on OCaml 5, you usually want to use a scheduler that
+      implements an effect handler directly, because that is likely to perform
+      better. *)
 end
-
-val set_picos_implementation : (module Implementation) -> unit
-(** [set_picos_implementation (module Implementation)] sets the implementation
-    of the primitive effects based operations of Picos.
-
-    ⚠️ This must be called exactly once in a program only on OCaml 4 before using
-    any effects based operations of Picos. *)
diff --git a/lib/picos/ptmc/dune b/lib/picos/ptmc/dune
new file mode 100644
index 0000000..373343b
--- /dev/null
+++ b/lib/picos/ptmc/dune
@@ -0,0 +1,10 @@
+(library
+ (name picos_ptmc)
+ (public_name picos.ptmc)
+ (libraries
+  picos.tls
+  (select
+   picos_ptmc.ml
+   from
+   (threads.posix -> picos_ptmc.common.ml)
+   (-> picos_ptmc.common.ml))))
diff --git a/lib/picos/ptmc/picos_ptmc.mli b/lib/picos/ptmc/picos_ptmc.mli
new file mode 100644
index 0000000..303fdeb
--- /dev/null
+++ b/lib/picos/ptmc/picos_ptmc.mli
@@ -0,0 +1,21 @@
+(** A per thread {!Mutex} and {!Condition} for {!Picos}. *)
+
+type t
+(** Pair of a {!Mutex} and a {!Condition} for a thread. *)
+
+val get : unit -> t
+(** [get ()] obtains the pair of a {!Mutex} and a {!Condition} for the current
+    thread. *)
+
+val lock : t -> unit
+(** [lock ptmc] locks the {!Mutex} of the [ptmc]. *)
+
+val unlock : t -> unit
+(** [unlock ptmc] unlocks the {!Mutex} of the [ptmc]. *)
+
+val broadcast : t -> unit
+(** [broadcast ptmc] locks and unlocks the {!Mutex} of the [ptmc] and then
+    broadcasts on the {!Condition} of the [ptmc]. *)
+
+val wait : t -> unit
+(** [wait ptmc] waits on the {!Condition} of the [ptmc]. *)
diff --git a/lib/picos/threaded/dune b/lib/picos/threaded/dune
index f4efbf4..e51edc6 100644
--- a/lib/picos/threaded/dune
+++ b/lib/picos/threaded/dune
@@ -6,8 +6,8 @@
  (libraries
   (re_export picos)
   (re_export picos.exn_bt)
-  picos.tls
   picos.domain
+  picos.ptmc
   (select
    select.ml
    from
diff --git a/lib/picos/threaded/picos_threaded.ml b/lib/picos/threaded/picos_threaded.ml
index a8f365a..d3ea9d2 100644
--- a/lib/picos/threaded/picos_threaded.ml
+++ b/lib/picos/threaded/picos_threaded.ml
@@ -1,82 +1,70 @@
 open Picos
 
-let implementation =
-  (module struct
-    module Per_thread = struct
-      type t = { mutex : Mutex.t; condition : Condition.t; fiber : Fiber.t }
+type t = Fiber.t
 
-      open struct
-        let create fiber =
-          let mutex = Mutex.create () and condition = Condition.create () in
-          { mutex; condition; fiber }
+let create = Fiber.create
 
-        let per_thread_key =
-          Picos_tls.new_key @@ fun () ->
-          create (Fiber.create ~forbid:false (Computation.create ()))
-      end
+let block trigger ptmc fiber =
+  Picos_ptmc.lock ptmc;
+  match
+    while not (Trigger.is_signaled trigger) do
+      Picos_ptmc.wait ptmc
+    done
+  with
+  | () ->
+      Picos_ptmc.unlock ptmc;
+      Fiber.canceled fiber
+  | exception exn ->
+      (* Condition.wait may be interrupted by asynchronous exceptions and we
+         must make sure to unlock even in that case. *)
+      Picos_ptmc.unlock ptmc;
+      raise exn
 
-      let get () = Picos_tls.get per_thread_key
-      let set fiber = Picos_tls.set per_thread_key (create fiber)
-    end
-
-    module Trigger = struct
-      let release _ _ (per_thread : Per_thread.t) =
-        Mutex.lock per_thread.mutex;
-        Mutex.unlock per_thread.mutex;
-        Condition.broadcast per_thread.condition
-
-      let block trigger (per_thread : Per_thread.t) =
-        Mutex.lock per_thread.mutex;
-        match
-          while not (Trigger.is_signaled trigger) do
-            Condition.wait per_thread.condition per_thread.mutex
-          done
-        with
-        | () ->
-            Mutex.unlock per_thread.mutex;
-            Fiber.canceled per_thread.fiber
-        | exception exn ->
-            (* Condition.wait may be interrupted by asynchronous exceptions and we
-               must make sure to unlock even in that case. *)
-            Mutex.unlock per_thread.mutex;
-            raise exn
+let release _ _ ptmc = Picos_ptmc.broadcast ptmc
 
-      let[@alert "-handler"] await trigger =
-        let per_thread = Per_thread.get () in
-        if Fiber.has_forbidden per_thread.fiber then
-          if Trigger.on_signal trigger () per_thread release then
-            block trigger per_thread
-          else None
-        else if Fiber.try_attach per_thread.fiber trigger then
-          if Trigger.on_signal trigger () per_thread release then
-            block trigger per_thread
-          else begin
-            Fiber.detach per_thread.fiber trigger;
-            Fiber.canceled per_thread.fiber
-          end
-        else begin
-          Trigger.dispose trigger;
-          Fiber.canceled per_thread.fiber
-        end
+let[@alert "-handler"] rec await fiber trigger =
+  if Fiber.has_forbidden fiber then
+    (* We could also have stored the per thread mutex and condition in the
+       context and avoid getting it here, but this is likely cheap enough at
+       this point anyway and makes the context trivial. *)
+    let ptmc = Picos_ptmc.get () in
+    if Trigger.on_signal trigger () ptmc release then block trigger ptmc fiber
+    else None
+  else if Fiber.try_attach fiber trigger then
+    let ptmc = Picos_ptmc.get () in
+    if Trigger.on_signal trigger () ptmc release then block trigger ptmc fiber
+    else begin
+      Fiber.detach fiber trigger;
+      Fiber.canceled fiber
     end
+  else begin
+    Trigger.dispose trigger;
+    Fiber.canceled fiber
+  end
 
-    module Computation = Select
+and cancel_after :
+    type a. t -> a Computation.t -> seconds:float -> Exn_bt.t -> unit =
+ fun fiber computation ~seconds exn_bt ->
+  Fiber.check fiber;
+  Select.cancel_after computation ~seconds exn_bt
 
-    module Fiber = struct
-      let current () =
-        let fiber = (Per_thread.get ()).fiber in
-        Fiber.check fiber;
-        fiber
+and current fiber =
+  Fiber.check fiber;
+  fiber
 
-      let spawn ~forbid computation mains =
-        mains
-        |> List.iter @@ fun main ->
-           Systhreads.create
-             (fun () ->
-               main (Per_thread.set (Fiber.create ~forbid computation)))
-             ()
-           |> ignore
+and yield fiber =
+  Fiber.check fiber;
+  Systhreads.yield ()
 
-      let yield () = Systhreads.yield ()
-    end
-  end : Implementation)
+and spawn :
+    type a. t -> forbid:bool -> a Computation.t -> (unit -> unit) list -> unit =
+ fun fiber ~forbid computation mains ->
+  Fiber.check fiber;
+  mains
+  |> List.iter @@ fun main ->
+     Systhreads.create
+       (fun () -> Picos.Handler.using handler (create ~forbid computation) main)
+       ()
+     |> ignore
+
+and handler = Handler.{ current; spawn; yield; cancel_after; await }
diff --git a/lib/picos/threaded/picos_threaded.mli b/lib/picos/threaded/picos_threaded.mli
index 51a6b1e..b4582bb 100644
--- a/lib/picos/threaded/picos_threaded.mli
+++ b/lib/picos/threaded/picos_threaded.mli
@@ -3,27 +3,30 @@
 
     Briefly:
 
-    - [Fiber.current] uses {{!Picos_tls} [TLS]} to store the current fiber.
-    - [Fiber.spawn] creates a new [Thread] for each fiber.
-    - [Fiber.yield] just calls [Thread.yield].
-    - [Computation.cancel_after] uses {{!Picos_domain.DLS} [DLS]} to store a
-      priority queue of timeouts and a per-domain {{!Picos_select} background
-      timeout thread} that runs a [Unix.select] loop to cancel computations.
-    - [Trigger.await] uses {{!Picos_tls} [TLS]} to store a [Mutex] and
-      [Condition] to suspend the thread.
-
-    The operations initialize their resources, per-thread, and per-domain state
-    and {{!Picos_select} background timeout thread}, only when actually used.
-    If the default {{!Picos.Computation.cancel_after} [cancel_after]} is not
-    used, no background timeout thread will be created and no per-domain state
-    will be used.  If none of the defaults are used, no per-thread state will be
-    used.
+    - [current] returns the current fiber from the per thread state.
+    - [spawn] creates a new [Thread] for each fiber.
+    - [yield] just calls [Thread.yield].
+    - [cancel_after] uses {{!Picos_domain.DLS} [DLS]} to store a priority queue
+      of timeouts and a per-domain {{!Picos_select} background timeout thread}
+      that runs a [Unix.select] loop to cancel computations.
+    - [await] uses {{!Picos_ptmc} a per thread mutex and condition} to suspend
+      the thread.
 
     ⚠️ This implementation is probably suitable for simple applications that do
-    not spawn a lot.  If an application uses a lot of short lived fibers, then a
-    more sophisticated implementation using some sort of thread pool will likely
-    perform significantly better. *)
+    not spawn a lot of fibers.  If an application uses a lot of short lived
+    fibers, then a more sophisticated implementation using some sort of thread
+    pool will likely perform significantly better. *)
+
+open Picos
+
+type t
+(** The per thread state for a {!handler}. *)
+
+val create : forbid:bool -> 'a Computation.t -> t
+(** [create ~forbid computation] creates a new per thread state for a fiber with
+    the given [computation] and propagation of cancelation [forbid]den as
+    given. *)
 
-val implementation : (module Picos.Implementation)
-(** A basic implementation of the effects based operations of {!Picos} using
-    [Thread]s and a [Unix.select] based IO event loop. *)
+val handler : t Handler.t
+(** A basic handler of the effects based operations of {!Picos} using [Thread]s
+    and a [Unix.select] based IO event loop. *)


Test patch:
diff --git a/test/lib/foundation/mutex_and_condition.ml b/lib/picos/ptmc/picos_ptmc.common.ml
similarity index 69%
rename from test/lib/foundation/mutex_and_condition.ml
rename to lib/picos/ptmc/picos_ptmc.common.ml
index 6ef6262..343fbad 100644
--- a/test/lib/foundation/mutex_and_condition.ml
+++ b/lib/picos/ptmc/picos_ptmc.common.ml
@@ -1,10 +1,10 @@
 type t = { mutex : Mutex.t; condition : Condition.t }
 
-let mutex_and_condition_key =
-  Thread_local_storage.new_key @@ fun () ->
+let key =
+  Picos_tls.new_key @@ fun () ->
   { mutex = Mutex.create (); condition = Condition.create () }
 
-let get () = Thread_local_storage.get mutex_and_condition_key
+let get () = Picos_tls.get key
 let lock t = Mutex.lock t.mutex
 let unlock t = Mutex.unlock t.mutex
 
diff --git a/test/lib/foundation/dune b/test/lib/foundation/dune
index 7952a50..7c5557b 100644
--- a/test/lib/foundation/dune
+++ b/test/lib/foundation/dune
@@ -1,3 +1,3 @@
 (library
  (name foundation)
- (libraries picos backoff threads.posix multicore-magic thread-local-storage))
+ (libraries picos backoff multicore-magic))
diff --git a/test/lib/foundation/mutex_and_condition.mli b/test/lib/foundation/mutex_and_condition.mli
deleted file mode 100644
index 3304fbb..0000000
--- a/test/lib/foundation/mutex_and_condition.mli
+++ /dev/null
@@ -1,17 +0,0 @@
-type t
-(** *)
-
-val get : unit -> t
-(** *)
-
-val lock : t -> unit
-(** *)
-
-val unlock : t -> unit
-(** *)
-
-val broadcast : t -> unit
-(** *)
-
-val wait : t -> unit
-(** *)
diff --git a/test/lib/schedulers/dune b/test/lib/schedulers/dune
index e4538d6..c8a2448 100644
--- a/test/lib/schedulers/dune
+++ b/test/lib/schedulers/dune
@@ -2,4 +2,4 @@
  (name schedulers)
  (enabled_if
   (>= %{ocaml_version} 5.0.0))
- (libraries picos foundation))
+ (libraries picos picos.ptmc foundation))
diff --git a/test/lib/schedulers/fifos.ml b/test/lib/schedulers/fifos.ml
index 58a3c38..0cb08b2 100644
--- a/test/lib/schedulers/fifos.ml
+++ b/test/lib/schedulers/fifos.ml
@@ -5,7 +5,7 @@ type t = {
   ready : (unit -> unit) Mpsc_queue.t;
   needs_wakeup : bool Atomic.t;
   num_alive_fibers : int ref;
-  mc : Mutex_and_condition.t;
+  mc : Picos_ptmc.t;
   on_unhandled : unit Try_handle.t option;
 }
 
@@ -15,13 +15,11 @@ let rec next t =
   | exception Mpsc_queue.Empty ->
       if !(t.num_alive_fibers) <> 0 then begin
         if Atomic.get t.needs_wakeup then begin
-          Mutex_and_condition.lock t.mc;
-          match
-            if Atomic.get t.needs_wakeup then Mutex_and_condition.wait t.mc
-          with
-          | () -> Mutex_and_condition.unlock t.mc
+          Picos_ptmc.lock t.mc;
+          match if Atomic.get t.needs_wakeup then Picos_ptmc.wait t.mc with
+          | () -> Picos_ptmc.unlock t.mc
           | exception exn ->
-              Mutex_and_condition.unlock t.mc;
+              Picos_ptmc.unlock t.mc;
               raise exn
         end
         else Atomic.set t.needs_wakeup true;
@@ -33,7 +31,7 @@ let run ?on_unhandled ~forbid main =
     let ready = Mpsc_queue.create () in
     let needs_wakeup = Atomic.make false in
     let num_alive_fibers = ref 1 in
-    let mc = Mutex_and_condition.get () in
+    let mc = Picos_ptmc.get () in
     { ready; needs_wakeup; num_alive_fibers; mc; on_unhandled }
   in
   let resume trigger fiber k =
@@ -58,7 +56,7 @@ let run ?on_unhandled ~forbid main =
     if
       Atomic.get t.needs_wakeup
       && Atomic.compare_and_set t.needs_wakeup true false
-    then Mutex_and_condition.broadcast t.mc
+    then Picos_ptmc.broadcast t.mc
   in
   let rec fork fiber main =
     let current = Some (fun k -> Fiber.continue fiber k fiber)
diff --git a/test/test_scheduler.ocaml4.ml b/test/test_scheduler.ocaml4.ml
index 12d10e7..bf43bbb 100644
--- a/test/test_scheduler.ocaml4.ml
+++ b/test/test_scheduler.ocaml4.ml
@@ -1,2 +1,6 @@
-let () = Picos.set_picos_implementation Picos_threaded.implementation
-let run main = main ()
+open Picos
+
+let run main =
+  Handler.using Picos_threaded.handler
+    (Picos_threaded.create ~forbid:false (Computation.create ()))
+    main