feat: Mutations Batcher shim (#1309)

**Changes Made:**

- Replaced mutations batcher implementation with one based off of the
data client.
- Reworked unit tests.
- Added additional system tests.

Author: gkevinzheng

google/cloud/bigtable/batcher.py

@@ -13,14 +13,12 @@
 # limitations under the License.
 
 """User friendly container for Google Cloud Bigtable MutationBatcher."""
-import threading
 import queue
-import concurrent.futures
 import atexit
 
 
-from google.api_core.exceptions import from_grpc_status
-from dataclasses import dataclass
+from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup
+from google.cloud.bigtable.data.mutations import RowMutationEntry
 
 
 FLUSH_COUNT = 100  # after this many elements, send out the batch
@@ -41,131 +39,6 @@ def __init__(self, message, exc):
         super().__init__(self.message)
 
 
-class _MutationsBatchQueue(object):
-    """Private Threadsafe Queue to hold rows for batching."""
-
-    def __init__(self, max_mutation_bytes=MAX_MUTATION_SIZE, flush_count=FLUSH_COUNT):
-        """Specify the queue constraints"""
-        self._queue = queue.Queue()
-        self.total_mutation_count = 0
-        self.total_size = 0
-        self.max_mutation_bytes = max_mutation_bytes
-        self.flush_count = flush_count
-
-    def get(self):
-        """
-        Retrieve an item from the queue. Recalculate queue size.
-
-        If the queue is empty, return None.
-        """
-        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 put(self, item):
-        """Insert an item to the queue. Recalculate queue size."""
-
-        mutation_count = len(item._get_mutations())
-
-        self._queue.put(item)
-
-        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
-
-
-@dataclass
-class _BatchInfo:
-    """Keeping track of size of a batch"""
-
-    mutations_count: int = 0
-    rows_count: int = 0
-    mutations_size: int = 0
-
-
-class _FlowControl(object):
-    def __init__(
-        self,
-        max_mutations=MAX_OUTSTANDING_ELEMENTS,
-        max_mutation_bytes=MAX_OUTSTANDING_BYTES,
-    ):
-        """Control the inflight requests. Keep track of the mutations, row bytes and row counts.
-        As requests to backend are being made, adjust the number of mutations being processed.
-
-        If threshold is reached, block the flow.
-        Reopen the flow as requests are finished.
-        """
-        self.max_mutations = max_mutations
-        self.max_mutation_bytes = max_mutation_bytes
-        self.inflight_mutations = 0
-        self.inflight_size = 0
-        self.event = threading.Event()
-        self.event.set()
-        self._lock = threading.Lock()
-
-    def is_blocked(self):
-        """Returns True if:
-
-        - inflight mutations >= max_mutations, or
-        - inflight bytes size >= max_mutation_bytes, or
-        """
-
-        return (
-            self.inflight_mutations >= self.max_mutations
-            or self.inflight_size >= self.max_mutation_bytes
-        )
-
-    def control_flow(self, batch_info):
-        """
-        Calculate the resources used by this batch
-        """
-
-        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):
-        """
-        Wait until flow control pushback has been released.
-        It awakens as soon as `event` is set.
-        """
-        self.event.wait()
-
-    def set_flow_control_status(self):
-        """Check the inflight mutations and size.
-
-        If values exceed the allowed threshold, block the event.
-        """
-        if self.is_blocked():
-            self.event.clear()  # sleep
-        else:
-            self.event.set()  # awaken the threads
-
-    def release(self, batch_info):
-        """
-        Release the resources.
-        Decrement the row size to allow enqueued mutations to be run.
-        """
-        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):
     """A MutationsBatcher is used in batch cases where the number of mutations
     is large or unknown. It will store :class:`DirectRow` in memory until one of the
@@ -214,29 +87,41 @@ def __init__(
         flush_interval=1,
         batch_completed_callback=None,
     ):
-        self._rows = _MutationsBatchQueue(
-            max_mutation_bytes=max_row_bytes, flush_count=flush_count
-        )
         self.table = table
-        self._executor = concurrent.futures.ThreadPoolExecutor()
-        atexit.register(self.close)
-        self._timer = threading.Timer(flush_interval, self.flush)
-        self._timer.start()
-        self.flow_control = _FlowControl(
-            max_mutations=MAX_OUTSTANDING_ELEMENTS,
-            max_mutation_bytes=MAX_OUTSTANDING_BYTES,
-        )
-        self.futures_mapping = {}
-        self.exceptions = queue.Queue()
+        self._batcher_kwargs = {
+            "flush_interval": flush_interval,
+            "flush_limit_mutation_count": flush_count,
+            "flush_limit_bytes": max_row_bytes,
+            "flow_control_max_mutation_count": MAX_OUTSTANDING_ELEMENTS,
+            "flow_control_max_bytes": MAX_OUTSTANDING_BYTES,
+        }
         self._user_batch_completed_callback = batch_completed_callback
+        self._init_batcher()
+        atexit.register(self.close)
+        self._exceptions = queue.Queue()
 
     @property
     def flush_count(self):
-        return self._rows.flush_count
+        return self._flush_count
 
     @property
     def max_row_bytes(self):
-        return self._rows.max_mutation_bytes
+        return self._max_row_bytes
+
+    def _init_batcher(self):
+        self._batcher = self.table._table_impl.mutations_batcher(**self._batcher_kwargs)
+        self._batcher._user_batch_completed_callback = (
+            self._user_batch_completed_callback
+        )
+
+    def _close_batcher(self):
+        try:
+            self._batcher.close()
+        except MutationsExceptionGroup as exc_group:
+            for error in exc_group.exceptions:
+                # Unpack the root cause of the FailedMutationEntryError
+                # and return that error to the user.
+                self._exceptions.put(error.__cause__)
 
     def __enter__(self):
         """Starting the MutationsBatcher as a context manager"""
@@ -260,10 +145,7 @@ def mutate(self, row):
             * :exc:`~.table._BigtableRetryableError` if any row returned a transient error.
             * :exc:`RuntimeError` if the number of responses doesn't match the number of rows that were retried
         """
