Add Capsule.protect

otherlibs/stdlib_alpha/capsule.ml

@@ -121,6 +121,9 @@ end
    it never returns is also [portable] *)
 external reraise : exn -> 'a @ portable @@ portable = "%reraise"
 
+external raise_with_backtrace :
+  exn -> Printexc.raw_backtrace -> 'a @ portable @@ portable = "%raise_with_backtrace"
+
 module Data = struct
   type ('a, 'k) t : value mod portable uncontended
 
@@ -349,10 +352,39 @@ module Rwlock = struct
 
 end
 
+module Condition = struct
+
+  type 'k t : value mod portable uncontended
+
+  external create : unit -> 'k t @@ portable = "caml_ml_condition_new"
+  external wait : 'k t -> M.t -> unit @@ portable = "caml_ml_condition_wait"
+  external signal : 'k t -> unit @@ portable = "caml_ml_condition_signal"
+  external broadcast : 'k t -> unit @@ portable = "caml_ml_condition_broadcast"
+
+  let wait t (mut : 'k Mutex.t) _password =
+    (* [mut] is locked, so we know it is not poisoned. *)
+    wait t mut.mutex
+
+end
+
 let create_with_mutex () =
   let (P name) = Name.make () in
   Mutex.P { name; mutex = M.create (); poisoned = false }
 
 let create_with_rwlock () =
   let (P name) = Name.make () in
   Rwlock.P { name; rwlock = Rw.create (); poisoned = false }
+
+exception Protected : 'k Mutex.t * (exn, 'k) Data.t -> exn
+
+(* CR-soon mslater: replace with portable stdlib *)
+let get_raw_backtrace : unit -> Printexc.raw_backtrace @@ portable =
+  O.magic O.magic Printexc.get_raw_backtrace
+
+let protect f =
+  try f () with
+  | exn ->
+    let (P mut) = create_with_mutex () in
+    raise_with_backtrace (Protected (mut, Data.unsafe_mk exn)) (get_raw_backtrace ())
+  ;;
+

otherlibs/stdlib_alpha/capsule.mli

@@ -147,9 +147,6 @@ val access_shared :
     with a shared {!Access.t} for ['k]. The result is within ['k] so it
     must be [portable] and it is marked [contended]. *)
 
-(** Mutual exclusion primtives for controlling uncontended access to a capsule.
-
-    Requires OCaml 5 runtime. *)
 module Mutex : sig
 
     type 'k t : value mod portable uncontended
@@ -236,6 +233,40 @@ module Rwlock : sig
         lock as poisoned. *)
 end
 
