feat(bigtable): route read_row/mutate_row through the accelerator with native fallback

Change-Id: I363bc3393ff13033e6c37ef9df4be13571837932

Author: mutianf

packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py

@@ -25,12 +25,18 @@
 
 import os
 import shutil
+import signal
 import socket
 import subprocess
 import tempfile
 import time
 from typing import Sequence
 
+# Environment variable that overrides the bundled binary location. Primarily
+# for development against a locally-built daemon, and for tests pointing at a
+# fake binary.
+_BIN_ENV_VAR = "BIGTABLE_ACCELERATOR_BIN"
+
 # Wheels ship the binary at this path relative to the `_accelerator/` package.
 _DEFAULT_BIN_RELATIVE_PATH = "bin/accelerator"
 
@@ -44,12 +50,34 @@
 _SIGKILL_GRACE_SECONDS = 2.0
 
 
-def _resolve_binary_path() -> str:
-    """Return the path of the bundled daemon binary, raising if it is absent."""
+def _resolve_binary_path(explicit_path: str | None = None) -> str:
+    """Resolve the daemon binary path, validating that it is a regular file.
+
+    Precedence: an explicit ``binary_path`` argument, then the
+    ``BIGTABLE_ACCELERATOR_BIN`` env var, then the binary bundled in the wheel.
+    An explicit path or env override that does not point at a regular file is a
+    hard error (a caller who named a path meant it); a missing bundled binary
+    reports how to supply one.
+    """
+    if explicit_path is not None:
+        if not os.path.isfile(explicit_path):
+            raise FileNotFoundError(
+                f"binary_path={explicit_path!r} does not point at a regular file"
+            )
+        return explicit_path
+    override = os.environ.get(_BIN_ENV_VAR)
+    if override:
+        if not os.path.isfile(override):
+            raise FileNotFoundError(
+                f"{_BIN_ENV_VAR}={override!r} does not point at a regular file"
+            )
+        return override
     bundled = os.path.join(os.path.dirname(__file__), _DEFAULT_BIN_RELATIVE_PATH)
     if not os.path.isfile(bundled):
         raise FileNotFoundError(
-            "Accelerator binary not found. Install a wheel that bundles the binary."
+            "No accelerator binary found. Set the "
+            f"{_BIN_ENV_VAR} env var to a daemon binary path, or install a "
+            "wheel that bundles the binary."
         )
     return bundled
 
@@ -76,20 +104,25 @@ def __init__(
         self,
         cli_flags: Sequence[str] = (),
         *,
+        binary_path: str | None = None,
         startup_timeout: float = _DEFAULT_STARTUP_TIMEOUT,
     ):
         """Resolve the binary and pick the UDS path (does not spawn anything).
 
         Args:
             cli_flags: extra arguments appended after ``--uds-path`` when
                 spawning the daemon (e.g. ``--project``/``--instance``).
+            binary_path: explicit path to the daemon binary. When omitted, the
+                path is resolved from the ``BIGTABLE_ACCELERATOR_BIN`` env var
+                and then the binary bundled in the wheel.
             startup_timeout: seconds ``start()`` waits for the daemon's UDS to
                 become connectable before raising.
 
         Raises:
-            FileNotFoundError: the bundled binary is not present in this wheel.
+            FileNotFoundError: no binary could be resolved, or an explicit
+                ``binary_path``/env override does not point at a regular file.
         """
-        self._binary_path = _resolve_binary_path()
+        self._binary_path = _resolve_binary_path(binary_path)
         self._cli_flags = list(cli_flags)
         self._startup_timeout = startup_timeout
         self._tempdir: str | None = None
