test(bigtable): add accelerator shared test harness and concurrency system tests

packages/google-cloud-bigtable/conftest.py

@@ -0,0 +1,48 @@
+# 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.
+"""Root pytest configuration for the package.
+
+The only job here is to keep ``pytest.ini`` portable across the range of
+``pytest-asyncio`` versions this package is tested with. ``pytest.ini`` sets
+``asyncio_default_fixture_loop_scope``/``asyncio_default_test_loop_scope`` so the
+session-scoped async fixtures in the system suites share one event loop on
+modern ``pytest-asyncio`` (>=0.23). Older ``pytest-asyncio`` (e.g. the 0.21.2 the
+nox ``system`` session pins) does not register those ini options, so it would
+emit a ``PytestConfigWarning: Unknown config option`` -- harmless on its own, but
+fatal in any environment that runs pytest with ``-W error``.
+
+Registering the options defensively here makes them "known" on older
+``pytest-asyncio`` (silencing that warning) while staying a no-op on newer
+versions, which already register them (the duplicate ``addini`` raises
+``ValueError``, which we swallow). Both the value and pytest-asyncio's own
+handling are otherwise untouched.
+"""
+
+
+def pytest_addoption(parser):
+    for name, help_text in (
+        (
+            "asyncio_default_fixture_loop_scope",
+            "default event-loop scope for asyncio fixtures (pytest-asyncio)",
+        ),
+        (
+            "asyncio_default_test_loop_scope",
+            "default event-loop scope for asyncio tests (pytest-asyncio)",
+        ),
+    ):
+        try:
+            parser.addini(name, help_text, default=None)
+        except ValueError:
+            # Already registered by a pytest-asyncio new enough to own it.
+            pass

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

@@ -20,6 +20,10 @@
 stops being dialed at all. A daemon whose subprocess has died mid-flight trips
 the breaker immediately — it will never recover.
 
+The same breaker also carries a reversible "degraded" flag, set by the health
+monitor when the daemon is up but too starved to answer its own health check;
+see :class:`AcceleratorBreaker`.
+
 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
@@ -36,9 +40,8 @@
 import threading
 from typing import TYPE_CHECKING
 
-from grpc import RpcError, StatusCode
-
 from google.api_core import exceptions as core_exceptions
+from grpc import RpcError, StatusCode
 
 if TYPE_CHECKING:
     from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon
@@ -57,26 +60,61 @@ class _AcceleratorFallback(Exception):
 class AcceleratorBreaker:
     """Tracks accelerator health and decides when to stop using it.
 
-    One instance per data target. Thread-safe so the generated sync client can
-    share a target 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.
+    One instance per Table. Thread-safe so the generated sync client can share a
+    Table across threads. There are two independent reasons to bypass, and they
+    differ in whether they can be undone:
+
+    * **Tripped** — permanent. The first ``UNIMPLEMENTED`` reply (the daemon
+      understands the RPC shape but has no working sessions), or an explicit
+      :meth:`trip` when the daemon subprocess is found dead. Neither condition
+      recovers, so the accelerator is abandoned for the life of the Table.
+    * **Degraded** — reversible. Set by the health monitor when the daemon stops
+      answering its own health check promptly (see
+      ``_async/_accelerator_health.py``). A daemon starved of CPU recovers when
+      the load that starved it goes away, so this must clear again — routing
+      permanently away from the accelerator on one bad minute would be a far
+      worse outcome than the slow calls it avoided.
+
+    Keeping the two flags separate is what makes that safe: :meth:`set_degraded`
+    can never resurrect an accelerator that :meth:`trip` has given up on. Each
+    flag is its own :class:`threading.Event`, so :meth:`bypass` — which runs on
+    every routed RPC — stays a pair of lock-free reads. The two are never read
+    as a pair: either one being set already means bypass.
     """
 
     def __init__(self):
         self._tripped = threading.Event()
+        self._degraded = threading.Event()
 
     def bypass(self) -> bool:
-        """Whether the accelerator should be skipped entirely from now on."""
-        return self._tripped.is_set()
+        """Whether the accelerator should be skipped for calls made right now."""
+        return self._tripped.is_set() or self._degraded.is_set()
 
     def trip(self) -> None:
         """Permanently bypass the accelerator (e.g. the daemon process died)."""
         self._tripped.set()
 