+module Condition : sig
+
+    type 'k t : value mod portable uncontended
+    (** ['k t] is the type of a condition variable associated with the capsule ['k].
+        This condition may only be used with the matching ['k Mutex.t]. *)
+
+  val create : unit -> 'k t @@ portable
+  (** [create ()] creates and returns a new condition variable.
+      This condition variable is associated with the matching ['k Mutex.t]
+      and with a certain property {i P} that is protected by the mutex. *)
+
+  val wait : 'k t -> 'k Mutex.t -> 'k Password.t @ local -> unit @@ portable
+  (** [wait c m] atomically unlocks the mutex [m] and suspends the
+      current thread on the condition variable [c]. This thread can
+      later be woken up after the condition variable [c] has been signaled
+      via {!signal} or {!broadcast}; however, it can also be woken up for
+      no reason. The mutex [m] is locked again before [wait] returns. One
+      cannot assume that the property {i P} associated with the condition
+      variable [c] holds when [wait] returns; one must explicitly test
+      whether {i P} holds after calling [wait]. *)
+
+  val signal : 'k t -> unit @@ portable
+  (** [signal c] wakes up one of the threads waiting on the condition
+      variable [c], if there is one. If there is none, this call has no effect.
+      It is recommended to call [signal c] inside a critical section, that is,
+      while the mutex [m] associated with [c] is locked. *)
+
+  val broadcast : 'k t -> unit @@ portable
+  (** [broadcast c] wakes up all threads waiting on the condition
+      variable [c]. If there are none, this call has no effect.
+      It is recommended to call [broadcast c] inside a critical section,
+      that is, while the mutex [m] associated with [c] is locked. *)
+end
+
 val create_with_mutex : unit -> Mutex.packed @@ portable
 (** [create_with_mutex ()] creates a new capsule with an associated mutex. *)
 
@@ -373,11 +404,23 @@ module Data : sig
         so it must be [portable] and it is marked [contended]. Since [nonportable]
         functions may enclose [uncontended] (and thus write) access to data,
         ['a] must cross [portability] *)
-
 end
 
 exception Encapsulated : 'k Name.t * (exn, 'k) Data.t -> exn
 (** If a function accessing the contents of the capsule raises an
     exception, it is wrapped in [Encapsulated] to avoid leaking access to
     the data. The [Name.t] can be used to associate the [Data.t] with a
     particular [Password.t] or [Mutex.t]. *)
+
+exception Protected : 'k Mutex.t * (exn, 'k) Data.t -> exn
+(** If a function passed to [protect] raises an exception, it is wrapped
+    in [Protected] to provide access to the capsule in which the function ran. *)
+(* CR-soon mslater: this should return a key, not a mutex. *)
+
+val protect
+  :  (unit -> 'a @ portable contended) @ local portable
+  -> 'a @ portable contended
+  @@ portable
+(** [protect f] runs [f] in a fresh capsule. If [f] returns normally, [protect]
+    merges this capsule into the caller's capsule. If [f] raises, [protect]
+    raises [Protected], giving the caller access to the encapsulated exception. *)

testsuite/tests/capsule-api/condition.ml

@@ -0,0 +1,53 @@
+(* TEST
+ include stdlib_alpha;
+ flags = "-extension-universe alpha";
+ runtime5;
+ { bytecode; }
+ { native; }
+*)
+
+[@@@ocaml.alert "-unsafe_parallelism"]
+
+module Capsule = Stdlib_alpha.Capsule
+
+external ref : 'a -> 'a ref @@ portable = "%makemutable"
+external ( ! ) : 'a ref -> 'a @@ portable = "%field0"
+external ( := ) : 'a ref -> 'a -> unit @@ portable = "%setfield0"
+
+let () = (* Signal *)
+  let (P mut) = Capsule.create_with_mutex () in
+  let cond = Capsule.Condition.create () in
+  let go = Capsule.Data.create (fun () -> ref true) in
+  let wait = Atomic.make true in
+  let domain = Domain.spawn (fun () ->
+    Capsule.Mutex.with_lock mut (fun password ->
+      Atomic.set wait false;
+      while Capsule.Data.extract password (fun go : bool -> !go) go do
+        Capsule.Condition.wait cond mut password
+      done))
+  in
+  while Atomic.get wait do () done;
+  Capsule.Mutex.with_lock mut (fun password ->
+    Capsule.Data.iter password (fun go -> go := false) go;
+  Capsule.Condition.signal cond);
+  Domain.join domain
+;;
+
+let () = (* Broadcast *)
+  let (P mut) = Capsule.create_with_mutex () in
+  let cond = Capsule.Condition.create () in
+  let go = Capsule.Data.create (fun () -> ref true) in
+  let ready = Atomic.make 0 in
+  let domains = List.init 4 (fun _ -> Domain.spawn (fun () ->
+    Capsule.Mutex.with_lock mut (fun password ->
+      Atomic.incr ready;
+      while Capsule.Data.extract password (fun go : bool -> !go) go do
+        Capsule.Condition.wait cond mut password
+      done)))
+  in
+  while Atomic.get ready < 4 do () done;
+  Capsule.Mutex.with_lock mut (fun password ->
+    Capsule.Data.iter password (fun go -> go := false) go;
+  Capsule.Condition.broadcast cond);
+  List.iter Domain.join domains
+;;

testsuite/tests/capsule-api/data.ml

@@ -176,3 +176,31 @@ let ptr' : (int, lost_capsule) Capsule.Data.t =
 let () =
   assert (Capsule.Data.project ptr' = 111)
 ;;
+
+
+(* [protect]. *)
+exception Exn of string
+
+let () =
+  match Capsule.protect (fun () -> "ok") with
+  | s -> assert (s = "ok")
+  | exception _ -> assert false
+;;
+
+let () =
+  match Capsule.protect (fun () -> Exn "ok") with
+  | Exn s -> assert (s = "ok")
+  | _ -> assert false
+;;
+
+let () =
+  match Capsule.protect (fun () -> reraise (Exn "fail")) with
+  | exception (Capsule.Protected (mut, exn)) ->
+    let s = Capsule.Mutex.with_lock mut (fun password ->
+      Capsule.Data.extract password (fun exn ->
+        match exn with
+        | Exn s -> s
+        | _ -> assert false) exn) in
+    assert (s = "fail")
+  | _ -> assert false
+;;