-
Notifications
You must be signed in to change notification settings - Fork 0
feat(bigtable): route read_row/mutate_row through the accelerator with native fallback #2
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.
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
Changes from all commits
a698ef1
bf53562
050258f
6d4ace5
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,126 @@ | ||
| # 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. | ||
|
|
||
| A daemon that cannot open any sessions replies ``UNIMPLEMENTED``, and the routing | ||
| 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 corresponding ``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 | ||
|
|
||
|
|
||
| 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: | ||
|
|
||
| * 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. | ||
| """ | ||
|
|
||
| def __init__(self): | ||
| 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).""" | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This seems like something we'd want to log, right? (See here for how to get a reference to a standard logger) |
||
| with self._lock: | ||
| self._tripped = True | ||
|
|
||
|
|
||
| def _grpc_code(exc: BaseException) -> StatusCode | None: | ||
| """Best-effort extraction of a gRPC status code from an exception.""" | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This would work for grpc.RpcError exceptions. But there's also api_core.exceptions.GoogleAPICallError errors, which store this as grpc_status_code. I don't know if that'll be relevant here (We have similar code here. Maybe we should make this a shared helper?) |
||
| 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`` -> 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 | ||
| # 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. | ||
| 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: | ||
| # 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 | ||
| raise core_exceptions.from_grpc_error(exc) from exc | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This could also live in google.cloud.bigtable.data.exceptions, depending on the usage