Add add_nonblocking, capacity and is_full for streams

Author: xavierm02

lib_eio/stream.ml

@@ -1,3 +1,5 @@
+type drop_priority = Newest | Oldest
+
 module Locking = struct
   type 'a t = {
     mutex : Mutex.t;
@@ -64,6 +66,26 @@ module Locking = struct
           )
       )
 
+  let add_nonblocking ~drop_priority t item =
+    Mutex.lock t.mutex;
+    match Waiters.wake_one t.readers item with
+    | `Ok -> Mutex.unlock t.mutex; None
+    | `Queue_empty ->
+      (* No-one is waiting for an item. Queue it. *)
+      if Queue.length t.items < t.capacity then (
+        Queue.add item t.items;
+        Mutex.unlock t.mutex;
+        None
+      ) else (
+        match drop_priority with
+        | Newest -> Mutex.unlock t.mutex; Some item
+        | Oldest ->
+          let dropped_item = Queue.take t.items in
+          Queue.add item t.items;
+          Mutex.unlock t.mutex;
+          Some dropped_item
+      )
+
   let take t =
     Mutex.lock t.mutex;
     match Queue.take_opt t.items with
@@ -101,6 +123,8 @@ module Locking = struct
     let len = Queue.length t.items in
     Mutex.unlock t.mutex;
     len
+  
+  let capacity t = t.capacity
 
   let dump f t =
     Fmt.pf f "<Locking stream: %d/%d items>" (length t) t.capacity
@@ -123,6 +147,11 @@ let take = function
   | Sync x -> Sync.take x |> Result.get_ok     (* todo: allow closing streams *)
   | Locking x -> Locking.take x
 
+let add_nonblocking ~drop_priority t v =
+  match t with
+  | Sync x -> if Sync.put_nonblocking x v then None else (Some v)
+  | Locking x -> Locking.add_nonblocking ~drop_priority x v
+
 let take_nonblocking = function
   | Locking x -> Locking.take_nonblocking x
   | Sync x ->
@@ -134,8 +163,14 @@ let length = function
   | Sync _ -> 0
   | Locking x -> Locking.length x
 
+let capacity = function
+  | Sync _ -> 0
+  | Locking x -> Locking.capacity x
+
 let is_empty t = (length t = 0)
 
+let is_full t = (length t = capacity t)
+
 let dump f = function
   | Sync x -> Sync.dump f x
   | Locking x -> Locking.dump f x

lib_eio/stream.mli

@@ -33,6 +33,21 @@ val take : 'a t -> 'a
 
     If no items are available, it waits until one becomes available. *)
 
+type drop_priority = Newest | Oldest
+
+val add_nonblocking : drop_priority: drop_priority -> 'a t -> 'a -> 'a option
+(** [add_nonblocking ~drop_priority t item] is like [(add t item); None] except that
+    it returns [Some dropped_item] if the stream is full rather than waiting, where
+    [dropped_item] is [item] if [drop_priority = Newest], and the first element of the
+    stream if [drop_priority = Oldest].
+        
+    In other words, if the stream is full then:
+    - [add_nonblocking ~drop_priority:Newest t item] is like [Some item]; and
+    - [add_nonblocking ~drop_priority:Oldest t item] is like
+      [let dropped_item = take t in add t item; Some dropped_item]
+      except that no other stream operation can happen (even in other threads)
+      between the [take] and the [add]. *)
+
 val take_nonblocking : 'a t -> 'a option
 (** [take_nonblocking t] is like [Some (take t)] except that
     it returns [None] if the stream is empty rather than waiting.
@@ -43,8 +58,14 @@ val take_nonblocking : 'a t -> 'a option
 val length : 'a t -> int
 (** [length t] returns the number of items currently in [t]. *)
 
+val capacity : 'a t -> int
+(** [capacity t] returns the number of items [t] can hold without blocking writers. *)
+
 val is_empty : 'a t -> bool
 (** [is_empty t] is [length t = 0]. *)
 
+val is_full : 'a t -> bool
+(** [is_full t] is [length t = capacity t]. *)
+
 val dump : 'a t Fmt.t
 (** For debugging. *)

lib_eio/sync.ml

@@ -441,6 +441,22 @@ let rec put (t : _ t) v =
     ~suspend:(put_suspend t v)
     ~closed:(fun () -> raise put_closed_err)
 
+let reject = Slot (fun _ -> false)
+
+let put_nonblocking (t : _ t) v =
+  match Balance.incr_if_negative t.balance with
+  | Balance_closed -> raise put_closed_err
+  | Update_refused -> false
+  | Updated ->
+    let rec aux cell =
+    producer_resume_cell t cell
+      ~success:(fun kc -> (if kc (Ok v) then () else put t v); true)
+      ~in_transition:(fun cell ->
+        Domain.cpu_relax ();
+        if Atomic.compare_and_set cell In_transition reject then false
+        else aux cell)
+    in aux (Q.next_resume t.consumers)
+
 (* Taking. *)
 
 (* Mirror of [producer_resume_cell]. *)
@@ -489,8 +505,6 @@ let take t =
    :  (_, [  `Closed ]) result 
    :> (_, [> `Closed ]) result)
 