-        self._rows.put(row)
-
-        if self._rows.full():
-            self._flush_async()
+        self._batcher.append(RowMutationEntry(row.row_key, row._get_mutations()))
 
     def mutate_rows(self, rows):
         """Add multiple rows to the batch. If the current batch meets one of the size
@@ -296,104 +178,10 @@ def flush(self):
             :dedent: 4
 
         :raises:
-            * :exc:`.batcherMutationsBatchError` if there's any error in the mutations.
-        """
-        rows_to_flush = []
-        row = self._rows.get()
-        while row is not None:
-            rows_to_flush.append(row)
-            row = self._rows.get()
-        response = self._flush_rows(rows_to_flush)
-        return response
-
-    def _flush_async(self):
-        """Sends the current batch to Cloud Bigtable asynchronously.
-
-        :raises:
-            * :exc:`.batcherMutationsBatchError` if there's any error in the mutations.
-        """
-        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
-            ):
-                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
-            self.futures_mapping[future] = batch_info
-            future.add_done_callback(self._batch_completed_callback)
-
-    def _batch_completed_callback(self, 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.
-        """
-        processed_rows = self.futures_mapping[future]
-        self.flow_control.release(processed_rows)
-        del self.futures_mapping[future]
-
-    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`.
-
-        :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.
+            * :exc:`~batcher.MutationsBatchError` if there's any error in the mutations.
         """
-        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.
-
-        :raises:
-            * :exc:`.batcherMutationsBatchError` if there's any error in the mutations.
-        """
-        responses = []
-        if len(rows_to_flush) > 0:
-            response = self.table.mutate_rows(rows_to_flush)
-
-            if self._user_batch_completed_callback:
-                self._user_batch_completed_callback(response)
-
-            for result in response:
-                if result.code != 0:
-                    exc = from_grpc_status(result.code, result.message)
-                    self.exceptions.put(exc)
-                responses.append(result)
-
-        return responses
+        self._close_batcher()
+        self._init_batcher()
 
     def __exit__(self, exc_type, exc_value, exc_traceback):
         """Clean up resources. Flush and shutdown the ThreadPoolExecutor."""
@@ -404,11 +192,10 @@ def close(self):
         Any errors will be raised.
 
         :raises:
-            * :exc:`.batcherMutationsBatchError` if there's any error in the mutations.
+            * :exc:`~batcher.MutationsBatchError` if there's any error in the mutations.
         """
-        self.flush()
-        self._executor.shutdown(wait=True)
+        self._close_batcher()
         atexit.unregister(self.close)
-        if self.exceptions.qsize() > 0:
-            exc = list(self.exceptions.queue)
+        if self._exceptions.qsize() > 0:
+            exc = list(self._exceptions.queue)
             raise MutationsBatchError("Errors in batch mutations.", exc=exc)

google/cloud/bigtable/data/_async/_mutate_rows.py

@@ -131,7 +131,7 @@ async def start(self):
                 self._handle_entry_error(idx, exc)
         finally:
             # raise exception detailing incomplete mutations
-            all_errors: list[Exception] = []
+            all_errors: list[bt_exceptions.FailedMutationEntryError] = []
             for idx, exc_list in self.errors.items():
                 if len(exc_list) == 0:
                     raise core_exceptions.ClientError(

google/cloud/bigtable/data/_async/client.py

@@ -16,7 +16,6 @@
 from __future__ import annotations
 
 from typing import (
-    Callable,
     cast,
     Any,
     AsyncIterable,
@@ -116,7 +115,6 @@
 if TYPE_CHECKING:
     from google.cloud.bigtable.data._helpers import RowKeySamples
     from google.cloud.bigtable.data._helpers import ShardedQuery
-    from google.rpc import status_pb2
 
     if CrossSync.is_async:
         from google.cloud.bigtable.data._async.mutations_batcher import (