mirror of
https://github.com/c-cube/moonpool.git
synced 2025-12-06 03:05:30 -05:00
88 lines
2.4 KiB
OCaml
88 lines
2.4 KiB
OCaml
module A = Atomic_
|
|
|
|
type suspension = unit Exn_bt.result -> unit
|
|
type task = unit -> unit
|
|
|
|
[@@@ifge 5.0]
|
|
|
|
type suspension_handler = {
|
|
handle:
|
|
run:(task -> unit) ->
|
|
resume:(suspension -> unit Exn_bt.result -> unit) ->
|
|
suspension ->
|
|
unit;
|
|
}
|
|
[@@unboxed]
|
|
|
|
[@@@ocaml.alert "-unstable"]
|
|
|
|
type _ Effect.t +=
|
|
| Suspend : suspension_handler -> unit Effect.t
|
|
| Yield : unit Effect.t
|
|
|
|
let[@inline] yield () = Effect.perform Yield
|
|
let[@inline] suspend h = Effect.perform (Suspend h)
|
|
|
|
type with_suspend_handler =
|
|
| WSH : {
|
|
on_suspend: unit -> 'state;
|
|
(** on_suspend called when [f()] suspends itself. *)
|
|
run: 'state -> task -> unit; (** run used to schedule new tasks *)
|
|
resume: 'state -> suspension -> unit Exn_bt.result -> unit;
|
|
(** resume run the suspension. Must be called exactly once. *)
|
|
}
|
|
-> with_suspend_handler
|
|
|
|
let with_suspend (WSH { on_suspend; run; resume }) (f : unit -> unit) : unit =
|
|
let module E = Effect.Deep in
|
|
(* effect handler *)
|
|
let effc : type e. e Effect.t -> ((e, _) E.continuation -> _) option =
|
|
function
|
|
| Suspend h ->
|
|
(* TODO: discontinue [k] if current fiber (if any) is cancelled? *)
|
|
Some
|
|
(fun k ->
|
|
let state = on_suspend () in
|
|
let k' : suspension = function
|
|
| Ok () -> E.continue k ()
|
|
| Error (exn, bt) -> E.discontinue_with_backtrace k exn bt
|
|
in
|
|
h.handle ~run:(run state) ~resume:(resume state) k')
|
|
| Yield ->
|
|
(* TODO: discontinue [k] if current fiber (if any) is cancelled? *)
|
|
Some
|
|
(fun k ->
|
|
let state = on_suspend () in
|
|
let k' : suspension = function
|
|
| Ok () -> E.continue k ()
|
|
| Error (exn, bt) -> E.discontinue_with_backtrace k exn bt
|
|
in
|
|
resume state k' @@ Ok ())
|
|
| _ -> None
|
|
in
|
|
|
|
E.try_with f () { E.effc }
|
|
|
|
(* DLA interop *)
|
|
let prepare_for_await () : Dla_.t =
|
|
(* current state *)
|
|
let st : (_ * suspension) option A.t = A.make None in
|
|
|
|
let release () : unit =
|
|
match A.exchange st None with
|
|
| None -> ()
|
|
| Some (resume, k) -> resume k @@ Ok ()
|
|
and await () : unit =
|
|
suspend { handle = (fun ~run:_ ~resume k -> A.set st (Some (resume, k))) }
|
|
in
|
|
|
|
let t = { Dla_.release; await } in
|
|
t
|
|
|
|
[@@@ocaml.alert "+unstable"]
|
|
[@@@else_]
|
|
|
|
let[@inline] with_suspend ~on_suspend:_ ~run:_ ~resume:_ f = f ()
|
|
let[@inline] prepare_for_await () = { Dla_.release = ignore; await = ignore }
|
|
|
|
[@@@endif]
|