Repository: ocaml-multicore/picos
Pull request #210 (2024-08-09 08:36:42)

Redesign `Spawn` and introduce concept of fatal errors
This PR essentially redesigns the `Fiber.spawn` mechanism to allow `FLS` to be populated before a fiber is spawned. Being able to populate the `FLS` before spawn can be used to reduce the amount of allocations required (i.e. this results in smaller `main` closure for the fiber) and also allows some use cases to be implemented more easily.

This also removes the early idea of giving a list of `main` functions, which allowed multiple fibers to be spawned with a single effect.  It seems that moving the allocation of the fiber record to happen before `spawn` also speeds up a single spawn.  This is nice, because, in most cases, one spawns one fiber at a time.

This also introduces the concept of _fatal errors_.  Previously, if the `main` function of a fiber raised an exception, it was considered to be undefined behavior.  This is a subtle way to say that `main` had to always handle all exceptions, because raising an exception had no behavior that one could rely upon.  This is now changed such that raising an exception from the `main` function is considered a _fatal error_, which is guaranteed to stop the entire program.

On the other hand, these changes move the requirement that fiber records are unique to be maintained outside of the scheduler.


Patch:
diff --git a/bench/bench_fib.ocaml5.ml b/bench/bench_fib.ocaml5.ml
index d880078..6712af4 100644
--- a/bench/bench_fib.ocaml5.ml
+++ b/bench/bench_fib.ocaml5.ml
@@ -8,8 +8,9 @@ let rec exp_fib i =
   if i < 2 then i
   else
     let computation = Computation.create ~mode:`LIFO () in
-    let main () = Computation.return computation (exp_fib (i - 2)) in
-    Fiber.spawn ~forbid:false computation [ main ];
+    let fiber = Fiber.create ~forbid:false computation in
+    let main _ = Computation.capture computation exp_fib (i - 2) in
+    Fiber.spawn fiber main;
     let f1 = exp_fib (i - 1) in
     let f2 = Computation.await computation in
     f1 + f2
diff --git a/bench/bench_memory.ml b/bench/bench_memory.ml
index fc862fc..325e1f7 100644
--- a/bench/bench_memory.ml
+++ b/bench/bench_memory.ml
@@ -112,7 +112,7 @@ let run_suite ~budgetf:_ =
         let latch = Latch.create 1 in
         let bytes =
           measure_live_bytes @@ fun () ->
-          let main () =
+          let main _ =
             try
               Latch.incr latch;
               while true do
@@ -121,7 +121,7 @@ let run_suite ~budgetf:_ =
             with Exit -> Latch.decr latch
           in
           for _ = 1 to n do
-            Fiber.spawn ~forbid:false computation [ main ]
+            Fiber.spawn (Fiber.create ~forbid:false computation) main
           done;
           Fiber.yield ()
         in
diff --git a/bench/bench_spawn.ml b/bench/bench_spawn.ml
index d82601f..c169899 100644
--- a/bench/bench_spawn.ml
+++ b/bench/bench_spawn.ml
@@ -5,7 +5,7 @@ let factor =
   Util.iter_factor
   * if String.starts_with ~prefix:"4." Sys.ocaml_version then 1 else 10
 
-let run_one ~budgetf ~at_a_time () =
+let run_one ~budgetf () =
   let n_spawns = 10 * factor in
 
   let init _ = () in
@@ -13,25 +13,22 @@ let run_one ~budgetf ~at_a_time () =
   let work _ () =
     let counter = ref n_spawns in
     let computation = Computation.create ~mode:`LIFO () in
-    for _ = 1 to n_spawns / at_a_time do
-      let main () =
-        let n = !counter - 1 in
-        counter := n;
-        if n = 0 then Computation.finish computation
-      in
-      let mains = List.init at_a_time @@ fun _ -> main in
-      Fiber.spawn ~forbid:false computation mains
+    let computation_packed = Computation.Packed computation in
+    let main _ =
+      let n = !counter - 1 in
+      counter := n;
+      if n = 0 then Computation.finish computation
+    in
+    for _ = 1 to n_spawns do
+      let fiber = Fiber.create_packed ~forbid:false computation_packed in
+      Fiber.spawn fiber main
     done;
     Computation.await computation
   in
 
-  let config = Printf.sprintf "%d at a time" at_a_time in
+  let config = "with packed computation" in
   Times.record ~budgetf ~n_domains:1 ~n_warmups:1 ~n_runs_min:1 ~init ~wrap
     ~work ()
   |> Times.to_thruput_metrics ~n:n_spawns ~singular:"spawn" ~config
 
-let run_suite ~budgetf =
-  if Sys.int_size <= 32 then []
-  else
-    [ 1; 2; 4; 8 ]
-    |> List.concat_map @@ fun at_a_time -> run_one ~budgetf ~at_a_time ()
+let run_suite ~budgetf = if Sys.int_size <= 32 then [] else run_one ~budgetf ()
diff --git a/lib/picos/bootstrap/picos_bootstrap.ml b/lib/picos/bootstrap/picos_bootstrap.ml
index 510aec7..870fc84 100644
--- a/lib/picos/bootstrap/picos_bootstrap.ml
+++ b/lib/picos/bootstrap/picos_bootstrap.ml
@@ -385,8 +385,7 @@ end
 module Handler = struct
   type 'c t = {
     current : 'c -> Fiber.t;
-    spawn :
-      'a. 'c -> forbid:bool -> 'a Computation.t -> (unit -> unit) list -> unit;
+    spawn : 'c -> Fiber.t -> (Fiber.t -> unit) -> unit;
     yield : 'c -> unit;
     cancel_after :
       'a. 'c -> 'a Computation.t -> seconds:float -> Exn_bt.t -> unit;
diff --git a/lib/picos/intf.ocaml5.ml b/lib/picos/intf.ocaml5.ml
index 83b5cbf..0ad43d2 100644
--- a/lib/picos/intf.ocaml5.ml
+++ b/lib/picos/intf.ocaml5.ml
@@ -132,13 +132,13 @@ module type Fiber = sig
       associated with the fiber has been canceled the scheduler is free to
       discontinue the fiber immediately before spawning new fibers.
 
-      The scheduler is free to run the newly created fibers on any domain and
+      The scheduler is free to run the newly created fiber on any domain and
       decide which fiber to give priority to.
 
-      ⚠️ The scheduler should guarantee that, when [Spawn] returns normally, all
-      of the [mains] will eventually be called by the scheduler and, when
-      [Spawn] raises an exception, none of the [mains] will be called.  In other
-      words, [Spawn] should check cancelation just once and be all or nothing.
+      ⚠️ The scheduler should guarantee that, when [Spawn] returns normally, the
+      given [main] will eventually be called by the scheduler and, when [Spawn]
+      raises an exception, the [main] will not be called.  In other words,
+      [Spawn] should check cancelation just once and be all or nothing.
       Furthermore, in case a newly spawned fiber is canceled before its main is
       called, the scheduler must still call the main.  This allows a program to
       ensure, i.e. keep track of, that all fibers it spawns are terminated
@@ -146,10 +146,5 @@ module type Fiber = sig
       properly. *)
   type _ Effect.t +=
     private
-    | Spawn : {
-        forbid : bool;
-        computation : 'a computation;
-        mains : (unit -> unit) list;
-      }
-        -> unit Effect.t
+    | Spawn : { fiber : t; main : t -> unit } -> unit Effect.t
 end
diff --git a/lib/picos/ocaml4/picos_ocaml.ml b/lib/picos/ocaml4/picos_ocaml.ml
index ad17437..cc15eeb 100644
--- a/lib/picos/ocaml4/picos_ocaml.ml
+++ b/lib/picos/ocaml4/picos_ocaml.ml
@@ -1,26 +1,26 @@
 open Picos_bootstrap
 
-let error () =
+let[@inline never] error _ =
   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
+    let current = error
+    and spawn _ _ = 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_thread.TLS.new_key @@ fun () -> default
   let get () = Picos_thread.TLS.get key
 
-  let using handler context thunk =
+  let using handler context main =
     let old = Picos_thread.TLS.get key in
     Picos_thread.TLS.set key (E { context; handler });
