add lock to flow control

googleapis · Dec 12, 2023 · 8a63ce4 · 8a63ce4
1 parent d218f4e
commit 8a63ce4
Showing 1 changed file with 9 additions and 6 deletions.
diff --git a/google/cloud/bigtable/batcher.py b/google/cloud/bigtable/batcher.py
@@ -110,6 +110,7 @@ def __init__(
         self.inflight_size = 0
         self.event = threading.Event()
         self.event.set()
+        self._lock = threading.Lock()
 
     def is_blocked(self):
         """Returns True if:
@@ -128,9 +129,10 @@ def control_flow(self, batch_info):
         Calculate the resources used by this batch
         """
 
-        self.inflight_mutations += batch_info.mutations_count
-        self.inflight_size += batch_info.mutations_size
-        self.set_flow_control_status()
+        with self._lock:
+            self.inflight_mutations += batch_info.mutations_count
+            self.inflight_size += batch_info.mutations_size
+            self.set_flow_control_status()
 
     def wait(self):
         """
@@ -154,9 +156,10 @@ def release(self, batch_info):
         Release the resources.
         Decrement the row size to allow enqueued mutations to be run.
         """
-        self.inflight_mutations -= batch_info.mutations_count
-        self.inflight_size -= batch_info.mutations_size
-        self.set_flow_control_status()
+        with self._lock:
+            self.inflight_mutations -= batch_info.mutations_count
+            self.inflight_size -= batch_info.mutations_size
+            self.set_flow_control_status()
 
 
 class MutationsBatcher(object):