fix: mutations batcher race condition

google/cloud/bigtable/batcher.py

@@ -18,6 +18,8 @@
 import concurrent.futures
 import atexit
 
+from functools import partial
+from copy import copy
 
 from google.api_core.exceptions import from_grpc_status
 from dataclasses import dataclass
@@ -51,36 +53,54 @@ def __init__(self, max_mutation_bytes=MAX_MUTATION_SIZE, flush_count=FLUSH_COUNT
         self.total_size = 0
         self.max_mutation_bytes = max_mutation_bytes
         self.flush_count = flush_count
+        self._lock = threading.Lock()
 
     def get(self):
-        """Retrieve an item from the queue. Recalculate queue size."""
-        row = self._queue.get()
-        mutation_size = row.get_mutations_size()
-        self.total_mutation_count -= len(row._get_mutations())
-        self.total_size -= mutation_size
-        return row
+        """
+        Retrieve an item from the queue. Recalculate queue size.
+
+        If the queue is empty, return None.
+        """
+        with self._lock:
+            try:
+                row = self._queue.get_nowait()
+                mutation_size = row.get_mutations_size()
+                self.total_mutation_count -= len(row._get_mutations())
+                self.total_size -= mutation_size
+                return row
+            except queue.Empty:
+                return None
+
+    def get_all(self):
+        """Get all items from the queue."""
+        with self._lock:
+            items = []
+            while not self._queue.empty():
+                items.append(self._queue.get_nowait())
+            self.total_mutation_count = 0
+            self.total_size = 0
+            return items
 
     def put(self, item):
         """Insert an item to the queue. Recalculate queue size."""
 
         mutation_count = len(item._get_mutations())
 
-        self._queue.put(item)
+        with self._lock:
+            self._queue.put(item)
 
-        self.total_size += item.get_mutations_size()
-        self.total_mutation_count += mutation_count
+            self.total_size += item.get_mutations_size()
+            self.total_mutation_count += mutation_count
 
     def full(self):
         """Check if the queue is full."""
-        if (
-            self.total_mutation_count >= self.flush_count
-            or self.total_size >= self.max_mutation_bytes
-        ):
-            return True
-        return False
-
-    def empty(self):
-        return self._queue.empty()
+        with self._lock:
+            if (
+                self.total_mutation_count >= self.flush_count
+                or self.total_size >= self.max_mutation_bytes
+            ):
+                return True
+            return False
 
 
 @dataclass
@@ -219,7 +239,6 @@ def __init__(
             max_mutations=MAX_OUTSTANDING_ELEMENTS,
             max_mutation_bytes=MAX_OUTSTANDING_BYTES,
         )
-        self.futures_mapping = {}
         self.exceptions = queue.Queue()
         self._user_batch_completed_callback = batch_completed_callback
 
@@ -291,9 +310,7 @@ def flush(self):
         :raises:
             * :exc:`.batcherMutationsBatchError` if there's any error in the mutations.
         """
-        rows_to_flush = []
-        while not self._rows.empty():
-            rows_to_flush.append(self._rows.get())
+        rows_to_flush = self._rows.get_all()
         response = self._flush_rows(rows_to_flush)
         return response
 
@@ -303,57 +320,66 @@ def _flush_async(self):
         :raises:
             * :exc:`.batcherMutationsBatchError` if there's any error in the mutations.
         """
-
-        rows_to_flush = []
-        mutations_count = 0
-        mutations_size = 0
-        rows_count = 0
-        batch_info = _BatchInfo()
-
-        while not self._rows.empty():
-            row = self._rows.get()
-            mutations_count += len(row._get_mutations())
-            mutations_size += row.get_mutations_size()
-            rows_count += 1
-            rows_to_flush.append(row)
-            batch_info.mutations_count = mutations_count
-            batch_info.rows_count = rows_count
-            batch_info.mutations_size = mutations_size
-
-            if (
-                rows_count >= self.flush_count
-                or mutations_size >= self.max_row_bytes
-                or mutations_count >= self.flow_control.max_mutations
-                or mutations_size >= self.flow_control.max_mutation_bytes
-                or self._rows.empty()  # submit when it reached the end of the queue
+        next_row = self._rows.get()
+        while next_row is not None:
+            # start a new batch
+            rows_to_flush = [next_row]
+            batch_info = _BatchInfo(
+                mutations_count=len(next_row._get_mutations()),
+                rows_count=1,
+                mutations_size=next_row.get_mutations_size(),
+            )
+            # fill up batch with rows
+            next_row = self._rows.get()
+            while next_row is not None and self._row_fits_in_batch(
+                next_row, batch_info
             ):
-                # wait for resources to become available, before submitting any new batch
-                self.flow_control.wait()
-                # once unblocked, submit a batch
-                # event flag will be set by control_flow to block subsequent thread, but not blocking this one
-                self.flow_control.control_flow(batch_info)
-                future = self._executor.submit(self._flush_rows, rows_to_flush)
-                self.futures_mapping[future] = batch_info
-                future.add_done_callback(self._batch_completed_callback)
-
-                # reset and start a new batch
-                rows_to_flush = []
-                mutations_size = 0
-                rows_count = 0
-                mutations_count = 0
-                batch_info = _BatchInfo()
-
-    def _batch_completed_callback(self, future):
+                rows_to_flush.append(next_row)
+                batch_info.mutations_count += len(next_row._get_mutations())
+                batch_info.rows_count += 1
+                batch_info.mutations_size += next_row.get_mutations_size()
+                next_row = self._rows.get()
+            # send batch over network
+            # wait for resources to become available
+            self.flow_control.wait()
+            # once unblocked, submit the batch
+            # event flag will be set by control_flow to block subsequent thread, but not blocking this one
+            self.flow_control.control_flow(batch_info)
+            future = self._executor.submit(self._flush_rows, rows_to_flush)
+            # schedule release of resources from flow control
+            future.add_done_callback(
+                partial(self._batch_completed_callback, copy(batch_info))
+            )
+
+    def _batch_completed_callback(self, batch_info, future):
         """Callback for when the mutation has finished to clean up the current batch
         and release items from the flow controller.
-
         Raise exceptions if there's any.
         Release the resources locked by the flow control and allow enqueued tasks to be run.
         """
+        self.flow_control.release(batch_info)
+
+    def _row_fits_in_batch(self, row, batch_info):
+        """Checks if a row can fit in the current batch.
+
+        :type row: class
+        :param row: :class:`~google.cloud.bigtable.row.DirectRow`.
 
-        processed_rows = self.futures_mapping[future]
-        self.flow_control.release(processed_rows)
-        del self.futures_mapping[future]
+        :type batch_info: :class:`_BatchInfo`
+        :param batch_info: Information about the current batch.
+
+        :rtype: bool
+        :returns: True if the row can fit in the current batch.
+        """
+        new_rows_count = batch_info.rows_count + 1
+        new_mutations_count = batch_info.mutations_count + len(row._get_mutations())
+        new_mutations_size = batch_info.mutations_size + row.get_mutations_size()
+        return (
+            new_rows_count <= self.flush_count
+            and new_mutations_size <= self.max_row_bytes
+            and new_mutations_count <= self.flow_control.max_mutations
+            and new_mutations_size <= self.flow_control.max_mutation_bytes
+        )
 
     def _flush_rows(self, rows_to_flush):
         """Mutate the specified rows.

tests/system/conftest.py

@@ -58,7 +58,7 @@ def location_id():
 
 @pytest.fixture(scope="session")
 def serve_nodes():
-    return 3
+    return 1
 
 
 @pytest.fixture(scope="session")

tests/system/test_data_api.py

@@ -381,3 +381,44 @@ def test_access_with_non_admin_client(data_client, data_instance_id, data_table_
     instance = data_client.instance(data_instance_id)
     table = instance.table(data_table_id)
     assert table.read_row("nonesuch") is None  # no raise
+
+
+def test_mutations_batcher_threading(data_table, rows_to_delete):
+    """
+    Test the mutations batcher by sending a bunch of mutations using different
+    flush methods
+    """
+    import mock
+    import time
+    from google.cloud.bigtable.batcher import MutationsBatcher
+
+    num_sent = 20
+    all_results = []
+
+    max_elements_per_batch = 2
+
+    def callback(results):
+        all_results.extend(results)
+        assert results <= max_elements_per_batch
+
+    # override flow control max elements
+    with mock.patch(
+        "google.cloud.bigtable.batcher.MAX_OUTSTANDING_ELEMENTS", max_elements_per_batch
+    ):
+        with MutationsBatcher(
+            data_table,
+            flush_count=5,
+            flush_interval=0.07,
+            batch_completed_callback=callback,
+        ) as batcher:
+            # send mutations in a way that timed flushes and count flushes interleave
+            for i in range(num_sent):
+                row = data_table.direct_row("row{}".format(i))
+                row.set_cell(
+                    COLUMN_FAMILY_ID1, COL_NAME1, "val{}".format(i).encode("utf-8")
+                )
+                rows_to_delete.append(row)
+                batcher.mutate(row)
+                time.sleep(0.01)
+    # ensure all mutations were sent
+    assert len(all_results) == num_sent