From 5571751f3eb4a50ca87a0fd0b16518dd1eb00c99 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Wed, 27 Dec 2023 21:33:04 -0500
Subject: [PATCH] missed some uses of `Atomic`

---
 src/util/b_queue.ml  | 14 ++++++++------
 src/util/mpsc_bag.ml | 12 +++++++-----
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/src/util/b_queue.ml b/src/util/b_queue.ml
index 1a77aa3..f5ee5f3 100644
--- a/src/util/b_queue.ml
+++ b/src/util/b_queue.ml
@@ -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)
diff --git a/src/util/mpsc_bag.ml b/src/util/mpsc_bag.ml
index a8f49aa..02aeadf 100644
--- a/src/util/mpsc_bag.ml
+++ b/src/util/mpsc_bag.ml
@@ -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)