Refine the Awaitable abstraction

The use of `Trigger` is exposed to allow awaiting for multiple things.

The internal cleanup is made more robust to make sure awaitables will not be
leaked.

Author: polytypic

lib/picos_std.awaitable/dune

@@ -1,7 +1,11 @@
 (library
  (name picos_std_awaitable)
  (public_name picos_std.awaitable)
- (libraries picos picos_aux.htbl backoff multicore-magic))
+ (libraries
+  (re_export picos)
+  picos_aux.htbl
+  backoff
+  multicore-magic))
 
 (mdx
  (package picos_meta)

lib/picos_std.awaitable/picos_std_awaitable.ml

@@ -355,13 +355,14 @@ module Awaitable = struct
   let update t ~signal ~count =
     try
       let signal = ref signal in
+      let count = ref count in
       let backoff = ref Backoff.default in
       while
         not
           (let before = Htbl.find_exn awaiters t in
            match
-             if !signal then Awaiters.signal before ~count
-             else Awaiters.cleanup before ~count
+             if !signal then Awaiters.signal before ~count:!count
+             else Awaiters.cleanup before ~count:!count
            with
            | Zero -> Htbl.try_compare_and_remove awaiters t before
            | One r ->
@@ -373,58 +374,70 @@ module Awaitable = struct
                before == after
                || Htbl.try_compare_and_set awaiters t before after)
       do
+        (* Even if the hash table update after signal fails, the trigger(s) have
+           been signaled. *)
         signal := false;
+        (* If a single awaiter and multi awaiter cleanup are attempted in
+           parallel it might be that a multi awaiter cleanup "succeeds" and yet
+           some awaiters are left in the queue.  For this reason we perform a
+           multi awaiter cleanup after failure.  It might be possible to improve
+           upon this with some more clever approach. *)
+        count := Int.max_int;
         backoff := Backoff.once !backoff
       done
     with Not_found -> ()
 
-  let add_as (type a) (t : a awaitable) value =
-    let trigger = Trigger.create () in
-    let one : Awaiters.is1 =
-      One { awaitable = t; value; trigger; counter = 0; next = Min0 Zero }
-    in
-    let backoff = ref Backoff.default in
-    while
-      not
-        (match Htbl.find_exn awaiters (Packed t) with
-        | before ->
-            let many = Awaiters.snoc before one in
-            Htbl.try_compare_and_set awaiters (Packed t) before (Min1 many)
-        | exception Not_found -> Htbl.try_add awaiters (Packed t) (Min1 one))
-    do
-      backoff := Backoff.once !backoff
-    done;
-    one
-
   module Awaiter = struct
     type t = Awaiters.is1
 
-    let add (type a) (t : a awaitable) =
-      add_as t (Sys.opaque_identity (Obj.magic awaiters : a))
+    let add_as (type a) (t : a awaitable) trigger value =
+      let one : Awaiters.is1 =
+        One { awaitable = t; value; trigger; counter = 0; next = Min0 Zero }
+      in
+      let backoff = ref Backoff.default in
+      while
+        not
+          (match Htbl.find_exn awaiters (Packed t) with
+          | before ->
+              let many = Awaiters.snoc before one in
+              Htbl.try_compare_and_set awaiters (Packed t) before (Min1 many)
+          | exception Not_found -> Htbl.try_add awaiters (Packed t) (Min1 one))
+      do
+        backoff := Backoff.once !backoff
+      done;
+      one
+
+    let add (type a) (t : a awaitable) trigger =
+      let unique_value = Sys.opaque_identity (Obj.magic awaiters : a) in
+      add_as t trigger unique_value
 
     let remove one =
       Awaiters.signal_and_clear one;
       update (Awaiters.awaitable_of one) ~signal:false ~count:1
+  end
 
-    let await one =
+  let await t value =
+    let trigger = Trigger.create () in
+    let one = Awaiter.add_as t trigger value in
+    if Awaiters.is_signalable one then Awaiter.remove one
+    else
       match Awaiters.await one with
       | None -> ()
       | Some exn_bt ->
           Awaiters.clear one;
           update (Awaiters.awaitable_of one) ~signal:true ~count:1;
           Printexc.raise_with_backtrace (fst exn_bt) (snd exn_bt)
-  end
-
-  let await t value =
-    let one = add_as t value in
-    if Awaiters.is_signalable one then Awaiter.remove one else Awaiter.await one
 
   let[@inline] broadcast t = update (Packed t) ~signal:true ~count:Int.max_int
   let[@inline] signal t = update (Packed t) ~signal:true ~count:1
 
   let () =
     Stdlib.at_exit @@ fun () ->
     match Htbl.find_random_exn awaiters with
-    | _ -> failwith "leaked awaitable"
+    | _ ->
+        (* This should not normally happen, but might happen due to the program
+           being forced to exit without proper cleanup.  Otherwise this may
+           indicate a bug in the cleanup of awaiters. *)
+        Printf.eprintf "Awaitable leaked\n%!"
     | exception Not_found -> ()
 end

lib/picos_std.awaitable/picos_std_awaitable.mli

@@ -1,6 +1,8 @@
 (** Basic {{:https://en.wikipedia.org/wiki/Futex} futex}-like awaitable atomic
     location for {!Picos}. *)
 
+open Picos
+
 (** {1 Modules} *)
 
 module Awaitable : sig
@@ -18,7 +20,7 @@ module Awaitable : sig
 
   (** {1 Atomic API} *)
 
-  type 'a t
+  type !'a t
   (** Represents an awaitable atomic location. *)
 
   val make : ?padded:bool -> 'a -> 'a t
@@ -80,7 +82,7 @@ module Awaitable : sig
       explicitly {!signal}ed and has a value other than [before].
 
       ⚠️ This operation is subject to the
-      {{:https://en.wikipedia.org/wiki/ABA_problem} ABA} problems.  An [await]
+      {{:https://en.wikipedia.org/wiki/ABA_problem} ABA} problem.  An [await]
       for value other than [A] may not return after the awaitable is signaled
       while having the value [B], because at a later point the awaitable has
       again the value [A].  Furthermore, by the time an [await] for value other
@@ -90,34 +92,26 @@ module Awaitable : sig
       implicitly wake up awaiters. *)
 
   module Awaiter : sig
-    (** Ability to await for a signal from the past.
-
-        {!Awaitable.await} only receives a signal at or after the point of
-        calling it.  This API allows the awaiting process to be broken into two
-        steps, {!add} and {!await}, such that a signal after {!add} can be
-        received by {!await}. *)
+    (** Low level interface for more flexible waiting. *)
 
     type 'a awaitable := 'a t
     (** An erased type alias for {!Awaitable.t}. *)
 
     type t
     (** Represents a single use awaiter of a signal to an {!awaitable}. *)
 
-    val add : 'a awaitable -> t
-    (** [add awaitable] create a single use awaiter, adds it to the FIFO
-        associated with the awaitable, and returns the awaiter. *)
-
-    val await : t -> unit
-    (** [await awaiter] awaits for the association awaitable to be signaled. *)
+    val add : 'a awaitable -> Trigger.t -> t
+    (** [add awaitable trigger] creates a single use awaiter, adds it to the
+        FIFO associated with the awaitable, and returns the awaiter. *)
 
     val remove : t -> unit
     (** [remove awaiter] marks the awaiter as having been signaled and removes it
         from the FIFO associated with the awaitable.
 
-        ⚠️ An explicit call of [remove] is needed when an {!add}ed awaiter is not
-        {!await}ed for.  In such a case, from the point of view of lost signals,
-        the caller of [remove] should be considered to have received or consumed
-        a signal before the call of [remove]. *)
+        ℹ️ If the associated trigger is used with only one awaiter and the
+        {!Trigger.await await} on the trigger returns [None], there is no need
+        to explicitly remove the awaiter, because it has already been
+        removed. *)
   end
 end
 
@@ -164,7 +158,7 @@ end
     {2 [Condition]}
 
     Let's also implement a condition variable.  For that we'll also make use of
-    low level operations in the {!Picos} core library:
+    low level abstractions and operations from the {!Picos} core library:
 
     {[
       # open Picos
@@ -180,20 +174,21 @@ end
         let create () = Awaitable.make ()
 
         let wait t mutex =
-          let awaiter = Awaitable.Awaiter.add t in
+          let trigger = Trigger.create () in
+          let awaiter = Awaitable.Awaiter.add t trigger in
           Mutex.unlock mutex;
           let lock_forbidden mutex =
             let fiber = Fiber.current () in
             let forbid = Fiber.exchange fiber ~forbid:true in
             Mutex.lock mutex;
             Fiber.set fiber ~forbid
           in
-          match Awaitable.Awaiter.await awaiter with
-          | () -> lock_forbidden mutex
-          | exception exn ->
-              let bt = Printexc.get_raw_backtrace () in
+          match Trigger.await trigger with
+          | None -> lock_forbidden mutex
+          | Some exn_bt ->
+              Awaitable.Awaiter.remove awaiter;
               lock_forbidden mutex;
-              Printexc.raise_with_backtrace exn bt
+              Printexc.raise_with_backtrace (fst exn_bt) (snd exn_bt)
 
         let signal = Awaitable.signal
         let broadcast = Awaitable.broadcast