From a698ef1633a2d9bfc33559e16ef161efc0cf391f Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Fri, 10 Jul 2026 18:45:54 +0000 Subject: [PATCH 1/4] feat(bigtable): route read_row/mutate_row through the accelerator with native fallback Change-Id: I6fc1d2563ca3185b12efb87b96eff2b7be438405 --- .../bigtable/data/_accelerator/_fallback.py | 144 ++++++++++++ .../bigtable/data/_accelerator/_routing.py | 27 +++ .../data/_async/_accelerator_client.py | 94 ++++++++ .../cloud/bigtable/data/_async/client.py | 219 +++++++++++++++++- .../data/_sync_autogen/_accelerator_client.py | 81 +++++++ .../bigtable/data/_sync_autogen/client.py | 176 +++++++++++++- .../unit/data/test_accelerator_enablement.py | 101 ++++++++ 7 files changed, 826 insertions(+), 16 deletions(-) create mode 100644 packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py create mode 100644 packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_routing.py create mode 100644 packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_client.py create mode 100644 packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_accelerator_client.py create mode 100644 packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py new file mode 100644 index 000000000000..d50a74cf2099 --- /dev/null +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py @@ -0,0 +1,144 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Client-side fallback policy for accelerator-routed RPCs. + +Mirrors the Go client's ``session.UnimplementedErrorInterceptor``: a daemon that +cannot open any sessions replies ``UNIMPLEMENTED``, and the routing layer +transparently retries the call on the native client. A sticky breaker trips +after enough consecutive ``UNIMPLEMENTED`` replies so a persistently-degraded +daemon stops being dialed at all. A daemon whose subprocess has died mid-flight +trips the breaker immediately — it will never recover. + +Any other gRPC error is a real, daemon-served result the native client would +reproduce (the daemon owns retries, so it has already exhausted them), so it is +translated to the matching ``google.api_core`` exception and raised without +falling back. + +This module is plain sync-only logic shared verbatim by the async and generated +sync clients; ``grpc.RpcError`` is the common base of both ``grpc.RpcError`` and +``grpc.aio.AioRpcError``, so no CrossSync branching is needed here. +""" + +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING + +from grpc import RpcError, StatusCode + +from google.api_core import exceptions as core_exceptions + +if TYPE_CHECKING: + from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon + +# Consecutive ``UNIMPLEMENTED`` replies that trip the sticky breaker, after which +# the accelerator is bypassed for the lifetime of the Table. Matches the Go +# client's ``session.DefaultUnimplementedThreshold``. +DEFAULT_UNIMPLEMENTED_THRESHOLD = 30 + + +class _AcceleratorFallback(Exception): + """Internal signal that an accelerator attempt should be retried natively. + + Never escapes the Table method that raises it: the method catches it and + falls through to the native code path. + """ + + +class AcceleratorBreaker: + """Tracks accelerator health and decides when to stop using it. + + One instance per Table. Thread-safe so the generated sync client can share a + Table across threads. Two triggers permanently bypass the accelerator: + + * ``threshold`` consecutive ``UNIMPLEMENTED`` replies (the daemon understands + the RPC shape but has no working sessions), and + * an explicit :meth:`trip` when the daemon subprocess is found dead. + + Any non-``UNIMPLEMENTED`` outcome resets the consecutive count, matching the + Go interceptor: a normal reply proves the daemon is healthy again. + """ + + def __init__(self, threshold: int = DEFAULT_UNIMPLEMENTED_THRESHOLD): + self._threshold = threshold + self._consecutive = 0 + self._tripped = False + self._lock = threading.Lock() + + def bypass(self) -> bool: + """Whether the accelerator should be skipped entirely from now on.""" + return self._tripped + + def trip(self) -> None: + """Permanently bypass the accelerator (e.g. the daemon process died).""" + with self._lock: + self._tripped = True + + def record_unimplemented(self) -> None: + """Note an ``UNIMPLEMENTED`` reply; trip the breaker at the threshold.""" + with self._lock: + self._consecutive += 1 + if self._consecutive >= self._threshold: + self._tripped = True + + def record_ok(self) -> None: + """Note any non-``UNIMPLEMENTED`` outcome; resets the consecutive count.""" + with self._lock: + self._consecutive = 0 + + +def _grpc_code(exc: BaseException) -> StatusCode | None: + """Best-effort extraction of a gRPC status code from an exception.""" + code = getattr(exc, "code", None) + if not callable(code): + return None + try: + return code() + except Exception: + return None + + +def handle_accelerator_error( + exc: BaseException, + *, + daemon: "AcceleratorDaemon | None", + breaker: AcceleratorBreaker, +) -> None: + """Classify an exception raised by an accelerator-routed RPC. + + Always raises. Either raises :class:`_AcceleratorFallback` to tell the caller + to retry on the native path, or raises the translated ``google.api_core`` + exception for the caller to propagate: + + * daemon subprocess dead -> trip the breaker, fall back (it will not recover) + * ``UNIMPLEMENTED`` -> count toward the breaker, fall back for this call + * any other gRPC error -> reset the counter, translate and raise + * a non-gRPC exception -> re-raise unchanged (never masked as a fallback) + """ + # A dead subprocess can surface as a channel error under any status code, so + # check liveness first: the "daemon died mid-flight" case always wins and is + # never recoverable. + if daemon is not None and not daemon.is_running: + breaker.trip() + raise _AcceleratorFallback() from exc + if not isinstance(exc, RpcError): + # A bug in our own merge machinery, not a daemon result. Do not mask it + # as a fallback; let it propagate unchanged. + raise exc + if _grpc_code(exc) == StatusCode.UNIMPLEMENTED: + breaker.record_unimplemented() + raise _AcceleratorFallback() from exc + breaker.record_ok() + raise core_exceptions.from_grpc_error(exc) from exc diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_routing.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_routing.py new file mode 100644 index 000000000000..a6ccd68029af --- /dev/null +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_routing.py @@ -0,0 +1,27 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Single source of truth for which RPCs are routed through the accelerator.""" + +from __future__ import annotations + +# Names mirror the `_DataApiTarget` method names, not the gRPC method names. +# Adding an entry here is not enough on its own: the corresponding method must +# also include a top-of-function branch that dispatches to the accelerator +# service. Keep this set in lockstep with the bundled daemon's capabilities. +_ACCELERATOR_SUPPORTED: frozenset[str] = frozenset({"read_row", "mutate_row"}) + + +def is_supported(method_name: str) -> bool: + return method_name in _ACCELERATOR_SUPPORTED diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_client.py new file mode 100644 index 000000000000..a343f9ec7176 --- /dev/null +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_client.py @@ -0,0 +1,94 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""gRPC client surface against the accelerator daemon's UDS server. + +The daemon registers the standard ``google.bigtable.v2.Bigtable`` service on +its Unix domain socket, so we register the same stubs the gapic transport +uses and send V2 protos verbatim. No translation in either direction. +""" + +from __future__ import annotations + +from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable_v2.types import ( + MutateRowRequest, + MutateRowResponse, + ReadRowsRequest, + ReadRowsResponse, +) + +if CrossSync.is_async: + from grpc.aio import insecure_channel +else: + from grpc import insecure_channel + +__CROSS_SYNC_OUTPUT__ = "google.cloud.bigtable.data._sync_autogen._accelerator_client" + +_MUTATE_ROW_METHOD = "/google.bigtable.v2.Bigtable/MutateRow" +_READ_ROWS_METHOD = "/google.bigtable.v2.Bigtable/ReadRows" + + +@CrossSync.convert_class(sync_name="_AcceleratorClient") +class _AsyncAcceleratorClient: + """Thin gRPC client bound to the daemon's UDS. + + Owns the channel and the per-RPC stubs. The set of registered RPCs is the + same set the daemon supports today; the routing layer + (``_accelerator/_routing.py``) decides which calls reach this object. + """ + + def __init__(self, uds_path: str): + self._uds_path = uds_path + self._channel = insecure_channel(f"unix://{uds_path}") + self._mutate_row_stub = self._channel.unary_unary( + _MUTATE_ROW_METHOD, + request_serializer=MutateRowRequest.serialize, + response_deserializer=MutateRowResponse.deserialize, + ) + self._read_rows_stub = self._channel.unary_stream( + _READ_ROWS_METHOD, + request_serializer=ReadRowsRequest.serialize, + response_deserializer=ReadRowsResponse.deserialize, + ) + + @property + def uds_path(self) -> str: + return self._uds_path + + @CrossSync.convert + async def mutate_row( + self, request: MutateRowRequest, *, timeout: float | None = None + ) -> MutateRowResponse: + return await self._mutate_row_stub(request, timeout=timeout) + + @CrossSync.convert + async def read_rows( + self, request: ReadRowsRequest, *, timeout: float | None = None + ): + """Open the server-streaming ReadRows RPC against the daemon. + + Returns the streaming call object — callers iterate it (``async for`` + in async, ``for`` in sync) to consume ``ReadRowsResponse`` messages. + Shape matches what ``_gapic_client.read_rows`` returns so the existing + chunk-merging machinery in ``_read_rows.py`` works unchanged. + """ + return self._read_rows_stub(request, timeout=timeout) + + @CrossSync.convert + async def close(self) -> None: + if CrossSync.is_async: + await self._channel.close() + else: + self._channel.close() diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py index 409038ec9f6c..c588c6b34dd8 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py @@ -64,6 +64,7 @@ _WarmedInstanceKey, ) from google.cloud.bigtable.data._metrics import ( + ActiveOperationMetric, BigtableClientSideMetricsController, OperationType, tracked_retry, @@ -103,9 +104,20 @@ SampleRowKeysRequest, ) +from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon +from google.cloud.bigtable.data._accelerator._fallback import ( + AcceleratorBreaker, + _AcceleratorFallback, + handle_accelerator_error, +) +from google.cloud.bigtable.data._accelerator._routing import is_supported + if CrossSync.is_async: from grpc.aio import insecure_channel + from google.cloud.bigtable.data._async._accelerator_client import ( + _AsyncAcceleratorClient as AcceleratorClientType, + ) from google.cloud.bigtable.data._async._swappable_channel import ( AsyncSwappableChannel as SwappableChannelType, ) @@ -124,6 +136,9 @@ from grpc import insecure_channel, intercept_channel + from google.cloud.bigtable.data._sync_autogen._accelerator_client import ( # noqa: F401 + _AcceleratorClient as AcceleratorClientType, + ) from google.cloud.bigtable.data._sync_autogen._swappable_channel import ( # noqa: F401 SwappableChannel as SwappableChannelType, ) @@ -185,6 +200,7 @@ def __init__( client_options: dict[str, Any] | "google.api_core.client_options.ClientOptions" | None = None, + use_accelerator: bool | None = None, **kwargs, ): """ @@ -204,11 +220,20 @@ def __init__( client_options: Client options used to set user options on the client. API Endpoint should be set through client_options. + use_accelerator: + Whether to route supported RPCs through the in-process + accelerator daemon. Enabled by default. When the accelerator + cannot run — the emulator is set, or the daemon is unavailable + for this platform — it is automatically disabled (with a + warning) and the native client is used. Pass ``False`` to + disable it explicitly, or ``True`` to require it (which raises + if the emulator is set). Raises: {RAISE_NO_LOOP} """ if "pool_size" in kwargs: warnings.warn("pool_size no longer supported") + self._use_accelerator = use_accelerator # set up client info headers for veneer library self.client_info = DEFAULT_CLIENT_INFO self.client_info.client_library_version = self._client_version() @@ -1124,6 +1149,90 @@ def __init__( f"{self.__class__.__name__} must be created within an async event loop context." ) from e + # Optional in-process accelerator daemon, scoped to this Table's + # (project, instance_id, app_profile_id) tuple. Enabled by default; + # controlled by the client's ``use_accelerator`` option. + self._accelerator_daemon: AcceleratorDaemon | None = None + self._accelerator_client: AcceleratorClientType | None = None + # Sticky fallback policy: a daemon that cannot serve a routed RPC + # (UNIMPLEMENTED / dead subprocess) is transparently bypassed in favor + # of the native client. See ``_accelerator/_fallback.py``. + self._accelerator_breaker = AcceleratorBreaker() + if self.client._use_accelerator is not False: + # None (default) or True: attempt to start. ``True`` is an explicit + # request, so a conflict with the emulator is a hard error. + self._maybe_start_accelerator(explicit=self.client._use_accelerator is True) + + def _maybe_start_accelerator(self, *, explicit: bool) -> None: + """Start the accelerator daemon unless the environment prevents it. + + The accelerator is on by default, so it must never break callers when + it can't run. When the emulator is set, or the daemon fails to start + (e.g. the binary isn't bundled for this platform), fall back to the + native client and warn. The one hard error is an explicit + ``use_accelerator=True`` combined with the emulator, which is a genuine + misconfiguration. + """ + if self.client._emulator_host is not None: + if explicit: + raise RuntimeError( + "use_accelerator=True is not supported when " + "BIGTABLE_EMULATOR_HOST is set; unset the emulator, or pass " + "use_accelerator=False to use the emulator." + ) + warnings.warn( + "Accelerator disabled because BIGTABLE_EMULATOR_HOST is set; " + "using the native client.", + RuntimeWarning, + stacklevel=2, + ) + return + try: + self._start_accelerator() + except Exception as exc: + warnings.warn( + "Failed to start the Bigtable accelerator daemon; falling back " + f"to the native client: {exc}", + RuntimeWarning, + stacklevel=2, + ) + self._accelerator_daemon = None + self._accelerator_client = None + + def _start_accelerator(self) -> None: + """Spawn the daemon for this Table and connect to its UDS. + + On any failure the daemon is torn down before the error propagates. + """ + flags = [ + "--project", + self.client.project, + "--instance", + self.instance_id, + ] + if self.app_profile_id: + flags.extend(["--app-profile", self.app_profile_id]) + server = AcceleratorDaemon(cli_flags=flags) + try: + server.start() + self._accelerator_client = AcceleratorClientType(server.uds_path) + except BaseException: + server.close() + raise + self._accelerator_daemon = server + + def _use_accelerator(self, method_name: str) -> bool: + """Whether this call should be routed through the accelerator daemon. + + False once the fallback breaker has tripped, so a daemon that has proven + unable to serve routed RPCs is skipped without another round trip. + """ + return ( + self._accelerator_client is not None + and is_supported(method_name) + and not self._accelerator_breaker.bypass() + ) + @property @abc.abstractmethod def _request_path(self) -> dict[str, str]: @@ -1244,6 +1353,48 @@ async def read_rows( ) return [row async for row in row_generator] + @CrossSync.convert + async def _read_row_via_accelerator( + self, query, operation_timeout, attempt_timeout + ): + """Drive a single ReadRows attempt through the accelerator daemon. + + The daemon owns retry, so we run one attempt through the existing merger + machinery and skip start_operation/tracked_retry. The daemon also owns + metrics for accelerated RPCs, so the merger gets a handler-less metric: + it satisfies the merger's state machine but never exports anything here. + + Raises ``_AcceleratorFallback`` if the caller should retry on the native + client; other gRPC errors are translated to ``google.api_core`` + exceptions and raised. + """ + row_merger = CrossSync._ReadRowsOperation( + query, + self, + operation_timeout=operation_timeout, + attempt_timeout=attempt_timeout, + metric=ActiveOperationMetric(OperationType.READ_ROWS, is_streaming=False), + retryable_exceptions=(), + ) + row_merger._operation_metric.start_attempt() + try: + stream = self._accelerator_client.read_rows( + row_merger.request, timeout=operation_timeout + ) + chunked_stream = row_merger.chunk_stream(stream) + results = [a async for a in row_merger.merge_rows(chunked_stream)] + except Exception as exc: + handle_accelerator_error( + exc, + daemon=self._accelerator_daemon, + breaker=self._accelerator_breaker, + ) + raise # unreachable: handle_accelerator_error always raises + try: + return results[0] + except IndexError: + return None + @CrossSync.convert async def read_row( self, @@ -1288,6 +1439,16 @@ async def read_row( operation_timeout, attempt_timeout = _get_timeouts( operation_timeout, attempt_timeout, self ) + + if self._use_accelerator("read_row"): + try: + return await self._read_row_via_accelerator( + query, operation_timeout, attempt_timeout + ) + except _AcceleratorFallback: + # Daemon can't serve this call; fall through to the native path. + pass + retryable_excs = _get_retryable_errors(retryable_errors, self) row_merger = CrossSync._ReadRowsOperation( @@ -1587,6 +1748,29 @@ def mutations_batcher( batch_retryable_errors=batch_retryable_errors, ) + @CrossSync.convert + async def _mutate_row_via_accelerator(self, request, operation_timeout): + """Drive a single MutateRow attempt through the accelerator daemon. + + The daemon owns retry and metrics for accelerated RPCs, so we skip + tracked_retry, the predicate, and operation-level metrics entirely. + + Raises ``_AcceleratorFallback`` if the caller should retry on the native + client; other gRPC errors are translated to ``google.api_core`` + exceptions and raised. + """ + try: + return await self._accelerator_client.mutate_row( + request, timeout=operation_timeout + ) + except Exception as exc: + handle_accelerator_error( + exc, + daemon=self._accelerator_daemon, + breaker=self._accelerator_breaker, + ) + raise # unreachable: handle_accelerator_error always raises + @CrossSync.convert async def mutate_row( self, @@ -1637,6 +1821,22 @@ async def mutate_row( raise ValueError("No mutations provided") mutations_list = mutations if isinstance(mutations, list) else [mutations] + request = MutateRowRequest( + row_key=row_key.encode("utf-8") if isinstance(row_key, str) else row_key, + mutations=[mutation._to_pb() for mutation in mutations_list], + app_profile_id=self.app_profile_id, + **self._request_path, + ) + + if self._use_accelerator("mutate_row"): + try: + return await self._mutate_row_via_accelerator( + request, operation_timeout + ) + except _AcceleratorFallback: + # Daemon can't serve this call; fall through to the native path. + pass + if all(mutation.is_idempotent() for mutation in mutations_list): # mutations are all idempotent and safe to retry predicate = retries.if_exception_type( @@ -1651,14 +1851,7 @@ async def mutate_row( ) as operation_metric: target = partial( self.client._gapic_client.mutate_row, - request=MutateRowRequest( - row_key=row_key.encode("utf-8") - if isinstance(row_key, str) - else row_key, - mutations=[mutation._to_pb() for mutation in mutations_list], - app_profile_id=self.app_profile_id, - **self._request_path, - ), + request=request, timeout=attempt_timeout, retry=None, ) @@ -1866,6 +2059,16 @@ async def close(self): self.client._remove_instance_registration( self.instance_id, self.app_profile_id, id(self) ) + if self._accelerator_client is not None: + try: + await self._accelerator_client.close() + finally: + self._accelerator_client = None + if self._accelerator_daemon is not None: + try: + self._accelerator_daemon.close() + finally: + self._accelerator_daemon = None @CrossSync.convert(sync_name="__enter__") async def __aenter__(self): diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_accelerator_client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_accelerator_client.py new file mode 100644 index 000000000000..3b5edd274060 --- /dev/null +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_accelerator_client.py @@ -0,0 +1,81 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +# This file is automatically generated by CrossSync. Do not edit manually. + +"""gRPC client surface against the accelerator daemon's UDS server. + +The daemon registers the standard ``google.bigtable.v2.Bigtable`` service on +its Unix domain socket, so we register the same stubs the gapic transport +uses and send V2 protos verbatim. No translation in either direction. +""" + +from __future__ import annotations + +from grpc import insecure_channel + +from google.cloud.bigtable_v2.types import ( + MutateRowRequest, + MutateRowResponse, + ReadRowsRequest, + ReadRowsResponse, +) + +_MUTATE_ROW_METHOD = "/google.bigtable.v2.Bigtable/MutateRow" +_READ_ROWS_METHOD = "/google.bigtable.v2.Bigtable/ReadRows" + + +class _AcceleratorClient: + """Thin gRPC client bound to the daemon's UDS. + + Owns the channel and the per-RPC stubs. The set of registered RPCs is the + same set the daemon supports today; the routing layer + (``_accelerator/_routing.py``) decides which calls reach this object. + """ + + def __init__(self, uds_path: str): + self._uds_path = uds_path + self._channel = insecure_channel(f"unix://{uds_path}") + self._mutate_row_stub = self._channel.unary_unary( + _MUTATE_ROW_METHOD, + request_serializer=MutateRowRequest.serialize, + response_deserializer=MutateRowResponse.deserialize, + ) + self._read_rows_stub = self._channel.unary_stream( + _READ_ROWS_METHOD, + request_serializer=ReadRowsRequest.serialize, + response_deserializer=ReadRowsResponse.deserialize, + ) + + @property + def uds_path(self) -> str: + return self._uds_path + + def mutate_row( + self, request: MutateRowRequest, *, timeout: float | None = None + ) -> MutateRowResponse: + return self._mutate_row_stub(request, timeout=timeout) + + def read_rows(self, request: ReadRowsRequest, *, timeout: float | None = None): + """Open the server-streaming ReadRows RPC against the daemon. + + Returns the streaming call object — callers iterate it (``async for`` + in async, ``for`` in sync) to consume ``ReadRowsResponse`` messages. + Shape matches what ``_gapic_client.read_rows`` returns so the existing + chunk-merging machinery in ``_read_rows.py`` works unchanged.""" + return self._read_rows_stub(request, timeout=timeout) + + def close(self) -> None: + self._channel.close() diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py index 6029d7f75f7c..f74b31fdeff0 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py @@ -44,6 +44,13 @@ from grpc import Channel, insecure_channel, intercept_channel from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT +from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon +from google.cloud.bigtable.data._accelerator._fallback import ( + AcceleratorBreaker, + _AcceleratorFallback, + handle_accelerator_error, +) +from google.cloud.bigtable.data._accelerator._routing import is_supported from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._helpers import ( _CONCURRENCY_LIMIT, @@ -58,10 +65,14 @@ _WarmedInstanceKey, ) from google.cloud.bigtable.data._metrics import ( + ActiveOperationMetric, BigtableClientSideMetricsController, OperationType, tracked_retry, ) +from google.cloud.bigtable.data._sync_autogen._accelerator_client import ( + _AcceleratorClient as AcceleratorClientType, +) from google.cloud.bigtable.data._sync_autogen._swappable_channel import ( SwappableChannel as SwappableChannelType, ) @@ -128,6 +139,7 @@ def __init__( client_options: dict[str, Any] | "google.api_core.client_options.ClientOptions" | None = None, + use_accelerator: bool | None = None, **kwargs, ): """Create a client instance for the Bigtable Data API @@ -146,10 +158,19 @@ def __init__( client_options: Client options used to set user options on the client. API Endpoint should be set through client_options. + use_accelerator: + Whether to route supported RPCs through the in-process + accelerator daemon. Enabled by default. When the accelerator + cannot run — the emulator is set, or the daemon is unavailable + for this platform — it is automatically disabled (with a + warning) and the native client is used. Pass ``False`` to + disable it explicitly, or ``True`` to require it (which raises + if the emulator is set). Raises: """ if "pool_size" in kwargs: warnings.warn("pool_size no longer supported") + self._use_accelerator = use_accelerator self.client_info = DEFAULT_CLIENT_INFO self.client_info.client_library_version = self._client_version() if type(client_options) is dict: @@ -887,6 +908,69 @@ def __init__( raise RuntimeError( f"{self.__class__.__name__} must be created within an async event loop context." ) from e + self._accelerator_daemon: AcceleratorDaemon | None = None + self._accelerator_client: AcceleratorClientType | None = None + self._accelerator_breaker = AcceleratorBreaker() + if self.client._use_accelerator is not False: + self._maybe_start_accelerator(explicit=self.client._use_accelerator is True) + + def _maybe_start_accelerator(self, *, explicit: bool) -> None: + """Start the accelerator daemon unless the environment prevents it. + + The accelerator is on by default, so it must never break callers when + it can't run. When the emulator is set, or the daemon fails to start + (e.g. the binary isn't bundled for this platform), fall back to the + native client and warn. The one hard error is an explicit + ``use_accelerator=True`` combined with the emulator, which is a genuine + misconfiguration.""" + if self.client._emulator_host is not None: + if explicit: + raise RuntimeError( + "use_accelerator=True is not supported when BIGTABLE_EMULATOR_HOST is set; unset the emulator, or pass use_accelerator=False to use the emulator." + ) + warnings.warn( + "Accelerator disabled because BIGTABLE_EMULATOR_HOST is set; using the native client.", + RuntimeWarning, + stacklevel=2, + ) + return + try: + self._start_accelerator() + except Exception as exc: + warnings.warn( + f"Failed to start the Bigtable accelerator daemon; falling back to the native client: {exc}", + RuntimeWarning, + stacklevel=2, + ) + self._accelerator_daemon = None + self._accelerator_client = None + + def _start_accelerator(self) -> None: + """Spawn the daemon for this Table and connect to its UDS. + + On any failure the daemon is torn down before the error propagates.""" + flags = ["--project", self.client.project, "--instance", self.instance_id] + if self.app_profile_id: + flags.extend(["--app-profile", self.app_profile_id]) + server = AcceleratorDaemon(cli_flags=flags) + try: + server.start() + self._accelerator_client = AcceleratorClientType(server.uds_path) + except BaseException: + server.close() + raise + self._accelerator_daemon = server + + def _use_accelerator(self, method_name: str) -> bool: + """Whether this call should be routed through the accelerator daemon. + + False once the fallback breaker has tripped, so a daemon that has proven + unable to serve routed RPCs is skipped without another round trip.""" + return ( + self._accelerator_client is not None + and is_supported(method_name) + and (not self._accelerator_breaker.bypass()) + ) @property @abc.abstractmethod @@ -999,6 +1083,42 @@ def read_rows( ) return [row for row in row_generator] + def _read_row_via_accelerator(self, query, operation_timeout, attempt_timeout): + """Drive a single ReadRows attempt through the accelerator daemon. + + The daemon owns retry, so we run one attempt through the existing merger + machinery and skip start_operation/tracked_retry. The daemon also owns + metrics for accelerated RPCs, so the merger gets a handler-less metric: + it satisfies the merger's state machine but never exports anything here. + + Raises ``_AcceleratorFallback`` if the caller should retry on the native + client; other gRPC errors are translated to ``google.api_core`` + exceptions and raised.""" + row_merger = CrossSync._Sync_Impl._ReadRowsOperation( + query, + self, + operation_timeout=operation_timeout, + attempt_timeout=attempt_timeout, + metric=ActiveOperationMetric(OperationType.READ_ROWS, is_streaming=False), + retryable_exceptions=(), + ) + row_merger._operation_metric.start_attempt() + try: + stream = self._accelerator_client.read_rows( + row_merger.request, timeout=operation_timeout + ) + chunked_stream = row_merger.chunk_stream(stream) + results = [a for a in row_merger.merge_rows(chunked_stream)] + except Exception as exc: + handle_accelerator_error( + exc, daemon=self._accelerator_daemon, breaker=self._accelerator_breaker + ) + raise + try: + return results[0] + except IndexError: + return None + def read_row( self, row_key: str | bytes, @@ -1039,6 +1159,13 @@ def read_row( operation_timeout, attempt_timeout = _get_timeouts( operation_timeout, attempt_timeout, self ) + if self._use_accelerator("read_row"): + try: + return self._read_row_via_accelerator( + query, operation_timeout, attempt_timeout + ) + except _AcceleratorFallback: + pass retryable_excs = _get_retryable_errors(retryable_errors, self) row_merger = CrossSync._Sync_Impl._ReadRowsOperation( query, @@ -1309,6 +1436,25 @@ def mutations_batcher( batch_retryable_errors=batch_retryable_errors, ) + def _mutate_row_via_accelerator(self, request, operation_timeout): + """Drive a single MutateRow attempt through the accelerator daemon. + + The daemon owns retry and metrics for accelerated RPCs, so we skip + tracked_retry, the predicate, and operation-level metrics entirely. + + Raises ``_AcceleratorFallback`` if the caller should retry on the native + client; other gRPC errors are translated to ``google.api_core`` + exceptions and raised.""" + try: + return self._accelerator_client.mutate_row( + request, timeout=operation_timeout + ) + except Exception as exc: + handle_accelerator_error( + exc, daemon=self._accelerator_daemon, breaker=self._accelerator_breaker + ) + raise + def mutate_row( self, row_key: str | bytes, @@ -1354,6 +1500,17 @@ def mutate_row( if not mutations: raise ValueError("No mutations provided") mutations_list = mutations if isinstance(mutations, list) else [mutations] + request = MutateRowRequest( + row_key=row_key.encode("utf-8") if isinstance(row_key, str) else row_key, + mutations=[mutation._to_pb() for mutation in mutations_list], + app_profile_id=self.app_profile_id, + **self._request_path, + ) + if self._use_accelerator("mutate_row"): + try: + return self._mutate_row_via_accelerator(request, operation_timeout) + except _AcceleratorFallback: + pass if all((mutation.is_idempotent() for mutation in mutations_list)): predicate = retries.if_exception_type( *_get_retryable_errors(retryable_errors, self) @@ -1365,14 +1522,7 @@ def mutate_row( ) as operation_metric: target = partial( self.client._gapic_client.mutate_row, - request=MutateRowRequest( - row_key=row_key.encode("utf-8") - if isinstance(row_key, str) - else row_key, - mutations=[mutation._to_pb() for mutation in mutations_list], - app_profile_id=self.app_profile_id, - **self._request_path, - ), + request=request, timeout=attempt_timeout, retry=None, ) @@ -1564,6 +1714,16 @@ def close(self): self.client._remove_instance_registration( self.instance_id, self.app_profile_id, id(self) ) + if self._accelerator_client is not None: + try: + self._accelerator_client.close() + finally: + self._accelerator_client = None + if self._accelerator_daemon is not None: + try: + self._accelerator_daemon.close() + finally: + self._accelerator_daemon = None def __enter__(self): """Implement async context manager protocol diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py new file mode 100644 index 000000000000..36c58257bcfd --- /dev/null +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py @@ -0,0 +1,101 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for how the client decides whether to start the accelerator +daemon. The accelerator is on by default and must degrade gracefully: it +disables itself (with a warning) when the emulator is set or the daemon can't +start, and only errors out on an explicit ``use_accelerator=True`` that +conflicts with the emulator. These exercise the pure enablement helpers on the +async target without spinning up a full client or event loop.""" + +from types import SimpleNamespace + +import pytest + +from google.cloud.bigtable.data._async.client import _DataApiTargetAsync + + +class _ConcreteTarget(_DataApiTargetAsync): + """Minimal concrete subclass so the abstract base can be instantiated.""" + + @property + def _request_path(self): + return {} + + +def _bare_target(start, emulator=None): + """A target that skips __init__, wired with a fake client and a fake + ``_start_accelerator`` so we can observe the enablement decision.""" + t = object.__new__(_ConcreteTarget) + t.client = SimpleNamespace(_emulator_host=emulator, project="p") + t.instance_id = "i" + t.app_profile_id = None + t._accelerator_daemon = "SENTINEL" + t._accelerator_client = "SENTINEL" + t._start_accelerator = start + return t + + +def _ok_start(target): + def _start(): + target._accelerator_client = "STARTED" + target._accelerator_daemon = "STARTED" + + return _start + + +class TestMaybeStartAccelerator: + def test_emulator_auto_disables_with_warning(self): + """Default (non-explicit) + emulator: warn and use the native client.""" + called = [] + t = _bare_target(emulator="localhost:8086", start=lambda: called.append(1)) + with pytest.warns(RuntimeWarning, match="Accelerator disabled"): + t._maybe_start_accelerator(explicit=False) + assert called == [] + + def test_emulator_explicit_raises(self): + """Explicit use_accelerator=True + emulator is a hard misconfiguration.""" + t = _bare_target(emulator="localhost:8086", start=lambda: None) + with pytest.raises(RuntimeError, match="use_accelerator=True is not supported"): + t._maybe_start_accelerator(explicit=True) + + @pytest.mark.parametrize("explicit", [False, True]) + def test_start_failure_falls_back_to_native(self, explicit): + """A daemon start failure (e.g. binary missing) never propagates; it + warns and leaves the target on the native client.""" + + def _boom(): + raise FileNotFoundError("no bundled binary") + + t = _bare_target(emulator=None, start=_boom) + with pytest.warns(RuntimeWarning, match="Failed to start"): + t._maybe_start_accelerator(explicit=explicit) + assert t._accelerator_client is None + assert t._accelerator_daemon is None + + def test_successful_start(self): + t = _bare_target(emulator=None, start=None) + t._start_accelerator = _ok_start(t) + t._maybe_start_accelerator(explicit=False) + assert t._accelerator_client == "STARTED" + assert t._accelerator_daemon == "STARTED" + + def test_keyboard_interrupt_propagates(self): + """Only Exception is swallowed; BaseException (Ctrl-C) must propagate.""" + + def _interrupt(): + raise KeyboardInterrupt() + + t = _bare_target(emulator=None, start=_interrupt) + with pytest.raises(KeyboardInterrupt): + t._maybe_start_accelerator(explicit=False) From bf535624437f4e3211aff8437d625e1ad93e67ea Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 11 Aug 2026 21:31:22 +0000 Subject: [PATCH 2/4] docs(bigtable): drop stale Go-client references from accelerator fallback comments Change-Id: I191c52745e0cd10e5ab99e5e06135cd2e518bb9f --- .../bigtable/data/_accelerator/_fallback.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py index d50a74cf2099..d02662edf736 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py @@ -14,12 +14,11 @@ # """Client-side fallback policy for accelerator-routed RPCs. -Mirrors the Go client's ``session.UnimplementedErrorInterceptor``: a daemon that -cannot open any sessions replies ``UNIMPLEMENTED``, and the routing layer -transparently retries the call on the native client. A sticky breaker trips -after enough consecutive ``UNIMPLEMENTED`` replies so a persistently-degraded -daemon stops being dialed at all. A daemon whose subprocess has died mid-flight -trips the breaker immediately — it will never recover. +A daemon that cannot open any sessions replies ``UNIMPLEMENTED``, and the routing +layer transparently retries the call on the native client. A sticky breaker +trips after enough consecutive ``UNIMPLEMENTED`` replies so a persistently- +degraded daemon stops being dialed at all. A daemon whose subprocess has died +mid-flight trips the breaker immediately — it will never recover. Any other gRPC error is a real, daemon-served result the native client would reproduce (the daemon owns retries, so it has already exhausted them), so it is @@ -44,8 +43,7 @@ from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon # Consecutive ``UNIMPLEMENTED`` replies that trip the sticky breaker, after which -# the accelerator is bypassed for the lifetime of the Table. Matches the Go -# client's ``session.DefaultUnimplementedThreshold``. +# the accelerator is bypassed for the lifetime of the Table. DEFAULT_UNIMPLEMENTED_THRESHOLD = 30 @@ -67,8 +65,8 @@ class AcceleratorBreaker: the RPC shape but has no working sessions), and * an explicit :meth:`trip` when the daemon subprocess is found dead. - Any non-``UNIMPLEMENTED`` outcome resets the consecutive count, matching the - Go interceptor: a normal reply proves the daemon is healthy again. + Any non-``UNIMPLEMENTED`` outcome resets the consecutive count: a normal + reply proves the daemon is healthy again. """ def __init__(self, threshold: int = DEFAULT_UNIMPLEMENTED_THRESHOLD): From 050258fbb75b066b5ed3632833169aa1c2fb6b59 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 11 Aug 2026 21:31:38 +0000 Subject: [PATCH 3/4] chore(bigtable): add TODO to emit a metric in accelerator error handling Change-Id: I72d5e456cdf722498ace021cc1f976cc1469a4dc --- .../google/cloud/bigtable/data/_accelerator/_fallback.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py index d02662edf736..f38a10b4fb06 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py @@ -125,6 +125,9 @@ def handle_accelerator_error( * any other gRPC error -> reset the counter, translate and raise * a non-gRPC exception -> re-raise unchanged (never masked as a fallback) """ + # TODO(accelerator): emit a metric here (e.g. a fallback/error counter keyed + # by reason: dead-daemon / unimplemented / translated-error) once client-side + # accelerator metrics are wired up. # A dead subprocess can surface as a channel error under any status code, so # check liveness first: the "daemon died mid-flight" case always wins and is # never recoverable. From 6d4ace5f091eaa4962c8d530890093a97d9d4e0f Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 11 Aug 2026 21:32:43 +0000 Subject: [PATCH 4/4] fix(bigtable): fall back to native immediately on first accelerator UNIMPLEMENTED Change-Id: Ifad0d9eb9d7b4d03fff07de2549a6d3294cf2f16 --- .../bigtable/data/_accelerator/_fallback.py | 47 ++++++------------- 1 file changed, 14 insertions(+), 33 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py index f38a10b4fb06..eb0cafa05232 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py @@ -15,14 +15,14 @@ """Client-side fallback policy for accelerator-routed RPCs. A daemon that cannot open any sessions replies ``UNIMPLEMENTED``, and the routing -layer transparently retries the call on the native client. A sticky breaker -trips after enough consecutive ``UNIMPLEMENTED`` replies so a persistently- -degraded daemon stops being dialed at all. A daemon whose subprocess has died -mid-flight trips the breaker immediately — it will never recover. +layer transparently retries the call on the native client. The first +``UNIMPLEMENTED`` reply trips a sticky breaker so a persistently-degraded daemon +stops being dialed at all. A daemon whose subprocess has died mid-flight trips +the breaker immediately — it will never recover. Any other gRPC error is a real, daemon-served result the native client would reproduce (the daemon owns retries, so it has already exhausted them), so it is -translated to the matching ``google.api_core`` exception and raised without +translated to the corresponding ``google.api_core`` exception and raised without falling back. This module is plain sync-only logic shared verbatim by the async and generated @@ -42,10 +42,6 @@ if TYPE_CHECKING: from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon -# Consecutive ``UNIMPLEMENTED`` replies that trip the sticky breaker, after which -# the accelerator is bypassed for the lifetime of the Table. -DEFAULT_UNIMPLEMENTED_THRESHOLD = 30 - class _AcceleratorFallback(Exception): """Internal signal that an accelerator attempt should be retried natively. @@ -61,17 +57,12 @@ class AcceleratorBreaker: One instance per Table. Thread-safe so the generated sync client can share a Table across threads. Two triggers permanently bypass the accelerator: - * ``threshold`` consecutive ``UNIMPLEMENTED`` replies (the daemon understands - the RPC shape but has no working sessions), and + * the first ``UNIMPLEMENTED`` reply (the daemon understands the RPC shape but + has no working sessions), and * an explicit :meth:`trip` when the daemon subprocess is found dead. - - Any non-``UNIMPLEMENTED`` outcome resets the consecutive count: a normal - reply proves the daemon is healthy again. """ - def __init__(self, threshold: int = DEFAULT_UNIMPLEMENTED_THRESHOLD): - self._threshold = threshold - self._consecutive = 0 + def __init__(self): self._tripped = False self._lock = threading.Lock() @@ -84,18 +75,6 @@ def trip(self) -> None: with self._lock: self._tripped = True - def record_unimplemented(self) -> None: - """Note an ``UNIMPLEMENTED`` reply; trip the breaker at the threshold.""" - with self._lock: - self._consecutive += 1 - if self._consecutive >= self._threshold: - self._tripped = True - - def record_ok(self) -> None: - """Note any non-``UNIMPLEMENTED`` outcome; resets the consecutive count.""" - with self._lock: - self._consecutive = 0 - def _grpc_code(exc: BaseException) -> StatusCode | None: """Best-effort extraction of a gRPC status code from an exception.""" @@ -121,8 +100,8 @@ def handle_accelerator_error( exception for the caller to propagate: * daemon subprocess dead -> trip the breaker, fall back (it will not recover) - * ``UNIMPLEMENTED`` -> count toward the breaker, fall back for this call - * any other gRPC error -> reset the counter, translate and raise + * ``UNIMPLEMENTED`` -> trip the breaker, fall back immediately + * any other gRPC error -> translate and raise * a non-gRPC exception -> re-raise unchanged (never masked as a fallback) """ # TODO(accelerator): emit a metric here (e.g. a fallback/error counter keyed @@ -139,7 +118,9 @@ def handle_accelerator_error( # as a fallback; let it propagate unchanged. raise exc if _grpc_code(exc) == StatusCode.UNIMPLEMENTED: - breaker.record_unimplemented() + # The daemon only replies UNIMPLEMENTED once it has no working sessions, + # a persistent condition, so trip the breaker and fall back immediately + # rather than re-dialing on every subsequent call. + breaker.trip() raise _AcceleratorFallback() from exc - breaker.record_ok() raise core_exceptions.from_grpc_error(exc) from exc