+    @property
+    def is_tripped(self) -> bool:
+        """Whether the accelerator has been permanently abandoned."""
+        return self._tripped.is_set()
+
+    @property
+    def is_degraded(self) -> bool:
+        """Whether the accelerator is currently being skipped for poor health."""
+        return self._degraded.is_set()
+
+    def set_degraded(self, degraded: bool) -> None:
+        """Start or stop bypassing the accelerator for poor health.
+
+        Reversible, and orthogonal to :meth:`trip`: clearing it does not undo a
+        permanent trip.
+        """
+        if degraded:
+            self._degraded.set()
+        else:
+            self._degraded.clear()
+
 
 def _grpc_code(exc: BaseException) -> StatusCode | None:
     """Best-effort extraction of a gRPC status code from an exception.

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

@@ -0,0 +1,89 @@
+# 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.
+#
+"""Wire codec for ``grpc.health.v1.Health/Check`` against the daemon.
+
+The daemon serves the standard gRPC health service. The generated Python stubs
+for it live in ``grpcio-health-checking``, a separate distribution this package
+does not depend on, and health.proto is far too small to justify adding one: the
+request we send is the empty message, and the response is a single varint.
+
+    HealthCheckRequest{service: ""}  ->  b""
+    HealthCheckResponse{status: S}   ->  b"\\x08" + varint(S)
+                                         b""          -> UNKNOWN (proto3 omits 0)
+                                         b"\\x08\\x01"  -> SERVING
+                                         b"\\x08\\x02"  -> NOT_SERVING
+
+So the two functions below are the whole codec. Note that we always probe the
+empty service name (""), which is the daemon's overall-process status.
+
+This module is plain sync-only logic, shared verbatim by the async and generated
+sync clients; no CrossSync branching is needed.
+"""
+
+from __future__ import annotations
+
+import enum
+
+#: Full method name of the health check RPC, as registered by the daemon.
+HEALTH_CHECK_METHOD = "/grpc.health.v1.Health/Check"
+
+#: Tag byte for ``HealthCheckResponse.status``: field 1, wire type 0 (varint).
+_STATUS_TAG = 0x08
+
+
+class ServingStatus(enum.IntEnum):
+    """Mirror of ``grpc.health.v1.HealthCheckResponse.ServingStatus``."""
+
+    UNKNOWN = 0
+    SERVING = 1
+    NOT_SERVING = 2
+    SERVICE_UNKNOWN = 3
+
+
+def serialize_health_request(_request: object = None) -> bytes:
+    """Serialize ``HealthCheckRequest{service: ""}``.
+
+    Every field is at its proto3 default, so the encoding is empty. Takes and
+    ignores an argument because gRPC calls it with the request object.
+    """
+    return b""
+
+
+def parse_health_response(payload: bytes) -> ServingStatus:
+    """Decode a ``HealthCheckResponse`` into its status.
+
+    Anything that isn't a recognizable status field decodes to ``UNKNOWN``,
+    which callers treat as a failed probe. That is the right default: the only
+    field the message has is ``status`` at field 1, and proto3 serializes fields
+    in ascending field-number order, so a payload that doesn't start with tag 1
+    is not a response we understand.
+    """
+    if not payload or payload[0] != _STATUS_TAG:
+        return ServingStatus.UNKNOWN
+    value = 0
+    shift = 0
+    for byte in payload[1:]:
+        value |= (byte & 0x7F) << shift
+        if not byte & 0x80:
+            break
+        shift += 7
+    else:
+        # Ran off the end of a continued varint: truncated message.
+        return ServingStatus.UNKNOWN
+    try:
+        return ServingStatus(value)
+    except ValueError:
+        # A status this client predates. Not SERVING as far as we know.
+        return ServingStatus.UNKNOWN

packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_client.py

@@ -21,6 +21,12 @@
 
 from __future__ import annotations
 
+from google.cloud.bigtable.data._accelerator._health import (
+    HEALTH_CHECK_METHOD,
+    ServingStatus,
+    parse_health_response,
+    serialize_health_request,
+)
 from google.cloud.bigtable.data._cross_sync import CrossSync
 from google.cloud.bigtable_v2.types import (
     MutateRowRequest,
@@ -39,6 +45,28 @@
 _MUTATE_ROW_METHOD = "/google.bigtable.v2.Bigtable/MutateRow"
 _READ_ROWS_METHOD = "/google.bigtable.v2.Bigtable/ReadRows"
 
+# Inbound message cap on the UDS channel. gRPC defaults this to 4 MiB, a limit
+# meant for a channel that crosses a network. This one doesn't leave the
+# machine: the daemon has already paid the cost of pulling the response off the
+# wire, and all the cap does on this hop is turn an otherwise-fine read into
+# RESOURCE_EXHAUSTED. Raise it well past anything the service itself produces so
+# the UDS leg is never the narrower of the two.
+_MAX_RECEIVE_MESSAGE_LENGTH = 256 * 1024 * 1024
+
+_CHANNEL_OPTIONS = (("grpc.max_receive_message_length", _MAX_RECEIVE_MESSAGE_LENGTH),)
+
+
+def _uds_target(uds_path: str) -> str:
+    """Build the gRPC target for the daemon's socket at ``uds_path``.
+
+    Must be the single-colon form. ``unix://`` opens an *authority* component,
+    which gRPC only accepts when the rest is an absolute POSIX path; hand it a
+    Windows path and ``C:`` parses as the authority, so channel construction
+    fails outright with "authority-based URIs not supported by the unix
+    scheme". ``unix:`` takes the remainder as a plain path on every platform.
+    """
+    return f"unix:{uds_path}"
+
 
 @CrossSync.convert_class(sync_name="_AcceleratorClient")
 class _AsyncAcceleratorClient:
@@ -49,9 +77,12 @@ class _AsyncAcceleratorClient:
     (``_accelerator/_routing.py``) decides which calls reach this object.
     """
 
