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 = {
mutex: Mutex.t;
cond: Condition.t;
q: 'a Mpsc_bag.t;
mutable closed: bool;
consumer_waiting: bool Atomic.t;
consumer_waiting: bool A.t;
}
exception Closed
@ -14,7 +16,7 @@ let create () : _ t =
cond = Condition.create ();
q = Mpsc_bag.create ();
closed = false;
consumer_waiting = Atomic.make false;
consumer_waiting = A.make false;
}
let close (self : _ t) =
@ -29,7 +31,7 @@ let push (self : _ t) x : unit =
if self.closed then raise Closed;
Mpsc_bag.add self.q x;
if self.closed then raise Closed;
if Atomic.get self.consumer_waiting then (
if A.get self.consumer_waiting then (
(* wakeup consumer *)
Mutex.lock self.mutex;
Condition.broadcast self.cond;
@ -42,14 +44,14 @@ let rec pop_all (self : 'a t) : 'a list =
| None ->
if self.closed then raise Closed;
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
last checked. However if we still find
nothing, because this comes after [consumer_waiting:=true],
any producer arriving after that will know to wake us up. *)
(match Mpsc_bag.pop_all self.q with
| Some l ->
Atomic.set self.consumer_waiting false;
A.set self.consumer_waiting false;
Mutex.unlock self.mutex;
l
| None ->
@ -58,6 +60,6 @@ let rec pop_all (self : 'a t) : 'a list =
raise Closed
);
Condition.wait self.cond self.mutex;
Atomic.set self.consumer_waiting false;
A.set self.consumer_waiting false;
Mutex.unlock self.mutex;
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 bag = Atomic.make [] in
let bag = A.make [] in
{ bag }
module Backoff = struct
@ -17,14 +19,14 @@ module Backoff = struct
end
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
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
let[@inline] add t x = add Backoff.default t x
let[@inline] pop_all t : _ list option =
match Atomic.exchange t.bag [] with
match A.exchange t.bag [] with
| [] -> None
| l -> Some (List.rev l)