diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py index 2f4dc3c3121c..2ee11ad9834b 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -16,8 +16,10 @@ import atexit import concurrent.futures +import contextlib import inspect import logging +import threading import time import warnings from collections import deque @@ -263,6 +265,9 @@ def __init__( ) self._closed = CrossSync.Event() + self._stage_lock: contextlib.AbstractContextManager[Any] = ( + threading.RLock() if not CrossSync.is_async else contextlib.nullcontext() + ) self._target = table self._staged_entries: list[RowMutationEntry] = [] self._staged_count, self._staged_bytes = 0, 0 @@ -287,9 +292,6 @@ def __init__( if not CrossSync.is_async else None ) - self._flush_timer = CrossSync.create_daemon_task( - self._timer_routine, flush_interval - ) self._flush_jobs: set[CrossSync.Future[None]] = set() # MutationExceptionGroup reports number of successful entries along with failures self._entries_processed_since_last_raise: int = 0 @@ -304,6 +306,9 @@ def __init__( self._user_batch_completed_callback: ( Callable[[list[status_pb2.Status]], Any] | None ) = None + self._flush_timer = CrossSync.create_daemon_task( + self._timer_routine, flush_interval + ) # clean up on program exit atexit.register(self._on_exit) @@ -325,8 +330,9 @@ async def _timer_routine(self, interval: float | None) -> None: await CrossSync.event_wait( self._closed, timeout=interval, async_break_early=False ) - if not self._closed.is_set() and self._staged_entries: - self._schedule_flush() + with self._stage_lock: + if not self._closed.is_set() and self._staged_entries: + self._schedule_flush() @CrossSync.convert async def append(self, mutation_entry: RowMutationEntry): @@ -340,21 +346,25 @@ async def append(self, mutation_entry: RowMutationEntry): ValueError: if an invalid mutation type is added """ # TODO: return a future to track completion of this entry - if self._closed.is_set(): - raise RuntimeError("Cannot append to closed MutationsBatcher") - if isinstance(cast(Mutation, mutation_entry), Mutation): - raise ValueError( - f"invalid mutation type: {type(mutation_entry).__name__}. Only RowMutationEntry objects are supported by batcher" - ) - self._staged_entries.append(mutation_entry) - # start a new flush task if limits exceeded - self._staged_count += len(mutation_entry.mutations) - self._staged_bytes += mutation_entry.size() - if ( - self._staged_count >= self._flush_limit_count - or self._staged_bytes >= self._flush_limit_bytes - ): - self._schedule_flush() + should_flush = False + with self._stage_lock: + if self._closed.is_set(): + raise RuntimeError("Cannot append to closed MutationsBatcher") + if isinstance(cast(Mutation, mutation_entry), Mutation): + raise ValueError( + f"invalid mutation type: {type(mutation_entry).__name__}. Only RowMutationEntry objects are supported by batcher" + ) + self._staged_entries.append(mutation_entry) + # start a new flush task if limits exceeded + self._staged_count += len(mutation_entry.mutations) + self._staged_bytes += mutation_entry.size() + if ( + self._staged_count >= self._flush_limit_count + or self._staged_bytes >= self._flush_limit_bytes + ): + self._schedule_flush() + should_flush = True + if should_flush: # yield to the event loop to allow flush to run await CrossSync.yield_to_event_loop() @@ -366,16 +376,25 @@ def _schedule_flush(self) -> CrossSync.Future[None] | None: Future[None] | None: future representing the background task, if started """ - if self._staged_entries: - entries, self._staged_entries = self._staged_entries, [] - self._staged_count, self._staged_bytes = 0, 0 - new_task = CrossSync.create_task( - self._flush_internal, entries, sync_executor=self._sync_flush_executor - ) - if not new_task.done(): - self._flush_jobs.add(new_task) - new_task.add_done_callback(self._flush_jobs.remove) - return new_task + with self._stage_lock: + if self._staged_entries: + entries, self._staged_entries = self._staged_entries, [] + self._staged_count, self._staged_bytes = 0, 0 + try: + new_task = CrossSync.create_task( + self._flush_internal, + entries, + sync_executor=self._sync_flush_executor, + ) + except RuntimeError: + self._staged_entries = entries + self._staged_entries + self._staged_count += sum(len(e.mutations) for e in entries) + self._staged_bytes += sum(e.size() for e in entries) + raise + if not new_task.done(): + self._flush_jobs.add(new_task) + new_task.add_done_callback(self._flush_jobs.discard) + return new_task return None @CrossSync.convert @@ -412,8 +431,9 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]): await self._wait_for_batch_results(*in_process_requests) ) # update exception data to reflect any new errors - self._entries_processed_since_last_raise += len(new_entries) - self._add_exceptions(found_exceptions) + with self._stage_lock: + self._entries_processed_since_last_raise += len(new_entries) + self._add_exceptions(found_exceptions) @CrossSync.convert async def _execute_mutate_rows( @@ -504,24 +524,25 @@ def _raise_exceptions(self): Raises: MutationsExceptionGroup: exception group with all unreported exceptions """ - if self._oldest_exceptions or self._newest_exceptions: - oldest, self._oldest_exceptions = self._oldest_exceptions, [] - newest = list(self._newest_exceptions) - self._newest_exceptions.clear() - entry_count, self._entries_processed_since_last_raise = ( - self._entries_processed_since_last_raise, - 0, - ) - exc_count, self._exceptions_since_last_raise = ( - self._exceptions_since_last_raise, - 0, - ) - raise MutationsExceptionGroup.from_truncated_lists( - first_list=oldest, - last_list=newest, - total_excs=exc_count, - entry_count=entry_count, - ) + with self._stage_lock: + if self._oldest_exceptions or self._newest_exceptions: + oldest, self._oldest_exceptions = self._oldest_exceptions, [] + newest = list(self._newest_exceptions) + self._newest_exceptions.clear() + entry_count, self._entries_processed_since_last_raise = ( + self._entries_processed_since_last_raise, + 0, + ) + exc_count, self._exceptions_since_last_raise = ( + self._exceptions_since_last_raise, + 0, + ) + raise MutationsExceptionGroup.from_truncated_lists( + first_list=oldest, + last_list=newest, + total_excs=exc_count, + entry_count=entry_count, + ) @CrossSync.convert(sync_name="__enter__") async def __aenter__(self): @@ -550,13 +571,17 @@ async def close(self): """ Flush queue and clean up resources """ - self._closed.set() - self._flush_timer.cancel() - entries = self._staged_entries - try: - self._schedule_flush() - except RuntimeError: - # Executor is already shut down (e.g. during atexit); flush synchronously. + with self._stage_lock: + self._closed.set() + self._flush_timer.cancel() + try: + self._schedule_flush() + entries: list[RowMutationEntry] = [] + except RuntimeError: + # Executor is already shut down (e.g. during atexit); flush synchronously. + entries, self._staged_entries = self._staged_entries, [] + self._staged_count, self._staged_bytes = 0, 0 + if entries: await self._flush_internal(entries) # shut down executors if self._sync_flush_executor: @@ -574,11 +599,12 @@ def _on_exit(self): """ Called when program is exited. Raises warning if unflushed mutations remain """ - if not self._closed.is_set() and self._staged_entries: - warnings.warn( - f"MutationsBatcher for target {self._target!r} was not closed. " - f"{len(self._staged_entries)} Unflushed mutations will not be sent to the server." - ) + with self._stage_lock: + if not self._closed.is_set() and self._staged_entries: + warnings.warn( + f"MutationsBatcher for target {self._target!r} was not closed. " + f"{len(self._staged_entries)} Unflushed mutations will not be sent to the server." + ) @staticmethod @CrossSync.convert diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py index 8a13c9da02c5..83ffcdecb33a 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py @@ -19,8 +19,10 @@ import atexit import concurrent.futures +import contextlib import inspect import logging +import threading import time import warnings from collections import deque @@ -230,6 +232,11 @@ def __init__( batch_retryable_errors, table ) self._closed = CrossSync._Sync_Impl.Event() + self._stage_lock: contextlib.AbstractContextManager[Any] = ( + threading.RLock() + if not CrossSync._Sync_Impl.is_async + else contextlib.nullcontext() + ) self._target = table self._staged_entries: list[RowMutationEntry] = [] self._staged_count, self._staged_bytes = (0, 0) @@ -252,9 +259,6 @@ def __init__( if not CrossSync._Sync_Impl.is_async else None ) - self._flush_timer = CrossSync._Sync_Impl.create_daemon_task( - self._timer_routine, flush_interval - ) self._flush_jobs: set[CrossSync._Sync_Impl.Future[None]] = set() self._entries_processed_since_last_raise: int = 0 self._exceptions_since_last_raise: int = 0 @@ -266,6 +270,9 @@ def __init__( self._user_batch_completed_callback: ( Callable[[list[status_pb2.Status]], Any] | None ) = None + self._flush_timer = CrossSync._Sync_Impl.create_daemon_task( + self._timer_routine, flush_interval + ) atexit.register(self._on_exit) def _timer_routine(self, interval: float | None) -> None: @@ -282,8 +289,9 @@ def _timer_routine(self, interval: float | None) -> None: CrossSync._Sync_Impl.event_wait( self._closed, timeout=interval, async_break_early=False ) - if not self._closed.is_set() and self._staged_entries: - self._schedule_flush() + with self._stage_lock: + if not self._closed.is_set() and self._staged_entries: + self._schedule_flush() def append(self, mutation_entry: RowMutationEntry): """Add a new set of mutations to the internal queue @@ -293,20 +301,24 @@ def append(self, mutation_entry: RowMutationEntry): Raises: RuntimeError: if batcher is closed ValueError: if an invalid mutation type is added""" - if self._closed.is_set(): - raise RuntimeError("Cannot append to closed MutationsBatcher") - if isinstance(cast(Mutation, mutation_entry), Mutation): - raise ValueError( - f"invalid mutation type: {type(mutation_entry).__name__}. Only RowMutationEntry objects are supported by batcher" - ) - self._staged_entries.append(mutation_entry) - self._staged_count += len(mutation_entry.mutations) - self._staged_bytes += mutation_entry.size() - if ( - self._staged_count >= self._flush_limit_count - or self._staged_bytes >= self._flush_limit_bytes - ): - self._schedule_flush() + should_flush = False + with self._stage_lock: + if self._closed.is_set(): + raise RuntimeError("Cannot append to closed MutationsBatcher") + if isinstance(cast(Mutation, mutation_entry), Mutation): + raise ValueError( + f"invalid mutation type: {type(mutation_entry).__name__}. Only RowMutationEntry objects are supported by batcher" + ) + self._staged_entries.append(mutation_entry) + self._staged_count += len(mutation_entry.mutations) + self._staged_bytes += mutation_entry.size() + if ( + self._staged_count >= self._flush_limit_count + or self._staged_bytes >= self._flush_limit_bytes + ): + self._schedule_flush() + should_flush = True + if should_flush: CrossSync._Sync_Impl.yield_to_event_loop() def _schedule_flush(self) -> CrossSync._Sync_Impl.Future[None] | None: @@ -315,16 +327,25 @@ def _schedule_flush(self) -> CrossSync._Sync_Impl.Future[None] | None: Returns: Future[None] | None: future representing the background task, if started""" - if self._staged_entries: - entries, self._staged_entries = (self._staged_entries, []) - self._staged_count, self._staged_bytes = (0, 0) - new_task = CrossSync._Sync_Impl.create_task( - self._flush_internal, entries, sync_executor=self._sync_flush_executor - ) - if not new_task.done(): - self._flush_jobs.add(new_task) - new_task.add_done_callback(self._flush_jobs.remove) - return new_task + with self._stage_lock: + if self._staged_entries: + entries, self._staged_entries = (self._staged_entries, []) + self._staged_count, self._staged_bytes = (0, 0) + try: + new_task = CrossSync._Sync_Impl.create_task( + self._flush_internal, + entries, + sync_executor=self._sync_flush_executor, + ) + except RuntimeError: + self._staged_entries = entries + self._staged_entries + self._staged_count += sum((len(e.mutations) for e in entries)) + self._staged_bytes += sum((e.size() for e in entries)) + raise + if not new_task.done(): + self._flush_jobs.add(new_task) + new_task.add_done_callback(self._flush_jobs.discard) + return new_task return None def _flush_internal(self, new_entries: list[RowMutationEntry]): @@ -353,8 +374,9 @@ def _flush_internal(self, new_entries: list[RowMutationEntry]): except RuntimeError: found_exceptions.extend(self._execute_mutate_rows(batch, metric)) found_exceptions.extend(self._wait_for_batch_results(*in_process_requests)) - self._entries_processed_since_last_raise += len(new_entries) - self._add_exceptions(found_exceptions) + with self._stage_lock: + self._entries_processed_since_last_raise += len(new_entries) + self._add_exceptions(found_exceptions) def _execute_mutate_rows( self, batch: list[RowMutationEntry], metric: ActiveOperationMetric @@ -425,24 +447,25 @@ def _raise_exceptions(self): Raises: MutationsExceptionGroup: exception group with all unreported exceptions""" - if self._oldest_exceptions or self._newest_exceptions: - oldest, self._oldest_exceptions = (self._oldest_exceptions, []) - newest = list(self._newest_exceptions) - self._newest_exceptions.clear() - entry_count, self._entries_processed_since_last_raise = ( - self._entries_processed_since_last_raise, - 0, - ) - exc_count, self._exceptions_since_last_raise = ( - self._exceptions_since_last_raise, - 0, - ) - raise MutationsExceptionGroup.from_truncated_lists( - first_list=oldest, - last_list=newest, - total_excs=exc_count, - entry_count=entry_count, - ) + with self._stage_lock: + if self._oldest_exceptions or self._newest_exceptions: + oldest, self._oldest_exceptions = (self._oldest_exceptions, []) + newest = list(self._newest_exceptions) + self._newest_exceptions.clear() + entry_count, self._entries_processed_since_last_raise = ( + self._entries_processed_since_last_raise, + 0, + ) + exc_count, self._exceptions_since_last_raise = ( + self._exceptions_since_last_raise, + 0, + ) + raise MutationsExceptionGroup.from_truncated_lists( + first_list=oldest, + last_list=newest, + total_excs=exc_count, + entry_count=entry_count, + ) def __enter__(self): """Allow use of context manager API""" @@ -462,12 +485,16 @@ def closed(self) -> bool: def close(self): """Flush queue and clean up resources""" - self._closed.set() - self._flush_timer.cancel() - entries = self._staged_entries - try: - self._schedule_flush() - except RuntimeError: + with self._stage_lock: + self._closed.set() + self._flush_timer.cancel() + try: + self._schedule_flush() + entries: list[RowMutationEntry] = [] + except RuntimeError: + entries, self._staged_entries = (self._staged_entries, []) + self._staged_count, self._staged_bytes = (0, 0) + if entries: self._flush_internal(entries) if self._sync_flush_executor: with self._sync_flush_executor: @@ -481,10 +508,11 @@ def close(self): def _on_exit(self): """Called when program is exited. Raises warning if unflushed mutations remain""" - if not self._closed.is_set() and self._staged_entries: - warnings.warn( - f"MutationsBatcher for target {self._target!r} was not closed. {len(self._staged_entries)} Unflushed mutations will not be sent to the server." - ) + with self._stage_lock: + if not self._closed.is_set() and self._staged_entries: + warnings.warn( + f"MutationsBatcher for target {self._target!r} was not closed. {len(self._staged_entries)} Unflushed mutations will not be sent to the server." + ) @staticmethod def _wait_for_batch_results( diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py index cba0ba5da7a0..7c9bb888e34c 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py @@ -1252,6 +1252,91 @@ async def test_close_executor_shutdown_fallback(self): assert exec_mock.call_count == 1 assert instance._staged_entries == [] + @CrossSync.pytest + async def test_schedule_flush_runtime_error_restores_staged_entries(self): + """If create_task raises RuntimeError in _schedule_flush, staged entries are preserved""" + async with self._make_one() as instance: + mutation = self._make_mutation(count=2, size=5) + await instance.append(mutation) + with mock.patch.object( + CrossSync, "create_task", side_effect=RuntimeError("shutdown") + ): + with pytest.raises(RuntimeError, match="shutdown"): + instance._schedule_flush() + assert instance._staged_entries == [mutation] + assert instance._staged_count == 2 + assert instance._staged_bytes == 5 + instance._staged_entries = [] + + @CrossSync.pytest + async def test_sync_concurrent_schedule_flush_and_close_thread_safety(self): + """ + Verify that in sync mode: + 1. Concurrent _schedule_flush() calls from user and timer threads never dispatch + duplicate batches (b/574637091). + 2. Concurrent timer flush and close() never shut down _sync_flush_executor mid-flush + or silently lose a batch (b/574637465). + """ + if CrossSync.is_async: + return + import threading + + # 1. Concurrent _schedule_flush() across threads dispatches batch exactly once + async with self._make_one(flush_interval=None) as instance: + flushed_batches = [] + flush_lock = threading.Lock() + barrier = threading.Barrier(4) + + def record_flush(entries): + with flush_lock: + flushed_batches.append(list(entries)) + + with mock.patch.object( + instance, "_flush_internal", side_effect=record_flush + ): + mutation = self._make_mutation(count=1, size=1) + instance._staged_entries = [mutation] + instance._staged_count = 1 + instance._staged_bytes = 1 + + def worker(): + barrier.wait() + instance._schedule_flush() + + threads = [threading.Thread(target=worker) for _ in range(4)] + for t in threads: + t.start() + for t in threads: + t.join() + + await instance.close() + assert len(flushed_batches) == 1 + assert flushed_batches[0] == [mutation] + + # 2. Concurrent timer wakeup and close() never loses batch or raises in timer thread + for _ in range(25): + flushed_entries = [] + flush_lock = threading.Lock() + instance = self._make_one( + flush_interval=0.001, flush_limit_mutation_count=100 + ) + with mock.patch.object( + instance, + "_execute_mutate_rows", + side_effect=lambda batch, metric: ( + flush_lock.acquire(), + flushed_entries.extend(batch), + flush_lock.release(), + [], + )[-1], + ): + mutation = self._make_mutation(count=1, size=1) + await instance.append(mutation) + time.sleep(0.001) + await instance.close() + assert instance._flush_timer.exception() is None + assert flushed_entries == [mutation] + @CrossSync.pytest async def test_close_w_exceptions(self): """Raise exceptions on close""" diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py index a7522c151418..51c28e95854e 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py @@ -1115,6 +1115,83 @@ def test_close_executor_shutdown_fallback(self): assert exec_mock.call_count == 1 assert instance._staged_entries == [] + def test_schedule_flush_runtime_error_restores_staged_entries(self): + """If create_task raises RuntimeError in _schedule_flush, staged entries are preserved""" + with self._make_one() as instance: + mutation = self._make_mutation(count=2, size=5) + instance.append(mutation) + with mock.patch.object( + CrossSync._Sync_Impl, + "create_task", + side_effect=RuntimeError("shutdown"), + ): + with pytest.raises(RuntimeError, match="shutdown"): + instance._schedule_flush() + assert instance._staged_entries == [mutation] + assert instance._staged_count == 2 + assert instance._staged_bytes == 5 + instance._staged_entries = [] + + def test_sync_concurrent_schedule_flush_and_close_thread_safety(self): + """Verify that in sync mode: + 1. Concurrent _schedule_flush() calls from user and timer threads never dispatch + duplicate batches (b/574637091). + 2. Concurrent timer flush and close() never shut down _sync_flush_executor mid-flush + or silently lose a batch (b/574637465).""" + import threading + + with self._make_one(flush_interval=None) as instance: + flushed_batches = [] + flush_lock = threading.Lock() + barrier = threading.Barrier(4) + + def record_flush(entries): + with flush_lock: + flushed_batches.append(list(entries)) + + with mock.patch.object( + instance, "_flush_internal", side_effect=record_flush + ): + mutation = self._make_mutation(count=1, size=1) + instance._staged_entries = [mutation] + instance._staged_count = 1 + instance._staged_bytes = 1 + + def worker(): + barrier.wait() + instance._schedule_flush() + + threads = [threading.Thread(target=worker) for _ in range(4)] + for t in threads: + t.start() + for t in threads: + t.join() + instance.close() + assert len(flushed_batches) == 1 + assert flushed_batches[0] == [mutation] + for _ in range(25): + flushed_entries = [] + flush_lock = threading.Lock() + instance = self._make_one( + flush_interval=0.001, flush_limit_mutation_count=100 + ) + with mock.patch.object( + instance, + "_execute_mutate_rows", + side_effect=lambda batch, metric: ( + flush_lock.acquire(), + flushed_entries.extend(batch), + flush_lock.release(), + [], + )[-1], + ): + mutation = self._make_mutation(count=1, size=1) + instance.append(mutation) + time.sleep(0.001) + instance.close() + assert instance._flush_timer.exception() is None + assert flushed_entries == [mutation] + def test_close_w_exceptions(self): """Raise exceptions on close""" from google.cloud.bigtable.data import exceptions