diff --git a/src/pool.ml b/src/pool.ml index b9f511b1..3518e33e 100644 --- a/src/pool.ml +++ b/src/pool.ml @@ -227,3 +227,90 @@ let shutdown_ ~wait (self : t) : unit = let shutdown_without_waiting (self : t) : unit = shutdown_ self ~wait:false let shutdown (self : t) : unit = shutdown_ self ~wait:true + +module Fork_join_ = struct + type 'a single_res = + | St_none + | St_some of 'a + | St_fail of exn * Printexc.raw_backtrace + + type ('a, 'b) t = { + mutable suspension: + ((unit, exn * Printexc.raw_backtrace) result -> unit) option; + (** suspended caller *) + left: 'a single_res; + right: 'b single_res; + } + + let get_exn (self : _ t A.t) = + match A.get self with + | { left = St_fail (e, bt); _ } | { right = St_fail (e, bt); _ } -> + Printexc.raise_with_backtrace e bt + | { left = St_some x; right = St_some y; _ } -> x, y + | _ -> assert false + + let check_if_state_complete_ (self : _ t) : unit = + match self.left, self.right, self.suspension with + | St_some _, St_some _, Some f -> f (Ok ()) + | St_fail (e, bt), _, Some f | _, St_fail (e, bt), Some f -> + f (Error (e, bt)) + | _ -> () + + let set_left_ (self : _ t A.t) (x : _ single_res) = + while + let old_st = A.get self in + let new_st = { old_st with left = x } in + if A.compare_and_set self old_st new_st then ( + check_if_state_complete_ new_st; + false + ) else + true + do + Domain_.relax () + done + + let set_right_ (self : _ t A.t) (y : _ single_res) = + while + let old_st = A.get self in + let new_st = { old_st with right = y } in + if A.compare_and_set self old_st new_st then ( + check_if_state_complete_ new_st; + false + ) else + true + do + Domain_.relax () + done +end + +let fork_join f g : _ * _ = + let open Fork_join_ in + let st = A.make { suspension = None; left = St_none; right = St_none } in + + let start_tasks ~run () : unit = + run (fun () -> + try + let res = f () in + set_left_ st (St_some res) + with e -> + let bt = Printexc.get_raw_backtrace () in + set_left_ st (St_fail (e, bt))); + + run (fun () -> + try + let res = g () in + set_right_ st (St_some res) + with e -> + let bt = Printexc.get_raw_backtrace () in + set_right_ st (St_fail (e, bt))) + in + + Suspend_.suspend + { + Suspend_types_.handle = + (fun ~run suspension -> + (* nothing else is started, no race condition possible *) + (A.get st).suspension <- Some suspension; + start_tasks ~run ()); + }; + get_exn st diff --git a/src/pool.mli b/src/pool.mli index 4df914e8..bde9aff6 100644 --- a/src/pool.mli +++ b/src/pool.mli @@ -75,7 +75,7 @@ val run_async : t -> (unit -> unit) -> unit (** [run_async pool f] schedules [f] for later execution on the pool in one of the threads. [f()] will run on one of the pool's worker threads. - @raise Shutdown if the pool was shut down before [run] was called. + @raise Shutdown if the pool was shut down before [run_async] was called. @since 0.3 *) val run : t -> (unit -> unit) -> unit @@ -91,3 +91,17 @@ val run_wait_block : t -> (unit -> 'a) -> 'a {b NOTE} be careful with deadlocks (see notes in {!Fut.wait_block}). @since 0.3 *) + +(** {2 Fork-join computations} *) + +val fork_join : (unit -> 'a) -> (unit -> 'b) -> 'a * 'b +(** [fork_join f g] runs [f()] and [g()], potentially in parallel, + and returns their result when both are done. + If any of [f()] and [g()] fails, then the whole computation fails. + + This must be run from within the pool, inside {!run} + (or inside a {!Fut.spawn} computation). + This is because it relies on an effect handler to be installed. + + @since 0.3 + {b NOTE} this is only available on OCaml 5. *)