-    match thunk () with
+    match main (handler.current context) with
     | value ->
         Picos_thread.TLS.set key old;
         value
@@ -42,9 +42,9 @@ module Fiber = struct
     let (E r) = Handler.get () in
     r.handler.current r.context
 
-  let spawn ~forbid computation mains =
+  let spawn fiber main =
     let (E r) = Handler.get () in
-    r.handler.spawn r.context ~forbid computation mains
+    r.handler.spawn r.context fiber main
 
   let yield () =
     let (E r) = Handler.get () in
diff --git a/lib/picos/ocaml5/picos_ocaml.ml b/lib/picos/ocaml5/picos_ocaml.ml
index 8e84513..6a8085a 100644
--- a/lib/picos/ocaml5/picos_ocaml.ml
+++ b/lib/picos/ocaml5/picos_ocaml.ml
@@ -25,15 +25,9 @@ module Fiber = struct
   let current () = Effect.perform Current
 
   type _ Effect.t +=
-    | Spawn : {
-        forbid : bool;
-        computation : 'a Computation.t;
-        mains : (unit -> unit) list;
-      }
-        -> unit Effect.t
+    | Spawn : { fiber : Fiber.t; main : Fiber.t -> unit } -> unit Effect.t
 
-  let spawn ~forbid computation mains =
-    Effect.perform @@ Spawn { forbid; computation; mains }
+  let spawn fiber main = Effect.perform @@ Spawn { fiber; main }
 
   type _ Effect.t += Yield : unit Effect.t
 
@@ -78,7 +72,7 @@ module Handler = struct
       | Fiber.Spawn r ->
           Some
             (fun k ->
-              match h.spawn c ~forbid:r.forbid r.computation r.mains with
+              match h.spawn c r.fiber r.main with
               | unit -> Effect.Deep.continue k unit
               | exception exn -> discontinue k exn)
       | Fiber.Yield -> yield
@@ -95,5 +89,5 @@ module Handler = struct
       | _ -> None
     in
     let handler = Effect.Deep.{ retc = Fun.id; exnc = raise; effc } in
-    fun main -> Effect.Deep.match_with main () handler
+    fun main -> Effect.Deep.match_with main (h.current c) handler
 end
diff --git a/lib/picos/picos.mli b/lib/picos/picos.mli
index 7b2ce18..4c61edf 100644
--- a/lib/picos/picos.mli
+++ b/lib/picos/picos.mli
@@ -758,30 +758,6 @@ module Fiber : sig
   (** [sleep ~seconds] suspends the current fiber for the specified number of
       [seconds]. *)
 
-  (** {2 Interface for spawning} *)
-
-  val spawn : forbid:bool -> 'a Computation.t -> (unit -> unit) list -> unit
-  (** [spawn ~forbid computation mains] starts new fibers by performing the
-      {!Spawn} effect.  The fibers will share the same [computation] and start
-      with {{!Fiber.has_forbidden} propagation of cancelation forbidden or
-      permitted} depending on the [forbid] flag.
-
-      ℹ️ Any {{!Computation} computation}, including the computation of the
-      current fiber, may be passed as the computation for new fibers.  Higher
-      level libraries are free to implement the desired structuring principles.
-
-      ⚠️ Behavior is undefined if any function in [mains] raises an exception.
-      For example, raising an exception might terminate the whole application
-      (recommended, but not required) or the exception might be ignored.  In
-      other words, the caller {i must} arrange for the computation to be
-      completed and errors reported in a desired manner.
-
-      ℹ️ The behavior is that
-
-      - on OCaml 5, [spawn] performs the {!Spawn} effect, and
-      - on OCaml 4, [spawn] will call the [spawn] operation of the {{!Handler}
-        current handler}. *)
-
   (** {2 Interface for current fiber} *)
 
   type t
@@ -966,6 +942,39 @@ module Fiber : sig
         ⚠️ It is only safe to call [set] from the fiber itself. *)
   end
 
+  (** {2 Interface for spawning} *)
+
+  val create_packed : forbid:bool -> Computation.packed -> t
+  (** [create_packed ~forbid packed] creates a new fiber record. *)
+
+  val create : forbid:bool -> 'a Computation.t -> t
+  (** [create ~forbid computation] is equivalent to
+      {{!create_packed} [create_packed ~forbid (Computation.Packed computation)]}. *)
+
+  val spawn : t -> (t -> unit) -> unit
+  (** [spawn fiber main] starts a new fiber by performing the {!Spawn} effect.
+
+      ⚠️ Fiber records must be unique and the caller of [spawn] must make sure
+      that a specific {{!fiber} fiber} record is not reused.  Failure to ensure
+      that fiber records are unique will break concurrent abstractions written
+      on top the the Picos interface.
+
+      ⚠️ If the [main] function raises an exception it is considered a {i fatal
+      error}.  A fatal error should effectively either directly exit the program
+      or stop the entire scheduler, without discontinuing existing fibers, and
+      force the invocations of the scheduler on all domains to exit.  What this
+      means is that the caller of [spawn] {i should} ideally arrange for any
+      exception to be handled by [main], but, in case that is not practical, it
+      is also possible to allow an exception to propagate out of [main], which
+      is then guaranteed to, one way or the other, to stop the entire program.
+      It is not possible to recover from a fatal error.
+
+      ℹ️ The behavior is that
+
+      - on OCaml 5, [spawn] performs the {!Spawn} effect, and
+      - on OCaml 4, [spawn] will call the [spawn] operation of the {{!Handler}
+        current handler}. *)
+
   (** {2 Interface for structuring} *)
 
   val get_computation : t -> Computation.packed
@@ -1057,13 +1066,6 @@ module Fiber : sig
 
   (** {2 Interface for schedulers} *)
 
