Source file bounded_queue_sync.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
module BQ = Bounded_queue

type push_res =
  | Closed
  | Pushed of { num_discarded: int }

(* a variant of {!Sync_queue} with more bespoke pushing behavior *)
module Q : sig
  type 'a t

  val create : unit -> 'a t

  val close : _ t -> unit

  val size : _ t -> int

  val closed : _ t -> bool

  val try_pop : 'a t -> 'a BQ.pop_result

  val push_while_not_full : high_watermark:int -> 'a t -> 'a list -> push_res
  (** [push_while_not_full q ~high_watermark xs] tries to push each item of [x]
      into [q].

      An item is not pushed if the queue is "full" (size >= high_watermark).

      This returns a pair [num_discarded, old_size] where [num_discarded] is the
      number of items that could not be pushed, and [old_size] is the size
      before anything was pushed. *)
end = struct
  module UM = Util_mutex

  type 'a t = {
    mutex: Mutex.t;
    q: 'a Queue.t;
    mutable closed: bool;
  }

  let create () : _ t =
    { mutex = Mutex.create (); q = Queue.create (); closed = false }

  (* NOTE: the race condition here is benign, assuming no tearing of
    a value of type [bool] which OCaml's memory model should guarantee. *)
  let[@inline] closed self = self.closed

  let[@inline] size self = UM.protect self.mutex (fun () -> Queue.length self.q)

  let close (self : _ t) =
    UM.protect self.mutex @@ fun () ->
    if not self.closed then self.closed <- true

  let try_pop (self : 'a t) : 'a BQ.pop_result =
    UM.protect self.mutex @@ fun () ->
    (* first, try to pop the queue. We want to drain it even if it's closed. *)
    try `Item (Queue.pop self.q)
    with Queue.Empty ->
      if self.closed then
        `Closed
      else
        `Empty

  let push_while_not_full ~high_watermark (self : 'a t) (xs : 'a list) :
      push_res =
    UM.protect self.mutex @@ fun () ->
    if self.closed then
      Closed
    else (
      let to_push = ref xs in

      let continue = ref true in
      while !continue && Queue.length self.q < high_watermark do
        match !to_push with
        | [] -> continue := false
        | x :: tl_xs ->
          to_push := tl_xs;
          Queue.push x self.q
      done;

      let num_discarded = List.length !to_push in
      (* Printf.eprintf "bq: pushed %d items (discarded: %d)\n%!" (List.length xs - num_discarded) num_discarded; *)

      Pushed { num_discarded }
    )
end

type 'a state = {
  n_discarded: int Atomic.t;
  high_watermark: int;
  q: 'a Q.t;
  on_non_empty: Cb_set.t;
  measure: 'a -> int;
}

let measure_all_ measure xs = List.fold_left (fun acc x -> acc + measure x) 0 xs

let push (self : _ state) x =
  if x <> [] then (
    match
      Q.push_while_not_full self.q ~high_watermark:self.high_watermark x
    with
    | Closed ->
      let n = measure_all_ self.measure x in
      ignore (Atomic.fetch_and_add self.n_discarded n : int)
    | Pushed { num_discarded } ->
      if num_discarded > 0 then (
        let n_signals = measure_all_ self.measure x in
        let total = Atomic.fetch_and_add self.n_discarded n_signals in
        Opentelemetry.Self_debug.log Warning (fun () ->
            Printf.sprintf
              "otel: dropped %d signals (queue full: %d/%d, total dropped: %d)"
              n_signals (Q.size self.q) self.high_watermark (total + n_signals))
      );
      (* wake up potentially asleep consumers *)
      Cb_set.trigger self.on_non_empty
  )

let[@inline] try_pop (self : _ state) : _ BQ.pop_result = Q.try_pop self.q

let to_bounded_queue (self : 'a state) : 'a BQ.t =
  let closed () = Q.closed self.q in
  let num_discarded () = Atomic.get self.n_discarded in
  let push x = push self x in
  let on_non_empty = Cb_set.register self.on_non_empty in
  let try_pop () = try_pop self in
  let size () = Q.size self.q in
  let high_watermark () = self.high_watermark in
  let close () =
    Q.close self.q;
    (* waiters will want to know *)
    Cb_set.trigger self.on_non_empty
  in
  let common = { BQ.Common.closed; num_discarded; size; high_watermark } in
  {
    BQ.send = { push; close; common };
    recv = { try_pop; on_non_empty; common };
  }

let create ?(measure = fun _ -> 1) ~high_watermark () : _ BQ.t =
  let st =
    {
      high_watermark;
      q = Q.create ();
      n_discarded = Atomic.make 0;
      on_non_empty = Cb_set.create ();
      measure;
    }
  in
  to_bounded_queue st