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 b69ddc304d1e..b02d9d5b236c 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 @@ -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 @@ -56,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. diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_health.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_health.py new file mode 100644 index 000000000000..33159a0d2a12 --- /dev/null +++ b/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 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 index b1ab1962695f..7b9b94993535 100644 --- 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 @@ -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, @@ -63,6 +69,13 @@ def __init__(self, uds_path: str, auth_secret: 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: @@ -89,6 +102,17 @@ async def read_rows( """ 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: if CrossSync.is_async: diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_health.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_health.py new file mode 100644 index 000000000000..44cba32331ba --- /dev/null +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_health.py @@ -0,0 +1,296 @@ +# 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. +# +"""Background health monitor for the accelerator daemon. + +Polls ``grpc.health.v1.Health/Check`` on the daemon's UDS and flips the +reversible "degraded" flag on the Table's :class:`AcceleratorBreaker`, which +routes calls to the native client while it is set. + +**The probe timeout is the actual signal.** The daemon answers a health check +entirely out of its own process — no session, no Channel, no network — so the +call is a few microseconds of work plus however long the daemon's Go scheduler +took to get to it. A handler that cheap missing its deadline means the daemon is +not being scheduled, which is exactly the failure mode this guards against (an +undersized GOMAXPROCS, CFS throttling against a cgroup quota, a long +stop-the-world pause). The daemon deliberately never self-reports NOT_SERVING +under load: it cannot pick that threshold from the inside, so the caller +measures instead. + +**This is a stall detector, not a latency monitor, and the difference matters.** +A measurably starved daemon runs on the order of tens of milliseconds behind on +each wakeup, which is well worth falling back from but sits uncomfortably close +to the noise floor of the measurement: the probe is issued from the caller's own +process, so a busy Python event loop inflates it just as a busy daemon does, and +a false positive there abandons the accelerator for a native path running in the +same busy process. A single-sample deadline cannot separate those two, so the +default is set to catch severe stalls confidently rather than mild degradation +unreliably. Catching the milder regime needs a different instrument — either the +daemon-side ``pacemaker_delays`` metric, which measures Go scheduling delay +directly and is already exported, or folding probe *latencies* into a +distribution here instead of thresholding each one. + +Both the trip and the recovery need several consecutive probes to agree. A +single slow probe is a GC pause, not a starved process, and flapping the route +on every one of those would be worse than either steady state. +""" + +from __future__ import annotations + +import logging +import threading +from typing import Any + +from grpc import StatusCode + +from google.cloud.bigtable.data._accelerator._fallback import ( + AcceleratorBreaker, + _grpc_code, +) +from google.cloud.bigtable.data._accelerator._health import ServingStatus +from google.cloud.bigtable.data._cross_sync import CrossSync + +if CrossSync.is_async: + from google.cloud.bigtable.data._async._accelerator_client import ( + _AsyncAcceleratorClient as AcceleratorClientType, + ) +else: + from google.cloud.bigtable.data._sync_autogen._accelerator_client import ( # noqa: F401 + _AcceleratorClient as AcceleratorClientType, + ) + +__CROSS_SYNC_OUTPUT__ = "google.cloud.bigtable.data._sync_autogen._accelerator_health" + +_LOGGER = logging.getLogger(__name__) + +#: Seconds between probes. Cheap enough to be frequent, sparse enough that the +#: monitor is never a meaningful share of the daemon's load. +DEFAULT_PROBE_INTERVAL = 5.0 + +#: Deadline for a single probe, and thereby the latency threshold that defines +#: "starved". A healthy round trip is well under a millisecond of actual work, +#: so this is enormous headroom -- deliberately. The probe is issued from the +#: caller's own process, so it measures that process's scheduling as well as the +#: daemon's, and on a busy event loop tens of milliseconds is normal for a +#: perfectly healthy daemon. Tightening this much further starts trading real +#: detection for false positives that fall back to a native path running in the +#: same busy process, which helps nobody. See the module docstring for what that +#: means for the failures this can and cannot see. +DEFAULT_PROBE_TIMEOUT = 0.25 + +#: Consecutive probes that must agree before the route changes: bad ones before +#: falling back, good ones before routing back. Three of them at the default +#: interval is roughly 10-15 seconds of sustained evidence in either direction — +#: a range rather than a figure because the loop waits before it probes, so a +#: change in the daemon is never seen sooner than the next scheduled probe. +DEFAULT_UNHEALTHY_THRESHOLD = 3 +DEFAULT_HEALTHY_THRESHOLD = 3 + +#: How long ``close()`` waits for the polling loop to finish. It is already +#: cancelled (async) or signalled (sync) by then, so the only thing this covers +#: is a probe in flight, bounded by the probe timeout. +_CLOSE_WAIT_SECONDS = 1.0 + + +@CrossSync.convert_class(sync_name="_AcceleratorHealthMonitor") +class _AsyncAcceleratorHealthMonitor: + """Polls the daemon's health and drives the breaker's degraded flag. + + One instance per Table, owned by the Table and stopped by ``Table.close``. + """ + + def __init__( + self, + accelerator_client: AcceleratorClientType, + breaker: AcceleratorBreaker, + *, + interval: float = DEFAULT_PROBE_INTERVAL, + timeout: float = DEFAULT_PROBE_TIMEOUT, + unhealthy_threshold: int = DEFAULT_UNHEALTHY_THRESHOLD, + healthy_threshold: int = DEFAULT_HEALTHY_THRESHOLD, + ): + self._client = accelerator_client + self._breaker = breaker + self._interval = interval + self._timeout = timeout + self._unhealthy_threshold = unhealthy_threshold + self._healthy_threshold = healthy_threshold + self._consecutive_failures = 0 + self._consecutive_successes = 0 + # A threading.Event in both variants, deliberately not CrossSync.Event: + # an asyncio.Event may only be set from the loop thread, and this one + # is also set by a weakref finalizer running on whatever thread the + # collection happened on. Nothing ever *awaits* it — the sync loop + # blocks on it, the async loop only reads it between sleeps — so it + # needs no loop affinity. See _run() and request_stop(). + self._closed = threading.Event() + # An asyncio.Task when async; a daemon thread when sync. See start(). + self._task: Any = None + + def start(self) -> None: + """Launch the polling task. Idempotent.""" + if self._task is not None: + return + if CrossSync.is_async: + self._task = CrossSync.create_task( + self._run, task_name="bigtable-accelerator-health" + ) + else: + # A dedicated daemon thread rather than a ThreadPoolExecutor + # worker: the loop runs for the life of the Table, and + # concurrent.futures *joins* its (non-daemon) workers at + # interpreter exit, so a Table that is never closed would hang the + # process on the way out. A dedicated thread also keeps the loop + # out of a pool slot other work is queued behind. + self._task = threading.Thread( + target=self._run, + name="bigtable-accelerator-health", + daemon=True, + ) + self._task.start() + + @property + def is_running(self) -> bool: + """Whether the polling loop is still alive.""" + if self._task is None: + return False + if CrossSync.is_async: + return not self._task.done() + return self._task.is_alive() + + def request_stop(self) -> None: + """Ask the loop to stop, without waiting for it to finish. + + For callers that cannot await ``close()`` — notably a weakref finalizer + on the owning Table, which is what stops a monitor whose Table was + dropped without being closed. Setting a threading.Event is safe from + any thread and during interpreter shutdown. The sync loop wakes on it + immediately; the async loop sees it at its next wakeup, since only + ``close()`` can cancel the task outright. + """ + self._closed.set() + + @CrossSync.convert + async def _run(self) -> None: + """Poll until closed, or until monitoring turns out to be pointless.""" + while not self._closed.is_set(): + if CrossSync.is_async: + # Sleep rather than wait on the event: it is a threading.Event + # (see __init__), so there is nothing here to await. close() + # cancels this task, which interrupts the sleep, and + # request_stop() is picked up by the check below. + await CrossSync.sleep(self._interval) + else: + self._closed.wait(timeout=self._interval) + if self._closed.is_set(): + return + if self._breaker.is_tripped: + # The accelerator has been permanently abandoned for some other + # reason; nothing this monitor observes can change that. + return + healthy = await self._probe() + if healthy is None: + # The daemon does not serve health checks at all. Stop probing + # and leave the breaker undegraded: an unmonitorable daemon is + # not a degraded one, and treating it as one would disable the + # accelerator against every daemon older than this feature. + # + # Clearing is not redundant. An earlier streak of probes can + # have degraded the breaker already — a daemon too starved to + # answer even UNIMPLEMENTED before the deadline looks exactly + # like an unhealthy one — and this return is the last thing + # that runs, so a flag left set here is never cleared by + # anyone. (A permanent trip is handled above and stays.) + self._breaker.set_degraded(False) + _LOGGER.info( + "Accelerator daemon does not implement the gRPC health " + "service; health-based fallback is disabled for this table." + ) + return + self._record(healthy) + + @CrossSync.convert + async def _probe(self) -> bool | None: + """Run one probe. + + Returns True if the daemon answered SERVING within the timeout, False if + it answered late, errored, or reported anything else, and None if it + does not implement the health service (see :meth:`_run`). + """ + # Deliberately silent per probe: a probe fails once every interval for as + # long as the daemon is unwell, so anything logged here is unbounded. The + # transitions in :meth:`_record` are the events worth reporting, and they + # fire once each. + try: + status = await self._client.check_health(timeout=self._timeout) + except Exception as exc: + if _grpc_code(exc) == StatusCode.UNIMPLEMENTED: + return None + return False + return status == ServingStatus.SERVING + + def _record(self, healthy: bool) -> None: + """Fold one probe result into the streak counters and act on it.""" + if healthy: + self._consecutive_failures = 0 + self._consecutive_successes += 1 + if ( + self._breaker.is_degraded + and self._consecutive_successes >= self._healthy_threshold + ): + _LOGGER.info( + "Accelerator daemon is healthy again after %d consecutive " + "successful health checks; resuming accelerated routing.", + self._consecutive_successes, + ) + self._breaker.set_degraded(False) + else: + self._consecutive_successes = 0 + self._consecutive_failures += 1 + if ( + not self._breaker.is_degraded + and self._consecutive_failures >= self._unhealthy_threshold + ): + _LOGGER.warning( + "Accelerator daemon failed %d consecutive health checks " + "(%gs timeout); falling back to the native Bigtable client " + "until it recovers.", + self._consecutive_failures, + self._timeout, + ) + self._breaker.set_degraded(True) + + @CrossSync.convert + async def close(self) -> None: + """Stop polling and wait for the loop to finish. + + Does not clear a degraded flag already set. Waiting is bounded by + ``_CLOSE_WAIT_SECONDS``: the caller is entitled to have the monitor + actually stopped when this returns, because the channel it probes over + is closed immediately afterwards, and a probe still running into that + would log an alarming "falling back" warning during orderly teardown. + """ + self._closed.set() + task, self._task = self._task, None + if task is None: + return + if CrossSync.is_async: + task.cancel() + # Retrieve any exception so asyncio doesn't log "Task exception was + # never retrieved" — as a callback too, in case the wait times out + # and the task completes after we stop looking. + task.add_done_callback(lambda t: None if t.cancelled() else t.exception()) + await CrossSync.wait([task], timeout=_CLOSE_WAIT_SECONDS) + else: + task.join(timeout=_CLOSE_WAIT_SECONDS) 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 8360e1a60ff2..ea9168abae6b 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 @@ -23,6 +23,7 @@ import sys import time import warnings +import weakref from functools import partial from typing import ( TYPE_CHECKING, @@ -127,6 +128,9 @@ from google.cloud.bigtable.data._async._accelerator_client import ( _AsyncAcceleratorClient as AcceleratorClientType, ) + from google.cloud.bigtable.data._async._accelerator_health import ( + _AsyncAcceleratorHealthMonitor as AcceleratorHealthMonitorType, + ) from google.cloud.bigtable.data._async._swappable_channel import ( AsyncSwappableChannel as SwappableChannelType, ) @@ -148,6 +152,9 @@ from google.cloud.bigtable.data._sync_autogen._accelerator_client import ( # noqa: F401 _AcceleratorClient as AcceleratorClientType, ) + from google.cloud.bigtable.data._sync_autogen._accelerator_health import ( # noqa: F401 + _AcceleratorHealthMonitor as AcceleratorHealthMonitorType, + ) from google.cloud.bigtable.data._sync_autogen._swappable_channel import ( # noqa: F401 SwappableChannel as SwappableChannelType, ) @@ -1354,9 +1361,12 @@ def __init__( # 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_health: AcceleratorHealthMonitorType | None = None + # Fallback policy: a daemon that cannot serve a routed RPC + # (UNIMPLEMENTED / dead subprocess) is permanently bypassed in favor of + # the native client, and one that is merely too starved to answer its + # own health check is bypassed until it recovers. See + # ``_accelerator/_fallback.py`` and ``_async/_accelerator_health.py``. self._accelerator_breaker = AcceleratorBreaker() if self.client._use_accelerator is not False: # None (default) or True: attempt to start. ``True`` is an explicit @@ -1462,6 +1472,29 @@ def _start_accelerator(self) -> None: except BaseException: server.close() raise + try: + self._accelerator_health = AcceleratorHealthMonitorType( + self._accelerator_client, self._accelerator_breaker + ) + self._accelerator_health.start() + # close() is the orderly stop, but a Table dropped without it would + # otherwise leave the loop probing for the life of the event loop, + # holding the monitor and through it the accelerator channel. The + # finalizer is bound to the monitor, not to self, so it does not + # keep this Table alive; it only flips a flag the loop reads. + weakref.finalize(self, self._accelerator_health.request_stop) + except Exception as exc: + # Monitoring is an optimization on top of a working accelerator, so + # failing to start it must not take the accelerator down with it — + # and must not leave this method raising after the daemon has been + # adopted, which would strand the subprocess. + warnings.warn( + "Could not start the accelerator health monitor; the " + f"accelerator will run without health-based fallback: {exc}", + RuntimeWarning, + stacklevel=2, + ) + self._accelerator_health = None def _use_accelerator(self, method_name: str) -> bool: """Whether this call should be routed through the accelerator daemon. @@ -2345,6 +2378,16 @@ async def close(self): """ Called to close the Table instance and release any resources held by it. """ + # Accelerator teardown goes first, so that nothing that can raise runs + # ahead of it and strands the daemon subprocess. Within it, the monitor + # stops before the channel it probes over goes away. + if self._accelerator_health is not None: + try: + await self._accelerator_health.close() + except Exception as exc: + _LOGGER.warning("Failed to close accelerator health monitor: %s", exc) + finally: + self._accelerator_health = None if self._accelerator_client is not None: try: await self._accelerator_client.close() 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 index 239904956dcb..bd0965843620 100644 --- 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 @@ -26,6 +26,12 @@ from grpc import insecure_channel +from google.cloud.bigtable.data._accelerator._health import ( + HEALTH_CHECK_METHOD, + ServingStatus, + parse_health_response, + serialize_health_request, +) from google.cloud.bigtable_v2.types import ( MutateRowRequest, MutateRowResponse, @@ -59,6 +65,11 @@ def __init__(self, uds_path: str, auth_secret: str): request_serializer=ReadRowsRequest.serialize, response_deserializer=ReadRowsResponse.deserialize, ) + 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: @@ -78,5 +89,14 @@ def read_rows(self, request: ReadRowsRequest, *, timeout: float | None = None): chunk-merging machinery in ``_read_rows.py`` works unchanged.""" return self._read_rows_stub(request, timeout=timeout, metadata=self._metadata) + 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 self._health_stub(None, timeout=timeout, metadata=self._metadata) + def close(self) -> None: self._channel.close() diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_accelerator_health.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_accelerator_health.py new file mode 100644 index 000000000000..1306bc333002 --- /dev/null +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_accelerator_health.py @@ -0,0 +1,202 @@ +# 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. + +"""Background health monitor for the accelerator daemon. + +Polls ``grpc.health.v1.Health/Check`` on the daemon's UDS and flips the +reversible "degraded" flag on the Table's :class:`AcceleratorBreaker`, which +routes calls to the native client while it is set. + +**The probe timeout is the actual signal.** The daemon answers a health check +entirely out of its own process — no session, no Channel, no network — so the +call is a few microseconds of work plus however long the daemon's Go scheduler +took to get to it. A handler that cheap missing its deadline means the daemon is +not being scheduled, which is exactly the failure mode this guards against (an +undersized GOMAXPROCS, CFS throttling against a cgroup quota, a long +stop-the-world pause). The daemon deliberately never self-reports NOT_SERVING +under load: it cannot pick that threshold from the inside, so the caller +measures instead. + +**This is a stall detector, not a latency monitor, and the difference matters.** +A measurably starved daemon runs on the order of tens of milliseconds behind on +each wakeup, which is well worth falling back from but sits uncomfortably close +to the noise floor of the measurement: the probe is issued from the caller's own +process, so a busy Python event loop inflates it just as a busy daemon does, and +a false positive there abandons the accelerator for a native path running in the +same busy process. A single-sample deadline cannot separate those two, so the +default is set to catch severe stalls confidently rather than mild degradation +unreliably. Catching the milder regime needs a different instrument — either the +daemon-side ``pacemaker_delays`` metric, which measures Go scheduling delay +directly and is already exported, or folding probe *latencies* into a +distribution here instead of thresholding each one. + +Both the trip and the recovery need several consecutive probes to agree. A +single slow probe is a GC pause, not a starved process, and flapping the route +on every one of those would be worse than either steady state. +""" + +from __future__ import annotations + +import logging +import threading +from typing import Any + +from grpc import StatusCode + +from google.cloud.bigtable.data._accelerator._fallback import ( + AcceleratorBreaker, + _grpc_code, +) +from google.cloud.bigtable.data._accelerator._health import ServingStatus +from google.cloud.bigtable.data._sync_autogen._accelerator_client import ( + _AcceleratorClient as AcceleratorClientType, +) + +_LOGGER = logging.getLogger(__name__) +DEFAULT_PROBE_INTERVAL = 5.0 +DEFAULT_PROBE_TIMEOUT = 0.25 +DEFAULT_UNHEALTHY_THRESHOLD = 3 +DEFAULT_HEALTHY_THRESHOLD = 3 +_CLOSE_WAIT_SECONDS = 1.0 + + +class _AcceleratorHealthMonitor: + """Polls the daemon's health and drives the breaker's degraded flag. + + One instance per Table, owned by the Table and stopped by ``Table.close``. + """ + + def __init__( + self, + accelerator_client: AcceleratorClientType, + breaker: AcceleratorBreaker, + *, + interval: float = DEFAULT_PROBE_INTERVAL, + timeout: float = DEFAULT_PROBE_TIMEOUT, + unhealthy_threshold: int = DEFAULT_UNHEALTHY_THRESHOLD, + healthy_threshold: int = DEFAULT_HEALTHY_THRESHOLD, + ): + self._client = accelerator_client + self._breaker = breaker + self._interval = interval + self._timeout = timeout + self._unhealthy_threshold = unhealthy_threshold + self._healthy_threshold = healthy_threshold + self._consecutive_failures = 0 + self._consecutive_successes = 0 + self._closed = threading.Event() + self._task: Any = None + + def start(self) -> None: + """Launch the polling task. Idempotent.""" + if self._task is not None: + return + self._task = threading.Thread( + target=self._run, name="bigtable-accelerator-health", daemon=True + ) + self._task.start() + + @property + def is_running(self) -> bool: + """Whether the polling loop is still alive.""" + if self._task is None: + return False + return self._task.is_alive() + + def request_stop(self) -> None: + """Ask the loop to stop, without waiting for it to finish. + + For callers that cannot await ``close()`` — notably a weakref finalizer + on the owning Table, which is what stops a monitor whose Table was + dropped without being closed. Setting a threading.Event is safe from + any thread and during interpreter shutdown. The sync loop wakes on it + immediately; the async loop sees it at its next wakeup, since only + ``close()`` can cancel the task outright.""" + self._closed.set() + + def _run(self) -> None: + """Poll until closed, or until monitoring turns out to be pointless.""" + while not self._closed.is_set(): + self._closed.wait(timeout=self._interval) + if self._closed.is_set(): + return + if self._breaker.is_tripped: + return + healthy = self._probe() + if healthy is None: + self._breaker.set_degraded(False) + _LOGGER.info( + "Accelerator daemon does not implement the gRPC health service; health-based fallback is disabled for this table." + ) + return + self._record(healthy) + + def _probe(self) -> bool | None: + """Run one probe. + + Returns True if the daemon answered SERVING within the timeout, False if + it answered late, errored, or reported anything else, and None if it + does not implement the health service (see :meth:`_run`).""" + try: + status = self._client.check_health(timeout=self._timeout) + except Exception as exc: + if _grpc_code(exc) == StatusCode.UNIMPLEMENTED: + return None + return False + return status == ServingStatus.SERVING + + def _record(self, healthy: bool) -> None: + """Fold one probe result into the streak counters and act on it.""" + if healthy: + self._consecutive_failures = 0 + self._consecutive_successes += 1 + if ( + self._breaker.is_degraded + and self._consecutive_successes >= self._healthy_threshold + ): + _LOGGER.info( + "Accelerator daemon is healthy again after %d consecutive successful health checks; resuming accelerated routing.", + self._consecutive_successes, + ) + self._breaker.set_degraded(False) + else: + self._consecutive_successes = 0 + self._consecutive_failures += 1 + if ( + not self._breaker.is_degraded + and self._consecutive_failures >= self._unhealthy_threshold + ): + _LOGGER.warning( + "Accelerator daemon failed %d consecutive health checks (%gs timeout); falling back to the native Bigtable client until it recovers.", + self._consecutive_failures, + self._timeout, + ) + self._breaker.set_degraded(True) + + def close(self) -> None: + """Stop polling and wait for the loop to finish. + + Does not clear a degraded flag already set. Waiting is bounded by + ``_CLOSE_WAIT_SECONDS``: the caller is entitled to have the monitor + actually stopped when this returns, because the channel it probes over + is closed immediately afterwards, and a probe still running into that + would log an alarming "falling back" warning during orderly teardown.""" + self._closed.set() + task, self._task = (self._task, None) + if task is None: + return + task.join(timeout=_CLOSE_WAIT_SECONDS) 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 e85d69b9f873..67a8b3f12175 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 @@ -26,6 +26,7 @@ import sys import time import warnings +import weakref from functools import partial from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional, Sequence, Set, cast @@ -82,6 +83,9 @@ from google.cloud.bigtable.data._sync_autogen._accelerator_client import ( _AcceleratorClient as AcceleratorClientType, ) +from google.cloud.bigtable.data._sync_autogen._accelerator_health import ( + _AcceleratorHealthMonitor as AcceleratorHealthMonitorType, +) from google.cloud.bigtable.data._sync_autogen._swappable_channel import ( SwappableChannel as SwappableChannelType, ) @@ -1060,6 +1064,7 @@ def __init__( ) from e self._accelerator_daemon: AcceleratorDaemon | None = None self._accelerator_client: AcceleratorClientType | None = None + self._accelerator_health: AcceleratorHealthMonitorType | 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) @@ -1138,6 +1143,19 @@ def _start_accelerator(self) -> None: except BaseException: server.close() raise + try: + self._accelerator_health = AcceleratorHealthMonitorType( + self._accelerator_client, self._accelerator_breaker + ) + self._accelerator_health.start() + weakref.finalize(self, self._accelerator_health.request_stop) + except Exception as exc: + warnings.warn( + f"Could not start the accelerator health monitor; the accelerator will run without health-based fallback: {exc}", + RuntimeWarning, + stacklevel=2, + ) + self._accelerator_health = None def _use_accelerator(self, method_name: str) -> bool: """Whether this call should be routed through the accelerator daemon. @@ -1923,6 +1941,13 @@ def read_modify_write_row( def close(self): """Called to close the Table instance and release any resources held by it.""" + if self._accelerator_health is not None: + try: + self._accelerator_health.close() + except Exception as exc: + _LOGGER.warning("Failed to close accelerator health monitor: %s", exc) + finally: + self._accelerator_health = None if self._accelerator_client is not None: try: self._accelerator_client.close() diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test__accelerator_health.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test__accelerator_health.py new file mode 100644 index 000000000000..458f80857ffa --- /dev/null +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test__accelerator_health.py @@ -0,0 +1,394 @@ +# 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. + +"""Probe and polling-loop tests for the accelerator health monitor. + +These run against a real gRPC server on a real Unix domain socket, serving the +real ``/grpc.health.v1.Health/Check`` method name, reached through the real +``_AcceleratorClient`` and its hand-rolled codec. Nothing about the transport is +faked, so the slow-daemon case below genuinely exercises a gRPC deadline rather +than a raised sentinel — which matters, because a probe timeout is the whole +signal this monitor is built on. + +The server replies with literal ``HealthCheckResponse`` bytes; that those +literals are what protobuf itself produces is established in +``tests/unit/data/test_accelerator_health.py``. +""" + +import threading +import time +from concurrent import futures + +import grpc +import pytest + +from google.cloud.bigtable.data._accelerator._fallback import AcceleratorBreaker +from google.cloud.bigtable.data._accelerator._health import ( + HEALTH_CHECK_METHOD, + ServingStatus, +) +from google.cloud.bigtable.data._cross_sync import CrossSync + +if CrossSync.is_async: + from google.cloud.bigtable.data._async._accelerator_client import ( + _AsyncAcceleratorClient as ClientType, + ) + from google.cloud.bigtable.data._async._accelerator_health import ( + _AsyncAcceleratorHealthMonitor as MonitorType, + ) +else: + from google.cloud.bigtable.data._sync_autogen._accelerator_client import ( + _AcceleratorClient as ClientType, + ) + from google.cloud.bigtable.data._sync_autogen._accelerator_health import ( + _AcceleratorHealthMonitor as MonitorType, + ) + +__CROSS_SYNC_OUTPUT__ = "tests.unit.data._sync_autogen.test__accelerator_health" + +_AUTH_SECRET = "test-secret" + +# HealthCheckResponse{status: S}: field 1, varint. See the codec test. +_WIRE_BY_STATUS = { + ServingStatus.UNKNOWN: b"", + ServingStatus.SERVING: b"\x08\x01", + ServingStatus.NOT_SERVING: b"\x08\x02", + ServingStatus.SERVICE_UNKNOWN: b"\x08\x03", +} + + +class _FakeDaemonHealthServer: + """A real gRPC server standing in for the daemon's health service. + + Mirrors the daemon in the two ways that matter to the monitor: it serves the + health method locally on the UDS, and (like the daemon's auth interceptor, + which runs before the ``isProxied`` gate) it sees the accelerator token on + every request. ``serve_health=False`` registers nothing, so gRPC itself + answers ``UNIMPLEMENTED`` — the shape an older daemon presents. Setting + ``abort_unimplemented`` later produces the same status from a server that + was answering normally, which is how a daemon that was merely too starved + to reply in time looks once it recovers. + """ + + def __init__(self, uds_path, *, status=ServingStatus.SERVING, serve_health=True): + self.status = status + self.abort_unimplemented = False + self.delay = 0.0 + self.request_count = 0 + self.last_metadata = None + self._lock = threading.Lock() + self._server = grpc.server(futures.ThreadPoolExecutor(max_workers=4)) + if serve_health: + handler = grpc.unary_unary_rpc_method_handler( + self._check, + request_deserializer=lambda payload: payload, + response_serializer=lambda payload: payload, + ) + service, method = HEALTH_CHECK_METHOD.lstrip("/").split("/") + self._server.add_generic_rpc_handlers( + (grpc.method_handlers_generic_handler(service, {method: handler}),) + ) + self._server.add_insecure_port(f"unix://{uds_path}") + self._server.start() + + def _check(self, request, context): + with self._lock: + self.request_count += 1 + self.last_metadata = dict(context.invocation_metadata()) + delay, status = self.delay, self.status + unimplemented = self.abort_unimplemented + if delay: + time.sleep(delay) + if unimplemented: + context.abort(grpc.StatusCode.UNIMPLEMENTED, "no health service") + return _WIRE_BY_STATUS[status] + + def stop(self): + self._server.stop(None) + + +@CrossSync.convert_class(sync_name="TestAcceleratorHealthMonitor") +class TestAsyncAcceleratorHealthMonitor: + @staticmethod + def _start_server(tmp_path, **kwargs): + return _FakeDaemonHealthServer(str(tmp_path / "bt_proxy.sock"), **kwargs) + + @staticmethod + @CrossSync.convert + def _make_client(tmp_path): + return ClientType(str(tmp_path / "bt_proxy.sock"), _AUTH_SECRET) + + @staticmethod + def _make_monitor(client, breaker, **kwargs): + kwargs.setdefault("interval", 0.02) + kwargs.setdefault("timeout", 5.0) + kwargs.setdefault("unhealthy_threshold", 2) + kwargs.setdefault("healthy_threshold", 2) + return MonitorType(client, breaker, **kwargs) + + @staticmethod + @CrossSync.convert + async def _wait_for(predicate, timeout=15.0): + """Poll ``predicate`` until true or the deadline passes.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + await CrossSync.sleep(0.01) + return predicate() + + @CrossSync.pytest + @pytest.mark.parametrize( + "status,expected", + [ + (ServingStatus.SERVING, True), + (ServingStatus.NOT_SERVING, False), + (ServingStatus.UNKNOWN, False), + (ServingStatus.SERVICE_UNKNOWN, False), + ], + ) + async def test_probe_reads_status_off_the_wire(self, tmp_path, status, expected): + server = self._start_server(tmp_path, status=status) + client = self._make_client(tmp_path) + try: + monitor = self._make_monitor(client, AcceleratorBreaker()) + assert await monitor._probe() is expected + finally: + await client.close() + server.stop() + + @CrossSync.pytest + async def test_probe_sends_the_accelerator_token(self, tmp_path): + # The daemon's auth interceptor is chained onto the whole server and + # runs ahead of the Bigtable-only proxy gate, so an unauthenticated + # probe would be rejected and every daemon would look degraded. + server = self._start_server(tmp_path) + client = self._make_client(tmp_path) + try: + monitor = self._make_monitor(client, AcceleratorBreaker()) + assert await monitor._probe() is True + assert server.last_metadata["x-accelerator-token"] == _AUTH_SECRET + finally: + await client.close() + server.stop() + + @CrossSync.pytest + async def test_probe_times_out_on_a_starved_daemon(self, tmp_path): + # The case the whole feature exists for: the daemon is up and would + # answer SERVING, but cannot get scheduled in time to say so. + server = self._start_server(tmp_path) + server.delay = 1.0 + client = self._make_client(tmp_path) + try: + monitor = self._make_monitor(client, AcceleratorBreaker(), timeout=0.05) + started = time.monotonic() + assert await monitor._probe() is False + # Failed on the deadline, not by waiting out the slow handler. + assert time.monotonic() - started < 0.9 + finally: + await client.close() + server.stop() + + @CrossSync.pytest + async def test_probe_reports_unimplemented_distinctly(self, tmp_path): + server = self._start_server(tmp_path, serve_health=False) + client = self._make_client(tmp_path) + try: + monitor = self._make_monitor(client, AcceleratorBreaker()) + # None, not False: an unmonitorable daemon is not a degraded one. + assert await monitor._probe() is None + finally: + await client.close() + server.stop() + + @CrossSync.pytest + async def test_probe_fails_when_nothing_is_listening(self, tmp_path): + client = self._make_client(tmp_path) + try: + monitor = self._make_monitor(client, AcceleratorBreaker(), timeout=1.0) + assert await monitor._probe() is False + finally: + await client.close() + + @CrossSync.pytest + async def test_loop_degrades_then_recovers(self, tmp_path): + server = self._start_server(tmp_path) + client = self._make_client(tmp_path) + breaker = AcceleratorBreaker() + monitor = self._make_monitor(client, breaker) + try: + monitor.start() + assert await self._wait_for(lambda: server.request_count >= 2) + assert not breaker.bypass() + + server.status = ServingStatus.NOT_SERVING + assert await self._wait_for(lambda: breaker.is_degraded) + assert breaker.bypass() + # Reversible, and never a permanent trip. + assert not breaker.is_tripped + + server.status = ServingStatus.SERVING + assert await self._wait_for(lambda: not breaker.is_degraded) + assert not breaker.bypass() + finally: + await monitor.close() + await client.close() + server.stop() + + @CrossSync.pytest + async def test_loop_stops_against_a_daemon_without_health(self, tmp_path): + # An older daemon must keep serving accelerated traffic, not be + # abandoned because it can't be monitored. + server = self._start_server(tmp_path, serve_health=False) + client = self._make_client(tmp_path) + breaker = AcceleratorBreaker() + monitor = self._make_monitor(client, breaker, unhealthy_threshold=1) + try: + monitor.start() + # The loop exits outright rather than backing off, and — with an + # unhealthy_threshold of 1, which would otherwise degrade on the + # very first probe — leaves the breaker alone. + assert await self._wait_for(lambda: not monitor.is_running) + assert not breaker.bypass() + finally: + await monitor.close() + await client.close() + server.stop() + + @CrossSync.pytest + async def test_unmonitorable_daemon_clears_an_earlier_degraded_flag(self, tmp_path): + # A daemon starved enough to miss the probe deadline degrades the + # breaker first; if it then recovers enough to answer UNIMPLEMENTED, + # the loop exits on the "can't be monitored" path. That return is the + # last thing the monitor ever does, so a degraded flag left set there + # would route this Table to the native client for the rest of its life + # against a daemon that is actually fine. + server = self._start_server(tmp_path, status=ServingStatus.NOT_SERVING) + client = self._make_client(tmp_path) + breaker = AcceleratorBreaker() + monitor = self._make_monitor(client, breaker, unhealthy_threshold=1) + try: + monitor.start() + assert await self._wait_for(lambda: breaker.is_degraded) + server.abort_unimplemented = True + assert await self._wait_for(lambda: not monitor.is_running) + assert not breaker.is_degraded + assert not breaker.bypass() + finally: + await monitor.close() + await client.close() + server.stop() + + @CrossSync.pytest + async def test_close_waits_for_the_loop_to_finish(self, tmp_path): + # Table.close() closes the accelerator channel immediately after this + # returns, so a probe still in flight would fail against a half-torn- + # down client and log a "falling back" warning during orderly teardown. + server = self._start_server(tmp_path) + client = self._make_client(tmp_path) + monitor = self._make_monitor(client, AcceleratorBreaker()) + try: + monitor.start() + assert await self._wait_for(lambda: server.request_count >= 1) + # Held across close(), which clears the attribute — asserting on + # monitor.is_running alone would pass even if nothing stopped. + handle = monitor._task + await monitor.close() + if CrossSync.is_async: + assert handle.done() + else: + assert not handle.is_alive() + finally: + await client.close() + server.stop() + + @CrossSync.pytest + async def test_request_stop_ends_the_loop_without_close(self, tmp_path): + # What the weakref finalizer on a dropped Table calls: no await, no + # lock, just a flag the loop reads at its next wakeup. + server = self._start_server(tmp_path) + client = self._make_client(tmp_path) + monitor = self._make_monitor(client, AcceleratorBreaker()) + try: + monitor.start() + assert await self._wait_for(lambda: server.request_count >= 1) + monitor.request_stop() + assert await self._wait_for(lambda: not monitor.is_running) + finally: + await monitor.close() + await client.close() + server.stop() + + @CrossSync.pytest + async def test_loop_stops_once_the_breaker_is_tripped(self, tmp_path): + server = self._start_server(tmp_path) + client = self._make_client(tmp_path) + breaker = AcceleratorBreaker() + monitor = self._make_monitor(client, breaker) + try: + monitor.start() + assert await self._wait_for(lambda: server.request_count >= 1) + breaker.trip() + assert await self._wait_for(lambda: not monitor.is_running) + finally: + await monitor.close() + await client.close() + server.stop() + + @CrossSync.pytest + async def test_close_stops_polling(self, tmp_path): + server = self._start_server(tmp_path) + client = self._make_client(tmp_path) + monitor = self._make_monitor(client, AcceleratorBreaker()) + try: + monitor.start() + assert await self._wait_for(lambda: server.request_count >= 1) + await monitor.close() + await CrossSync.sleep(0.2) + counted = server.request_count + await CrossSync.sleep(0.2) + assert server.request_count == counted + finally: + await client.close() + server.stop() + + @CrossSync.pytest + async def test_close_leaves_a_degraded_flag_set(self, tmp_path): + # Closing the monitor is a teardown step, not a verdict that the daemon + # got better; it must not silently re-enable accelerated routing. + server = self._start_server(tmp_path, status=ServingStatus.NOT_SERVING) + client = self._make_client(tmp_path) + breaker = AcceleratorBreaker() + monitor = self._make_monitor(client, breaker) + try: + monitor.start() + assert await self._wait_for(lambda: breaker.is_degraded) + await monitor.close() + assert breaker.bypass() + finally: + await client.close() + server.stop() + + @CrossSync.pytest + async def test_start_is_idempotent(self): + # A long interval keeps the task parked in its first wait, so no probe + # is attempted and the client is never touched. + monitor = self._make_monitor(None, AcceleratorBreaker(), interval=60.0) + try: + monitor.start() + first = monitor._task + monitor.start() + assert monitor._task is first + finally: + await monitor.close() diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__accelerator_health.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__accelerator_health.py new file mode 100644 index 000000000000..e2619cb9d523 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__accelerator_health.py @@ -0,0 +1,332 @@ +# 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. + +"""Probe and polling-loop tests for the accelerator health monitor. + +These run against a real gRPC server on a real Unix domain socket, serving the +real ``/grpc.health.v1.Health/Check`` method name, reached through the real +``_AcceleratorClient`` and its hand-rolled codec. Nothing about the transport is +faked, so the slow-daemon case below genuinely exercises a gRPC deadline rather +than a raised sentinel — which matters, because a probe timeout is the whole +signal this monitor is built on. + +The server replies with literal ``HealthCheckResponse`` bytes; that those +literals are what protobuf itself produces is established in +``tests/unit/data/test_accelerator_health.py``. +""" + +import threading +import time +from concurrent import futures + +import grpc +import pytest + +from google.cloud.bigtable.data._accelerator._fallback import AcceleratorBreaker +from google.cloud.bigtable.data._accelerator._health import ( + HEALTH_CHECK_METHOD, + ServingStatus, +) +from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable.data._sync_autogen._accelerator_client import ( + _AcceleratorClient as ClientType, +) +from google.cloud.bigtable.data._sync_autogen._accelerator_health import ( + _AcceleratorHealthMonitor as MonitorType, +) + +_AUTH_SECRET = "test-secret" +_WIRE_BY_STATUS = { + ServingStatus.UNKNOWN: b"", + ServingStatus.SERVING: b"\x08\x01", + ServingStatus.NOT_SERVING: b"\x08\x02", + ServingStatus.SERVICE_UNKNOWN: b"\x08\x03", +} + + +class _FakeDaemonHealthServer: + """A real gRPC server standing in for the daemon's health service. + + Mirrors the daemon in the two ways that matter to the monitor: it serves the + health method locally on the UDS, and (like the daemon's auth interceptor, + which runs before the ``isProxied`` gate) it sees the accelerator token on + every request. ``serve_health=False`` registers nothing, so gRPC itself + answers ``UNIMPLEMENTED`` — the shape an older daemon presents. Setting + ``abort_unimplemented`` later produces the same status from a server that + was answering normally, which is how a daemon that was merely too starved + to reply in time looks once it recovers. + """ + + def __init__(self, uds_path, *, status=ServingStatus.SERVING, serve_health=True): + self.status = status + self.abort_unimplemented = False + self.delay = 0.0 + self.request_count = 0 + self.last_metadata = None + self._lock = threading.Lock() + self._server = grpc.server(futures.ThreadPoolExecutor(max_workers=4)) + if serve_health: + handler = grpc.unary_unary_rpc_method_handler( + self._check, + request_deserializer=lambda payload: payload, + response_serializer=lambda payload: payload, + ) + service, method = HEALTH_CHECK_METHOD.lstrip("/").split("/") + self._server.add_generic_rpc_handlers( + (grpc.method_handlers_generic_handler(service, {method: handler}),) + ) + self._server.add_insecure_port(f"unix://{uds_path}") + self._server.start() + + def _check(self, request, context): + with self._lock: + self.request_count += 1 + self.last_metadata = dict(context.invocation_metadata()) + delay, status = (self.delay, self.status) + unimplemented = self.abort_unimplemented + if delay: + time.sleep(delay) + if unimplemented: + context.abort(grpc.StatusCode.UNIMPLEMENTED, "no health service") + return _WIRE_BY_STATUS[status] + + def stop(self): + self._server.stop(None) + + +class TestAcceleratorHealthMonitor: + @staticmethod + def _start_server(tmp_path, **kwargs): + return _FakeDaemonHealthServer(str(tmp_path / "bt_proxy.sock"), **kwargs) + + @staticmethod + def _make_client(tmp_path): + return ClientType(str(tmp_path / "bt_proxy.sock"), _AUTH_SECRET) + + @staticmethod + def _make_monitor(client, breaker, **kwargs): + kwargs.setdefault("interval", 0.02) + kwargs.setdefault("timeout", 5.0) + kwargs.setdefault("unhealthy_threshold", 2) + kwargs.setdefault("healthy_threshold", 2) + return MonitorType(client, breaker, **kwargs) + + @staticmethod + def _wait_for(predicate, timeout=15.0): + """Poll ``predicate`` until true or the deadline passes.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + CrossSync._Sync_Impl.sleep(0.01) + return predicate() + + @pytest.mark.parametrize( + "status,expected", + [ + (ServingStatus.SERVING, True), + (ServingStatus.NOT_SERVING, False), + (ServingStatus.UNKNOWN, False), + (ServingStatus.SERVICE_UNKNOWN, False), + ], + ) + def test_probe_reads_status_off_the_wire(self, tmp_path, status, expected): + server = self._start_server(tmp_path, status=status) + client = self._make_client(tmp_path) + try: + monitor = self._make_monitor(client, AcceleratorBreaker()) + assert monitor._probe() is expected + finally: + client.close() + server.stop() + + def test_probe_sends_the_accelerator_token(self, tmp_path): + server = self._start_server(tmp_path) + client = self._make_client(tmp_path) + try: + monitor = self._make_monitor(client, AcceleratorBreaker()) + assert monitor._probe() is True + assert server.last_metadata["x-accelerator-token"] == _AUTH_SECRET + finally: + client.close() + server.stop() + + def test_probe_times_out_on_a_starved_daemon(self, tmp_path): + server = self._start_server(tmp_path) + server.delay = 1.0 + client = self._make_client(tmp_path) + try: + monitor = self._make_monitor(client, AcceleratorBreaker(), timeout=0.05) + started = time.monotonic() + assert monitor._probe() is False + assert time.monotonic() - started < 0.9 + finally: + client.close() + server.stop() + + def test_probe_reports_unimplemented_distinctly(self, tmp_path): + server = self._start_server(tmp_path, serve_health=False) + client = self._make_client(tmp_path) + try: + monitor = self._make_monitor(client, AcceleratorBreaker()) + assert monitor._probe() is None + finally: + client.close() + server.stop() + + def test_probe_fails_when_nothing_is_listening(self, tmp_path): + client = self._make_client(tmp_path) + try: + monitor = self._make_monitor(client, AcceleratorBreaker(), timeout=1.0) + assert monitor._probe() is False + finally: + client.close() + + def test_loop_degrades_then_recovers(self, tmp_path): + server = self._start_server(tmp_path) + client = self._make_client(tmp_path) + breaker = AcceleratorBreaker() + monitor = self._make_monitor(client, breaker) + try: + monitor.start() + assert self._wait_for(lambda: server.request_count >= 2) + assert not breaker.bypass() + server.status = ServingStatus.NOT_SERVING + assert self._wait_for(lambda: breaker.is_degraded) + assert breaker.bypass() + assert not breaker.is_tripped + server.status = ServingStatus.SERVING + assert self._wait_for(lambda: not breaker.is_degraded) + assert not breaker.bypass() + finally: + monitor.close() + client.close() + server.stop() + + def test_loop_stops_against_a_daemon_without_health(self, tmp_path): + server = self._start_server(tmp_path, serve_health=False) + client = self._make_client(tmp_path) + breaker = AcceleratorBreaker() + monitor = self._make_monitor(client, breaker, unhealthy_threshold=1) + try: + monitor.start() + assert self._wait_for(lambda: not monitor.is_running) + assert not breaker.bypass() + finally: + monitor.close() + client.close() + server.stop() + + def test_unmonitorable_daemon_clears_an_earlier_degraded_flag(self, tmp_path): + server = self._start_server(tmp_path, status=ServingStatus.NOT_SERVING) + client = self._make_client(tmp_path) + breaker = AcceleratorBreaker() + monitor = self._make_monitor(client, breaker, unhealthy_threshold=1) + try: + monitor.start() + assert self._wait_for(lambda: breaker.is_degraded) + server.abort_unimplemented = True + assert self._wait_for(lambda: not monitor.is_running) + assert not breaker.is_degraded + assert not breaker.bypass() + finally: + monitor.close() + client.close() + server.stop() + + def test_close_waits_for_the_loop_to_finish(self, tmp_path): + server = self._start_server(tmp_path) + client = self._make_client(tmp_path) + monitor = self._make_monitor(client, AcceleratorBreaker()) + try: + monitor.start() + assert self._wait_for(lambda: server.request_count >= 1) + handle = monitor._task + monitor.close() + assert not handle.is_alive() + finally: + client.close() + server.stop() + + def test_request_stop_ends_the_loop_without_close(self, tmp_path): + server = self._start_server(tmp_path) + client = self._make_client(tmp_path) + monitor = self._make_monitor(client, AcceleratorBreaker()) + try: + monitor.start() + assert self._wait_for(lambda: server.request_count >= 1) + monitor.request_stop() + assert self._wait_for(lambda: not monitor.is_running) + finally: + monitor.close() + client.close() + server.stop() + + def test_loop_stops_once_the_breaker_is_tripped(self, tmp_path): + server = self._start_server(tmp_path) + client = self._make_client(tmp_path) + breaker = AcceleratorBreaker() + monitor = self._make_monitor(client, breaker) + try: + monitor.start() + assert self._wait_for(lambda: server.request_count >= 1) + breaker.trip() + assert self._wait_for(lambda: not monitor.is_running) + finally: + monitor.close() + client.close() + server.stop() + + def test_close_stops_polling(self, tmp_path): + server = self._start_server(tmp_path) + client = self._make_client(tmp_path) + monitor = self._make_monitor(client, AcceleratorBreaker()) + try: + monitor.start() + assert self._wait_for(lambda: server.request_count >= 1) + monitor.close() + CrossSync._Sync_Impl.sleep(0.2) + counted = server.request_count + CrossSync._Sync_Impl.sleep(0.2) + assert server.request_count == counted + finally: + client.close() + server.stop() + + def test_close_leaves_a_degraded_flag_set(self, tmp_path): + server = self._start_server(tmp_path, status=ServingStatus.NOT_SERVING) + client = self._make_client(tmp_path) + breaker = AcceleratorBreaker() + monitor = self._make_monitor(client, breaker) + try: + monitor.start() + assert self._wait_for(lambda: breaker.is_degraded) + monitor.close() + assert breaker.bypass() + finally: + client.close() + server.stop() + + def test_start_is_idempotent(self): + monitor = self._make_monitor(None, AcceleratorBreaker(), interval=60.0) + try: + monitor.start() + first = monitor._task + monitor.start() + assert monitor._task is first + finally: + monitor.close() diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_health.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_health.py new file mode 100644 index 000000000000..b011586bbe25 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_health.py @@ -0,0 +1,219 @@ +# 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 the hand-rolled health codec and the breaker it drives. + +The codec in ``_accelerator/_health.py`` exists so this package doesn't have to +depend on ``grpcio-health-checking`` for two trivial messages. That trade is only +sound if the bytes it produces and accepts are exactly what a real protobuf +implementation produces and accepts, so the tests here check it against +``google.protobuf.proto_builder`` — protobuf's own encoder, building a message +with the same field number and wire type as ``HealthCheckResponse.status`` +rather than re-stating the byte strings the codec was written from. + +The probe and polling loop are exercised over a real UDS gRPC server in +``tests/unit/data/_async/test__accelerator_health.py``. +""" + +import pytest +from google.protobuf import descriptor_pb2, proto_builder + +from google.cloud.bigtable.data._accelerator._fallback import AcceleratorBreaker +from google.cloud.bigtable.data._accelerator._health import ( + HEALTH_CHECK_METHOD, + ServingStatus, + parse_health_response, + serialize_health_request, +) +from google.cloud.bigtable.data._sync_autogen._accelerator_health import ( + _AcceleratorHealthMonitor, +) + +# A message shaped like HealthCheckResponse: one field, number 1, varint. An +# enum and an int32 share a wire type, so protobuf encodes these identically to +# the real generated class. +_ResponseLike = proto_builder.MakeSimpleProtoClass( + {"status": descriptor_pb2.FieldDescriptorProto.TYPE_INT32}, + full_name="tests.bigtable.accelerator.HealthCheckResponseLike", +) + + +def _encode(status: int) -> bytes: + return _ResponseLike(status=status).SerializeToString() + + +class TestHealthCodec: + def test_request_is_the_empty_message(self): + # HealthCheckRequest{service: ""} has every field at its proto3 default. + assert serialize_health_request(None) == b"" + assert serialize_health_request(object()) == b"" + + def test_method_name(self): + assert HEALTH_CHECK_METHOD == "/grpc.health.v1.Health/Check" + + @pytest.mark.parametrize("status", list(ServingStatus)) + def test_parses_what_protobuf_encodes(self, status): + assert parse_health_response(_encode(status.value)) == status + + def test_empty_payload_is_unknown(self): + # proto3 omits a field at its default, so a real SERVING_STATUS of + # UNKNOWN arrives as zero bytes. proto_builder emits proto2, which + # writes the explicit zero instead; both must decode the same way. + assert parse_health_response(b"") == ServingStatus.UNKNOWN + assert parse_health_response(_encode(0)) == ServingStatus.UNKNOWN + + @pytest.mark.parametrize( + "payload,reason", + [ + (b"\x12\x01", "field 2, not the status field"), + (b"\x09\x01", "field 1 but the wrong wire type"), + (b"\x08", "tag with no varint after it"), + (b"\x08\x80", "varint truncated mid-continuation"), + (b"\x08\xac\x02", "a status value this client has never heard of"), + ], + ) + def test_unreadable_payloads_are_unknown(self, payload, reason): + # UNKNOWN is what the monitor counts as a failed probe, which is the + # safe reading of a response we cannot interpret. + assert parse_health_response(payload) == ServingStatus.UNKNOWN, reason + + +class TestBreakerDegraded: + def test_degraded_bypasses_and_clears(self): + breaker = AcceleratorBreaker() + assert not breaker.bypass() + breaker.set_degraded(True) + assert breaker.bypass() and breaker.is_degraded and not breaker.is_tripped + breaker.set_degraded(False) + assert not breaker.bypass() + + def test_clearing_degraded_cannot_undo_a_trip(self): + # The whole point of keeping the two flags apart: a health recovery must + # never resurrect an accelerator that replied UNIMPLEMENTED or died. + breaker = AcceleratorBreaker() + breaker.trip() + breaker.set_degraded(True) + breaker.set_degraded(False) + assert breaker.bypass() and breaker.is_tripped and not breaker.is_degraded + + +class TestMonitorStateMachine: + """Drives ``_record`` directly, so thresholds are checked without timing.""" + + @staticmethod + def _make_one(breaker, *, unhealthy=3, healthy=3): + return _AcceleratorHealthMonitor( + accelerator_client=None, + breaker=breaker, + unhealthy_threshold=unhealthy, + healthy_threshold=healthy, + ) + + def test_degrades_only_on_the_full_streak(self): + breaker = AcceleratorBreaker() + monitor = self._make_one(breaker) + for _ in range(2): + monitor._record(False) + assert not breaker.is_degraded + monitor._record(False) + assert breaker.is_degraded + + def test_one_good_probe_resets_the_failure_streak(self): + # A daemon that is merely occasionally slow never trips: this is the + # anti-flap property, and the reason a single GC pause is survivable. + breaker = AcceleratorBreaker() + monitor = self._make_one(breaker) + for _ in range(10): + monitor._record(False) + monitor._record(False) + monitor._record(True) + assert not breaker.is_degraded + + def test_recovers_only_on_the_full_streak(self): + breaker = AcceleratorBreaker() + monitor = self._make_one(breaker) + for _ in range(3): + monitor._record(False) + assert breaker.is_degraded + for _ in range(2): + monitor._record(True) + assert breaker.is_degraded + monitor._record(True) + assert not breaker.is_degraded + + def test_one_bad_probe_resets_the_recovery_streak(self): + breaker = AcceleratorBreaker() + monitor = self._make_one(breaker) + for _ in range(3): + monitor._record(False) + for _ in range(10): + monitor._record(True) + monitor._record(True) + monitor._record(False) + assert breaker.is_degraded + + def test_thresholds_of_one_flip_immediately(self): + breaker = AcceleratorBreaker() + monitor = self._make_one(breaker, unhealthy=1, healthy=1) + monitor._record(False) + assert breaker.is_degraded + monitor._record(True) + assert not breaker.is_degraded + + def test_never_clears_a_permanent_trip(self): + breaker = AcceleratorBreaker() + breaker.trip() + monitor = self._make_one(breaker, healthy=1) + monitor._record(True) + assert breaker.bypass() + + +class TestSyncMonitorThread: + """The sync monitor's thread, which has no counterpart in the async one.""" + + @staticmethod + def _make_one(**kwargs): + kwargs.setdefault("interval", 0.01) + return _AcceleratorHealthMonitor( + accelerator_client=None, breaker=AcceleratorBreaker(), **kwargs + ) + + def test_loop_runs_on_a_daemon_thread(self): + # Not a ThreadPoolExecutor worker: concurrent.futures joins its + # non-daemon workers at interpreter exit, and this loop only ends on + # close(), so a Table left unclosed would hang the process on the way + # out. A daemon thread cannot hold exit up. + monitor = self._make_one(interval=60.0) + thread = None + try: + monitor.start() + thread = monitor._task + assert thread.daemon + assert monitor.is_running + finally: + monitor.close() + assert not thread.is_alive() + + def test_request_stop_ends_the_loop(self): + # The weakref-finalizer path: sets a flag, waits for nothing, and the + # loop notices at its next wakeup. + monitor = self._make_one() + monitor.start() + thread = monitor._task + try: + monitor.request_stop() + thread.join(timeout=10.0) + assert not thread.is_alive() + assert not monitor.is_running + finally: + monitor.close()