From 3d3e98b2bb7abc9ee6fd097700043d29c42688bd Mon Sep 17 00:00:00 2001 From: Tomo Suzuki Date: Fri, 23 Oct 2026 14:33:34 -0400 Subject: [PATCH 1/4] feat(google/cloud/observability/v1): add google-cloud-observability (#18807) Onboard google/cloud/observability/v1 API (google-cloud-observability). Internal tracking bug http://b/575514351 PiperOrigin-RevId: 1002511692 ## Debug Info ### System Details - Python version: Python 3.13.15 - OS: Linux 7.1.6-1rodete1-amd64 x86_64 - Librarian version: v0.49.0 ### Commands Executed - `python3 -m venv venv && source venv/bin/activate` - `pip install --upgrade pip` - `go run github.com/googleapis/librarian/cmd/librarian@v0.49.0 install` - `go run github.com/googleapis/librarian/cmd/librarian@v0.49.0 add google/cloud/observability/v1` - `go run github.com/googleapis/librarian/cmd/librarian@v0.49.0 generate -v google-cloud-observability` - `go run github.com/googleapis/librarian/cmd/librarian@v0.49.0 tidy` - `git add .` - `git commit -m "feat(google/cloud/observability/v1): add google-cloud-observability"` From 73052bc39e5aacc6bc67605897aa3a5ae8bd6e5e Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Fri, 23 Oct 2026 14:59:31 +0000 Subject: [PATCH 2/4] fix(bigtable): synchronize MutationsBatcher staging and flush scheduling --- .../bigtable/data/_async/mutations_batcher.py | 100 +++++++++++------- .../data/_sync_autogen/mutations_batcher.py | 94 ++++++++++------ .../data/_async/test_mutations_batcher.py | 85 +++++++++++++++ .../_sync_autogen/test_mutations_batcher.py | 77 ++++++++++++++ 4 files changed, 283 insertions(+), 73 deletions(-) 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..f5180af8e317 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 @@ -18,6 +18,7 @@ import concurrent.futures import inspect import logging +import threading import time import warnings from collections import deque @@ -263,6 +264,7 @@ def __init__( ) self._closed = CrossSync.Event() + self._stage_lock = threading.RLock() self._target = table self._staged_entries: list[RowMutationEntry] = [] self._staged_count, self._staged_bytes = 0, 0 @@ -287,9 +289,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 +303,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 +327,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): @@ -346,15 +349,21 @@ async def append(self, mutation_entry: RowMutationEntry): 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") + 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 +375,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 +430,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( @@ -550,13 +569,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 +597,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..9dab02ed30a4 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 @@ -21,6 +21,7 @@ import concurrent.futures import inspect import logging +import threading import time import warnings from collections import deque @@ -230,6 +231,7 @@ def __init__( batch_retryable_errors, table ) self._closed = CrossSync._Sync_Impl.Event() + self._stage_lock = threading.RLock() self._target = table self._staged_entries: list[RowMutationEntry] = [] self._staged_count, self._staged_bytes = (0, 0) @@ -252,9 +254,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 +265,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 +284,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 @@ -299,14 +302,20 @@ def append(self, mutation_entry: RowMutationEntry): 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") + 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 +324,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 +371,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 @@ -462,12 +481,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 +504,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 From 3de64573bb505b7201b7cb6745feb761f1d1dbd3 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Fri, 23 Oct 2026 15:16:06 +0000 Subject: [PATCH 3/4] perf(bigtable): use nullcontext for _stage_lock in MutationsBatcherAsync --- .../google/cloud/bigtable/data/_async/mutations_batcher.py | 5 ++++- .../cloud/bigtable/data/_sync_autogen/mutations_batcher.py | 7 ++++++- 2 files changed, 10 insertions(+), 2 deletions(-) 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 f5180af8e317..19ce2f5552c6 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,6 +16,7 @@ import atexit import concurrent.futures +import contextlib import inspect import logging import threading @@ -264,7 +265,9 @@ def __init__( ) self._closed = CrossSync.Event() - self._stage_lock = threading.RLock() + 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 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 9dab02ed30a4..39035ba9cd68 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,6 +19,7 @@ import atexit import concurrent.futures +import contextlib import inspect import logging import threading @@ -231,7 +232,11 @@ def __init__( batch_retryable_errors, table ) self._closed = CrossSync._Sync_Impl.Event() - self._stage_lock = threading.RLock() + 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) From 1831d04e325f23754fd1a1fc3f4382fe1ddcabf9 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Fri, 23 Oct 2026 18:31:58 +0000 Subject: [PATCH 4/4] fix(bigtable): remove duplicate closed check in append and lock _raise_exceptions --- .../bigtable/data/_async/mutations_batcher.py | 47 +++++++++---------- .../data/_sync_autogen/mutations_batcher.py | 47 +++++++++---------- 2 files changed, 46 insertions(+), 48 deletions(-) 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 19ce2f5552c6..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 @@ -346,16 +346,14 @@ 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" - ) 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) @@ -526,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): 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 39035ba9cd68..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 @@ -301,16 +301,14 @@ 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" - ) 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() @@ -449,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"""