-    def __init__(self, uds_path: str):
+    def __init__(self, uds_path: str, auth_secret: str):
         self._uds_path = uds_path
-        self._channel = insecure_channel(f"unix://{uds_path}")
+        self._metadata = (("x-accelerator-token", auth_secret),)
+        self._channel = insecure_channel(
+            _uds_target(uds_path), options=_CHANNEL_OPTIONS
+        )
         self._mutate_row_stub = self._channel.unary_unary(
             _MUTATE_ROW_METHOD,
             request_serializer=MutateRowRequest.serialize,
@@ -62,6 +93,13 @@ def __init__(self, uds_path: str):
             request_serializer=ReadRowsRequest.serialize,
             response_deserializer=ReadRowsResponse.deserialize,
         )
+        # Hand-rolled codec instead of the generated health stubs, which live in
+        # a distribution we don't depend on. See ``_accelerator/_health.py``.
+        self._health_stub = self._channel.unary_unary(
+            HEALTH_CHECK_METHOD,
+            request_serializer=serialize_health_request,
+            response_deserializer=parse_health_response,
+        )
 
     @property
     def uds_path(self) -> str:
@@ -71,7 +109,9 @@ def uds_path(self) -> str:
     async def mutate_row(
         self, request: MutateRowRequest, *, timeout: float | None = None
     ) -> MutateRowResponse:
-        return await self._mutate_row_stub(request, timeout=timeout)
+        return await self._mutate_row_stub(
+            request, timeout=timeout, metadata=self._metadata
+        )
 
     @CrossSync.convert
     async def read_rows(
@@ -84,7 +124,18 @@ async def read_rows(
         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)
+        return self._read_rows_stub(request, timeout=timeout, metadata=self._metadata)
+
+    @CrossSync.convert
+    async def check_health(self, *, timeout: float | None = None) -> ServingStatus:
+        """Probe the daemon's overall serving status.
+
+        The daemon answers this out of its own process without touching a
+        session, the Channel, or the network, so how long it takes is almost
+        purely a measure of how contended that process is. That makes the
+        ``timeout`` the real signal here, more than the returned status.
+        """
+        return await self._health_stub(None, timeout=timeout, metadata=self._metadata)
 
     @CrossSync.convert
     async def close(self) -> None: