Add Eio.Lazy

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.

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

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

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
 

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

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. *)

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 = ()
+```