@@ -139,10 +172,6 @@ def start(self) -> None:
         """
         if self._proc is not None:
             raise RuntimeError("AcceleratorDaemon.start() called twice")
-        if not hasattr(socket, "AF_UNIX"):
-            raise OSError(
-                "Unix domain sockets (AF_UNIX) are not supported on this platform."
-            )
         self._tempdir = tempfile.mkdtemp(prefix="bt-accel-")
         self._uds_path = os.path.join(self._tempdir, "sock")
         # Redirect the daemon's stdout/stderr to a log file rather than
@@ -157,10 +186,9 @@ def start(self) -> None:
         # The log file handle only needs to live long enough for Popen to dup
         # it into the child, so it stays local to start() rather than being an
         # attribute. Startup failures read the tail back from the path.
+        log_file = open(self._log_path, "wb")
         argv = [self._binary_path, "--uds-path", self._uds_path, *self._cli_flags]
-        log_file = None
         try:
-            log_file = open(self._log_path, "wb")
             self._proc = subprocess.Popen(
                 argv,
                 stdin=subprocess.PIPE,
@@ -177,11 +205,10 @@ def start(self) -> None:
             # Whether or not the spawn succeeded, the parent no longer needs its
             # copy of the log fd: on success the child holds its own dup, and on
             # failure there is nothing to keep open.
-            if log_file is not None:
-                try:
-                    log_file.close()
-                except OSError:
-                    pass
+            try:
+                log_file.close()
+            except OSError:
+                pass
         try:
             self._wait_until_ready(self._startup_timeout)
         except BaseException:
@@ -205,17 +232,11 @@ def close(self) -> None:
                         pass
                 if not self._wait_for_exit(_STDIN_GRACE_SECONDS):
                     # Step 2: SIGTERM.
-                    try:
-                        proc.terminate()
-                    except OSError:
-                        pass
+                    proc.terminate()
                     if not self._wait_for_exit(_SIGTERM_GRACE_SECONDS):
                         # Step 3: SIGKILL. Bounded wait so teardown can't hang
                         # forever if the process is stuck unreapable.
-                        try:
-                            proc.kill()
-                        except OSError:
-                            pass
+                        proc.kill()
                         self._wait_for_exit(_SIGKILL_GRACE_SECONDS)
         finally:
             self._proc = None
@@ -295,7 +316,7 @@ def _force_kill(self) -> None:
         if self._proc is None or self._proc.poll() is not None:
             return
         try:
-            self._proc.kill()
+            self._proc.send_signal(signal.SIGKILL)
         except (OSError, ProcessLookupError):
             pass
         try:
@@ -309,4 +330,3 @@ def _cleanup_tempdir(self) -> None:
             shutil.rmtree(self._tempdir, ignore_errors=True)
         self._tempdir = None
         self._uds_path = None
-        self._log_path = None

packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py

@@ -0,0 +1,146 @@
+# 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 logging
+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
+
+_LOGGER = logging.getLogger(__name__)
+
+
+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)."""
+        with self._lock:
+            self._tripped = True
+
+
+def _grpc_code(exc: BaseException) -> StatusCode | None:
+    """Best-effort extraction of a gRPC status code from an exception.
+
+    Handles both ``grpc.RpcError`` (status via a ``code()`` method) and
+    ``google.api_core.exceptions.GoogleAPICallError`` (status stored on the
+    ``grpc_status_code`` attribute), returning ``None`` for anything else.
+    """
+    code = getattr(exc, "code", None)
+    if callable(code):
+        try:
+            return code()
+        except Exception:
+            return None
+    grpc_status = getattr(exc, "grpc_status_code", None)
+    if isinstance(grpc_status, StatusCode):
+        return grpc_status
+    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:
+        _LOGGER.warning(
+            "Accelerator daemon is no longer running; permanently falling back "
+            "to the native Bigtable client for this table.",
+            exc_info=exc,
+        )
+        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.
+        _LOGGER.warning(
+            "Accelerator daemon replied UNIMPLEMENTED; permanently falling back "
+            "to the native Bigtable client for this table."
+        )
+        breaker.trip()
+        raise _AcceleratorFallback() from exc
+    raise core_exceptions.from_grpc_error(exc) from exc

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