diff --git a/lib/picos_std.structured/run.ml b/lib/picos_std.structured/run.ml index bc9122f5..79ebe292 100644 --- a/lib/picos_std.structured/run.ml +++ b/lib/picos_std.structured/run.ml @@ -1,39 +1,37 @@ open Picos -let wrap_all t main _ = - if Bundle.is_running t then begin - try main () with exn -> Bundle.error t exn (Printexc.get_raw_backtrace ()) - end; - Bundle.decr t +let[@inline never] wrap_all t main = + match main () with + | () -> Bundle.decr t + | exception exn -> Bundle.raised exn 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 (Printexc.get_raw_backtrace ()) - end; - Bundle.decr t +let[@inline never] wrap_any t main = + match main () with + | () -> + Bundle.terminate t; + Bundle.decr t + | exception exn -> Bundle.raised exn t -let rec spawn (Bundle r as t : Bundle.t) wrap = function +let rec spawn (Bundle r as t : Bundle.t) ~all = function | [] -> () | [ main ] -> Bundle.unsafe_incr t; - let unused_fake_fiber = Obj.magic () in - wrap t main unused_fake_fiber + if Bundle.is_running t then + if all then wrap_all t main else wrap_any t main + else Bundle.decr t | main :: mains -> Bundle.unsafe_incr t; let fiber = Fiber.create_packed ~forbid:false r.bundle in (* Note that [Fiber.spawn] checks the cancelation status of the bundle. *) - Fiber.spawn fiber (wrap t main); - spawn t wrap mains + Fiber.spawn fiber (fun _ -> + if Bundle.is_running t then + if all then wrap_all t main else wrap_any t main + else Bundle.decr t); + spawn t ~all mains -let run actions wrap = - Bundle.join_after @@ fun (Bundle _ as t : Bundle.t) -> - try spawn t wrap actions - with exn -> - let bt = Printexc.get_raw_backtrace () in - Bundle.decr t; - Bundle.error t exn bt +let run actions ~all = + Bundle.join_after @@ fun (t : Bundle.t) -> + try spawn t ~all actions with exn -> Bundle.raised exn t -let all actions = run actions wrap_all -let any actions = run actions wrap_any +let all actions = run actions ~all:true +let any actions = run actions ~all:false