Skip to content

fix(bigtable): synchronize MutationsBatcher staging and flush scheduling #18827

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mutianf merged 4 commits into from
Oct 23, 2026
Merged

fix(bigtable): synchronize MutationsBatcher staging and flush scheduling #18827

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3d3e98b
feat(google/cloud/observability/v1): add google-cloud-observability (…
suztomo Oct 23, 2026
73052bc
fix(bigtable): synchronize MutationsBatcher staging and flush scheduling
mutianf Oct 23, 2026
3de6457
perf(bigtable): use nullcontext for _stage_lock in MutationsBatcherAsync
mutianf Oct 23, 2026
1831d04
fix(bigtable): remove duplicate closed check in append and lock _rais…
mutianf Oct 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
150 changes: 88 additions & 62 deletions packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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)

Expand All @@ -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):
Expand All @@ -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")

@daniel-sanche daniel-sanche Oct 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We also have one of these at the start of the function. Are they both needed?

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()

Expand All @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down
Loading
Loading