diff --git a/lib_eio/core/eio__core.mli b/lib_eio/core/eio__core.mli index 830680a86..98b0aa6a9 100644 --- a/lib_eio/core/eio__core.mli +++ b/lib_eio/core/eio__core.mli @@ -284,6 +284,9 @@ module Fiber : sig Many operations automatically check this before starting. @raise Cancel.Cancelled if the fiber's context has been cancelled. *) + val is_cancelled : unit -> bool + (** [is_cancelled ()] is [true] iff {!check} would raise an exception. *) + val yield : unit -> unit (** [yield ()] asks the scheduler to switch to the next runnable task. The current task remains runnable, but goes to the back of the queue. diff --git a/lib_eio/core/fiber.ml b/lib_eio/core/fiber.ml index 14e40f44d..617590499 100644 --- a/lib_eio/core/fiber.ml +++ b/lib_eio/core/fiber.ml @@ -144,6 +144,10 @@ let any fs = let first f g = any [f; g] +let is_cancelled () = + let ctx = Effect.perform Cancel.Get_context in + not (Cancel.is_on ctx.cancel_context) + let check () = let ctx = Effect.perform Cancel.Get_context in Cancel.check ctx.cancel_context diff --git a/lib_eio/eio.ml b/lib_eio/eio.ml index 909ca5894..4eb68911b 100644 --- a/lib_eio/eio.ml +++ b/lib_eio/eio.ml @@ -8,6 +8,7 @@ module Semaphore = Semaphore module Mutex = Eio_mutex module Condition = Condition module Stream = Stream +module Lazy = Lazy module Exn = Exn module Resource = Resource module Flow = Flow diff --git a/lib_eio/eio.mli b/lib_eio/eio.mli index 3f26dcb04..349ad0585 100644 --- a/lib_eio/eio.mli +++ b/lib_eio/eio.mli @@ -36,6 +36,9 @@ module Condition = Condition (** A stream/queue. *) module Stream = Stream +(** Delayed evaluation. *) +module Lazy = Lazy + (** Cancelling fibers. *) module Cancel = Eio__core.Cancel diff --git a/lib_eio/lazy.ml b/lib_eio/lazy.ml new file mode 100644 index 000000000..0917ed1f2 --- /dev/null +++ b/lib_eio/lazy.ml @@ -0,0 +1,47 @@ +open Std + +type 'a state = + | Value of 'a + | Waiting of (unit Promise.u -> unit) + | Running of unit Promise.t (* Wait until resolved and check again *) + | Failed of Exn.with_bt + +type 'a t = 'a state Atomic.t + +let init = Waiting (fun _ -> assert false) + +let from_fun ~cancel fn = + let state = Atomic.make init in + let rec force r = + match + if cancel = `Protect then Cancel.protect fn else fn () + with + | x -> + Atomic.set state (Value x); + Promise.resolve r () + | exception ex -> + let bt = Printexc.get_raw_backtrace () in + match ex with + | Cancel.Cancelled _ when cancel = `Restart && Fiber.is_cancelled () -> + Atomic.set state (Waiting force); + Promise.resolve r (); + Fiber.check () + | _ -> + Atomic.set state (Failed (ex, bt)); + Promise.resolve r (); + Printexc.raise_with_backtrace ex bt + in + Atomic.set state @@ Waiting force; + state + +let from_val v = Atomic.make (Value v) + +let rec force t = + match Atomic.get t with + | Value v -> v + | Failed (ex, bt) -> Printexc.raise_with_backtrace ex bt + | Running p -> Promise.await p; force t + | Waiting fn as prev -> + let p, r = Promise.create () in + if Atomic.compare_and_set t prev (Running p) then fn r; + force t diff --git a/lib_eio/lazy.mli b/lib_eio/lazy.mli new file mode 100644 index 000000000..46eaaf996 --- /dev/null +++ b/lib_eio/lazy.mli @@ -0,0 +1,27 @@ +(** This is like [Stdlib.Lazy], but multiple fibers or domains can force at once. *) + +type 'a t +(** A lazy value that produces a value of type ['a]. *) + +val from_fun : + cancel:[`Restart | `Record | `Protect] -> + (unit -> 'a) -> 'a t +(** [from_fun ~cancel fn] is a lazy value that runs [fn ()] the first time it is forced. + + [cancel] determines how cancellation is handled while forcing: + + - [`Restart] : if the forcing fiber is cancelled, the next waiting fiber runs [fn] again. + - [`Record] : the failure is recorded and the lazy value will always report cancelled if used. + - [`Protect] : the forcing fiber is protected from cancellation while running. *) + +val from_val : 'a -> 'a t +(** [from_val v] is a lazy value that is already forced. + + It is equivalent to [from_fun (fun () -> v)]. *) + +val force : 'a t -> 'a +(** [force t] returns the result of running the function passed to {!from_fun}. + + If the function is currently running, this waits for it to finish and then retries. + If the function has already completed then it returns the saved result. + If the function returned an exception then [force] re-raises it. *) diff --git a/tests/lazy.md b/tests/lazy.md new file mode 100644 index 000000000..93f23ea65 --- /dev/null +++ b/tests/lazy.md @@ -0,0 +1,131 @@ +# Setting up the environment + +```ocaml +# #require "eio.mock";; +``` +```ocaml +open Eio.Std + +let test label v = + traceln "%s: forcing..." label; + match Eio.Lazy.force v with + | v -> + Fiber.check (); + traceln "%s: %d" label v + | exception ex -> + traceln "%s: %a" label Fmt.exn ex; + Fiber.check () +``` + +# Tests + +Two fibers request the value. It's only computed once: + +```ocaml +# Eio_mock.Backend.run @@ fun () -> + let v = Eio.Lazy.from_fun ~cancel:`Restart (fun () -> + traceln "calculating..."; + Fiber.yield (); + traceln "complete"; + 42 + ) in + Fiber.both + (fun () -> test "a" v) + (fun () -> test "b" v) + ;; ++a: forcing... ++calculating... ++b: forcing... ++complete ++a: 42 ++b: 42 +- : unit = () +``` + +The calculation fails. It's still only performed once: + +```ocaml +# Eio_mock.Backend.run @@ fun () -> + let v = Eio.Lazy.from_fun ~cancel:`Restart (fun () -> + traceln "calculating..."; + Fiber.yield (); + failwith "failed"; + ) in + Fiber.both + (fun () -> test "a" v) + (fun () -> test "b" v) + ;; ++a: forcing... ++calculating... ++b: forcing... ++a: Failure("failed") ++b: Failure("failed") +- : unit = () +``` + +## Cancellation + +The first fiber cancels. What happens depends on the cancel mode: + +```ocaml +let test_cancel cancel = + Eio_mock.Backend.run @@ fun () -> + let v = Eio.Lazy.from_fun ~cancel (fun () -> + traceln "calculating..."; + Fiber.yield (); + traceln "complete"; + 42 + ) in + Fiber.both + (fun () -> + let x = + Fiber.first + (fun () -> test "a" v; assert false) + (fun () -> 5) + in + traceln "a: %d" x + ) + (fun () -> test "b" v) + ;; +``` + +In record mode, the second fiber sees the cancelled exception: + +```ocaml +# test_cancel `Record;; ++a: forcing... ++calculating... ++b: forcing... ++a: Cancelled: Eio__core__Fiber.Not_first ++b: Cancelled: Eio__core__Fiber.Not_first ++a: 5 +- : unit = () +``` + +In protect mode, the first calculation succeeds: + +```ocaml +# test_cancel `Protect;; ++a: forcing... ++calculating... ++b: forcing... ++complete ++b: 42 ++a: 5 +- : unit = () +``` + +In restart mode, the second fiber restarts the calculation: + +```ocaml +# test_cancel `Restart;; ++a: forcing... ++calculating... ++b: forcing... ++a: Cancelled: Eio__core__Fiber.Not_first ++calculating... ++a: 5 ++complete ++b: 42 +- : unit = () +```