-let reject = Slot (fun _ -> false)
-
 let take_nonblocking (t : _ t) =
   match Balance.decr_if_positive t.balance with
   | Balance_closed -> Error `Closed

lib_eio/sync.mli

@@ -22,6 +22,12 @@ val put : 'a t -> 'a -> unit
 
     @raise Invalid_argument if [t] was closed before [x] was added. *)
 
+val put_nonblocking : 'a t -> 'a -> bool
+(** [put_nonblocking t x] is like [put t x; true], except that it returns [false]
+    instead of waiting if no consumer is immediately available.
+
+    @raise Invalid_argument if [t] was closed before [x] was added. *)
+
 val take : 'a t -> ('a, [> `Closed]) result
 (** [take t] waits until a producer is available with an item and then returns it.
 

tests/stream.md

@@ -20,6 +20,16 @@ let add t v =
   S.add t v;
   traceln "Added %d to stream" v
 
+let add_nonblocking ~drop_priority t v =
+  traceln "Adding %d to stream" v;
+  match S.add_nonblocking ~drop_priority t v with
+  | None -> traceln "Added %d to stream" v
+  | Some d ->
+    match drop_priority, S.capacity t with
+    | Newest, _ | _, 0 -> assert (d = v); traceln "Dropped %i instead of adding it to stream" v
+    | Oldest, _ -> traceln "Dropped %i from stream and added %i to stream" d v
+  
+
 let take t =
   traceln "Reading from stream";
   traceln "Got %d from stream" (S.take t)
@@ -320,6 +330,48 @@ Cancelling writing to a stream:
 - : unit = ()
 ```
 
+Non-blocking add (capacity = 1):
+
+```ocaml
+# run @@ fun () ->
+  let t = S.create 1 in
+  add t 0;
+  add_nonblocking ~drop_priority:Newest t 1;
+  add_nonblocking ~drop_priority:Oldest t 2;;
++Adding 0 to stream
++Added 0 to stream
++Adding 1 to stream
++Dropped 1 instead of adding it to stream
++Adding 2 to stream
++Dropped 0 from stream and added 2 to stream
+- : unit = ()
+```
+
+Non-blocking add (capacity = 0):
+
+```ocaml
+# run @@ fun () ->
+  let t = S.create 0 in
+  Fiber.both
+    (fun () -> ignore @@ take t; ignore @@ take t)
+    (fun () ->
+      add_nonblocking ~drop_priority:Newest t 1;
+      add_nonblocking ~drop_priority:Newest t 2;
+      Fiber.yield ();
+      add_nonblocking ~drop_priority:Newest t 3);;
++Reading from stream
++Adding 1 to stream
++Added 1 to stream
++Adding 2 to stream
++Dropped 2 instead of adding it to stream
++Got 1 from stream
++Reading from stream
++Adding 3 to stream
++Added 3 to stream
++Got 3 from stream
+- : unit = ()
+```
+
 Non-blocking take:
 
 ```ocaml