-  val create_packed : forbid:bool -> Computation.packed -> t
-  (** [create_packed ~forbid packed] creates a new fiber. *)
-
-  val create : forbid:bool -> 'a Computation.t -> t
-  (** [create ~forbid computation] is equivalent to
-      {{!create_packed} [create_packed ~forbid (Computation.Packed computation)]}. *)
-
   val try_suspend :
     t -> Trigger.t -> 'x -> 'y -> (Trigger.t -> 'x -> 'y -> unit) -> bool
   (** [try_suspend fiber trigger x y resume] tries to suspend the [fiber] to
@@ -1121,8 +1123,7 @@ module Handler : sig
 
   type 'c t = {
     current : 'c -> Fiber.t;  (** See {!Picos.Fiber.current}. *)
-    spawn :
-      'a. 'c -> forbid:bool -> 'a Computation.t -> (unit -> unit) list -> unit;
+    spawn : 'c -> Fiber.t -> (Fiber.t -> unit) -> unit;
         (** See {!Picos.Fiber.spawn}. *)
     yield : 'c -> unit;  (** See {!Picos.Fiber.yield}. *)
     cancel_after :
@@ -1134,10 +1135,10 @@ module Handler : sig
   (** 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
+  val using : 'c t -> 'c -> (Fiber.t -> unit) -> unit
+  (** [using handler context main] sets the [handler] and the [context] for the
       handler of the primitive effects based operations of Picos while running
-      [thunk].
+      [main].
 
       ℹ️ The behavior is that
 
diff --git a/lib/picos_fifos/picos_fifos.ml b/lib/picos_fifos/picos_fifos.ml
index dc5641a..88874f0 100644
--- a/lib/picos_fifos/picos_fifos.ml
+++ b/lib/picos_fifos/picos_fifos.ml
@@ -5,7 +5,7 @@ let[@inline never] quota_non_positive () = invalid_arg "quota must be positive"
 (* As a minor optimization, we avoid allocating closures, which take slightly
    more memory than values of this type. *)
 type ready =
-  | Spawn of Fiber.t * (unit -> unit)
+  | Spawn of Fiber.t * (Fiber.t -> unit)
   | Current of Fiber.t * (Fiber.t, unit) Effect.Deep.continuation
   | Continue of Fiber.t * (unit, unit) Effect.Deep.continuation
   | Resume of Fiber.t * (Exn_bt.t option, unit) Effect.Deep.continuation
@@ -32,19 +32,12 @@ type t = {
   mutable remaining_quota : int;
 }
 
-let rec spawn t n forbid packed = function
-  | [] -> Atomic.fetch_and_add t.num_alive_fibers n |> ignore
-  | main :: mains ->
-      let fiber = Fiber.create_packed ~forbid packed in
-      Picos_mpscq.push t.ready (Spawn (fiber, main));
-      spawn t (n + 1) forbid packed mains
-
 let rec next t =
   match Picos_mpscq.pop_exn t.ready with
   | Spawn (fiber, main) ->
       t.fiber <- Fiber.Maybe.of_fiber fiber;
       t.remaining_quota <- t.quota;
-      Effect.Deep.match_with main () t.handler
+      Effect.Deep.match_with main fiber t.handler
   | Current (fiber, k) ->
       t.fiber <- Fiber.Maybe.of_fiber fiber;
       t.remaining_quota <- t.quota;
@@ -82,7 +75,7 @@ let rec next t =
         next t
       end
 
-let run ?quota ?(forbid = false) main =
+let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
   let quota =
     match quota with
     | None -> Int.max_int
@@ -153,7 +146,7 @@ let run ?quota ?(forbid = false) main =
       (fun k ->
         let fiber = Fiber.Maybe.to_fiber t.fiber in
         Fiber.continue fiber k ())
-  and handler = { retc; exnc = raise; effc }
+  and handler = { retc; exnc; effc }
   and[@alert "-handler"] effc :
       type a. a Effect.t -> ((a, _) Effect.Deep.continuation -> _) option =
     function
@@ -164,10 +157,10 @@ let run ?quota ?(forbid = false) main =
     | Fiber.Spawn r ->
         (* We check cancelation status once and then either perform the
            whole operation or discontinue the fiber. *)
-        let fiber = Fiber.Maybe.to_fiber t.fiber in
         if Fiber.is_canceled fiber then t.discontinue
         else begin
-          spawn t 0 r.forbid (Packed r.computation) r.mains;
+          Atomic.incr t.num_alive_fibers;
+          Picos_mpscq.push t.ready (Spawn (r.fiber, r.main));
           t.return
         end
     | Fiber.Yield -> t.yield
@@ -233,9 +226,12 @@ let run ?quota ?(forbid = false) main =
       Condition.broadcast t.condition
     end
   in
+  Picos_mpscq.push t.ready (Spawn (fiber, main));
+  next t
+
+let run ?quota ?fatal_exn_handler ?(forbid = false) main =
   let computation = Computation.create ~mode:`LIFO () in
   let fiber = Fiber.create ~forbid computation in
-  let main = Computation.capture computation main in
-  Picos_mpscq.push t.ready (Spawn (fiber, main));
-  next t;
+  let main _ = Computation.capture computation main () in
+  run_fiber ?quota ?fatal_exn_handler fiber main;
   Computation.await computation
diff --git a/lib/picos_fifos/picos_fifos.mli b/lib/picos_fifos/picos_fifos.mli
index a374f54..a0d6651 100644
--- a/lib/picos_fifos/picos_fifos.mli
+++ b/lib/picos_fifos/picos_fifos.mli
@@ -22,13 +22,29 @@
     This scheduler also gives priority to fibers woken up from
     {{!Picos.Trigger.await} [await]} due to being canceled. *)
 
-val run : ?quota:int -> ?forbid:bool -> (unit -> 'a) -> 'a
-(** [run main] runs the [main] thunk with the scheduler.  Returns after [main]
-    and all of the fibers spawned by [main] have returned.
+open Picos
 
-    The optional [forbid] argument defaults to [false] and determines whether
-    propagation of cancelation is initially allowed.
+val run_fiber :
+  ?quota:int ->
+  ?fatal_exn_handler:(exn -> unit) ->
+  Fiber.t ->
+  (Fiber.t -> unit) ->
+  unit
+(** [run_fiber fiber main] runs the [main] program as the specified [fiber] and
+    returns after [main] and all of the fibers spawned by [main] have returned.
 
     The optional [quota] argument defaults to [Int.max_int] and determines the
     number of effects a fiber is allowed to perform before it is forced to
     yield. *)
+
+val run :
+  ?quota:int ->
+  ?fatal_exn_handler:(exn -> unit) ->
+  ?forbid:bool ->
+  (unit -> 'a) ->
+  'a
+(** [run main] is equivalent to calling {!run_fiber} with a freshly created
+    fiber and [main] wrapped to capture the result of [main].
+
+    The optional [forbid] argument defaults to [false] and determines whether
+    propagation of cancelation is initially allowed. *)
diff --git a/lib/picos_lwt/picos_lwt.ml b/lib/picos_lwt/picos_lwt.ml
index a438ee5..917f6ac 100644
--- a/lib/picos_lwt/picos_lwt.ml
+++ b/lib/picos_lwt/picos_lwt.ml
@@ -49,13 +49,8 @@ let[@alert "-handler"] rec go :
           (fun k ->
             match Fiber.canceled fiber with
             | None ->
-                let packed = Computation.Packed r.computation in
-                List.iter
-                  (fun main ->
-                    let fiber = Fiber.create_packed ~forbid:r.forbid packed in
-                    Lwt.async @@ fun () ->
-                    go fiber system (Effect.Shallow.fiber main) (Ok ()))
-                  r.mains;
+                Lwt.async (fun () ->
+                    go r.fiber system (Effect.Shallow.fiber r.main) (Ok r.fiber));
                 go fiber system k (Ok ())
             | Some exn_bt -> go fiber system k (Error exn_bt))
     | Fiber.Yield ->
@@ -105,12 +100,13 @@ let[@alert "-handler"] rec go :
   | Ok v -> Effect.Shallow.continue_with k v handler
   | Error exn_bt -> Exn_bt.discontinue_with k exn_bt handler
 
-let run ?(forbid = false) system main =
+let run_fiber system fiber main =
   if not (Picos_thread.is_main_thread ()) then not_main_thread ();
+  go fiber system (Effect.Shallow.fiber main) (Ok fiber)
+
+let run ?(forbid = false) system main =
   let computation = Computation.create ~mode:`LIFO () in
   let fiber = Fiber.create ~forbid computation in
-  let main () =
-    Computation.capture computation main ();
-    Computation.await computation
-  in
-  go fiber system (Effect.Shallow.fiber main) (Ok ())
+  let main _ = Computation.capture computation main () in
+  run_fiber system fiber main
+  |> Lwt.map @@ fun () -> Computation.await computation
diff --git a/lib/picos_lwt/picos_lwt.mli b/lib/picos_lwt/picos_lwt.mli
index 5510a81..d6a0d73 100644
--- a/lib/picos_lwt/picos_lwt.mli
+++ b/lib/picos_lwt/picos_lwt.mli
@@ -7,6 +7,8 @@
     ℹ️ This is a {{!System} system} independent interface to {!Lwt}.  See
     {!Picos_lwt_unix} for a {!Unix} specific interface. *)
 
+open Picos
+
 val await : 'a Lwt.t -> 'a
 (** [await promise] awaits for the [promise] to resolve and returns the result.
 
@@ -15,16 +17,21 @@ val await : 'a Lwt.t -> 'a
 
 include module type of Intf
 
-val run : ?forbid:bool -> (module System) -> (unit -> 'a) -> 'a Lwt.t
-(** [run (module System) main] runs the [main] program implemented in {!Picos}
-    as a promise with {!Lwt} as the scheduler using the given {!System} module.
-    In other words, the [main] program will be run as a {!Lwt} promise or fiber.
+val run_fiber : (module System) -> Fiber.t -> (Fiber.t -> unit) -> unit Lwt.t
+(** [run_fiber (module System) fiber main] runs the [main] program as the
+    specified [fiber] as a promise with {!Lwt} as the scheduler using the given
+    {!System} module.  In other words, the [main] program will be run as a
+    {!Lwt} promise or fiber.
 
     ℹ️ Inside [main] you can use anything implemented in Picos for concurrent
     programming.  In particular, you only need to call [run] with a {!System}
     module implementation at the entry point of your application.
 
-    The optional [forbid] argument defaults to [false] and determines whether
-    propagation of cancelation is initially allowed.
-
     ⚠️ This may only be called on the main thread on which {!Lwt} runs. *)
+
+val run : ?forbid:bool -> (module System) -> (unit -> 'a) -> 'a Lwt.t
+(** [run (module System) main] is equivalent to calling {!run_fiber} with a
+    freshly created fiber and [main] wrapped to capture the result of [main].
+
+    The optional [forbid] argument defaults to [false] and determines whether
+    propagation of cancelation is initially allowed. *)
diff --git a/lib/picos_lwt_unix/picos_lwt_unix.ml b/lib/picos_lwt_unix/picos_lwt_unix.ml
index e32aa13..9573a9a 100644
--- a/lib/picos_lwt_unix/picos_lwt_unix.ml
+++ b/lib/picos_lwt_unix/picos_lwt_unix.ml
@@ -1,3 +1,5 @@
+open Picos
+
 let[@inline never] not_main_thread () =
   invalid_arg "not called from the main thread"
 
@@ -61,7 +63,7 @@ let notification_decr _ =
   notification.ref_count <- ref_count;
   if ref_count = 0 then Lwt_unix.stop_notification notification.id
 
-let run ?forbid main =
+let run_fiber fiber main =
   if not (Picos_thread.is_main_thread ()) then not_main_thread ();
   begin
     let ref_count = notification.ref_count + 1 in
@@ -69,8 +71,14 @@ let run ?forbid main =
     if ref_count = 1 then
       notification.id <- Lwt_unix.make_notification notify_callback
   end;
-  let promise = Picos_lwt.run ?forbid system main in
+  let promise = Picos_lwt.run_fiber system fiber main in
   Lwt.on_any promise notification_decr notification_decr;
   promise
 
+let run ?(forbid = false) main =
+  let computation = Computation.create ~mode:`LIFO () in
+  let fiber = Fiber.create ~forbid computation in
+  let main _ = Computation.capture computation main () in
+  run_fiber fiber main |> Lwt.map @@ fun () -> Computation.await computation
+
 let () = Lwt_main.run (Lwt_unix.sleep 0.0)
diff --git a/lib/picos_lwt_unix/picos_lwt_unix.mli b/lib/picos_lwt_unix/picos_lwt_unix.mli
index 7b36a14..d658dd0 100644
--- a/lib/picos_lwt_unix/picos_lwt_unix.mli
+++ b/lib/picos_lwt_unix/picos_lwt_unix.mli
@@ -1,9 +1,19 @@
 (** Direct style {!Picos} compatible interface to {!Lwt} with {!Lwt_unix} for
     OCaml 5. *)
 
-val run : ?forbid:bool -> (unit -> 'a) -> 'a Lwt.t
-(** [run main] runs the [main] program implemented in {!Picos} as a promise with
-    {!Lwt} as the scheduler using a {!Lwt_unix} based {{!Picos_lwt.System}
-    [System]} module.
+open Picos
+
+val run_fiber : Fiber.t -> (Fiber.t -> unit) -> unit Lwt.t
+(** [run_fiber fiber main] runs the [main] program as the specified [fiber] as a
+    promise with {!Lwt} as the scheduler using a {!Lwt_unix} based
+    {{!Picos_lwt.System} [System]} module.  In other words, the [main] program
+    will be run as a {!Lwt} promise or fiber.
 
     ⚠️ This may only be called on the main thread on which {!Lwt} runs. *)
+
+val run : ?forbid:bool -> (unit -> 'a) -> 'a Lwt.t
+(** [run main] is equivalent to calling {!run_fiber} with a freshly created
+    fiber and [main] wrapped to capture the result of [main].
+
+    The optional [forbid] argument defaults to [false] and determines whether
+    propagation of cancelation is initially allowed. *)
diff --git a/lib/picos_randos/picos_randos.ml b/lib/picos_randos/picos_randos.ml
index 480e410..1f777ab 100644
--- a/lib/picos_randos/picos_randos.ml
+++ b/lib/picos_randos/picos_randos.ml
@@ -27,7 +27,7 @@ module Collection = struct
 end
 
 type ready =
-  | Spawn of Fiber.t * (unit -> unit)
+  | Spawn of Fiber.t * (Fiber.t -> unit)
   | Current of Fiber.t * (Fiber.t, unit) Effect.Deep.continuation
   | Continue of Fiber.t * (unit, unit) Effect.Deep.continuation
   | Resume of Fiber.t * (Exn_bt.t option, unit) Effect.Deep.continuation
@@ -44,21 +44,13 @@ type t = {
     (Exn_bt.t option, unit) Effect.Deep.continuation ->
     unit;
   retc : unit -> unit;
+  exnc : exn -> unit;
   num_waiters : int ref;
   condition : Condition.t;
   mutex : Mutex.t;
   mutable run : bool;
 }
 
-let rec spawn t forbid packed = function
-  | [] -> ()
-  | main :: mains ->
-      let fiber = Fiber.create_packed ~forbid packed in
-      Atomic.incr t.num_alive_fibers;
-      Collection.push t.ready (Spawn (fiber, main));
-      if !(t.num_waiters_non_zero) then Condition.signal t.condition;
-      spawn t forbid packed mains
-
 let rec next t =
   match Collection.pop_exn t.ready with
   | Spawn (fiber, main) ->
@@ -84,7 +76,9 @@ let rec next t =
         | Fiber.Spawn r ->
             if Fiber.is_canceled fiber then yield
             else begin
-              spawn t r.forbid (Packed r.computation) r.mains;
+              Atomic.incr t.num_alive_fibers;
+              Collection.push t.ready (Spawn (r.fiber, r.main));
+              if !(t.num_waiters_non_zero) then Condition.signal t.condition;
               return
             end
         | Fiber.Yield -> yield
@@ -112,7 +106,7 @@ let rec next t =
                 end)
         | _ -> None
       in
-      Effect.Deep.match_with main () { retc = t.retc; exnc = raise; effc }
+      Effect.Deep.match_with main fiber { retc = t.retc; exnc = t.exnc; effc }
   | Raise (k, exn_bt) -> Exn_bt.discontinue k exn_bt
   | Return k -> Effect.Deep.continue k ()
   | Current (fiber, k) -> Effect.Deep.continue k fiber
@@ -151,8 +145,24 @@ let rec next t =
         Condition.broadcast t.condition
       end
 
-let context () =
+let default_fatal_exn_handler exn =
+  prerr_string "Fatal error: exception ";
+  prerr_string (Printexc.to_string exn);
+  prerr_char '\n';
+  Printexc.print_backtrace stderr;
+  flush stderr;
+  exit 2
+
+let context ?fatal_exn_handler () =
   Select.check_configured ();
+  let exnc =
+    match fatal_exn_handler with
+    | None -> default_fatal_exn_handler
+    | Some handler ->
+        fun exn ->
+          handler exn;
+          raise exn
+  in
   let rec t =
     {
       ready = Collection.create ();
@@ -160,6 +170,7 @@ let context () =
       num_alive_fibers = Atomic.make 1 |> Multicore_magic.copy_as_padded;
       resume;
       retc;
+      exnc;
       num_waiters = ref 0 |> Multicore_magic.copy_as_padded;
       condition = Condition.create ();
       mutex = Mutex.create ();
@@ -188,20 +199,17 @@ let runner_on_this_thread t =
   Select.check_configured ();
   next t
 
-let rec await t computation =
+let rec await t =
   if !(t.num_waiters_non_zero) then begin
     match Condition.wait t.condition t.mutex with
-    | () -> await t computation
+    | () -> await t
     | exception async_exn ->
         Mutex.unlock t.mutex;
         raise async_exn
   end
-  else begin
-    Mutex.unlock t.mutex;
-    Computation.await computation
-  end
+  else Mutex.unlock t.mutex
 
-let run ?context:t_opt ?(forbid = false) main =
+let run_fiber ?context:t_opt fiber main =
   let t =
     match t_opt with
     | Some t ->
@@ -217,17 +225,21 @@ let run ?context:t_opt ?(forbid = false) main =
   else begin
     t.run <- true;
     Mutex.unlock t.mutex;
-    let computation = Computation.create ~mode:`LIFO () in
-    let fiber = Fiber.create ~forbid computation in
-    let main = Computation.capture computation main in
     Collection.push t.ready (Spawn (fiber, main));
     next t;
     Mutex.lock t.mutex;
-    await t computation
+    await t
   end
 
-let rec run_on n ?forbid main context =
-  if n <= 1 then run ~context ?forbid main
+let run ?context ?(forbid = false) main =
+  let computation = Computation.create ~mode:`LIFO () in
+  let fiber = Fiber.create ~forbid computation in
+  let main _ = Computation.capture computation main () in
+  run_fiber ?context fiber main;
+  Computation.await computation
+
+let rec run_fiber_on n fiber main context =
+  if n <= 1 then run_fiber ~context fiber main
   else
     let runner =
       try Domain.spawn @@ fun () -> runner_on_this_thread context
@@ -236,7 +248,7 @@ let rec run_on n ?forbid main context =
         run ~context Fun.id;
         Printexc.raise_with_backtrace exn bt
     in
-    match run_on (n - 1) ?forbid main context with
+    match run_fiber_on (n - 1) fiber main context with
     | result ->
         Domain.join runner;
         result
@@ -245,6 +257,13 @@ let rec run_on n ?forbid main context =
         Domain.join runner;
         Printexc.raise_with_backtrace exn bt
 
-let run_on ~n_domains ?forbid main =
+let run_fiber_on ?fatal_exn_handler ~n_domains fiber main =
   if n_domains < 1 then invalid_arg "n_domains must be positive";
-  run_on n_domains ?forbid main (context ())
+  run_fiber_on n_domains fiber main (context ?fatal_exn_handler ())
+
+let run_on ?fatal_exn_handler ~n_domains ?(forbid = false) main =
+  let computation = Computation.create ~mode:`LIFO () in
+  let fiber = Fiber.create ~forbid computation in
+  let main _ = Computation.capture computation main () in
+  run_fiber_on ?fatal_exn_handler ~n_domains fiber main;
+  Computation.await computation
diff --git a/lib/picos_randos/picos_randos.mli b/lib/picos_randos/picos_randos.mli
index 6aa0b7a..d23838d 100644
--- a/lib/picos_randos/picos_randos.mli
+++ b/lib/picos_randos/picos_randos.mli
@@ -10,10 +10,12 @@
     bugs in programs implemented in Picos that make invalid scheduling
     assumptions. *)
 
+open Picos
+
 type t
 (** Represents a shared context for randomized runners. *)
 
-val context : unit -> t
+val context : ?fatal_exn_handler:(exn -> unit) -> unit -> t
 (** [context ()] creates a new context for randomized runners.  The context
     should be consumed by a call of {{!run} [run ~context ...]}. *)
 
@@ -22,21 +24,42 @@ val runner_on_this_thread : t -> unit
     fibers on the context.  The runner returns when {{!run} [run ~context ...]}
     returns. *)
 
-val run : ?context:t -> ?forbid:bool -> (unit -> 'a) -> 'a
-(** [run main] runs the [main] thunk with the scheduler.  Returns after [main]
-    and all of the fibers spawned by [main] have returned.
+val run_fiber : ?context:t -> Fiber.t -> (Fiber.t -> unit) -> unit
+(** [run_fiber fiber main] runs the [main] program as the specified [fiber] and
+    returns [main] and all of the fibers spawned by [main] have returned.
 
     The optional [context] argument specifies a context in which to run the
-    [main] thunk.  If unspecified, a new context is automatically created and
+    [main] program.  If unspecified, a new context is automatically created and
     the scheduler will be single-threaded.  By {{!context} creating a context},
     spawning concurrent or parallel {{!runner_on_this_thread} runners} on to the
     context, and then explicitly passing the context to [run ~context ...] one
     can create a multi-threaded scheduler.  Only a single call of {!run} per
-    context is allowed.
+    context is allowed. *)
+
+val run : ?context:t -> ?forbid:bool -> (unit -> 'a) -> 'a
+(** [run main] is equivalent to calling {!run_fiber} with a freshly created
+    fiber and [main] wrapped to capture the result of [main].
 
     The optional [forbid] argument defaults to [false] and determines whether
     propagation of cancelation is initially allowed. *)
 
-val run_on : n_domains:int -> ?forbid:bool -> (unit -> 'a) -> 'a
-(** [run_on ~n_domains main] spawns [n_domains - 1] additional domains and runs
-    the [main] on the current domain and those additional domains. *)
+val run_fiber_on :
+  ?fatal_exn_handler:(exn -> unit) ->
+  n_domains:int ->
+  Fiber.t ->
+  (Fiber.t -> unit) ->
+  unit
+(** [run_fiber_on ~n_domains main] spawns [n_domains - 1] additional domains and
+    runs the [main] on the current domain and those additional domains. *)
+
+val run_on :
+  ?fatal_exn_handler:(exn -> unit) ->
+  n_domains:int ->
+  ?forbid:bool ->
+  (unit -> 'a) ->
+  'a
+(** [run_on ~n_domains main] is equivalent to calling {!run_fiber_on} with a
+    freshly created fiber and [main] wrapped to capture the result of [main].
+
+    The optional [forbid] argument defaults to [false] and determines whether
+    propagation of cancelation is initially allowed. *)
diff --git a/lib/picos_structured/bundle.ml b/lib/picos_structured/bundle.ml
index f7cdd08..565b807 100644
--- a/lib/picos_structured/bundle.ml
+++ b/lib/picos_structured/bundle.ml
@@ -6,7 +6,7 @@ type _ tdt =
   | Nothing : [> `Nothing ] tdt
   | Bundle : {
       num_fibers : int Atomic.t;
-      bundle : unit Computation.t;
+      bundle : Computation.packed;
       errors : Control.Errors.t;
       finished : Trigger.t;
     }
@@ -17,12 +17,12 @@ let flock_key : [ `Bundle | `Nothing ] tdt Fiber.FLS.key =
 
 type t = [ `Bundle ] tdt
 
-let terminate ?callstack (Bundle r : t) =
-  Computation.cancel r.bundle (Control.terminate_bt ?callstack ())
+let terminate ?callstack (Bundle { bundle = Packed bundle; _ } : t) =
+  Computation.cancel bundle (Control.terminate_bt ?callstack ())
 
-let terminate_after ?callstack (Bundle r : t) ~seconds =
-  Computation.cancel_after r.bundle ~seconds
-    (Control.terminate_bt ?callstack ())
+let terminate_after ?callstack (Bundle { bundle = Packed bundle; _ } : t)
+    ~seconds =
+  Computation.cancel_after bundle ~seconds (Control.terminate_bt ?callstack ())
 
 let error ?callstack (Bundle r as t : t) (exn_bt : Exn_bt.t) =
   if exn_bt.Exn_bt.exn != Control.Terminate then begin
@@ -33,7 +33,8 @@ let error ?callstack (Bundle r as t : t) (exn_bt : Exn_bt.t) =
 let decr (Bundle r : t) =
   let n = Atomic.fetch_and_add r.num_fibers (-1) in
   if n = 1 then begin
-    Computation.finish r.bundle;
+    let (Packed bundle) = r.bundle in
+    Computation.cancel bundle (Control.terminate_bt ());
     Trigger.signal r.finished
   end
 
@@ -65,7 +66,7 @@ let join_after_pass (type a) (fn : a -> _) (pass : a pass) =
   (* The sequence of operations below ensures that nothing is leaked. *)
   let (Bundle r as t : t) =
     let num_fibers = Atomic.make 1 in
-    let bundle = Computation.create ~mode:`LIFO () in
+    let bundle = Computation.Packed (Computation.create ~mode:`LIFO ()) in
     let errors = Control.Errors.create () in
     let finished = Trigger.create () in
     Bundle { num_fibers; bundle; errors; finished }
@@ -75,12 +76,12 @@ let join_after_pass (type a) (fn : a -> _) (pass : a pass) =
     match pass with FLS -> Fiber.FLS.get fiber flock_key | Arg -> Nothing
   in
   let (Packed parent as packed) = Fiber.get_computation fiber in
-  let bundle = Computation.Packed r.bundle in
-  let canceler = Computation.attach_canceler ~from:parent ~into:r.bundle 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 bundle;
+    Fiber.set_computation fiber r.bundle;
     fn (match pass with FLS -> Fiber.FLS.set fiber flock_key t | Arg -> t)
   with
   | value ->
@@ -103,13 +104,14 @@ let fork_as_promise_pass (type a) (Bundle r as t : t) thunk (pass : a pass) =
   incr t Backoff.default;
   try
     let child = Computation.create ~mode:`LIFO () in
-    let canceler = Computation.attach_canceler ~from:r.bundle ~into:child in
+    let fiber = Fiber.create ~forbid:false child in
+    let (Packed bundle) = r.bundle in
+    let canceler = Computation.attach_canceler ~from:bundle ~into:child in
     let main =
       match pass with
       | FLS ->
-          let main () =
-            let fiber = Fiber.current () in
-            Fiber.FLS.set fiber flock_key t;
+          Fiber.FLS.set fiber flock_key t;
+          fun fiber ->
             begin
               match thunk () with
               | value -> Computation.return child value
@@ -119,12 +121,11 @@ let fork_as_promise_pass (type a) (Bundle r as t : t) thunk (pass : a pass) =
                   error (get_flock fiber) exn_bt
             end;
             let (Bundle r as t : t) = get_flock fiber in
-            Computation.detach r.bundle canceler;
+            let (Packed bundle) = r.bundle in
+            Computation.detach bundle canceler;
             decr t
-          in
-          [ main ]
       | Arg ->
-          let main () =
+          fun _ ->
             begin
               match thunk () with
               | value -> Computation.return child value
@@ -133,12 +134,11 @@ let fork_as_promise_pass (type a) (Bundle r as t : t) thunk (pass : a pass) =
                   Computation.cancel child exn_bt;
                   error t exn_bt
             end;
-            Computation.detach r.bundle canceler;
+            let (Packed bundle) = r.bundle in
+            Computation.detach bundle canceler;
             decr t
-          in
-          [ main ]
     in
-    Fiber.spawn ~forbid:false child main;
+    Fiber.spawn fiber main;
     child
   with canceled_exn ->
     (* We don't worry about detaching the [canceler], because at this point we
@@ -151,38 +151,36 @@ let fork_pass (type a) (Bundle r as t : t) thunk (pass : a pass) =
   (* The sequence of operations below ensures that nothing is leaked. *)
   incr t Backoff.default;
   try
+    let fiber = Fiber.create_packed ~forbid:false r.bundle in
     let main =
       match pass with
       | FLS ->
-          let main () =
-            let fiber = Fiber.current () in
-            Fiber.FLS.set fiber flock_key t;
+          Fiber.FLS.set fiber flock_key t;
+          fun fiber ->
             begin
               try thunk ()
               with exn -> error (get_flock fiber) (Exn_bt.get exn)
             end;
             decr (get_flock fiber)
-          in
-          [ main ]
       | Arg ->
-          let main () =
+          fun _ ->
             begin
               try thunk () with exn -> error t (Exn_bt.get exn)
             end;
             decr t
-          in
-          [ main ]
     in
-    Fiber.spawn ~forbid:false r.bundle main
+    Fiber.spawn fiber main
   with canceled_exn ->
     decr t;
     raise canceled_exn
 
 (* *)
 
+let is_running (Bundle { bundle = Packed bundle; _ } : t) =
+  Computation.is_running bundle
+
 let join_after fn = join_after_pass fn Arg
 let fork t thunk = fork_pass t thunk Arg
 let fork_as_promise t thunk = fork_as_promise_pass t thunk Arg
-let is_running (Bundle r : t) = Computation.is_running r.bundle
 let unsafe_incr (Bundle r : t) = Atomic.incr r.num_fibers
 let unsafe_reset (Bundle r : t) = Atomic.set r.num_fibers 1
diff --git a/lib/picos_structured/run.ml b/lib/picos_structured/run.ml
index e7e7935..b17c6a1 100644
--- a/lib/picos_structured/run.ml
+++ b/lib/picos_structured/run.ml
@@ -1,32 +1,34 @@
 open Picos
 
-let wrap_all t main =
-  Bundle.unsafe_incr t;
-  fun () ->
-    if Bundle.is_running t then begin
-      try main () with exn -> Bundle.error t (Exn_bt.get exn)
-    end;
-    Bundle.decr t
+let wrap_all t main _ =
+  if Bundle.is_running t then begin
+    try main () with exn -> Bundle.error t (Exn_bt.get exn)
+  end;
+  Bundle.decr t
 
-let wrap_any t main =
-  Bundle.unsafe_incr t;
-  fun () ->
-    if Bundle.is_running t then begin
-      try
-        main ();
-        Bundle.terminate t
-      with exn -> Bundle.error t (Exn_bt.get exn)
-    end;
-    Bundle.decr t
+let wrap_any t main _ =
+  if Bundle.is_running t then begin
+    match main () with
+    | () -> Bundle.terminate t
+    | exception exn -> Bundle.error t (Exn_bt.get exn)
+  end;
+  Bundle.decr t
+
+let rec spawn (Bundle r as t : Bundle.t) wrap = function
+  | [] -> ()
+  | main :: mains ->
+      Bundle.unsafe_incr t;
+      let fiber = Fiber.create_packed ~forbid:false r.bundle in
+      Fiber.spawn fiber (wrap t main);
+      spawn t wrap mains
 
 let run actions wrap =
-  Bundle.join_after @@ fun (Bundle r as t : Bundle.t) ->
-  try
-    let mains = List.map (wrap t) actions in
-    Fiber.spawn ~forbid:false r.bundle mains
+  Bundle.join_after @@ fun (Bundle _ as t : Bundle.t) ->
+  try spawn t wrap actions
   with exn ->
-    Bundle.unsafe_reset t;
-    raise exn
+    let exn_bt = Exn_bt.get exn in
+    Bundle.decr t;
+    Bundle.error t exn_bt
 
 let all actions = run actions wrap_all
 let any actions = run actions wrap_any
diff --git a/lib/picos_threaded/picos_threaded.ml b/lib/picos_threaded/picos_threaded.ml
index 6b8a363..489fcaf 100644
--- a/lib/picos_threaded/picos_threaded.ml
+++ b/lib/picos_threaded/picos_threaded.ml
@@ -1,12 +1,15 @@
 open Picos
 
-type t = { fiber : Fiber.t; mutex : Mutex.t; condition : Condition.t }
+type t = {
+  fiber : Fiber.t;
+  mutex : Mutex.t;
+  condition : Condition.t;
+  fatal_exn_handler : exn -> unit;
+}
 
-let create_packed ~forbid packed =
-  let fiber = Fiber.create_packed ~forbid packed in
-  let mutex = Mutex.create () in
-  let condition = Condition.create () in
-  { fiber; mutex; condition }
+let create ~fatal_exn_handler fiber =
+  let mutex = Mutex.create () and condition = Condition.create () in
+  { fiber; mutex; condition; fatal_exn_handler }
 
 let rec block trigger t =
   if not (Trigger.is_signaled trigger) then begin
@@ -45,6 +48,14 @@ let resume trigger t _ =
   end;
   Condition.broadcast t.condition
 
+let default_fatal_exn_handler exn =
+  prerr_string "Fatal error: exception ";
+  prerr_string (Printexc.to_string exn);
+  prerr_char '\n';
+  Printexc.print_backtrace stderr;
+  flush stderr;
+  exit 2
+
 let[@alert "-handler"] rec await t trigger =
   if Fiber.try_suspend t.fiber trigger t t resume then block trigger t;
   Fiber.canceled t.fiber
@@ -67,79 +78,31 @@ and cancel_after : type a. _ -> a Computation.t -> _ =
   Fiber.check t.fiber;
   Select.cancel_after computation ~seconds exn_bt
 
-and spawn : type a. _ -> forbid:bool -> a Computation.t -> _ =
- fun t ~forbid computation mains ->
+and spawn t fiber main =
   Fiber.check t.fiber;
-  let packed = Computation.Packed computation in
-  match mains with
-  | [ main ] ->
-      Thread.create
-        (fun () ->
-          (* We need to (recursively) install the handler on each new thread
-             that we create. *)
-          Handler.using handler (create_packed ~forbid packed) main)
-        ()
-      |> ignore
-  | mains -> begin
-      (* We try to be careful to implement the all-or-nothing behaviour based on
-         the assumption that we may run out of threads well before we run out of
-         memory.  In a thread pool based scheduler this should actually not
-         require special treatment. *)
-      let all_or_nothing = ref `Wait in
-      match
-        mains
-        |> List.iter @@ fun main ->
-           Thread.create
-             (fun () ->
-               if !all_or_nothing == `Wait then begin
-                 Mutex.lock t.mutex;
-                 match
-                   while
-                     match !all_or_nothing with
-                     | `Wait ->
-                         Condition.wait t.condition t.mutex;
-                         true
-                     | `All | `Nothing -> false
-                   do
-                     ()
-                   done
-                 with
-                 | () -> Mutex.unlock t.mutex
-                 | exception async_exn ->
-                     (* Condition.wait may be interrupted by asynchronous
-                        exceptions and we must make sure to unlock even in that
-                        case. *)
-                     Mutex.unlock t.mutex;
-                     raise async_exn
-               end;
-               if !all_or_nothing == `All then
-                 (* We need to (recursively) install the handler on each new
-                    thread that we create. *)
-                 Handler.using handler (create_packed ~forbid packed) main)
-             ()
-           |> ignore
-      with
-      | () ->
-          Mutex.lock t.mutex;
-          all_or_nothing := `All;
-          Mutex.unlock t.mutex;
-          Condition.broadcast t.condition
-      | exception exn ->
-          Mutex.lock t.mutex;
-          all_or_nothing := `Nothing;
-          Mutex.unlock t.mutex;
-          Condition.broadcast t.condition;
-          raise exn
-    end
+  match Thread.create start (fiber, t.fatal_exn_handler, main) with
+  | _ -> ( (* We assume that [main] is now guaranteed to be called. *) )
+  | exception exn ->
+      let exn_bt = Exn_bt.get exn in
+      (* [main] wasn't called, so we need to finalize. *)
+      let (Packed computation) = Fiber.get_computation fiber in
+      Computation.cancel computation exn_bt;
+      Exn_bt.raise exn_bt
 
 and handler = Handler.{ current; spawn; yield; cancel_after; await }
 
-let run ?(forbid = false) main =
+and start (fiber, fatal_exn_handler, main) =
+  (* We need to install the handler on each new thread that we create. *)
+  try Handler.using handler (create ~fatal_exn_handler fiber) main
+  with exn -> fatal_exn_handler exn
+
+let run_fiber ?(fatal_exn_handler = default_fatal_exn_handler) fiber main =
   Select.check_configured ();
+  Handler.using handler (create ~fatal_exn_handler fiber) main
+
+let run ?(forbid = false) ?fatal_exn_handler main =
   let computation = Computation.create ~mode:`LIFO () in
-  let context = create_packed ~forbid (Packed computation) in
-  let main () =
-    Computation.capture computation main ();
-    Computation.await computation
-  in
-  Handler.using handler context main
+  let fiber = Fiber.create ~forbid computation in
+  let main _ = Computation.capture computation main () in
+  run_fiber ?fatal_exn_handler fiber main;
+  Computation.await computation
diff --git a/lib/picos_threaded/picos_threaded.mli b/lib/picos_threaded/picos_threaded.mli
index cdab398..215cecf 100644
--- a/lib/picos_threaded/picos_threaded.mli
+++ b/lib/picos_threaded/picos_threaded.mli
@@ -24,9 +24,16 @@
     OCaml 5 a scheduler that implements an effect handler directly is likely to
     perform better. *)
 
-val run : ?forbid:bool -> (unit -> 'a) -> 'a
-(** [run main] runs the [main] thunk with the scheduler.  Returns after [main]
-    and all of the fibers spawned by [main] have returned.
+open Picos
+
+val run_fiber :
+  ?fatal_exn_handler:(exn -> unit) -> Fiber.t -> (Fiber.t -> unit) -> unit
+(** [run_fiber fiber main] runs the [main] program as the specified [fiber] and
+    returns [main] and all of the fibers spawned by [main] have returned. *)
+
+val run : ?forbid:bool -> ?fatal_exn_handler:(exn -> unit) -> (unit -> 'a) -> 'a
+(** [run main] is equivalent to calling {!run_fiber} with a freshly created
+    fiber and [main] wrapped to capture the result of [main].
 
     The optional [forbid] argument defaults to [false] and determines whether
     propagation of cancelation is initially allowed. *)


Test patch:
diff --git a/test/test_picos.ml b/test/test_picos.ml
index 9763412..bf47970 100644
--- a/test/test_picos.ml
+++ b/test/test_picos.ml
@@ -3,7 +3,9 @@ open Picos_finally
 
 let run_in_fiber main =
   let computation = Computation.create ~mode:`LIFO () in
-  Fiber.spawn ~forbid:false computation [ Computation.capture computation main ];
+  let fiber = Fiber.create ~forbid:false computation in
+  let main _ = Computation.capture computation main () in
+  Fiber.spawn fiber main;
   Computation.await computation
 
 let test_fls_basics =
@@ -90,13 +92,15 @@ let test_thread_cancelation () =
   let@ _ =
     finally Computation.await @@ fun () ->
     let result = Computation.create () in
-    let main =
-      Computation.capture result @@ fun () ->
-      while true do
-        Fiber.yield ()
-      done
+    let main _ =
+      Computation.capture result
+        (fun () ->
+          while true do
+            Fiber.yield ()
+          done)
+        ()
     in
-    Fiber.spawn ~forbid:false computation [ main ];
+    Fiber.spawn (Fiber.create ~forbid:false computation) main;
     result
   in
   Computation.cancel computation (Exn_bt.get_callstack 0 Exit)
@@ -105,13 +109,15 @@ let test_cancel_after () =
   Alcotest.check_raises "should be canceled" Not_found @@ fun () ->
   Test_scheduler.run ~max_domains:2 @@ fun () ->
   let computation = Computation.create () in
-  let main =
-    Computation.capture computation @@ fun () ->
-    while true do
-      Fiber.yield ()
-    done
+  let main _ =
+    Computation.capture computation
+      (fun () ->
+        while true do
+          Fiber.yield ()
+        done)
+      ()
   in
-  Fiber.spawn ~forbid:false computation [ main ];
+  Fiber.spawn (Fiber.create ~forbid:false computation) main;
   Computation.cancel_after computation ~seconds:0.01
     (Exn_bt.get_callstack 0 Not_found);
   Computation.await computation
diff --git a/test/test_scheduler.ocaml4.ml b/test/test_scheduler.ocaml4.ml
index 63b6cf2..df04791 100644
--- a/test/test_scheduler.ocaml4.ml
+++ b/test/test_scheduler.ocaml4.ml
@@ -1,6 +1,8 @@
 let () = Printexc.record_backtrace true
 let () = Random.self_init ()
 
+open Picos
+
 let () =
   Picos_select.check_configured ();
 
@@ -15,5 +17,12 @@ let () =
   in
   propagate ()
 
-let run ?max_domains:_ ?allow_lwt:_ ?forbid main =
-  Picos_threaded.run ?forbid main
+let run_fiber ?max_domains:_ ?allow_lwt:_ ?fatal_exn_handler fiber main =
+  Picos_threaded.run_fiber ?fatal_exn_handler fiber main
+
+let run ?max_domains ?allow_lwt ?fatal_exn_handler ?(forbid = false) main =
+  let computation = Computation.create ~mode:`LIFO () in
+  let fiber = Fiber.create ~forbid computation in
+  let main _ = Computation.capture computation main () in
+  run_fiber ?max_domains ?allow_lwt ?fatal_exn_handler fiber main;
+  Computation.await computation
diff --git a/test/test_scheduler.ocaml5.ml b/test/test_scheduler.ocaml5.ml
index 2d3d19f..9c8450b 100644
--- a/test/test_scheduler.ocaml5.ml
+++ b/test/test_scheduler.ocaml5.ml
@@ -1,6 +1,8 @@
 let () = Printexc.record_backtrace true
 let () = Random.self_init ()
 
+open Picos
+
 let () =
   Picos_select.check_configured ();
 
@@ -15,18 +17,44 @@ let () =
   in
   propagate ()
 
-let rec run ?(max_domains = 1) ?(allow_lwt = true) ?forbid main =
+let rec run_fiber ?(max_domains = 1) ?(allow_lwt = true) ?fatal_exn_handler
+    fiber main =
   let scheduler =
     match Random.int 3 with 0 -> `Fifos | 1 -> `Randos | _ -> `Lwt
   in
+  ignore
+    (match scheduler with
+    | `Fifos -> "fifos"
+    | `Randos -> "randos"
+    | `Lwt -> "lwt");
   match scheduler with
   | `Lwt ->
-      if Picos_thread.is_main_thread () && allow_lwt then
-        Lwt_main.run (Picos_lwt_unix.run ?forbid main)
-      else run ~max_domains ~allow_lwt ?forbid main
+      if Picos_thread.is_main_thread () && allow_lwt then begin
+        let old_hook = !Lwt.async_exception_hook in
+        begin
+          match fatal_exn_handler with
+          | None -> ()
+          | Some hook -> Lwt.async_exception_hook := hook
+        end;
+        match Lwt_main.run (Picos_lwt_unix.run_fiber fiber main) with
+        | result ->
+            Lwt.async_exception_hook := old_hook;
+            result
+        | exception exn ->
+            Lwt.async_exception_hook := old_hook;
+            raise exn
+      end
+      else run_fiber ~max_domains ~allow_lwt fiber main
   | `Randos ->
       let n_domains =
         Int.min max_domains (Domain.recommended_domain_count ())
       in
-      Picos_randos.run_on ~n_domains ?forbid main
-  | `Fifos -> Picos_fifos.run ?forbid main
+      Picos_randos.run_fiber_on ?fatal_exn_handler ~n_domains fiber main
+  | `Fifos -> Picos_fifos.run_fiber ?fatal_exn_handler fiber main
+
+let run ?max_domains ?allow_lwt ?fatal_exn_handler ?(forbid = false) main =
+  let computation = Computation.create ~mode:`LIFO () in
+  let fiber = Fiber.create ~forbid computation in
+  let main _ = Computation.capture computation main () in
+  run_fiber ?max_domains ?allow_lwt ?fatal_exn_handler fiber main;
+  Computation.await computation
diff --git a/test/test_schedulers.ml b/test/test_schedulers.ml
index f646e91..3ed71df 100644
--- a/test/test_schedulers.ml
+++ b/test/test_schedulers.ml
@@ -55,6 +55,30 @@ let test_cancel_after_long_timeout () =
   | () -> Computation.finish computation
   | exception Invalid_argument _ -> ()
 
+let test_fatal () =
+  match
+    let computation = Computation.create () in
+    let fiber = Fiber.create ~forbid:false computation in
+    let fatal_exn_handler exn =
+      Computation.cancel computation (Exn_bt.get exn);
+      raise exn
+    in
+    Test_scheduler.run_fiber ~fatal_exn_handler ~max_domains:3 fiber @@ fun _ ->
+    for _ = 1 to 100 do
+      Fiber.spawn (Fiber.create ~forbid:false computation) @@ fun _ ->
+      while true do
+        Fiber.yield ()
+      done
+    done;
+    Fiber.spawn (Fiber.create ~forbid:false computation) (fun _ ->
+        failwith "fatal");
+    while true do
+      Fiber.yield ()
+    done
+  with
+  | _ -> assert false
+  | exception Failure msg -> assert (msg = "fatal")
+
 let () =
   [
     ("Trivial main returns", [ Alcotest.test_case "" `Quick test_returns ]);
@@ -66,5 +90,9 @@ let () =
         Alcotest.test_case "basic" `Quick test_cancel_after_basic;
         Alcotest.test_case "long timeout" `Quick test_cancel_after_long_timeout;
       ] );
+    (* The fatal exn test must be kept last, because it may leave some
+       schedulers in a bad state. *)
+    ( "Fatal exception terminates scheduler",
+      [ Alcotest.test_case "" `Quick test_fatal ] );
   ]
   |> Alcotest.run "Picos schedulers"
diff --git a/test/test_structured.ml b/test/test_structured.ml
index 8033804..999a0b2 100644
--- a/test/test_structured.ml
+++ b/test/test_structured.ml
@@ -108,13 +108,13 @@ let test_block_raises_sys_error () =
   let success = ref false in
   let finished = Trigger.create () in
   let computation = Computation.create () in
-  let main () =
+  let main _ =
     begin
       try Control.block () with Sys_error _ -> success := true
     end;
     Trigger.signal finished
   in
-  Fiber.spawn ~forbid:false computation [ main ];
+  Fiber.spawn (Fiber.create ~forbid:false computation) main;
   Control.sleep ~seconds:0.1;
   Computation.finish computation;
   Trigger.await finished |> ignore;
diff --git a/test/test_sync.ml b/test/test_sync.ml
index 7d799e0..3608d9a 100644
--- a/test/test_sync.ml
+++ b/test/test_sync.ml
@@ -7,10 +7,11 @@ let msgs = ref []
 module Fiber = struct
   include Fiber
 
-  let start thunk =
+  let start main =
     let computation = Computation.create ~mode:`LIFO () in
-    Fiber.spawn ~forbid:false computation
-      [ Computation.capture computation thunk ];
+    let fiber = Fiber.create ~forbid:false computation in
+    let main _ = Computation.capture computation main () in
+    Fiber.spawn fiber main;
     computation
 end
 
@@ -30,13 +31,14 @@ let test_mutex_and_condition_basics () =
            let n = Atomic.make 10 in
            let test = Computation.create () in
 
-           let main () =
+           let main _ =
              Mutex.protect ?checked mutex (fun () ->
                  Condition.wait condition mutex);
              if 1 = Atomic.fetch_and_add n (-1) then Computation.finish test
            in
-           Fiber.spawn ~forbid:false computation
-             (List.init (Atomic.get n) @@ fun _ -> main);
+           for _ = 1 to Atomic.get n do
+             Fiber.spawn (Fiber.create ~forbid:false computation) main
+           done;
 
            while Computation.is_running test do
              Fiber.yield ();
@@ -84,7 +86,7 @@ let test_mutex_and_condition_cancelation () =
     Array.init n @@ fun _ -> Computation.Packed (Computation.create ())
   in
 
-  let attempt i finished ?checked () =
+  let attempt i finished ?checked _ =
     computations.(i) <- Fiber.get_computation (Fiber.current ());
     match
       Atomic.incr step_1;
@@ -135,8 +137,9 @@ let test_mutex_and_condition_cancelation () =
     do
       let finished = Trigger.create () in
       let checked = if Random.State.bool state then None else some_false in
-      Fiber.spawn ~forbid:false (Computation.create ())
-        [ attempt i finished ?checked ];
+      Fiber.spawn
+        (Fiber.create ~forbid:false (Computation.create ()))
+        (attempt i finished ?checked);
       Trigger.await finished |> ignore
     done
   in
@@ -194,14 +197,14 @@ let test_lazy_cancelation () =
     let to_await = Trigger.create () in
     let s = susp ~to_await 42 in
     let tried = Atomic.make false in
-    let main () =
+    let main _ =
       match Lazy.force s with
       | _ -> assert false
       | exception Exit -> Atomic.set tried true
     in
     let computation = Computation.create () in
     Computation.cancel computation (Exn_bt.get_callstack 0 Exit);
-    Fiber.spawn ~forbid:false computation [ main ];
+    Fiber.spawn (Fiber.create ~forbid:false computation) main;
     while not (Atomic.get tried) do
       Fiber.yield ()
     done;