missed some uses of Atomic

This commit is contained in:
Simon Cruanes 2023-12-27 21:33:04 -05:00
parent 6aeb1ea007
commit 5571751f3e
No known key found for this signature in database
GPG key ID: EBFFF6F283F3A2B4
2 changed files with 15 additions and 11 deletions

View file

@ -1,9 +1,11 @@
module A = Trace_core.Internal_.Atomic_
type 'a t = { type 'a t = {
mutex: Mutex.t; mutex: Mutex.t;
cond: Condition.t; cond: Condition.t;
q: 'a Mpsc_bag.t; q: 'a Mpsc_bag.t;
mutable closed: bool; mutable closed: bool;
consumer_waiting: bool Atomic.t; consumer_waiting: bool A.t;
} }
exception Closed exception Closed
@ -14,7 +16,7 @@ let create () : _ t =
cond = Condition.create (); cond = Condition.create ();
q = Mpsc_bag.create (); q = Mpsc_bag.create ();
closed = false; closed = false;
consumer_waiting = Atomic.make false; consumer_waiting = A.make false;
} }
let close (self : _ t) = let close (self : _ t) =
@ -29,7 +31,7 @@ let push (self : _ t) x : unit =
if self.closed then raise Closed; if self.closed then raise Closed;
Mpsc_bag.add self.q x; Mpsc_bag.add self.q x;
if self.closed then raise Closed; if self.closed then raise Closed;
if Atomic.get self.consumer_waiting then ( if A.get self.consumer_waiting then (
(* wakeup consumer *) (* wakeup consumer *)
Mutex.lock self.mutex; Mutex.lock self.mutex;
Condition.broadcast self.cond; Condition.broadcast self.cond;
@ -42,14 +44,14 @@ let rec pop_all (self : 'a t) : 'a list =
| None -> | None ->
if self.closed then raise Closed; if self.closed then raise Closed;
Mutex.lock self.mutex; Mutex.lock self.mutex;
Atomic.set self.consumer_waiting true; A.set self.consumer_waiting true;
(* check again, a producer might have pushed an element since we (* check again, a producer might have pushed an element since we
last checked. However if we still find last checked. However if we still find
nothing, because this comes after [consumer_waiting:=true], nothing, because this comes after [consumer_waiting:=true],
any producer arriving after that will know to wake us up. *) any producer arriving after that will know to wake us up. *)
(match Mpsc_bag.pop_all self.q with (match Mpsc_bag.pop_all self.q with
| Some l -> | Some l ->
Atomic.set self.consumer_waiting false; A.set self.consumer_waiting false;
Mutex.unlock self.mutex; Mutex.unlock self.mutex;
l l
| None -> | None ->
@ -58,6 +60,6 @@ let rec pop_all (self : 'a t) : 'a list =
raise Closed raise Closed
); );
Condition.wait self.cond self.mutex; Condition.wait self.cond self.mutex;
Atomic.set self.consumer_waiting false; A.set self.consumer_waiting false;
Mutex.unlock self.mutex; Mutex.unlock self.mutex;
pop_all self) pop_all self)

View file

@ -1,7 +1,9 @@
type 'a t = { bag: 'a list Atomic.t } [@@unboxed] module A = Trace_core.Internal_.Atomic_
type 'a t = { bag: 'a list A.t } [@@unboxed]
let create () = let create () =
let bag = Atomic.make [] in let bag = A.make [] in
{ bag } { bag }
module Backoff = struct module Backoff = struct
@ -17,14 +19,14 @@ module Backoff = struct
end end
let rec add backoff t x = let rec add backoff t x =
let before = Atomic.get t.bag in let before = A.get t.bag in
let after = x :: before in let after = x :: before in
if not (Atomic.compare_and_set t.bag before after) then if not (A.compare_and_set t.bag before after) then
add (Backoff.once backoff) t x add (Backoff.once backoff) t x
let[@inline] add t x = add Backoff.default t x let[@inline] add t x = add Backoff.default t x
let[@inline] pop_all t : _ list option = let[@inline] pop_all t : _ list option =
match Atomic.exchange t.bag [] with match A.exchange t.bag [] with
| [] -> None | [] -> None
| l -> Some (List.rev l) | l -> Some (List.rev l)