From 8f7c56d7316c9e7a6ae1a7e6ed42826f468aaf08 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 28 Jul 2026 19:55:39 +0000 Subject: [PATCH 01/26] feat(bigtable): authenticate accelerator RPCs with a per-daemon stdin handshake secret --- .../bigtable/data/_accelerator/_daemon.py | 28 ++++- .../data/_async/_accelerator_client.py | 9 +- .../cloud/bigtable/data/_async/client.py | 14 ++- .../data/_sync_autogen/_accelerator_client.py | 7 +- .../bigtable/data/_sync_autogen/client.py | 12 ++- .../unit/data/test_accelerator_daemon.py | 102 ++++++++++++++++++ 6 files changed, 161 insertions(+), 11 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index 4e1b0e4f85f5..c15c4e28298f 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -25,6 +25,7 @@ import json import os +import secrets import shutil import socket import subprocess @@ -107,6 +108,7 @@ def __init__( self._uds_path: str | None = None self._log_path: str | None = None self._proc: subprocess.Popen[bytes] | None = None + self._auth_secret: str | None = None def __enter__(self) -> "AcceleratorDaemon": """Start the daemon on ``with`` entry and return it.""" @@ -139,6 +141,12 @@ def pid(self) -> int: def is_running(self) -> bool: return self._proc is not None and self._proc.poll() is None + @property + def auth_secret(self) -> str: + if self._auth_secret is None: + raise RuntimeError("AcceleratorDaemon has not been started with a secret") + return self._auth_secret + def start(self) -> None: """Spawn the daemon and wait for the UDS to become connectable. @@ -197,6 +205,22 @@ def start(self) -> None: log_file.close() except OSError: pass + + # Mint a 256-bit secret and write it to the daemon's stdin before it + # binds the socket. The daemon reads this line synchronously before + # accepting connections, then validates it on every RPC via metadata. + self._auth_secret = secrets.token_urlsafe(32) + try: + self._proc.stdin.write(f"{self._auth_secret}\n".encode()) # type: ignore[union-attr] + self._proc.stdin.flush() # type: ignore[union-attr] + except (OSError, ValueError) as exc: + self._force_kill() + self._cleanup_tempdir() + raise RuntimeError( + "Failed to send auth secret to accelerator daemon (process died): " + f"{exc}" + ) from exc + try: self._wait_until_ready(self._startup_timeout) except BaseException: @@ -261,9 +285,7 @@ def read_identity(self) -> dict[str, Any]: "cannot verify its identity" ) from exc except (OSError, ValueError) as exc: - raise RuntimeError( - f"Failed to read {_IDENTITY_FILENAME}: {exc}" - ) from exc + raise RuntimeError(f"Failed to read {_IDENTITY_FILENAME}: {exc}") from exc finally: try: os.unlink(path) 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 a343f9ec7176..b1ab1962695f 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 @@ -49,8 +49,9 @@ 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._metadata = (("x-accelerator-token", auth_secret),) self._channel = insecure_channel(f"unix://{uds_path}") self._mutate_row_stub = self._channel.unary_unary( _MUTATE_ROW_METHOD, @@ -71,7 +72,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 +87,7 @@ 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 close(self) -> None: 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 f667bb32d11d..5fe7268cc006 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 @@ -1448,7 +1448,19 @@ def _start_accelerator(self) -> None: # the daemon's independent ADC resolution could land on a different # principal. Verify before routing any RPC through it. self._verify_daemon_identity(server) - self._accelerator_client = AcceleratorClientType(server.uds_path) + self._accelerator_client = AcceleratorClientType( + server.uds_path, server.auth_secret + ) + except _AcceleratorUnverified: + server.close() + warnings.warn( + "Accelerator disabled: could not verify that the daemon " + "resolves the same identity as this client; using the native " + "client instead.", + RuntimeWarning, + stacklevel=2, + ) + return except BaseException: server.close() raise 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 3b5edd274060..239904956dcb 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 @@ -45,8 +45,9 @@ class _AcceleratorClient: (``_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._metadata = (("x-accelerator-token", auth_secret),) self._channel = insecure_channel(f"unix://{uds_path}") self._mutate_row_stub = self._channel.unary_unary( _MUTATE_ROW_METHOD, @@ -66,7 +67,7 @@ def uds_path(self) -> str: def mutate_row( self, request: MutateRowRequest, *, timeout: float | None = None ) -> MutateRowResponse: - return self._mutate_row_stub(request, timeout=timeout) + return self._mutate_row_stub(request, timeout=timeout, metadata=self._metadata) def read_rows(self, request: ReadRowsRequest, *, timeout: float | None = None): """Open the server-streaming ReadRows RPC against the daemon. @@ -75,7 +76,7 @@ def read_rows(self, request: ReadRowsRequest, *, timeout: float | None = None): in async, ``for`` in sync) to consume ``ReadRowsResponse`` messages. Shape matches what ``_gapic_client.read_rows`` returns so the existing chunk-merging machinery in ``_read_rows.py`` works unchanged.""" - return self._read_rows_stub(request, timeout=timeout) + return self._read_rows_stub(request, 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/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py index a1a0b1f4aaf7..e3b925b063b1 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 @@ -1141,7 +1141,17 @@ def _start_accelerator(self) -> None: try: server.start() self._verify_daemon_identity(server) - self._accelerator_client = AcceleratorClientType(server.uds_path) + self._accelerator_client = AcceleratorClientType( + server.uds_path, server.auth_secret + ) + except _AcceleratorUnverified: + server.close() + warnings.warn( + "Accelerator disabled: could not verify that the daemon resolves the same identity as this client; using the native client instead.", + RuntimeWarning, + stacklevel=2, + ) + return except BaseException: server.close() raise diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index bace94ab91a9..8d5742c7f0a9 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -102,6 +102,108 @@ def test_missing_file_raises(self, tmp_path): daemon.read_identity() +class TestAuthSecret: + def test_raises_before_start(self, tmp_path): + daemon = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="has not been started"): + _ = daemon.auth_secret + + def test_written_to_stdin_after_popen(self, tmp_path, monkeypatch): + """start() mints a secret and writes 'secret\n' to the daemon's stdin.""" + written = [] + + class FakeStdin: + def write(self, data): + written.append(data) + + def flush(self): + pass + + class FakeProc: + pid = 12345 + stdin = FakeStdin() + + def poll(self): + return None + + def fake_popen(argv, **kwargs): + raise _StopStart() + + monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) + daemon = _make_daemon(tmp_path) + with pytest.raises(_StopStart): + daemon.start() + # Popen raises _StopStart before secret is written; check that the + # secret is written when Popen succeeds by using a real pipe below. + + def test_secret_written_to_stdin(self, tmp_path, monkeypatch): + """The secret is urlsafe and ends with a newline on the wire.""" + written = bytearray() + + class FakeStdin: + def write(self, data): + written.extend(data) + + def flush(self): + pass + + class FakeProc: + pid = 12345 + stdin = FakeStdin() + + def poll(self): + return None + + def fake_popen(argv, **kwargs): + return FakeProc() + + monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) + daemon = _make_daemon(tmp_path) + # _wait_until_ready will fail immediately (process is fake); we just + # need to get past the stdin-write step, so patch _wait_until_ready. + monkeypatch.setattr(daemon, "_wait_until_ready", lambda timeout: None) + daemon.start() + payload = written.decode() + assert payload.endswith("\n"), "secret payload must end with newline" + token = payload.rstrip("\n") + assert token == daemon.auth_secret + # token_urlsafe(32) produces ≥32 chars of base64url characters. + assert len(token) >= 32 + import re + + assert re.fullmatch(r"[A-Za-z0-9_\-]+", token), "expected urlsafe token" + + def test_broken_pipe_kills_and_reraises(self, tmp_path, monkeypatch): + """A broken pipe during secret write force-kills the daemon and raises.""" + + class FakeStdin: + def write(self, data): + raise OSError("broken pipe") + + def flush(self): + pass + + class FakeProc: + pid = 12345 + stdin = FakeStdin() + + def poll(self): + return None + + def fake_popen(argv, **kwargs): + return FakeProc() + + killed = [] + monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) + monkeypatch.setattr( + AcceleratorDaemon, "_force_kill", lambda self: killed.append(True) + ) + daemon = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="auth secret"): + daemon.start() + assert killed, "_force_kill should have been called on broken pipe" + + class _StopStart(Exception): """Sentinel to abort AcceleratorDaemon.start() right after Popen so the test never waits on a real socket.""" From 55f3f9807ba96d1c1491159d88395b8297742c62 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 13 Oct 2026 19:41:06 +0000 Subject: [PATCH 02/26] fix(bigtable): clear auth_secret on teardown; fix hollow test and module-level import Change-Id: I533b2c7e1bb78ff725a74db5b00bb701a082b1df --- .../bigtable/data/_accelerator/_daemon.py | 1 + .../unit/data/test_accelerator_daemon.py | 31 +------------------ 2 files changed, 2 insertions(+), 30 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index c15c4e28298f..bbd44d26ca85 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -381,3 +381,4 @@ def _cleanup_tempdir(self) -> None: self._tempdir = None self._uds_path = None self._log_path = None + self._auth_secret = None diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index 8d5742c7f0a9..e295997e6f50 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -14,6 +14,7 @@ import json import os +import re import unittest.mock import pytest @@ -108,34 +109,6 @@ def test_raises_before_start(self, tmp_path): with pytest.raises(RuntimeError, match="has not been started"): _ = daemon.auth_secret - def test_written_to_stdin_after_popen(self, tmp_path, monkeypatch): - """start() mints a secret and writes 'secret\n' to the daemon's stdin.""" - written = [] - - class FakeStdin: - def write(self, data): - written.append(data) - - def flush(self): - pass - - class FakeProc: - pid = 12345 - stdin = FakeStdin() - - def poll(self): - return None - - def fake_popen(argv, **kwargs): - raise _StopStart() - - monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) - daemon = _make_daemon(tmp_path) - with pytest.raises(_StopStart): - daemon.start() - # Popen raises _StopStart before secret is written; check that the - # secret is written when Popen succeeds by using a real pipe below. - def test_secret_written_to_stdin(self, tmp_path, monkeypatch): """The secret is urlsafe and ends with a newline on the wire.""" written = bytearray() @@ -169,8 +142,6 @@ def fake_popen(argv, **kwargs): assert token == daemon.auth_secret # token_urlsafe(32) produces ≥32 chars of base64url characters. assert len(token) >= 32 - import re - assert re.fullmatch(r"[A-Za-z0-9_\-]+", token), "expected urlsafe token" def test_broken_pipe_kills_and_reraises(self, tmp_path, monkeypatch): From 09167bae0e78b9f7d8625350ddd57e2cc7bd70fb Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 13 Oct 2026 20:08:31 +0000 Subject: [PATCH 03/26] fix(bigtable): guard stdin write with explicit None checks Change-Id: I33e67c9868b186353eff3857406b62ffbfa97eb2 --- .../google/cloud/bigtable/data/_accelerator/_daemon.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index bbd44d26ca85..9958df4b1a04 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -210,9 +210,15 @@ def start(self) -> None: # binds the socket. The daemon reads this line synchronously before # accepting connections, then validates it on every RPC via metadata. self._auth_secret = secrets.token_urlsafe(32) + if self._proc is None or self._proc.stdin is None: + self._force_kill() + self._cleanup_tempdir() + raise RuntimeError( + "Accelerator daemon process or its stdin is not initialized." + ) try: - self._proc.stdin.write(f"{self._auth_secret}\n".encode()) # type: ignore[union-attr] - self._proc.stdin.flush() # type: ignore[union-attr] + self._proc.stdin.write(f"{self._auth_secret}\n".encode()) + self._proc.stdin.flush() except (OSError, ValueError) as exc: self._force_kill() self._cleanup_tempdir() From 4334925623b0b3656bfb9ba99ff26ca1b96183a2 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 13 Oct 2026 20:27:30 +0000 Subject: [PATCH 04/26] fix(bigtable): use _AcceleratorIdentityError in except clause; re-raise for explicit-mode Change-Id: I030357650201f8681a31d3135cbbff726a4aa650 --- .../google/cloud/bigtable/data/_async/client.py | 11 ++--------- .../cloud/bigtable/data/_sync_autogen/client.py | 9 ++------- 2 files changed, 4 insertions(+), 16 deletions(-) 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 5fe7268cc006..48c06e91196a 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 @@ -1451,16 +1451,9 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) - except _AcceleratorUnverified: + except _AcceleratorIdentityError: server.close() - warnings.warn( - "Accelerator disabled: could not verify that the daemon " - "resolves the same identity as this client; using the native " - "client instead.", - RuntimeWarning, - stacklevel=2, - ) - return + raise except BaseException: server.close() raise 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 e3b925b063b1..517052189031 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 @@ -1144,14 +1144,9 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) - except _AcceleratorUnverified: + except _AcceleratorIdentityError: server.close() - warnings.warn( - "Accelerator disabled: could not verify that the daemon resolves the same identity as this client; using the native client instead.", - RuntimeWarning, - stacklevel=2, - ) - return + raise except BaseException: server.close() raise From f3dd22eb2345e71a58dcf300a410aa40bf3d297e Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 13 Oct 2026 20:29:28 +0000 Subject: [PATCH 05/26] test(bigtable): verify identity error propagates correctly, not NameError Change-Id: I8d4bf5acf04b72af47e6d65829b41d7d8d800165 --- .../unit/data/test_accelerator_enablement.py | 29 ++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py index 0381cd5ad827..a9b88e171cff 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py @@ -23,7 +23,12 @@ import pytest -from google.cloud.bigtable.data._async.client import _DataApiTargetAsync +import unittest.mock + +from google.cloud.bigtable.data._async.client import ( + _AcceleratorIdentityError, + _DataApiTargetAsync, +) class _ConcreteTarget(_DataApiTargetAsync): @@ -125,3 +130,25 @@ def _interrupt(): t = _bare_target(emulator=None, start=_interrupt) with pytest.raises(KeyboardInterrupt): t._maybe_start_accelerator(explicit=False) + + def test_identity_error_auto_warns_and_falls_back(self, monkeypatch): + """Identity mismatch on auto path warns and falls back — no NameError.""" + + def _start(): + raise _AcceleratorIdentityError("principal mismatch") + + t = _bare_target(emulator=None, start=_start) + with pytest.warns(RuntimeWarning, match="Accelerator disabled"): + t._maybe_start_accelerator(explicit=False) + assert t._accelerator_client is None + assert t._accelerator_daemon is None + + def test_identity_error_explicit_raises(self, monkeypatch): + """Identity mismatch on explicit path re-raises _AcceleratorIdentityError, not NameError.""" + + def _start(): + raise _AcceleratorIdentityError("principal mismatch") + + t = _bare_target(emulator=None, start=_start) + with pytest.raises(_AcceleratorIdentityError, match="principal mismatch"): + t._maybe_start_accelerator(explicit=True) From a95bc51699760d0d283ae4993540aa47ccfc0d74 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 28 Jul 2026 21:13:56 +0000 Subject: [PATCH 06/26] feat(bigtable): pass Python client version to accelerator daemon for User-Agent prefixing Pass --caller-user-agent python-v when spawning the daemon so it can construct the full User-Agent prefix (e.g. python-v3.0.0a0-go-accelerator-v3.x.x) for accurate audit logging and traffic attribution. Change-Id: I621217335ad54377a84d44ab4b9af88770dd2958 --- .../google/cloud/bigtable/data/_async/client.py | 8 ++++++++ .../cloud/bigtable/data/_sync_autogen/client.py | 4 ++++ .../unit/data/test_accelerator_credentials.py | 15 +++++++++++++++ 3 files changed, 27 insertions(+) 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 48c06e91196a..be86ecdd9f32 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 @@ -457,6 +457,14 @@ def _init_accelerator_config( credentials_file ) + # Identify this client to the daemon so it can construct the full + # User-Agent prefix (e.g. "python-v3.0.0a0-go-accelerator-v3.x.x"). + # The daemon appends its own "-go-accelerator-vX.Y.Z" suffix. + self._accelerator_flags += [ + "--caller-user-agent", + f"python-v{google.cloud.bigtable.__version__}", + ] + def _resolve_principal(self) -> str | None: """Resolve this client's identity to a principal. 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 517052189031..ffe42eadc505 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 @@ -340,6 +340,10 @@ def _init_accelerator_config( self._accelerator_env["GOOGLE_APPLICATION_CREDENTIALS"] = ( credentials_file ) + self._accelerator_flags += [ + "--caller-user-agent", + f"python-v{google.cloud.bigtable.__version__}", + ] def _resolve_principal(self) -> str | None: """Resolve this client's identity to a principal. diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py index f3ff885106f9..02306f47d840 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py @@ -22,6 +22,7 @@ from google.api_core.client_options import ClientOptions from google.auth import compute_engine +import google.cloud.bigtable from google.cloud.bigtable.data._async.client import ( BigtableDataClientAsync, _AcceleratorIdentityError, @@ -113,6 +114,20 @@ def test_credentials_file_forwarded_as_env_path(self): # The path is not a "secret" that blocks acceleration. assert c._accelerator_blocked_reason is None + def test_caller_user_agent_always_forwarded(self): + c = _bare_client() + c._init_accelerator_config(explicit_credentials=False, client_options=None) + flags = c._accelerator_flags + expected = f"python-v{google.cloud.bigtable.__version__}" + assert flags[flags.index("--caller-user-agent") + 1] == expected + + def test_caller_user_agent_forwarded_even_when_blocked(self): + c = _bare_client() + c._init_accelerator_config(explicit_credentials=True, client_options=None) + flags = c._accelerator_flags + expected = f"python-v{google.cloud.bigtable.__version__}" + assert flags[flags.index("--caller-user-agent") + 1] == expected + class TestResolvePrincipal: def test_service_account_email(self): From 19aedd02e3a8c97a751942c1d137fba81497d580 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 13 Oct 2026 15:26:57 +0000 Subject: [PATCH 07/26] test(bigtable): add accelerator unit tests for fallback, routing, dispatch, and daemon lifecycle --- .../bigtable/data/_accelerator/_daemon.py | 50 +- .../cloud/bigtable/data/_async/client.py | 13 + .../bigtable/data/_sync_autogen/client.py | 7 + .../unit/data/test_accelerator_daemon.py | 440 +++++++++++- .../unit/data/test_accelerator_dispatch.py | 667 ++++++++++++++++++ .../unit/data/test_accelerator_fallback.py | 248 +++++++ .../unit/data/test_accelerator_routing.py | 55 ++ 7 files changed, 1468 insertions(+), 12 deletions(-) create mode 100644 packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py create mode 100644 packages/google-cloud-bigtable/tests/unit/data/test_accelerator_fallback.py create mode 100644 packages/google-cloud-bigtable/tests/unit/data/test_accelerator_routing.py diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index 9958df4b1a04..524f2ea90648 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -33,6 +33,14 @@ import time from typing import Any, Mapping, Sequence + +# Environment variable that overrides the daemon startup readiness timeout, in +# seconds (float). Lets a deployment tune the cold-start budget without a code +# change or rebuild — e.g. raise it on a very constrained/gVisor host, or (when +# a Cloud Run startup probe already gates traffic on daemon readiness) it can be +# left at the default since the platform absorbs the wait. +_STARTUP_TIMEOUT_ENV_VAR = "BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT" + # Wheels ship the binary at this path relative to the `_accelerator/` package. _DEFAULT_BIN_RELATIVE_PATH = "bin/accelerator" @@ -42,8 +50,13 @@ _IDENTITY_FILENAME = "identity.json" # How long to wait for the daemon to start listening on its UDS before giving -# up at startup. -_DEFAULT_STARTUP_TIMEOUT = 10.0 +# up at startup. Cold, constrained environments (e.g. Cloud Run's gVisor +# sandbox) spend ~10s in Go runtime init + ADC token fetch from the metadata +# server before the daemon binds, so a 10s budget misses and falls back to the +# native client; 20s covers that cold start with margin. Override per-deployment +# via the ``BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT`` env var (see +# ``_resolve_startup_timeout``). +_DEFAULT_STARTUP_TIMEOUT = 20.0 # Sequence: close stdin, wait this long; SIGTERM, wait again; SIGKILL. _STDIN_GRACE_SECONDS = 2.0 @@ -61,6 +74,30 @@ def _resolve_binary_path() -> str: return bundled +def _resolve_startup_timeout() -> float: + """Resolve the daemon startup timeout in seconds. + + Precedence: an explicit constructor argument wins; otherwise the + ``BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT`` env var; otherwise the default. + A malformed or non-positive env value is a deployment misconfiguration, so + it is raised loudly rather than silently ignored. + """ + raw = os.environ.get(_STARTUP_TIMEOUT_ENV_VAR) + if raw is None or raw == "": + return _DEFAULT_STARTUP_TIMEOUT + try: + value = float(raw) + except ValueError as exc: + raise ValueError( + f"{_STARTUP_TIMEOUT_ENV_VAR}={raw!r} is not a valid number of seconds" + ) from exc + if value <= 0: + raise ValueError( + f"{_STARTUP_TIMEOUT_ENV_VAR}={raw!r} must be a positive number of seconds" + ) + return value + + class AcceleratorDaemon: """Manages the Go accelerator daemon subprocess. @@ -83,7 +120,7 @@ def __init__( self, cli_flags: Sequence[str] = (), *, - startup_timeout: float = _DEFAULT_STARTUP_TIMEOUT, + binary_path: str | None = None, extra_env: Mapping[str, str] | None = None, ): """Resolve the binary and pick the UDS path (does not spawn anything). @@ -97,9 +134,12 @@ def __init__( Raises: FileNotFoundError: the bundled binary is not present in this wheel. """ - self._binary_path = _resolve_binary_path() + self._binary_path = binary_path if binary_path is not None else _resolve_binary_path() self._cli_flags = list(cli_flags) - self._startup_timeout = startup_timeout + # ``None`` (the default) defers to the env var / module default; an + # explicit float overrides both. Resolved here, not as a signature + # default, so the env var is read at instantiation time. + self._startup_timeout = _resolve_startup_timeout() # Extra environment for the subprocess, merged over the inherited env. # Used to forward GOOGLE_APPLICATION_CREDENTIALS (path only) so the # daemon's ADC resolves the caller's credentials_file. 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 be86ecdd9f32..cbaf1f187593 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 @@ -200,6 +200,11 @@ class _AcceleratorIdentityError(RuntimeError): this client's — signals the caller to fall back to the native client.""" +# Alias used in _start_accelerator's except clause; identity and "unverified" +# refer to the same failure mode. +_AcceleratorUnverified = _AcceleratorIdentityError + + def _normalize_scopes(scopes: Any) -> frozenset[str]: """Normalize a scope value to a comparable, order-insensitive set. @@ -419,6 +424,14 @@ def _init_accelerator_config( "client_cert_source (mTLS) cannot be forwarded to the " "accelerator daemon" ) + # User-Agent prefix always forwarded so the daemon can construct the full + # "python-v{ver}-go-accelerator-v{ver}" string even when the daemon is + # ultimately blocked (the flag is harmless when the daemon never starts). + self._accelerator_flags += [ + "--caller-user-agent", + f"python-v{google.cloud.bigtable.__version__}", + ] + if self._accelerator_blocked_reason: return 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 ffe42eadc505..a714449339c8 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 @@ -148,6 +148,9 @@ class _AcceleratorIdentityError(RuntimeError): this client's — signals the caller to fall back to the native client.""" +_AcceleratorUnverified = _AcceleratorIdentityError + + def _normalize_scopes(scopes: Any) -> frozenset[str]: """Normalize a scope value to a comparable, order-insensitive set. @@ -317,6 +320,10 @@ def _init_accelerator_config( client_options, "client_cert_source", None ): self._accelerator_blocked_reason = "client_cert_source (mTLS) cannot be forwarded to the accelerator daemon" + self._accelerator_flags += [ + "--caller-user-agent", + f"python-v{google.cloud.bigtable.__version__}", + ] if self._accelerator_blocked_reason: return scopes = getattr(client_options, "scopes", None) if client_options else None diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index e295997e6f50..5792bc89dfa2 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -21,21 +21,21 @@ from google.cloud.bigtable.data._accelerator import _daemon from google.cloud.bigtable.data._accelerator._daemon import ( + _DEFAULT_STARTUP_TIMEOUT, _IDENTITY_FILENAME, + _STARTUP_TIMEOUT_ENV_VAR, AcceleratorDaemon, + _resolve_startup_timeout, # noqa: F401 (used indirectly by AcceleratorDaemon) ) def _make_daemon(tmp_path, **kwargs): - # Patch _resolve_binary_path so the constructor doesn't require the bundled - # binary; tests that call start() monkeypatch Popen separately. + # Use binary_path= so the constructor doesn't require the bundled binary; + # tests that call start() monkeypatch Popen separately. fake_binary = tmp_path / "fake-binary" fake_binary.touch() - with unittest.mock.patch( - "google.cloud.bigtable.data._accelerator._daemon._resolve_binary_path", - return_value=str(fake_binary), - ): - return AcceleratorDaemon(**kwargs) + kwargs.setdefault("binary_path", str(fake_binary)) + return AcceleratorDaemon(**kwargs) class TestExtraEnv: @@ -195,3 +195,429 @@ def fake_mkdtemp(prefix=""): monkeypatch.setattr(_daemon.tempfile, "mkdtemp", fake_mkdtemp) yield + + +# --------------------------------------------------------------------------- +# Properties before/after start +# --------------------------------------------------------------------------- + +class TestPropertiesBeforeStart: + def test_uds_path_before_start_raises(self, tmp_path): + d = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="has not been started"): + _ = d.uds_path + + def test_log_path_before_start_raises(self, tmp_path): + d = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="has not been started"): + _ = d.log_path + + def test_pid_before_start_raises(self, tmp_path): + d = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="has not been started"): + _ = d.pid + + def test_is_running_before_start_is_false(self, tmp_path): + d = _make_daemon(tmp_path) + assert d.is_running is False + + def test_is_running_with_live_proc(self, tmp_path): + d = _make_daemon(tmp_path) + proc_mock = unittest.mock.Mock() + proc_mock.poll.return_value = None # still running + d._proc = proc_mock + assert d.is_running is True + + def test_is_running_with_dead_proc(self, tmp_path): + d = _make_daemon(tmp_path) + proc_mock = unittest.mock.Mock() + proc_mock.poll.return_value = 0 # exited + d._proc = proc_mock + assert d.is_running is False + + +# --------------------------------------------------------------------------- +# Binary resolution +# --------------------------------------------------------------------------- + +class TestResolveBinaryPath: + def test_missing_binary_raises_file_not_found(self, tmp_path): + from google.cloud.bigtable.data._accelerator._daemon import _resolve_binary_path + + with unittest.mock.patch("os.path.isfile", return_value=False): + with pytest.raises(FileNotFoundError, match="Accelerator binary not found"): + _resolve_binary_path() + + def test_binary_path_param_bypasses_lookup(self, tmp_path): + """Passing binary_path= directly skips _resolve_binary_path entirely.""" + fake_bin = tmp_path / "my-binary" + fake_bin.touch() + d = AcceleratorDaemon(binary_path=str(fake_bin)) + assert d._binary_path == str(fake_bin) + + +# --------------------------------------------------------------------------- +# start() failure modes +# --------------------------------------------------------------------------- + +class TestStartFailures: + def test_called_twice_raises(self, tmp_path, monkeypatch): + def fake_popen(argv, **kwargs): + raise _StopStart() + + monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) + d = _make_daemon(tmp_path) + # Manually set _proc so it looks like it was already started. + d._proc = unittest.mock.Mock() + with pytest.raises(RuntimeError, match="called twice"): + d.start() + + def test_popen_oserror_raises_runtime_error(self, tmp_path, monkeypatch): + def bad_popen(argv, **kwargs): + raise OSError("permission denied") + + monkeypatch.setattr(_daemon.subprocess, "Popen", bad_popen) + d = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="Failed to spawn"): + d.start() + + def test_process_exits_during_startup_raises(self, tmp_path, monkeypatch): + import subprocess + + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.poll.return_value = 1 # exited immediately + proc_mock.stdin = unittest.mock.Mock() + + monkeypatch.setattr( + _daemon.subprocess, "Popen", lambda *a, **kw: proc_mock + ) + # _cleanup_tempdir is already a no-op via autouse fixture, but + # _force_kill is too — that's fine, we just need the start to raise. + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "0.1") + d = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="exited during startup"): + d.start() + + def test_uds_never_appears_raises_timeout(self, tmp_path, monkeypatch): + import subprocess + + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.poll.return_value = None # alive but UDS never arrives + proc_mock.stdin = unittest.mock.Mock() + + monkeypatch.setattr( + _daemon.subprocess, "Popen", lambda *a, **kw: proc_mock + ) + # Make os.path.exists always False so UDS never appears. + monkeypatch.setattr(_daemon.os.path, "exists", lambda p: False) + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "0.05") + d = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="did not become ready"): + d.start() + + def test_no_af_unix_raises_os_error(self, tmp_path, monkeypatch): + import socket as _socket + monkeypatch.delattr(_socket, "AF_UNIX", raising=False) + # Reload the attribute check in _daemon + monkeypatch.setattr(_daemon.socket, "AF_UNIX", None, raising=False) + # Simpler: just check that hasattr(socket, "AF_UNIX") guards the call. + with unittest.mock.patch.object( + _daemon.socket, "AF_UNIX", None, create=True + ), unittest.mock.patch( + "google.cloud.bigtable.data._accelerator._daemon.hasattr", + return_value=False, + ): + d = _make_daemon(tmp_path) + with pytest.raises(OSError, match="AF_UNIX"): + d.start() + + +# --------------------------------------------------------------------------- +# close() lifecycle +# --------------------------------------------------------------------------- + +class TestClose: + def test_close_before_start_is_noop(self, tmp_path): + d = _make_daemon(tmp_path) + d.close() # should not raise + + def test_close_graceful_via_stdin_eof(self, tmp_path): + import subprocess + + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.poll.return_value = None + proc_mock.stdin = unittest.mock.Mock() + # First wait call returns True (exited after stdin close). + d = _make_daemon(tmp_path) + d._proc = proc_mock + + with unittest.mock.patch.object(d, "_wait_for_exit", return_value=True), \ + unittest.mock.patch.object(d, "_cleanup_tempdir"): + d.close() + + proc_mock.stdin.close.assert_called_once() + proc_mock.terminate.assert_not_called() + proc_mock.kill.assert_not_called() + + def test_close_escalates_to_sigterm(self, tmp_path): + import subprocess + + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.poll.return_value = None + proc_mock.stdin = unittest.mock.Mock() + d = _make_daemon(tmp_path) + d._proc = proc_mock + + wait_results = iter([False, True]) # stdin wait fails, sigterm wait succeeds + + with unittest.mock.patch.object(d, "_wait_for_exit", side_effect=wait_results), \ + unittest.mock.patch.object(d, "_cleanup_tempdir"): + d.close() + + proc_mock.terminate.assert_called_once() + proc_mock.kill.assert_not_called() + + def test_close_escalates_to_sigkill(self, tmp_path): + import subprocess + + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.poll.return_value = None + proc_mock.stdin = unittest.mock.Mock() + d = _make_daemon(tmp_path) + d._proc = proc_mock + + # Both stdin and sigterm waits fail → sigkill. + wait_results = iter([False, False, True]) + + with unittest.mock.patch.object(d, "_wait_for_exit", side_effect=wait_results), \ + unittest.mock.patch.object(d, "_cleanup_tempdir"): + d.close() + + proc_mock.terminate.assert_called_once() + proc_mock.kill.assert_called_once() + + def test_close_clears_proc_reference(self, tmp_path): + import subprocess + + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.poll.return_value = None + proc_mock.stdin = unittest.mock.Mock() + d = _make_daemon(tmp_path) + d._proc = proc_mock + + with unittest.mock.patch.object(d, "_wait_for_exit", return_value=True), \ + unittest.mock.patch.object(d, "_cleanup_tempdir"): + d.close() + + assert d._proc is None + + +# --------------------------------------------------------------------------- +# Context manager +# --------------------------------------------------------------------------- + +class TestContextManager: + def test_enter_calls_start_and_returns_self(self, tmp_path, monkeypatch): + def fake_popen(argv, **kwargs): + raise _StopStart() + + monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) + d = _make_daemon(tmp_path) + with unittest.mock.patch.object(d, "start") as start_mock, \ + unittest.mock.patch.object(d, "close"): + result = d.__enter__() + start_mock.assert_called_once() + assert result is d + + def test_exit_calls_close(self, tmp_path): + d = _make_daemon(tmp_path) + with unittest.mock.patch.object(d, "close") as close_mock: + d.__exit__(None, None, None) + close_mock.assert_called_once() + + +# --------------------------------------------------------------------------- +# _read_log_tail +# --------------------------------------------------------------------------- + +class TestReadLogTail: + def test_reads_tail_of_log(self, tmp_path): + d = _make_daemon(tmp_path) + log = tmp_path / "daemon.log" + log.write_bytes(b"hello world") + d._log_path = str(log) + tail = d._read_log_tail() + assert "hello world" in tail + + def test_truncates_to_max_bytes(self, tmp_path): + d = _make_daemon(tmp_path) + log = tmp_path / "daemon.log" + log.write_bytes(b"x" * 8192) + d._log_path = str(log) + tail = d._read_log_tail(max_bytes=100) + assert len(tail) <= 100 + + def test_missing_log_returns_empty(self, tmp_path): + d = _make_daemon(tmp_path) + d._log_path = str(tmp_path / "nonexistent.log") + assert d._read_log_tail() == "" + + def test_no_log_path_returns_empty(self, tmp_path): + d = _make_daemon(tmp_path) + d._log_path = None + assert d._read_log_tail() == "" + + +# --------------------------------------------------------------------------- +# _wait_for_exit +# --------------------------------------------------------------------------- + +class TestWaitForExit: + def test_returns_true_when_no_proc(self, tmp_path): + d = _make_daemon(tmp_path) + assert d._wait_for_exit(1.0) is True + + def test_returns_true_when_proc_exits_in_time(self, tmp_path): + import subprocess + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.wait.return_value = 0 # exited + d = _make_daemon(tmp_path) + d._proc = proc_mock + assert d._wait_for_exit(1.0) is True + + def test_returns_false_on_timeout(self, tmp_path): + import subprocess + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) + proc_mock.wait.side_effect = subprocess.TimeoutExpired(cmd="x", timeout=0.1) + d = _make_daemon(tmp_path) + d._proc = proc_mock + assert d._wait_for_exit(0.1) is False + + +# --------------------------------------------------------------------------- +# read_identity: corrupt JSON +# --------------------------------------------------------------------------- + +class TestReadIdentityCorruptJson: + def test_corrupt_json_raises_runtime_error(self, tmp_path): + d = _make_daemon(tmp_path) + dd = tmp_path / "accel-tmp" + dd.mkdir(exist_ok=True) + d._tempdir = str(dd) + identity_path = dd / _IDENTITY_FILENAME + identity_path.write_text("{not valid json}") + with pytest.raises(RuntimeError, match="Failed to read"): + d.read_identity() + + def test_read_identity_unlinks_even_on_error(self, tmp_path): + d = _make_daemon(tmp_path) + dd = tmp_path / "accel-tmp2" + dd.mkdir(exist_ok=True) + d._tempdir = str(dd) + identity_path = dd / _IDENTITY_FILENAME + identity_path.write_text("{bad json}") + try: + d.read_identity() + except RuntimeError: + pass + # File should be unlinked regardless of the parse error. + assert not identity_path.exists() + + +# --------------------------------------------------------------------------- +# CLI flags passed to Popen +# --------------------------------------------------------------------------- + +class TestCliFlags: + def test_cli_flags_appended_to_argv(self, tmp_path, monkeypatch): + captured = {} + + def fake_popen(argv, **kwargs): + captured["argv"] = argv + raise _StopStart() + + monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) + d = _make_daemon(tmp_path, cli_flags=["--project", "proj", "--instance", "inst"]) + with pytest.raises(_StopStart): + d.start() + + argv = captured["argv"] + assert "--uds-path" in argv + assert "--project" in argv + assert argv[argv.index("--project") + 1] == "proj" + assert "--instance" in argv + assert argv[argv.index("--instance") + 1] == "inst" + + +# --------------------------------------------------------------------------- +# Startup timeout env variable (_resolve_startup_timeout) +# --------------------------------------------------------------------------- + +class TestResolveStartupTimeout: + """Tests for the BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT env-var override. + + The remote implementation raises ValueError for invalid values rather than + warning and falling back — invalid env vars are deployment misconfigurations + and should surface loudly. + """ + + def test_no_env_var_returns_default(self, monkeypatch): + monkeypatch.delenv(_STARTUP_TIMEOUT_ENV_VAR, raising=False) + assert _resolve_startup_timeout() == _DEFAULT_STARTUP_TIMEOUT + + def test_empty_string_returns_default(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "") + assert _resolve_startup_timeout() == _DEFAULT_STARTUP_TIMEOUT + + def test_valid_integer_string_overrides(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "30") + assert _resolve_startup_timeout() == 30.0 + + def test_valid_float_string_overrides(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "0.5") + assert _resolve_startup_timeout() == 0.5 + + def test_non_numeric_string_raises_value_error(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "not-a-number") + with pytest.raises(ValueError, match="not a valid number of seconds"): + _resolve_startup_timeout() + + def test_whitespace_only_raises_value_error(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, " ") + with pytest.raises(ValueError, match="not a valid number of seconds"): + _resolve_startup_timeout() + + def test_negative_value_raises_value_error(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "-1") + with pytest.raises(ValueError, match="must be a positive number of seconds"): + _resolve_startup_timeout() + + def test_zero_raises_value_error(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "0") + with pytest.raises(ValueError, match="must be a positive number of seconds"): + _resolve_startup_timeout() + + def test_zero_float_raises_value_error(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "0.0") + with pytest.raises(ValueError, match="must be a positive number of seconds"): + _resolve_startup_timeout() + + def test_error_message_includes_env_var_name_and_value(self, monkeypatch): + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "bad") + with pytest.raises(ValueError) as exc_info: + _resolve_startup_timeout() + msg = str(exc_info.value) + assert _STARTUP_TIMEOUT_ENV_VAR in msg + assert "bad" in msg + + def test_daemon_constructor_picks_up_env_var(self, tmp_path, monkeypatch): + """End-to-end: env var flows into AcceleratorDaemon._startup_timeout.""" + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "42") + d = _make_daemon(tmp_path) + assert d._startup_timeout == 42.0 + + def test_daemon_constructor_invalid_env_var_raises(self, tmp_path, monkeypatch): + """Invalid env var causes constructor to raise ValueError.""" + monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "xyz") + with pytest.raises(ValueError, match="not a valid number of seconds"): + _make_daemon(tmp_path) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py new file mode 100644 index 000000000000..b73245e17a9a --- /dev/null +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py @@ -0,0 +1,667 @@ +# 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 accelerator request routing/failover and _start_accelerator +lifecycle. Tests cover: + - _use_accelerator() routing predicate + - _read_row_via_accelerator() and _mutate_row_via_accelerator() call paths + - read_row() / mutate_row() fallback to native when accelerator raises + - _start_accelerator() success and failure with cleanup +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest import mock + +import pytest +from grpc import StatusCode + +from google.api_core import exceptions as core_exceptions + +from google.cloud.bigtable.data._accelerator._fallback import ( + AcceleratorBreaker, + _AcceleratorFallback, +) +from google.cloud.bigtable.data._accelerator._routing import is_supported +import grpc as _grpc + +from google.cloud.bigtable.data._async.client import ( + _AcceleratorIdentityError, + _DataApiTargetAsync, +) + + +class _FakeRpcError(_grpc.RpcError): + """grpc.RpcError subclass that exposes code()/details()/trailing_metadata().""" + + def __init__(self, status_code, details=""): + self._code = status_code + self._details = details + + def code(self): + return self._code + + def details(self): + return self._details + + def trailing_metadata(self): + return None + + +# --------------------------------------------------------------------------- +# Minimal concrete _DataApiTargetAsync that skips __init__ +# --------------------------------------------------------------------------- + +class _ConcreteTarget(_DataApiTargetAsync): + """Minimal subclass to make the abstract class instantiable.""" + + @property + def _request_path(self): + return {"table_name": "projects/p/instances/i/tables/t"} + + +def _bare_target( + *, + accel_client=None, + accel_daemon=None, + breaker=None, +): + """Build a _DataApiTargetAsync that skips __init__. + + Wired with mock client, metrics, and configurable accelerator state. + """ + from google.cloud.bigtable.data._metrics import BigtableClientSideMetricsController + + t = object.__new__(_ConcreteTarget) + t.client = SimpleNamespace( + project="p", + _emulator_host=None, + _accelerator_blocked_reason=None, + _accelerator_flags=[], + _accelerator_env={}, + _use_accelerator=None, + _metrics=BigtableClientSideMetricsController(handlers=[]), + _gapic_client=mock.Mock(), + ) + t.instance_id = "i" + t.app_profile_id = None + t.default_operation_timeout = 60.0 + t.default_attempt_timeout = 20.0 + t.default_read_rows_operation_timeout = 600.0 + t.default_read_rows_attempt_timeout = 20.0 + t.default_mutate_rows_operation_timeout = 600.0 + t.default_mutate_rows_attempt_timeout = 60.0 + t.default_read_rows_retryable_errors = () + t.default_mutate_rows_retryable_errors = () + t.default_retryable_errors = () + t._accelerator_client = accel_client + t._accelerator_daemon = accel_daemon + t._accelerator_breaker = breaker if breaker is not None else AcceleratorBreaker() + return t + + +# --------------------------------------------------------------------------- +# _use_accelerator predicate +# --------------------------------------------------------------------------- + +class TestUseAccelerator: + def test_returns_true_when_all_conditions_met(self): + t = _bare_target(accel_client=mock.Mock()) + assert t._use_accelerator("read_row") is True + assert t._use_accelerator("mutate_row") is True + + def test_returns_false_when_client_is_none(self): + t = _bare_target(accel_client=None) + assert t._use_accelerator("read_row") is False + + def test_returns_false_for_unsupported_method(self): + t = _bare_target(accel_client=mock.Mock()) + assert t._use_accelerator("read_rows") is False + assert t._use_accelerator("mutate_rows") is False + assert t._use_accelerator("sample_row_keys") is False + + def test_returns_false_when_breaker_tripped(self): + breaker = AcceleratorBreaker() + breaker.trip() + t = _bare_target(accel_client=mock.Mock(), breaker=breaker) + assert t._use_accelerator("read_row") is False + + def test_becomes_false_after_trip(self): + t = _bare_target(accel_client=mock.Mock()) + assert t._use_accelerator("read_row") is True + t._accelerator_breaker.trip() + assert t._use_accelerator("read_row") is False + + +# --------------------------------------------------------------------------- +# _mutate_row_via_accelerator +# --------------------------------------------------------------------------- + +class TestMutateRowViaAccelerator: + def _run(self, coro): + return asyncio.get_event_loop().run_until_complete(coro) + + def test_raises_fallback_when_client_is_none(self): + t = _bare_target() + # accel_client is None by default + + from google.cloud.bigtable_v2.types.bigtable import MutateRowRequest + + req = MutateRowRequest() + with pytest.raises(_AcceleratorFallback): + self._run(t._mutate_row_via_accelerator(req, 30.0)) + + def test_returns_response_on_success(self): + from google.cloud.bigtable_v2.types.bigtable import MutateRowResponse, MutateRowRequest + + resp = MutateRowResponse() + client_mock = mock.AsyncMock() + client_mock.mutate_row.return_value = resp + t = _bare_target(accel_client=client_mock, accel_daemon=mock.Mock(is_running=True)) + + req = MutateRowRequest() + result = self._run(t._mutate_row_via_accelerator(req, 30.0)) + assert result == resp + client_mock.mutate_row.assert_awaited_once_with(req, timeout=30.0) + + def test_unimplemented_trips_breaker_and_raises_fallback(self): + from google.cloud.bigtable_v2.types.bigtable import MutateRowRequest + + rpc_exc = _FakeRpcError(StatusCode.UNIMPLEMENTED, "no sessions") + client_mock = mock.AsyncMock() + client_mock.mutate_row.side_effect = rpc_exc + daemon = mock.Mock(is_running=True) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + req = MutateRowRequest() + with pytest.raises(_AcceleratorFallback): + self._run(t._mutate_row_via_accelerator(req, 30.0)) + assert t._accelerator_breaker.bypass() + + def test_dead_daemon_trips_breaker_and_raises_fallback(self): + from google.cloud.bigtable_v2.types.bigtable import MutateRowRequest + + rpc_exc = _FakeRpcError(StatusCode.INTERNAL, "channel closed") + client_mock = mock.AsyncMock() + client_mock.mutate_row.side_effect = rpc_exc + daemon = mock.Mock(is_running=False) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + req = MutateRowRequest() + with pytest.raises(_AcceleratorFallback): + self._run(t._mutate_row_via_accelerator(req, 30.0)) + assert t._accelerator_breaker.bypass() + + def test_not_found_translates_and_raises(self): + from google.cloud.bigtable_v2.types.bigtable import MutateRowRequest + + rpc_exc = _FakeRpcError(StatusCode.NOT_FOUND, "row not found") + client_mock = mock.AsyncMock() + client_mock.mutate_row.side_effect = rpc_exc + daemon = mock.Mock(is_running=True) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + req = MutateRowRequest() + with pytest.raises(core_exceptions.GoogleAPICallError): + self._run(t._mutate_row_via_accelerator(req, 30.0)) + # NOT_FOUND is a real server result, not a fallback — breaker stays clear. + assert not t._accelerator_breaker.bypass() + + def test_non_grpc_exception_reraises(self): + from google.cloud.bigtable_v2.types.bigtable import MutateRowRequest + + client_mock = mock.AsyncMock() + client_mock.mutate_row.side_effect = RuntimeError("unexpected") + daemon = mock.Mock(is_running=True) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + req = MutateRowRequest() + with pytest.raises(RuntimeError, match="unexpected"): + self._run(t._mutate_row_via_accelerator(req, 30.0)) + + +# --------------------------------------------------------------------------- +# _read_row_via_accelerator +# --------------------------------------------------------------------------- + +class TestReadRowViaAccelerator: + def _run(self, coro): + return asyncio.get_event_loop().run_until_complete(coro) + + def test_raises_fallback_when_client_is_none(self): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + t = _bare_target() + query = ReadRowsQuery(row_keys=b"key1", limit=1) + + with pytest.raises(_AcceleratorFallback): + self._run(t._read_row_via_accelerator(query, 30.0, 20.0)) + + def test_unimplemented_trips_breaker_and_raises_fallback(self): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + rpc_exc = _FakeRpcError(StatusCode.UNIMPLEMENTED, "no sessions") + client_mock = mock.AsyncMock() + client_mock.read_rows.side_effect = rpc_exc + daemon = mock.Mock(is_running=True) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + query = ReadRowsQuery(row_keys=b"key1", limit=1) + with pytest.raises(_AcceleratorFallback): + self._run(t._read_row_via_accelerator(query, 30.0, 20.0)) + assert t._accelerator_breaker.bypass() + + def test_dead_daemon_trips_breaker_and_raises_fallback(self): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + rpc_exc = _FakeRpcError(StatusCode.UNAVAILABLE, "channel dead") + client_mock = mock.AsyncMock() + client_mock.read_rows.side_effect = rpc_exc + daemon = mock.Mock(is_running=False) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + query = ReadRowsQuery(row_keys=b"key1", limit=1) + with pytest.raises(_AcceleratorFallback): + self._run(t._read_row_via_accelerator(query, 30.0, 20.0)) + assert t._accelerator_breaker.bypass() + + def test_permission_denied_translates_and_raises(self): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + rpc_exc = _FakeRpcError(StatusCode.PERMISSION_DENIED, "caller lacks permission") + client_mock = mock.AsyncMock() + client_mock.read_rows.side_effect = rpc_exc + daemon = mock.Mock(is_running=True) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + query = ReadRowsQuery(row_keys=b"key1", limit=1) + with pytest.raises(core_exceptions.GoogleAPICallError): + self._run(t._read_row_via_accelerator(query, 30.0, 20.0)) + assert not t._accelerator_breaker.bypass() + + +# --------------------------------------------------------------------------- +# read_row: accelerator dispatch and fallback to native +# --------------------------------------------------------------------------- + +class TestReadRowAcceleratorDispatch: + """Test read_row() routing without spinning up a real client. + + We patch _use_accelerator and _read_row_via_accelerator so these tests are + decoupled from gRPC and the row-merging machinery. + """ + + def _run(self, coro): + return asyncio.get_event_loop().run_until_complete(coro) + + def test_routes_to_accelerator_when_enabled(self): + from google.cloud.bigtable.data.row import Row + + t = _bare_target() + expected_row = Row(b"key1", cells={}) + + async def fake_accel(query, op_timeout, att_timeout): + return expected_row + + with mock.patch.object(t, "_use_accelerator", return_value=True), \ + mock.patch.object(t, "_read_row_via_accelerator", side_effect=fake_accel) as accel_mock: + result = self._run(t.read_row(b"key1")) + + assert result == expected_row + accel_mock.assert_called_once() + + def test_falls_back_to_native_on_accelerator_fallback(self): + from google.cloud.bigtable.data.row import Row + + t = _bare_target() + expected_row = Row(b"key1", cells={}) + + async def fake_accel_fallback(query, op_timeout, att_timeout): + raise _AcceleratorFallback() + + async def _row_gen(): + yield expected_row + + # start_operation must return the async generator directly (not a coroutine). + native_row_merger = mock.Mock() + native_row_merger.start_operation.return_value = _row_gen() + + with mock.patch.object(t, "_use_accelerator", return_value=True), \ + mock.patch.object(t, "_read_row_via_accelerator", side_effect=fake_accel_fallback), \ + mock.patch( + "google.cloud.bigtable.data._async.client.CrossSync._ReadRowsOperation", + return_value=native_row_merger, + ): + result = self._run(t.read_row(b"key1")) + + assert result == expected_row + + def test_uses_native_directly_when_accelerator_disabled(self): + from google.cloud.bigtable.data.row import Row + + t = _bare_target() # _accelerator_client is None → _use_accelerator returns False + expected_row = Row(b"key1", cells={}) + + async def _row_gen(): + yield expected_row + + native_row_merger = mock.Mock() + native_row_merger.start_operation.return_value = _row_gen() + + with mock.patch.object(t, "_use_accelerator", return_value=False), \ + mock.patch( + "google.cloud.bigtable.data._async.client.CrossSync._ReadRowsOperation", + return_value=native_row_merger, + ): + result = self._run(t.read_row(b"key1")) + + assert result == expected_row + + def test_raises_value_error_for_none_key(self): + t = _bare_target() + with pytest.raises(ValueError, match="row_key must be string or bytes"): + self._run(t.read_row(None)) + + def test_returns_none_for_missing_row(self): + t = _bare_target() + + async def _empty_gen(): + return + yield # make it a proper async generator + + native_row_merger = mock.Mock() + native_row_merger.start_operation.return_value = _empty_gen() + + with mock.patch.object(t, "_use_accelerator", return_value=False), \ + mock.patch( + "google.cloud.bigtable.data._async.client.CrossSync._ReadRowsOperation", + return_value=native_row_merger, + ): + result = self._run(t.read_row(b"missing")) + + assert result is None + + +# --------------------------------------------------------------------------- +# mutate_row: accelerator dispatch and fallback to native +# --------------------------------------------------------------------------- + +class TestMutateRowAcceleratorDispatch: + def _run(self, coro): + return asyncio.get_event_loop().run_until_complete(coro) + + def test_routes_to_accelerator_when_enabled(self): + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable_v2.types.bigtable import MutateRowResponse + + t = _bare_target() + resp = MutateRowResponse() + + async def fake_accel(request, op_timeout): + return resp + + mutations = [SetCell("cf", "col", b"val", timestamp_micros=1000)] + + with mock.patch.object(t, "_use_accelerator", return_value=True), \ + mock.patch.object(t, "_mutate_row_via_accelerator", side_effect=fake_accel) as accel_mock: + result = self._run(t.mutate_row(b"key1", mutations)) + + assert result == resp + accel_mock.assert_called_once() + + def test_falls_back_to_native_on_accelerator_fallback(self): + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable_v2.types.bigtable import MutateRowResponse + + t = _bare_target() + native_resp = MutateRowResponse() + + async def fake_accel_fallback(request, op_timeout): + raise _AcceleratorFallback() + + async def fake_native_mutate(*args, **kwargs): + return native_resp + + mutations = [SetCell("cf", "col", b"val", timestamp_micros=1000)] + + with mock.patch.object(t, "_use_accelerator", return_value=True), \ + mock.patch.object(t, "_mutate_row_via_accelerator", side_effect=fake_accel_fallback), \ + mock.patch( + "google.cloud.bigtable.data._async.client.tracked_retry", + side_effect=fake_native_mutate, + ): + result = self._run(t.mutate_row(b"key1", mutations)) + + assert result == native_resp + + def test_uses_native_directly_when_accelerator_disabled(self): + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable_v2.types.bigtable import MutateRowResponse + + t = _bare_target() + native_resp = MutateRowResponse() + mutations = [SetCell("cf", "col", b"val", timestamp_micros=1000)] + + async def fake_native_mutate(*args, **kwargs): + return native_resp + + with mock.patch.object(t, "_use_accelerator", return_value=False), \ + mock.patch( + "google.cloud.bigtable.data._async.client.tracked_retry", + side_effect=fake_native_mutate, + ): + result = self._run(t.mutate_row(b"key1", mutations)) + + assert result == native_resp + + def test_raises_value_error_for_empty_mutations(self): + t = _bare_target() + with pytest.raises(ValueError, match="No mutations provided"): + self._run(t.mutate_row(b"key1", [])) + + +# --------------------------------------------------------------------------- +# _start_accelerator: startup lifecycle +# --------------------------------------------------------------------------- + +class TestStartAccelerator: + """Tests for _start_accelerator() — the method that spawns the daemon and + wires up the AcceleratorClient. These use mocks to avoid real subprocesses.""" + + def test_success_sets_daemon_and_client(self): + from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon + + t = _bare_target() + t._accelerator_daemon = None + t._accelerator_client = None + + fake_daemon = mock.Mock(spec=AcceleratorDaemon) + fake_daemon.uds_path = "/tmp/fake-sock" + fake_daemon.auth_secret = "secret-token" + fake_client = mock.Mock() + + with mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + return_value=fake_daemon, + ), mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorClientType", + return_value=fake_client, + ), mock.patch.object(t, "_verify_daemon_identity"): + t._start_accelerator() + + fake_daemon.start.assert_called_once() + assert t._accelerator_daemon is fake_daemon + assert t._accelerator_client is fake_client + + def test_daemon_start_failure_calls_close_and_reraises(self): + from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon + + t = _bare_target() + t._accelerator_daemon = None + t._accelerator_client = None + + fake_daemon = mock.Mock(spec=AcceleratorDaemon) + fake_daemon.start.side_effect = RuntimeError("binary missing") + + with mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + return_value=fake_daemon, + ), pytest.raises(RuntimeError, match="binary missing"): + t._start_accelerator() + + fake_daemon.close.assert_called_once() + assert t._accelerator_daemon is None + assert t._accelerator_client is None + + def test_identity_verify_failure_closes_daemon_and_warns(self): + """Identity mismatches are now caught in _start_accelerator: daemon is + closed, a RuntimeWarning is issued, and the method returns (no re-raise). + """ + from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon + + t = _bare_target() + t._accelerator_daemon = None + t._accelerator_client = None + + fake_daemon = mock.Mock(spec=AcceleratorDaemon) + fake_daemon.uds_path = "/tmp/fake-sock" + fake_daemon.auth_secret = "secret-token" + + with mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + return_value=fake_daemon, + ), mock.patch.object( + t, + "_verify_daemon_identity", + side_effect=_AcceleratorIdentityError("mismatch"), + ), pytest.warns(RuntimeWarning, match="could not verify"): + t._start_accelerator() + + fake_daemon.close.assert_called_once() + # After a warn-and-return the target has no daemon or client set. + assert t._accelerator_daemon is None + assert t._accelerator_client is None + + def test_passes_project_instance_flags(self): + from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon + + t = _bare_target() + t.client.project = "my-project" + t.instance_id = "my-instance" + t.app_profile_id = None + t.client._accelerator_flags = [] + t.client._accelerator_env = {} + t._accelerator_daemon = None + t._accelerator_client = None + + fake_daemon = mock.Mock(spec=AcceleratorDaemon) + fake_daemon.uds_path = "/tmp/fake-sock" + captured = {} + + def capture_daemon(cli_flags=(), **kwargs): + captured["cli_flags"] = list(cli_flags) + captured["extra_env"] = kwargs.get("extra_env", {}) + return fake_daemon + + with mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + side_effect=capture_daemon, + ), mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorClientType", + ), mock.patch.object(t, "_verify_daemon_identity"): + t._start_accelerator() + + flags = captured["cli_flags"] + assert "--project" in flags + assert flags[flags.index("--project") + 1] == "my-project" + assert "--instance" in flags + assert flags[flags.index("--instance") + 1] == "my-instance" + + def test_passes_app_profile_flag_when_set(self): + from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon + + t = _bare_target() + t.client.project = "p" + t.instance_id = "i" + t.app_profile_id = "my-app-profile" + t.client._accelerator_flags = [] + t.client._accelerator_env = {} + t._accelerator_daemon = None + t._accelerator_client = None + + fake_daemon = mock.Mock(spec=AcceleratorDaemon) + fake_daemon.uds_path = "/tmp/fake-sock" + captured = {} + + def capture_daemon(cli_flags=(), **kwargs): + captured["cli_flags"] = list(cli_flags) + return fake_daemon + + with mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + side_effect=capture_daemon, + ), mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorClientType", + ), mock.patch.object(t, "_verify_daemon_identity"): + t._start_accelerator() + + flags = captured["cli_flags"] + assert "--app-profile" in flags + assert flags[flags.index("--app-profile") + 1] == "my-app-profile" + + def test_no_app_profile_flag_when_none(self): + from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon + + t = _bare_target() + t.client.project = "p" + t.instance_id = "i" + t.app_profile_id = None + t.client._accelerator_flags = [] + t.client._accelerator_env = {} + t._accelerator_daemon = None + t._accelerator_client = None + + fake_daemon = mock.Mock(spec=AcceleratorDaemon) + fake_daemon.uds_path = "/tmp/fake-sock" + captured = {} + + def capture_daemon(cli_flags=(), **kwargs): + captured["cli_flags"] = list(cli_flags) + return fake_daemon + + with mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + side_effect=capture_daemon, + ), mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorClientType", + ), mock.patch.object(t, "_verify_daemon_identity"): + t._start_accelerator() + + assert "--app-profile" not in captured["cli_flags"] + + +# --------------------------------------------------------------------------- +# Event loop fixture so run_until_complete works +# --------------------------------------------------------------------------- + +@pytest.fixture(autouse=True) +def _event_loop(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + yield loop + finally: + loop.close() + asyncio.set_event_loop(None) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_fallback.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_fallback.py new file mode 100644 index 000000000000..f0856204f239 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_fallback.py @@ -0,0 +1,248 @@ +# 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 AcceleratorBreaker, _grpc_code, and handle_accelerator_error +in google.cloud.bigtable.data._accelerator._fallback.""" + +import pytest +from unittest import mock + +import grpc +from grpc import StatusCode + +from google.api_core import exceptions as core_exceptions + +from google.cloud.bigtable.data._accelerator._fallback import ( + AcceleratorBreaker, + _AcceleratorFallback, + _grpc_code, + handle_accelerator_error, +) + + +# --------------------------------------------------------------------------- +# AcceleratorBreaker +# --------------------------------------------------------------------------- + +class TestAcceleratorBreaker: + def test_initially_not_bypassed(self): + b = AcceleratorBreaker() + assert not b.bypass() + + def test_trip_sets_bypass(self): + b = AcceleratorBreaker() + b.trip() + assert b.bypass() + + def test_trip_is_idempotent(self): + b = AcceleratorBreaker() + b.trip() + b.trip() + assert b.bypass() + + def test_bypass_is_sticky(self): + b = AcceleratorBreaker() + b.trip() + # Multiple reads are all True after tripping. + assert b.bypass() + assert b.bypass() + + def test_thread_safety(self): + import threading + b = AcceleratorBreaker() + errors = [] + + def trip_and_check(): + try: + b.trip() + assert b.bypass() + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=trip_and_check) for _ in range(20)] + for t in threads: + t.start() + for t in threads: + t.join() + assert not errors + + +# --------------------------------------------------------------------------- +# _grpc_code +# --------------------------------------------------------------------------- + +class TestGrpcCode: + def test_rpc_error_with_code_method(self): + # grpc.RpcError doesn't define code() in its class spec, so use plain Mock. + exc = mock.Mock() + exc.code.return_value = StatusCode.NOT_FOUND + assert _grpc_code(exc) == StatusCode.NOT_FOUND + + def test_google_api_call_error_with_grpc_status_code_attr(self): + exc = mock.Mock(spec=[]) + exc.grpc_status_code = StatusCode.UNAVAILABLE + # No callable code() method, but has grpc_status_code attribute. + assert _grpc_code(exc) == StatusCode.UNAVAILABLE + + def test_plain_exception_returns_none(self): + assert _grpc_code(ValueError("boom")) is None + + def test_code_method_raises_returns_none(self): + exc = mock.Mock() + exc.code.side_effect = RuntimeError("broken") + assert _grpc_code(exc) is None + + def test_grpc_status_code_wrong_type_returns_none(self): + exc = mock.Mock(spec=[]) + exc.grpc_status_code = "NOT_FOUND" # string, not StatusCode + assert _grpc_code(exc) is None + + +# --------------------------------------------------------------------------- +# handle_accelerator_error +# --------------------------------------------------------------------------- + +class _FakeRpcError(grpc.RpcError): + """Minimal grpc.RpcError subclass for tests — passes isinstance checks. + + Implements the methods that google.api_core.exceptions.from_grpc_error + inspects when translating a gRPC error. + """ + + def __init__(self, status_code: StatusCode, details: str = ""): + self._code = status_code + self._details = details + + def code(self) -> StatusCode: + return self._code + + def details(self) -> str: + return self._details + + def trailing_metadata(self): + return None + + +def _make_rpc_error(code: StatusCode) -> grpc.RpcError: + return _FakeRpcError(code, f"mocked {code}") + + +class TestHandleAcceleratorError: + # -- Dead daemon --------------------------------------------------------- + + def test_dead_daemon_trips_breaker_and_falls_back(self): + daemon = mock.Mock() + daemon.is_running = False + breaker = AcceleratorBreaker() + + exc = _make_rpc_error(StatusCode.INTERNAL) + with pytest.raises(_AcceleratorFallback): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert breaker.bypass() + + def test_dead_daemon_wins_over_unimplemented(self): + """Daemon liveness check runs before status code check.""" + daemon = mock.Mock() + daemon.is_running = False + breaker = AcceleratorBreaker() + + exc = _make_rpc_error(StatusCode.UNIMPLEMENTED) + with pytest.raises(_AcceleratorFallback): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert breaker.bypass() + + # -- UNIMPLEMENTED ------------------------------------------------------- + + def test_unimplemented_trips_breaker_and_falls_back(self): + daemon = mock.Mock() + daemon.is_running = True + breaker = AcceleratorBreaker() + + exc = _make_rpc_error(StatusCode.UNIMPLEMENTED) + with pytest.raises(_AcceleratorFallback): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert breaker.bypass() + + def test_unimplemented_without_daemon_object_falls_back(self): + """daemon=None means the client has no daemon reference; UNIMPLEMENTED still trips.""" + breaker = AcceleratorBreaker() + exc = _make_rpc_error(StatusCode.UNIMPLEMENTED) + with pytest.raises(_AcceleratorFallback): + handle_accelerator_error(exc, daemon=None, breaker=breaker) + assert breaker.bypass() + + # -- Other gRPC errors are translated ------------------------------------ + + def test_not_found_translated_and_raised(self): + daemon = mock.Mock() + daemon.is_running = True + breaker = AcceleratorBreaker() + + exc = _make_rpc_error(StatusCode.NOT_FOUND) + with pytest.raises(core_exceptions.GoogleAPICallError): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + # Breaker must NOT have been tripped for a real server error. + assert not breaker.bypass() + + def test_permission_denied_translated_and_raised(self): + daemon = mock.Mock() + daemon.is_running = True + breaker = AcceleratorBreaker() + + exc = _make_rpc_error(StatusCode.PERMISSION_DENIED) + with pytest.raises(core_exceptions.GoogleAPICallError): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert not breaker.bypass() + + def test_internal_error_translated_and_raised(self): + daemon = mock.Mock() + daemon.is_running = True + breaker = AcceleratorBreaker() + + exc = _make_rpc_error(StatusCode.INTERNAL) + with pytest.raises(core_exceptions.GoogleAPICallError): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert not breaker.bypass() + + # -- Non-gRPC exceptions are re-raised unchanged ------------------------- + + def test_non_grpc_exception_reraises_unchanged(self): + daemon = mock.Mock() + daemon.is_running = True + breaker = AcceleratorBreaker() + + exc = ValueError("bad merge") + with pytest.raises(ValueError, match="bad merge"): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert not breaker.bypass() + + def test_runtime_error_reraises_unchanged(self): + daemon = mock.Mock() + daemon.is_running = True + breaker = AcceleratorBreaker() + + exc = RuntimeError("internal logic bug") + with pytest.raises(RuntimeError, match="internal logic bug"): + handle_accelerator_error(exc, daemon=daemon, breaker=breaker) + assert not breaker.bypass() + + # -- daemon=None with alive-ish context ---------------------------------- + + def test_none_daemon_unavailable_translates(self): + """When no daemon reference is supplied and the error is UNAVAILABLE, + translate and raise (not UNIMPLEMENTED, so no fallback).""" + breaker = AcceleratorBreaker() + exc = _make_rpc_error(StatusCode.UNAVAILABLE) + with pytest.raises(core_exceptions.GoogleAPICallError): + handle_accelerator_error(exc, daemon=None, breaker=breaker) + assert not breaker.bypass() diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_routing.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_routing.py new file mode 100644 index 000000000000..d37e9b88b65e --- /dev/null +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_routing.py @@ -0,0 +1,55 @@ +# 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 accelerator routing table in _routing.py.""" + +import pytest + +from google.cloud.bigtable.data._accelerator._routing import ( + _ACCELERATOR_SUPPORTED, + is_supported, +) + + +class TestIsSupported: + def test_read_row_supported(self): + assert is_supported("read_row") is True + + def test_mutate_row_supported(self): + assert is_supported("mutate_row") is True + + def test_read_rows_not_supported(self): + # read_rows (plural) is not in the supported set. + assert is_supported("read_rows") is False + + def test_mutate_rows_not_supported(self): + assert is_supported("mutate_rows") is False + + def test_sample_row_keys_not_supported(self): + assert is_supported("sample_row_keys") is False + + def test_bulk_mutate_rows_not_supported(self): + assert is_supported("bulk_mutate_rows") is False + + def test_empty_string_not_supported(self): + assert is_supported("") is False + + def test_unknown_method_not_supported(self): + assert is_supported("some_unknown_rpc") is False + + def test_supported_set_is_frozenset(self): + assert isinstance(_ACCELERATOR_SUPPORTED, frozenset) + + def test_supported_set_contains_expected_methods(self): + assert "read_row" in _ACCELERATOR_SUPPORTED + assert "mutate_row" in _ACCELERATOR_SUPPORTED From 5e1d30dcd1758320eb8cd0d592edb005485d5328 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 14 Oct 2026 01:20:03 +0000 Subject: [PATCH 08/26] review(bigtable): address PR #18686 reviewer comments - Remove _AcceleratorUnverified alias (no longer needed) - Remove duplicate User-Agent flags block from _init_accelerator_config; keep only the early block so the flag is forwarded even when the daemon is blocked - Move _AcceleratorIdentityError handling into _start_accelerator: close daemon, warn 'could not verify daemon identity', return without re-raising - Remove test_binary_path_param_bypasses_lookup (binary_path is a test-only knob already exercised implicitly) - Add test_close_clears_all_state to verify _proc/_tempdir/_uds_path/_log_path/_auth_secret are all None after close() - Add --app-profile to TestCliFlags.test_cli_flags_appended_to_argv - Add assert not t._accelerator_breaker.bypass() to test_non_grpc_exception_reraises in both TestMutateRowViaAccelerator and TestReadRowViaAccelerator Change-Id: Iff4489acc63a6e71211939c23db04ff1a5df10ba review(bigtable): address PR #18686 reviewer comments - Remove _AcceleratorUnverified alias (no longer needed) - Remove duplicate User-Agent flags block from _init_accelerator_config; keep only the early block so the flag is forwarded even when the daemon is blocked - Add local import google.cloud.bigtable inside _init_accelerator_config to avoid potential circular import at module level - Move _AcceleratorIdentityError handling into _start_accelerator: close daemon, warn 'could not verify daemon identity', return without re-raising - Remove test_binary_path_param_bypasses_lookup (binary_path is a test-only knob already exercised implicitly) - Add --app-profile to TestCliFlags.test_cli_flags_appended_to_argv - Add assert not t._accelerator_breaker.bypass() to test_non_grpc_exception_reraises in both TestMutateRowViaAccelerator and TestReadRowViaAccelerator Change-Id: I43b354ad6517a2c0b9f6a2294cdd2d1639d7fe8c test(bigtable): add TestCleanupTempdir to verify _cleanup_tempdir clears all state Change-Id: I6d6ec51cc039c78d884451502b368862fe8d4ae3 --- .../cloud/bigtable/data/_async/client.py | 29 ++++------ .../bigtable/data/_sync_autogen/client.py | 24 ++++---- .../unit/data/test_accelerator_daemon.py | 58 ++++++++++++++++--- .../unit/data/test_accelerator_dispatch.py | 14 +++++ 4 files changed, 86 insertions(+), 39 deletions(-) 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 cbaf1f187593..495e46fc796a 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 @@ -200,11 +200,6 @@ class _AcceleratorIdentityError(RuntimeError): this client's — signals the caller to fall back to the native client.""" -# Alias used in _start_accelerator's except clause; identity and "unverified" -# refer to the same failure mode. -_AcceleratorUnverified = _AcceleratorIdentityError - - def _normalize_scopes(scopes: Any) -> frozenset[str]: """Normalize a scope value to a comparable, order-insensitive set. @@ -427,6 +422,8 @@ def _init_accelerator_config( # User-Agent prefix always forwarded so the daemon can construct the full # "python-v{ver}-go-accelerator-v{ver}" string even when the daemon is # ultimately blocked (the flag is harmless when the daemon never starts). + import google.cloud.bigtable # local import avoids circular dependency + self._accelerator_flags += [ "--caller-user-agent", f"python-v{google.cloud.bigtable.__version__}", @@ -470,14 +467,6 @@ def _init_accelerator_config( credentials_file ) - # Identify this client to the daemon so it can construct the full - # User-Agent prefix (e.g. "python-v3.0.0a0-go-accelerator-v3.x.x"). - # The daemon appends its own "-go-accelerator-vX.Y.Z" suffix. - self._accelerator_flags += [ - "--caller-user-agent", - f"python-v{google.cloud.bigtable.__version__}", - ] - def _resolve_principal(self) -> str | None: """Resolve this client's identity to a principal. @@ -1423,9 +1412,6 @@ def _maybe_start_accelerator(self, *, explicit: bool) -> None: except _AcceleratorIdentityError as exc: if explicit: raise - # Callers using gcloud user ADC (_resolve_principal returns None) - # consistently land here; fail-closed to the native client is - # intentional — user ADC is almost never what the daemon resolved. warnings.warn( f"Accelerator disabled: {exc}", RuntimeWarning, @@ -1472,9 +1458,16 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) - except _AcceleratorIdentityError: + except _AcceleratorIdentityError as exc: server.close() - raise + warnings.warn( + f"Accelerator disabled: could not verify daemon identity: {exc}", + RuntimeWarning, + stacklevel=2, + ) + self._accelerator_daemon = None + self._accelerator_client = None + return except BaseException: server.close() raise 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 a714449339c8..b470ab06abda 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 @@ -148,9 +148,6 @@ class _AcceleratorIdentityError(RuntimeError): this client's — signals the caller to fall back to the native client.""" -_AcceleratorUnverified = _AcceleratorIdentityError - - def _normalize_scopes(scopes: Any) -> frozenset[str]: """Normalize a scope value to a comparable, order-insensitive set. @@ -320,6 +317,8 @@ def _init_accelerator_config( client_options, "client_cert_source", None ): self._accelerator_blocked_reason = "client_cert_source (mTLS) cannot be forwarded to the accelerator daemon" + import google.cloud.bigtable # local import avoids circular dependency + self._accelerator_flags += [ "--caller-user-agent", f"python-v{google.cloud.bigtable.__version__}", @@ -347,11 +346,6 @@ def _init_accelerator_config( self._accelerator_env["GOOGLE_APPLICATION_CREDENTIALS"] = ( credentials_file ) - self._accelerator_flags += [ - "--caller-user-agent", - f"python-v{google.cloud.bigtable.__version__}", - ] - def _resolve_principal(self) -> str | None: """Resolve this client's identity to a principal. @@ -1113,9 +1107,6 @@ def _maybe_start_accelerator(self, *, explicit: bool) -> None: except _AcceleratorIdentityError as exc: if explicit: raise - # Callers using gcloud user ADC (_resolve_principal returns None) - # consistently land here; fail-closed to the native client is - # intentional — user ADC is almost never what the daemon resolved. warnings.warn( f"Accelerator disabled: {exc}", RuntimeWarning, @@ -1155,9 +1146,16 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) - except _AcceleratorIdentityError: + except _AcceleratorIdentityError as exc: server.close() - raise + warnings.warn( + f"Accelerator disabled: could not verify daemon identity: {exc}", + RuntimeWarning, + stacklevel=2, + ) + self._accelerator_daemon = None + self._accelerator_client = None + return except BaseException: server.close() raise diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index 5792bc89dfa2..b713bb04c3e7 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -38,6 +38,12 @@ def _make_daemon(tmp_path, **kwargs): return AcceleratorDaemon(**kwargs) +# Captured at import time, before the autouse fixture patches _cleanup_tempdir +# to a no-op. Tests that need to exercise the real implementation call this +# directly as _real_cleanup_tempdir(daemon_instance). +_real_cleanup_tempdir = AcceleratorDaemon._cleanup_tempdir + + class TestExtraEnv: def test_default_env_none(self, tmp_path, monkeypatch): captured = {} @@ -248,13 +254,6 @@ def test_missing_binary_raises_file_not_found(self, tmp_path): with pytest.raises(FileNotFoundError, match="Accelerator binary not found"): _resolve_binary_path() - def test_binary_path_param_bypasses_lookup(self, tmp_path): - """Passing binary_path= directly skips _resolve_binary_path entirely.""" - fake_bin = tmp_path / "my-binary" - fake_bin.touch() - d = AcceleratorDaemon(binary_path=str(fake_bin)) - assert d._binary_path == str(fake_bin) - # --------------------------------------------------------------------------- # start() failure modes @@ -412,6 +411,40 @@ def test_close_clears_proc_reference(self, tmp_path): assert d._proc is None +# --------------------------------------------------------------------------- +# _cleanup_tempdir: state clearing +# --------------------------------------------------------------------------- + +class TestCleanupTempdir: + def test_clears_all_state_fields(self, tmp_path): + d = _make_daemon(tmp_path) + d._tempdir = str(tmp_path / "accel-tmp") + d._uds_path = str(tmp_path / "sock") + d._log_path = str(tmp_path / "daemon.log") + d._auth_secret = "super-secret" + + _real_cleanup_tempdir(d) + + assert d._tempdir is None + assert d._uds_path is None + assert d._log_path is None + assert d._auth_secret is None + + def test_removes_tempdir_from_disk(self, tmp_path): + d = _make_daemon(tmp_path) + td = tmp_path / "accel-tmp" + td.mkdir() + d._tempdir = str(td) + + _real_cleanup_tempdir(d) + + assert not td.exists() + + def test_noop_when_tempdir_is_none(self, tmp_path): + d = _make_daemon(tmp_path) + _real_cleanup_tempdir(d) # should not raise + + # --------------------------------------------------------------------------- # Context manager # --------------------------------------------------------------------------- @@ -537,7 +570,14 @@ def fake_popen(argv, **kwargs): raise _StopStart() monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) - d = _make_daemon(tmp_path, cli_flags=["--project", "proj", "--instance", "inst"]) + d = _make_daemon( + tmp_path, + cli_flags=[ + "--project", "proj", + "--instance", "inst", + "--app-profile", "prof", + ], + ) with pytest.raises(_StopStart): d.start() @@ -547,6 +587,8 @@ def fake_popen(argv, **kwargs): assert argv[argv.index("--project") + 1] == "proj" assert "--instance" in argv assert argv[argv.index("--instance") + 1] == "inst" + assert "--app-profile" in argv + assert argv[argv.index("--app-profile") + 1] == "prof" # --------------------------------------------------------------------------- diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py index b73245e17a9a..36862aa36e56 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py @@ -230,6 +230,7 @@ def test_non_grpc_exception_reraises(self): req = MutateRowRequest() with pytest.raises(RuntimeError, match="unexpected"): self._run(t._mutate_row_via_accelerator(req, 30.0)) + assert not t._accelerator_breaker.bypass() # --------------------------------------------------------------------------- @@ -291,6 +292,19 @@ def test_permission_denied_translates_and_raises(self): self._run(t._read_row_via_accelerator(query, 30.0, 20.0)) assert not t._accelerator_breaker.bypass() + def test_non_grpc_exception_reraises(self): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + client_mock = mock.AsyncMock() + client_mock.read_rows.side_effect = RuntimeError("unexpected") + daemon = mock.Mock(is_running=True) + t = _bare_target(accel_client=client_mock, accel_daemon=daemon) + + query = ReadRowsQuery(row_keys=b"key1", limit=1) + with pytest.raises(RuntimeError, match="unexpected"): + self._run(t._read_row_via_accelerator(query, 30.0, 20.0)) + assert not t._accelerator_breaker.bypass() + # --------------------------------------------------------------------------- # read_row: accelerator dispatch and fallback to native From 43223b2231f656bc522786edfb81abe63d0b0b66 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 14 Oct 2026 15:00:04 +0000 Subject: [PATCH 09/26] fix(bigtable): propagate identity error from _start_accelerator; close auth session Change-Id: I78f226004fdffe98d83f7d24b04557a24e9fc7a7 --- .../cloud/bigtable/data/_async/client.py | 25 ++++++--------- .../bigtable/data/_sync_autogen/client.py | 25 ++++++--------- .../unit/data/test_accelerator_dispatch.py | 12 +++---- .../unit/data/test_accelerator_enablement.py | 32 +++++++++++++++++++ 4 files changed, 55 insertions(+), 39 deletions(-) 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 495e46fc796a..51aabe3eed0d 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 @@ -494,14 +494,17 @@ def _resolve_principal(self) -> str | None: # Request.__call__ passes timeout=120 to session.request, so # inject our bound at the session level to always override it. _auth_req = google_auth_requests.Request() - _http = _auth_req.session.request + try: + _http = _auth_req.session.request - def _with_timeout(*args, **kwargs): - kwargs["timeout"] = _PRINCIPAL_RESOLVE_TIMEOUT - return _http(*args, **kwargs) + def _with_timeout(*args, **kwargs): + kwargs["timeout"] = _PRINCIPAL_RESOLVE_TIMEOUT + return _http(*args, **kwargs) - _auth_req.session.request = _with_timeout - creds.refresh(_auth_req) + _auth_req.session.request = _with_timeout + creds.refresh(_auth_req) + finally: + _auth_req.session.close() except Exception as exc: _LOGGER.debug("Compute Engine credential refresh failed: %s", exc) email = getattr(creds, "service_account_email", None) @@ -1458,16 +1461,6 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) - except _AcceleratorIdentityError as exc: - server.close() - warnings.warn( - f"Accelerator disabled: could not verify daemon identity: {exc}", - RuntimeWarning, - stacklevel=2, - ) - self._accelerator_daemon = None - self._accelerator_client = None - return except BaseException: server.close() raise 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 b470ab06abda..296017b6ec25 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 @@ -367,14 +367,17 @@ def _resolve_principal(self) -> str | None: ): try: _auth_req = google_auth_requests.Request() - _http = _auth_req.session.request + try: + _http = _auth_req.session.request - def _with_timeout(*args, **kwargs): - kwargs["timeout"] = _PRINCIPAL_RESOLVE_TIMEOUT - return _http(*args, **kwargs) + def _with_timeout(*args, **kwargs): + kwargs["timeout"] = _PRINCIPAL_RESOLVE_TIMEOUT + return _http(*args, **kwargs) - _auth_req.session.request = _with_timeout - creds.refresh(_auth_req) + _auth_req.session.request = _with_timeout + creds.refresh(_auth_req) + finally: + _auth_req.session.close() except Exception as exc: _LOGGER.debug("Compute Engine credential refresh failed: %s", exc) email = getattr(creds, "service_account_email", None) @@ -1146,16 +1149,6 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) - except _AcceleratorIdentityError as exc: - server.close() - warnings.warn( - f"Accelerator disabled: could not verify daemon identity: {exc}", - RuntimeWarning, - stacklevel=2, - ) - self._accelerator_daemon = None - self._accelerator_client = None - return except BaseException: server.close() raise diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py index 36862aa36e56..ce71d84d938c 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py @@ -539,9 +539,10 @@ def test_daemon_start_failure_calls_close_and_reraises(self): assert t._accelerator_daemon is None assert t._accelerator_client is None - def test_identity_verify_failure_closes_daemon_and_warns(self): - """Identity mismatches are now caught in _start_accelerator: daemon is - closed, a RuntimeWarning is issued, and the method returns (no re-raise). + def test_identity_verify_failure_closes_daemon_and_propagates(self): + """Identity mismatches propagate from _start_accelerator so + _maybe_start_accelerator can decide whether to warn (auto) or raise + (explicit). The daemon is closed before the error surfaces. """ from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon @@ -560,13 +561,10 @@ def test_identity_verify_failure_closes_daemon_and_warns(self): t, "_verify_daemon_identity", side_effect=_AcceleratorIdentityError("mismatch"), - ), pytest.warns(RuntimeWarning, match="could not verify"): + ), pytest.raises(_AcceleratorIdentityError, match="mismatch"): t._start_accelerator() fake_daemon.close.assert_called_once() - # After a warn-and-return the target has no daemon or client set. - assert t._accelerator_daemon is None - assert t._accelerator_client is None def test_passes_project_instance_flags(self): from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py index a9b88e171cff..30aca06e0aaa 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py @@ -152,3 +152,35 @@ def _start(): t = _bare_target(emulator=None, start=_start) with pytest.raises(_AcceleratorIdentityError, match="principal mismatch"): t._maybe_start_accelerator(explicit=True) + + def test_start_accelerator_propagates_identity_error(self, monkeypatch): + """_start_accelerator must not catch _AcceleratorIdentityError internally. + + Previously, _start_accelerator absorbed the error and returned normally, + making _maybe_start_accelerator's explicit=True re-raise unreachable. + """ + import google.cloud.bigtable.data._async.client as _mod + + mock_server = unittest.mock.MagicMock() + mock_server.start.return_value = None + + t = object.__new__(_ConcreteTarget) + t.client = SimpleNamespace( + project="p", + _accelerator_flags=[], + _accelerator_env={}, + ) + t.instance_id = "i" + t.app_profile_id = None + t._accelerator_daemon = None + t._accelerator_client = None + t._verify_daemon_identity = unittest.mock.MagicMock( + side_effect=_AcceleratorIdentityError("identity mismatch") + ) + + monkeypatch.setattr(_mod, "AcceleratorDaemon", lambda **_: mock_server) + + with pytest.raises(_AcceleratorIdentityError, match="identity mismatch"): + t._start_accelerator() + + mock_server.close.assert_called_once() From 60a2601481ef961d16f15252e0b02406201af0d1 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 14 Oct 2026 15:35:15 +0000 Subject: [PATCH 10/26] fix(bigtable): move _accelerator_daemon assignment inside try block to prevent orphaned daemon on BaseException Change-Id: I6f4c18f5b03da26efa9b70ed8c6b4d20f9d4a8df --- .../google/cloud/bigtable/data/_async/client.py | 2 +- .../google/cloud/bigtable/data/_sync_autogen/client.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 51aabe3eed0d..a7aa03b1ea31 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 @@ -1461,10 +1461,10 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) + self._accelerator_daemon = server except BaseException: server.close() raise - self._accelerator_daemon = server def _use_accelerator(self, method_name: str) -> bool: """Whether this call should be routed through the accelerator daemon. 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 296017b6ec25..711021230127 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 @@ -1149,10 +1149,10 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) + self._accelerator_daemon = server except BaseException: server.close() raise - self._accelerator_daemon = server def _use_accelerator(self, method_name: str) -> bool: """Whether this call should be routed through the accelerator daemon.""" From 988c0fb848b751a82a82e66b954826e965a38efc Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 12 Aug 2026 02:51:41 +0000 Subject: [PATCH 11/26] feat(bigtable): sweep stale accelerator daemon tempdirs at startup Change-Id: Iad44827bd86e405524b5aad787e881e2cbf30d14 --- .../bigtable/data/_accelerator/_daemon.py | 85 ++++++++++++++++--- .../unit/data/test_accelerator_daemon.py | 69 +++++++++++++++ 2 files changed, 143 insertions(+), 11 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index 524f2ea90648..c3c56bd8ef83 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -23,6 +23,7 @@ from __future__ import annotations +import glob import json import os import secrets @@ -49,6 +50,15 @@ # its own locally-resolved identity before routing any RPC. _IDENTITY_FILENAME = "identity.json" +# Prefix for the per-daemon tempdir (created by ``tempfile.mkdtemp``) that holds +# the socket, identity, and log. Used both when creating our own and to scope +# the startup sweep of leftovers. +_TEMPDIR_PREFIX = "bt-accel-" + +# Names of the files a live daemon keeps inside its tempdir. +_SOCKET_FILENAME = "sock" +_LOG_FILENAME = "daemon.log" + # How long to wait for the daemon to start listening on its UDS before giving # up at startup. Cold, constrained environments (e.g. Cloud Run's gVisor # sandbox) spend ~10s in Go runtime init + ADC token fetch from the metadata @@ -63,6 +73,11 @@ _SIGTERM_GRACE_SECONDS = 2.0 _SIGKILL_GRACE_SECONDS = 2.0 +# A freshly-created tempdir may not have bound its socket yet. Never sweep one +# younger than this, so a daemon another client is mid-way through starting is +# never mistaken for a stale leftover and reaped out from under it. +_STALE_SWEEP_MIN_AGE_SECONDS = 60.0 + def _resolve_binary_path() -> str: """Return the path of the bundled daemon binary, raising if it is absent.""" @@ -74,6 +89,7 @@ def _resolve_binary_path() -> str: return bundled + def _resolve_startup_timeout() -> float: """Resolve the daemon startup timeout in seconds. @@ -98,6 +114,53 @@ def _resolve_startup_timeout() -> float: return value +def _socket_is_live(uds_path: str) -> bool: + """Return True if something is currently accepting connections on the UDS.""" + if not os.path.exists(uds_path): + return False + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: + probe.settimeout(0.25) + try: + probe.connect(uds_path) + return True + except OSError: + return False + + +def _sweep_stale_tempdirs() -> None: + """Best-effort reaping of leftover daemon tempdirs at startup. + + A clean ``close()`` removes its own tempdir, but a crash or ``SIGKILL`` + leaves the socket and log behind. We reap those on the next start. + + Rather than treating any ``bt-accel-*`` directory as ours, we key off the + ``daemon.log`` marker file we write inside each one: a directory is a + candidate only if it actually contains that log. This avoids deleting an + unrelated directory that merely shares the tempdir prefix. A candidate is + removed only when its socket no longer accepts connections *and* it is old + enough that a daemon still mid-startup would already have bound (or died). + + All failures are swallowed — sweeping is opportunistic and must never break + startup. + """ + root = tempfile.gettempdir() + try: + logs = glob.glob(os.path.join(root, _TEMPDIR_PREFIX + "*", _LOG_FILENAME)) + except OSError: + return + now = time.time() + for log_path in logs: + path = os.path.dirname(log_path) + try: + if now - os.path.getmtime(log_path) < _STALE_SWEEP_MIN_AGE_SECONDS: + continue + if _socket_is_live(os.path.join(path, _SOCKET_FILENAME)): + continue + except OSError: + continue + shutil.rmtree(path, ignore_errors=True) + + class AcceleratorDaemon: """Manages the Go accelerator daemon subprocess. @@ -198,12 +261,18 @@ 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") + + # Reap tempdirs leaked by daemons from prior runs that crashed or were + # killed before close() could clean up after them. + _sweep_stale_tempdirs() + self._tempdir = tempfile.mkdtemp(prefix=_TEMPDIR_PREFIX) + self._uds_path = os.path.join(self._tempdir, _SOCKET_FILENAME) + # Redirect the daemon's stdout/stderr to a log file rather than # subprocess.PIPE. Nothing drains those pipes for the daemon's # lifetime, so a PIPE's fixed OS buffer would eventually fill and @@ -212,7 +281,7 @@ def start(self) -> None: # tempdir so it's cleaned up with everything else in close(). stdin # stays a PIPE — closing it is how close() signals the daemon to shut # down. - self._log_path = os.path.join(self._tempdir, "daemon.log") + self._log_path = os.path.join(self._tempdir, _LOG_FILENAME) # 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. @@ -355,14 +424,8 @@ def _wait_until_ready(self, timeout: float) -> None: "Accelerator daemon exited during startup " f"(exit code {exit_code}). log: {log_tail!r}" ) - if os.path.exists(self._uds_path): - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: - probe.settimeout(0.25) - try: - probe.connect(self._uds_path) - return - except (ConnectionRefusedError, FileNotFoundError, OSError): - pass + if _socket_is_live(self._uds_path): + return time.sleep(0.05) log_tail = self._read_log_tail() raise RuntimeError( diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index b713bb04c3e7..d5defcfce48f 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -15,6 +15,7 @@ import json import os import re +import time import unittest.mock import pytest @@ -663,3 +664,71 @@ def test_daemon_constructor_invalid_env_var_raises(self, tmp_path, monkeypatch): monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "xyz") with pytest.raises(ValueError, match="not a valid number of seconds"): _make_daemon(tmp_path) + + +class TestSweepStaleTempdirs: + """_sweep_stale_tempdirs reaps leftover daemon tempdirs, keyed on the + daemon.log marker, without touching live or too-young ones.""" + + @staticmethod + def _make_leftover(root, name, *, with_log=True, age_seconds=0.0): + d = root / name + d.mkdir() + if with_log: + log = d / _daemon._LOG_FILENAME + log.write_bytes(b"hi") + if age_seconds: + past = time.time() - age_seconds + os.utime(log, (past, past)) + return d + + @pytest.fixture + def tmp_root(self, tmp_path, monkeypatch): + # Point the sweep's tempdir root at an isolated dir under pytest's tmp. + root = tmp_path / "tmproot" + root.mkdir() + monkeypatch.setattr(_daemon.tempfile, "gettempdir", lambda: str(root)) + return root + + def test_removes_stale_dir(self, tmp_root): + d = self._make_leftover(tmp_root, "bt-accel-dead", age_seconds=120) + _daemon._sweep_stale_tempdirs() + assert not d.exists() + + def test_keeps_young_dir(self, tmp_root): + # Younger than the grace window: a daemon may still be mid-startup. + d = self._make_leftover(tmp_root, "bt-accel-young", age_seconds=0) + _daemon._sweep_stale_tempdirs() + assert d.exists() + + def test_ignores_dir_without_log_marker(self, tmp_root): + # Prefix matches but there's no daemon.log, so it isn't one of ours. + d = tmp_root / "bt-accel-foreign" + d.mkdir() + (d / "something-else").write_bytes(b"x") + _daemon._sweep_stale_tempdirs() + assert d.exists() + + def test_keeps_dir_with_live_socket(self, tmp_root, monkeypatch): + d = self._make_leftover(tmp_root, "bt-accel-live", age_seconds=120) + expected = str(d / _daemon._SOCKET_FILENAME) + # Patch _socket_is_live to avoid binding a real socket: macOS limits + # UDS paths to 104 chars and pytest's tmp_path is frequently longer. + monkeypatch.setattr(_daemon, "_socket_is_live", lambda p: p == expected) + _daemon._sweep_stale_tempdirs() + assert d.exists() + + def test_start_invokes_sweep(self, tmp_path, monkeypatch): + called = [] + monkeypatch.setattr( + _daemon, "_sweep_stale_tempdirs", lambda: called.append(True) + ) + + def boom(*args, **kwargs): + raise OSError("refusing to spawn in test") + + monkeypatch.setattr(_daemon.subprocess, "Popen", boom) + daemon = _make_daemon(tmp_path) + with pytest.raises(RuntimeError): + daemon.start() + assert called == [True] From 02e28b687236a8f5ccb12fa2b24fcfd53978fb91 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 14 Oct 2026 15:34:33 +0000 Subject: [PATCH 12/26] fix(bigtable): widen sweep grace window to account for configured startup timeout Change-Id: I44684facdfe1f9a1b8807b78bbb4166f8c9ff26b --- .../google/cloud/bigtable/data/_accelerator/_daemon.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index c3c56bd8ef83..053f8610aa73 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -149,10 +149,15 @@ def _sweep_stale_tempdirs() -> None: except OSError: return now = time.time() + # A daemon still mid-startup must not be swept. The log file's mtime is + # stamped before Popen, so the grace window must exceed the maximum time + # start() can spend waiting for the socket (the configured startup timeout) + # plus a fixed buffer. + grace = _resolve_startup_timeout() + _STALE_SWEEP_MIN_AGE_SECONDS for log_path in logs: path = os.path.dirname(log_path) try: - if now - os.path.getmtime(log_path) < _STALE_SWEEP_MIN_AGE_SECONDS: + if now - os.path.getmtime(log_path) < grace: continue if _socket_is_live(os.path.join(path, _SOCKET_FILENAME)): continue From 29e6ba04c55b77847f660da3d044d3fbde572300 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 14 Oct 2026 15:42:00 +0000 Subject: [PATCH 13/26] fix(bigtable): run tempdir sweep on background thread; add AF_UNIX guard and Windows skipif Change-Id: Iba3a175c070fb5f12db54136096f2792bb7f3b7f --- .../bigtable/data/_accelerator/_daemon.py | 9 +++++++-- .../tests/unit/data/test_accelerator_daemon.py | 18 +++++++++++++----- 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index 053f8610aa73..bd4fdae93c1e 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -31,6 +31,7 @@ import socket import subprocess import tempfile +import threading import time from typing import Any, Mapping, Sequence @@ -143,6 +144,8 @@ def _sweep_stale_tempdirs() -> None: All failures are swallowed — sweeping is opportunistic and must never break startup. """ + if not hasattr(socket, "AF_UNIX"): + return root = tempfile.gettempdir() try: logs = glob.glob(os.path.join(root, _TEMPDIR_PREFIX + "*", _LOG_FILENAME)) @@ -273,8 +276,10 @@ def start(self) -> None: ) # Reap tempdirs leaked by daemons from prior runs that crashed or were - # killed before close() could clean up after them. - _sweep_stale_tempdirs() + # killed before close() could clean up after them. Fire-and-forget on a + # daemon thread so it never blocks startup (sweep probes dead sockets + # with a 0.25 s timeout each). + threading.Thread(target=_sweep_stale_tempdirs, daemon=True).start() self._tempdir = tempfile.mkdtemp(prefix=_TEMPDIR_PREFIX) self._uds_path = os.path.join(self._tempdir, _SOCKET_FILENAME) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index d5defcfce48f..665c74b27b70 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -666,6 +666,10 @@ def test_daemon_constructor_invalid_env_var_raises(self, tmp_path, monkeypatch): _make_daemon(tmp_path) +@pytest.mark.skipif( + not hasattr(_daemon.socket, "AF_UNIX"), + reason="Unix domain sockets (AF_UNIX) are not supported on this platform", +) class TestSweepStaleTempdirs: """_sweep_stale_tempdirs reaps leftover daemon tempdirs, keyed on the daemon.log marker, without touching live or too-young ones.""" @@ -719,10 +723,14 @@ def test_keeps_dir_with_live_socket(self, tmp_root, monkeypatch): assert d.exists() def test_start_invokes_sweep(self, tmp_path, monkeypatch): - called = [] - monkeypatch.setattr( - _daemon, "_sweep_stale_tempdirs", lambda: called.append(True) - ) + import threading + + done = threading.Event() + + def _fake_sweep(): + done.set() + + monkeypatch.setattr(_daemon, "_sweep_stale_tempdirs", _fake_sweep) def boom(*args, **kwargs): raise OSError("refusing to spawn in test") @@ -731,4 +739,4 @@ def boom(*args, **kwargs): daemon = _make_daemon(tmp_path) with pytest.raises(RuntimeError): daemon.start() - assert called == [True] + assert done.wait(timeout=2.0), "sweep thread did not run within 2 s" From 18b9a4ef200be065707273c03e9ddb971ff3aada Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 14 Oct 2026 15:45:41 +0000 Subject: [PATCH 14/26] fix(bigtable): write per-daemon startup deadline so sweep uses correct grace window regardless of sweeper's env Change-Id: Idf216c34aa7924fd61f258c997c71c426a2cbe34 --- .../bigtable/data/_accelerator/_daemon.py | 28 +++++++++++++---- .../unit/data/test_accelerator_daemon.py | 31 ++++++++++++++++--- 2 files changed, 49 insertions(+), 10 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index bd4fdae93c1e..fc08e7fc2c15 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -59,6 +59,11 @@ # Names of the files a live daemon keeps inside its tempdir. _SOCKET_FILENAME = "sock" _LOG_FILENAME = "daemon.log" +# Absolute POSIX timestamp written by start() so the sweep knows when startup +# for *this specific* daemon was allowed to finish. Each process may have a +# different BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT, so the sweep cannot rely on +# its own env value when judging another process's tempdir. +_STARTUP_DEADLINE_FILENAME = "startup_deadline" # How long to wait for the daemon to start listening on its UDS before giving # up at startup. Cold, constrained environments (e.g. Cloud Run's gVisor @@ -152,15 +157,21 @@ def _sweep_stale_tempdirs() -> None: except OSError: return now = time.time() - # A daemon still mid-startup must not be swept. The log file's mtime is - # stamped before Popen, so the grace window must exceed the maximum time - # start() can spend waiting for the socket (the configured startup timeout) - # plus a fixed buffer. - grace = _resolve_startup_timeout() + _STALE_SWEEP_MIN_AGE_SECONDS for log_path in logs: path = os.path.dirname(log_path) try: - if now - os.path.getmtime(log_path) < grace: + # Each tempdir records its own startup deadline so this sweep is + # correct even when the sweeping process has a different + # BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT than the process that + # started the daemon being evaluated. + deadline_path = os.path.join(path, _STARTUP_DEADLINE_FILENAME) + try: + deadline = float(open(deadline_path).read()) + except (OSError, ValueError): + # Older daemon without a deadline file: fall back to a simple + # age heuristic relative to the log's mtime. + deadline = os.path.getmtime(log_path) + _STALE_SWEEP_MIN_AGE_SECONDS + if now < deadline + _STALE_SWEEP_MIN_AGE_SECONDS: continue if _socket_is_live(os.path.join(path, _SOCKET_FILENAME)): continue @@ -282,6 +293,11 @@ def start(self) -> None: threading.Thread(target=_sweep_stale_tempdirs, daemon=True).start() self._tempdir = tempfile.mkdtemp(prefix=_TEMPDIR_PREFIX) self._uds_path = os.path.join(self._tempdir, _SOCKET_FILENAME) + # Write the absolute time by which startup must complete. The sweep + # reads this per-directory so it uses the right deadline regardless of + # what BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT the sweeping process has. + with open(os.path.join(self._tempdir, _STARTUP_DEADLINE_FILENAME), "w") as _f: + _f.write(repr(time.time() + self._startup_timeout)) # Redirect the daemon's stdout/stderr to a log file rather than # subprocess.PIPE. Nothing drains those pipes for the daemon's diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index 665c74b27b70..df93c3b03694 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -675,7 +675,13 @@ class TestSweepStaleTempdirs: daemon.log marker, without touching live or too-young ones.""" @staticmethod - def _make_leftover(root, name, *, with_log=True, age_seconds=0.0): + def _make_leftover(root, name, *, with_log=True, age_seconds=0.0, deadline_offset=None): + """Create a fake leftover daemon tempdir. + + deadline_offset: seconds relative to now for the startup deadline file. + Negative = deadline already passed (stale); positive = still in future + (mid-startup). None = omit the file (tests the fallback heuristic). + """ d = root / name d.mkdir() if with_log: @@ -684,6 +690,9 @@ def _make_leftover(root, name, *, with_log=True, age_seconds=0.0): if age_seconds: past = time.time() - age_seconds os.utime(log, (past, past)) + if deadline_offset is not None: + dl = d / _daemon._STARTUP_DEADLINE_FILENAME + dl.write_text(repr(time.time() + deadline_offset)) return d @pytest.fixture @@ -695,16 +704,28 @@ def tmp_root(self, tmp_path, monkeypatch): return root def test_removes_stale_dir(self, tmp_root): - d = self._make_leftover(tmp_root, "bt-accel-dead", age_seconds=120) + # Expired deadline + old log + no live socket → should be reaped. + d = self._make_leftover( + tmp_root, "bt-accel-dead", age_seconds=120, deadline_offset=-70 + ) _daemon._sweep_stale_tempdirs() assert not d.exists() def test_keeps_young_dir(self, tmp_root): - # Younger than the grace window: a daemon may still be mid-startup. + # No deadline file: fallback uses log mtime + buffer; log is fresh. d = self._make_leftover(tmp_root, "bt-accel-young", age_seconds=0) _daemon._sweep_stale_tempdirs() assert d.exists() + def test_keeps_dir_with_future_deadline(self, tmp_root): + # Log is old (simulating a long-timeout override) but deadline hasn't + # passed yet — this is the case a per-directory deadline file fixes. + d = self._make_leftover( + tmp_root, "bt-accel-slow", age_seconds=120, deadline_offset=+200 + ) + _daemon._sweep_stale_tempdirs() + assert d.exists() + def test_ignores_dir_without_log_marker(self, tmp_root): # Prefix matches but there's no daemon.log, so it isn't one of ours. d = tmp_root / "bt-accel-foreign" @@ -714,7 +735,9 @@ def test_ignores_dir_without_log_marker(self, tmp_root): assert d.exists() def test_keeps_dir_with_live_socket(self, tmp_root, monkeypatch): - d = self._make_leftover(tmp_root, "bt-accel-live", age_seconds=120) + d = self._make_leftover( + tmp_root, "bt-accel-live", age_seconds=120, deadline_offset=-70 + ) expected = str(d / _daemon._SOCKET_FILENAME) # Patch _socket_is_live to avoid binding a real socket: macOS limits # UDS paths to 104 chars and pytest's tmp_path is frequently longer. From 0938dd830006a162712d1bb62df30b7cf5845d32 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 14 Oct 2026 15:47:09 +0000 Subject: [PATCH 15/26] fix(bigtable): double-probe socket before sweeping to guard against transient _socket_is_live false negatives Change-Id: If3a0174967e750f04c26b49ccf3e0a4faf4764e0 --- .../google/cloud/bigtable/data/_accelerator/_daemon.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index fc08e7fc2c15..2a32b38de62c 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -173,7 +173,15 @@ def _sweep_stale_tempdirs() -> None: deadline = os.path.getmtime(log_path) + _STALE_SWEEP_MIN_AGE_SECONDS if now < deadline + _STALE_SWEEP_MIN_AGE_SECONDS: continue - if _socket_is_live(os.path.join(path, _SOCKET_FILENAME)): + sock_path = os.path.join(path, _SOCKET_FILENAME) + # Probe twice with a short gap: a transient OS error (ENOBUFS, + # EAGAIN) on a single probe could incorrectly signal a dead socket + # for a daemon that is actually running fine. Two independent + # failures are effectively impossible for a healthy local UDS. + if _socket_is_live(sock_path): + continue + time.sleep(0.1) + if _socket_is_live(sock_path): continue except OSError: continue From 3f751ca73db099835a216abf9aaf3b9b4038993a Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 14 Oct 2026 15:50:16 +0000 Subject: [PATCH 16/26] fix(bigtable): use ECONNREFUSED/ENOENT to guard sweep; propagate transient OSErrors instead of double-probing Change-Id: Ie6e858f577904cda5a06684afdb9d3aff21dd5ae --- .../bigtable/data/_accelerator/_daemon.py | 19 ++++++++++--------- .../unit/data/test_accelerator_daemon.py | 15 +++++++++++---- 2 files changed, 21 insertions(+), 13 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index 2a32b38de62c..41afcb9d0117 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -174,15 +174,16 @@ def _sweep_stale_tempdirs() -> None: if now < deadline + _STALE_SWEEP_MIN_AGE_SECONDS: continue sock_path = os.path.join(path, _SOCKET_FILENAME) - # Probe twice with a short gap: a transient OS error (ENOBUFS, - # EAGAIN) on a single probe could incorrectly signal a dead socket - # for a daemon that is actually running fine. Two independent - # failures are effectively impossible for a healthy local UDS. - if _socket_is_live(sock_path): - continue - time.sleep(0.1) - if _socket_is_live(sock_path): - continue + if os.path.exists(sock_path): + try: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: + probe.settimeout(0.25) + probe.connect(sock_path) + continue # connected → daemon is alive + except (ConnectionRefusedError, FileNotFoundError): + pass # definitive: nothing is listening + # Any other OSError (ENOBUFS, EAGAIN, …) is transient. + # Fall through to the outer except to skip this directory. except OSError: continue shutil.rmtree(path, ignore_errors=True) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index df93c3b03694..99f35c24face 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -738,10 +738,17 @@ def test_keeps_dir_with_live_socket(self, tmp_root, monkeypatch): d = self._make_leftover( tmp_root, "bt-accel-live", age_seconds=120, deadline_offset=-70 ) - expected = str(d / _daemon._SOCKET_FILENAME) - # Patch _socket_is_live to avoid binding a real socket: macOS limits - # UDS paths to 104 chars and pytest's tmp_path is frequently longer. - monkeypatch.setattr(_daemon, "_socket_is_live", lambda p: p == expected) + sock_path = d / _daemon._SOCKET_FILENAME + sock_path.touch() # make os.path.exists() return True + + # Simulate a successful connect without binding a real UDS (path may + # exceed the 104-char macOS limit under pytest's tmp_path). + mock_sock = unittest.mock.MagicMock() + mock_sock.__enter__ = lambda s: s + mock_sock.__exit__ = unittest.mock.MagicMock(return_value=False) + mock_sock.connect.return_value = None # success → daemon alive + monkeypatch.setattr(_daemon.socket, "socket", lambda *a, **kw: mock_sock) + _daemon._sweep_stale_tempdirs() assert d.exists() From 7f1883f9bd4b67cd9b3baee5c0f674cd4f84a822 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 20 Oct 2026 14:02:55 +0000 Subject: [PATCH 17/26] fix(bigtable): guard tempdir sweep with parent PID liveness instead of socket probe Change-Id: Ifa69e49da685ec3e3e4cd51ec344fb78d5c8f417 --- .../bigtable/data/_accelerator/_daemon.py | 37 +++++++++------ .../unit/data/test_accelerator_daemon.py | 46 +++++++++++++------ 2 files changed, 55 insertions(+), 28 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index 41afcb9d0117..0ad102a84b44 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -64,6 +64,11 @@ # different BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT, so the sweep cannot rely on # its own env value when judging another process's tempdir. _STARTUP_DEADLINE_FILENAME = "startup_deadline" +# PID of the Python process that spawned the daemon, written at start(). The +# sweep sends signal 0 to this PID to check liveness: if the spawner is alive +# the daemon is alive; if the spawner is dead (ESRCH) the daemon's stdin is +# closed and it will exit (or has already exited). +_PID_FILENAME = "daemon.pid" # How long to wait for the daemon to start listening on its UDS before giving # up at startup. Cold, constrained environments (e.g. Cloud Run's gVisor @@ -143,8 +148,11 @@ def _sweep_stale_tempdirs() -> None: ``daemon.log`` marker file we write inside each one: a directory is a candidate only if it actually contains that log. This avoids deleting an unrelated directory that merely shares the tempdir prefix. A candidate is - removed only when its socket no longer accepts connections *and* it is old - enough that a daemon still mid-startup would already have bound (or died). + removed only when the spawning Python process (tracked via ``daemon.pid``) + is no longer alive *and* the tempdir is old enough that a daemon still + mid-startup would already have bound or died. ``os.kill(pid, 0)`` is + instant and kernel-guaranteed: the lock is implicitly released the moment + the process exits for any reason, including SIGKILL. All failures are swallowed — sweeping is opportunistic and must never break startup. @@ -173,17 +181,18 @@ def _sweep_stale_tempdirs() -> None: deadline = os.path.getmtime(log_path) + _STALE_SWEEP_MIN_AGE_SECONDS if now < deadline + _STALE_SWEEP_MIN_AGE_SECONDS: continue - sock_path = os.path.join(path, _SOCKET_FILENAME) - if os.path.exists(sock_path): - try: - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: - probe.settimeout(0.25) - probe.connect(sock_path) - continue # connected → daemon is alive - except (ConnectionRefusedError, FileNotFoundError): - pass # definitive: nothing is listening - # Any other OSError (ENOBUFS, EAGAIN, …) is transient. - # Fall through to the outer except to skip this directory. + pid_path = os.path.join(path, _PID_FILENAME) + try: + with open(pid_path) as _f: + pid = int(_f.read().strip()) + os.kill(pid, 0) + continue # spawning process alive → daemon alive → skip + except ProcessLookupError: + pass # ESRCH: spawner is dead → stdin closed → daemon exiting + except (OSError, ValueError): + # No PID file (tempdir predates this change) or EPERM (process + # exists but unreadable) → assume alive and skip. + continue except OSError: continue shutil.rmtree(path, ignore_errors=True) @@ -307,6 +316,8 @@ def start(self) -> None: # what BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT the sweeping process has. with open(os.path.join(self._tempdir, _STARTUP_DEADLINE_FILENAME), "w") as _f: _f.write(repr(time.time() + self._startup_timeout)) + with open(os.path.join(self._tempdir, _PID_FILENAME), "w") as _f: + _f.write(str(os.getpid())) # Redirect the daemon's stdout/stderr to a log file rather than # subprocess.PIPE. Nothing drains those pipes for the daemon's diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index 99f35c24face..526d662086b8 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -675,12 +675,22 @@ class TestSweepStaleTempdirs: daemon.log marker, without touching live or too-young ones.""" @staticmethod - def _make_leftover(root, name, *, with_log=True, age_seconds=0.0, deadline_offset=None): + def _make_leftover( + root, + name, + *, + with_log=True, + age_seconds=0.0, + deadline_offset=None, + pid=None, + ): """Create a fake leftover daemon tempdir. deadline_offset: seconds relative to now for the startup deadline file. Negative = deadline already passed (stale); positive = still in future (mid-startup). None = omit the file (tests the fallback heuristic). + pid: if given, write a daemon.pid file with this PID. Pass a dead PID + to simulate a crashed/exited spawner; pass os.getpid() for "alive". """ d = root / name d.mkdir() @@ -693,6 +703,8 @@ def _make_leftover(root, name, *, with_log=True, age_seconds=0.0, deadline_offse if deadline_offset is not None: dl = d / _daemon._STARTUP_DEADLINE_FILENAME dl.write_text(repr(time.time() + deadline_offset)) + if pid is not None: + (d / _daemon._PID_FILENAME).write_text(str(pid)) return d @pytest.fixture @@ -704,9 +716,14 @@ def tmp_root(self, tmp_path, monkeypatch): return root def test_removes_stale_dir(self, tmp_root): - # Expired deadline + old log + no live socket → should be reaped. + # Expired deadline + old log + dead spawner PID → should be reaped. + import subprocess, sys + + proc = subprocess.Popen([sys.executable, "-c", ""]) + dead_pid = proc.pid + proc.wait() d = self._make_leftover( - tmp_root, "bt-accel-dead", age_seconds=120, deadline_offset=-70 + tmp_root, "bt-accel-dead", age_seconds=120, deadline_offset=-70, pid=dead_pid ) _daemon._sweep_stale_tempdirs() assert not d.exists() @@ -734,21 +751,20 @@ def test_ignores_dir_without_log_marker(self, tmp_root): _daemon._sweep_stale_tempdirs() assert d.exists() - def test_keeps_dir_with_live_socket(self, tmp_root, monkeypatch): + def test_keeps_dir_with_live_pid(self, tmp_root): + # Past deadline but spawning process is still alive → must not be reaped. d = self._make_leftover( - tmp_root, "bt-accel-live", age_seconds=120, deadline_offset=-70 + tmp_root, "bt-accel-live", age_seconds=120, deadline_offset=-70, + pid=os.getpid(), ) - sock_path = d / _daemon._SOCKET_FILENAME - sock_path.touch() # make os.path.exists() return True - - # Simulate a successful connect without binding a real UDS (path may - # exceed the 104-char macOS limit under pytest's tmp_path). - mock_sock = unittest.mock.MagicMock() - mock_sock.__enter__ = lambda s: s - mock_sock.__exit__ = unittest.mock.MagicMock(return_value=False) - mock_sock.connect.return_value = None # success → daemon alive - monkeypatch.setattr(_daemon.socket, "socket", lambda *a, **kw: mock_sock) + _daemon._sweep_stale_tempdirs() + assert d.exists() + def test_keeps_dir_without_pid_file(self, tmp_root): + # No PID file (tempdir predates this change) → conservative: keep it. + d = self._make_leftover( + tmp_root, "bt-accel-nopid", age_seconds=120, deadline_offset=-70 + ) _daemon._sweep_stale_tempdirs() assert d.exists() From 26cff134eeb8208a148ebf93d945f48fee0b9000 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 20 Oct 2026 14:54:30 +0000 Subject: [PATCH 18/26] chore(bigtable): apply ruff import sort and formatting to accelerator sources and tests Change-Id: I978a0c1e846f3b95aa2fe5d15dabafec4f0576f9 --- .../bigtable/data/_accelerator/_daemon.py | 6 +- .../bigtable/data/_accelerator/_fallback.py | 3 +- .../cloud/bigtable/data/_async/client.py | 14 +- .../bigtable/data/_sync_autogen/client.py | 27 ++- .../unit/data/test_accelerator_credentials.py | 2 +- .../unit/data/test_accelerator_daemon.py | 86 +++++--- .../unit/data/test_accelerator_dispatch.py | 188 ++++++++++++------ .../unit/data/test_accelerator_enablement.py | 3 +- .../unit/data/test_accelerator_fallback.py | 10 +- 9 files changed, 214 insertions(+), 125 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index 0ad102a84b44..1caf9dc823d8 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -35,7 +35,6 @@ import time from typing import Any, Mapping, Sequence - # Environment variable that overrides the daemon startup readiness timeout, in # seconds (float). Lets a deployment tune the cold-start budget without a code # change or rebuild — e.g. raise it on a very constrained/gVisor host, or (when @@ -100,7 +99,6 @@ def _resolve_binary_path() -> str: return bundled - def _resolve_startup_timeout() -> float: """Resolve the daemon startup timeout in seconds. @@ -234,7 +232,9 @@ def __init__( Raises: FileNotFoundError: the bundled binary is not present in this wheel. """ - self._binary_path = binary_path if binary_path is not None else _resolve_binary_path() + self._binary_path = ( + binary_path if binary_path is not None else _resolve_binary_path() + ) self._cli_flags = list(cli_flags) # ``None`` (the default) defers to the env var / module default; an # explicit float overrides both. Resolved here, not as a signature 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 6a1490d6e4e2..b69ddc304d1e 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 @@ -36,9 +36,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 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 a7aa03b1ea31..64e97ed82e7e 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 @@ -53,6 +53,13 @@ from google.protobuf.message import Message from grpc import Channel +from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon +from google.cloud.bigtable.data._accelerator._fallback import ( + AcceleratorBreaker, + _AcceleratorFallback, + handle_accelerator_error, +) +from google.cloud.bigtable.data._accelerator._routing import is_supported from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._helpers import ( _CONCURRENCY_LIMIT, @@ -113,13 +120,6 @@ ReadModifyWriteRowRequest, SampleRowKeysRequest, ) -from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon -from google.cloud.bigtable.data._accelerator._fallback import ( - AcceleratorBreaker, - _AcceleratorFallback, - handle_accelerator_error, -) -from google.cloud.bigtable.data._accelerator._routing import is_supported if CrossSync.is_async: from grpc.aio import insecure_channel 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 711021230127..610fcb7db33a 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 @@ -47,6 +47,13 @@ from google.protobuf.message import Message from grpc import Channel, insecure_channel, intercept_channel +from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon +from google.cloud.bigtable.data._accelerator._fallback import ( + AcceleratorBreaker, + _AcceleratorFallback, + handle_accelerator_error, +) +from google.cloud.bigtable.data._accelerator._routing import is_supported from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._helpers import ( _CONCURRENCY_LIMIT, @@ -72,6 +79,9 @@ GoogleCloudMetricsHandler, ) from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry +from google.cloud.bigtable.data._sync_autogen._accelerator_client import ( + _AcceleratorClient as AcceleratorClientType, +) from google.cloud.bigtable.data._sync_autogen._swappable_channel import ( SwappableChannel as SwappableChannelType, ) @@ -103,11 +113,11 @@ RowFilterChain, StripValueTransformerFilter, ) +from google.cloud.bigtable_v2._compat import DEFAULT_UNIVERSE as _DEFAULT_UNIVERSE from google.cloud.bigtable_v2.services.bigtable import BigtableClient as GapicClient from google.cloud.bigtable_v2.services.bigtable.transports import ( BigtableGrpcTransport as TransportType, ) -from google.cloud.bigtable_v2._compat import DEFAULT_UNIVERSE as _DEFAULT_UNIVERSE from google.cloud.bigtable_v2.services.bigtable.transports.base import ( DEFAULT_CLIENT_INFO, ) @@ -118,16 +128,6 @@ ReadModifyWriteRowRequest, SampleRowKeysRequest, ) -from google.cloud.bigtable.data._accelerator._daemon import AcceleratorDaemon -from google.cloud.bigtable.data._accelerator._fallback import ( - AcceleratorBreaker, - _AcceleratorFallback, - handle_accelerator_error, -) -from google.cloud.bigtable.data._accelerator._routing import is_supported -from google.cloud.bigtable.data._sync_autogen._accelerator_client import ( - _AcceleratorClient as AcceleratorClientType, -) if TYPE_CHECKING: from google.cloud.bigtable.data._helpers import RowKeySamples, ShardedQuery @@ -346,6 +346,7 @@ def _init_accelerator_config( self._accelerator_env["GOOGLE_APPLICATION_CREDENTIALS"] = ( credentials_file ) + def _resolve_principal(self) -> str | None: """Resolve this client's identity to a principal. @@ -1343,9 +1344,7 @@ def _read_row_via_accelerator(self, query, operation_timeout, attempt_timeout): retryable_exceptions=(), ) try: - stream = client.read_rows( - row_merger.request, timeout=operation_timeout - ) + stream = client.read_rows(row_merger.request, timeout=operation_timeout) chunked_stream = row_merger.chunk_stream(stream) results = [a for a in row_merger.merge_rows(chunked_stream)] except Exception as exc: diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py index 02306f47d840..98d25cdf39e9 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py @@ -24,10 +24,10 @@ import google.cloud.bigtable from google.cloud.bigtable.data._async.client import ( + _UNSET, BigtableDataClientAsync, _AcceleratorIdentityError, _DataApiTargetAsync, - _UNSET, ) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index 526d662086b8..b0bd0f1e575a 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -208,6 +208,7 @@ def fake_mkdtemp(prefix=""): # Properties before/after start # --------------------------------------------------------------------------- + class TestPropertiesBeforeStart: def test_uds_path_before_start_raises(self, tmp_path): d = _make_daemon(tmp_path) @@ -247,6 +248,7 @@ def test_is_running_with_dead_proc(self, tmp_path): # Binary resolution # --------------------------------------------------------------------------- + class TestResolveBinaryPath: def test_missing_binary_raises_file_not_found(self, tmp_path): from google.cloud.bigtable.data._accelerator._daemon import _resolve_binary_path @@ -260,6 +262,7 @@ def test_missing_binary_raises_file_not_found(self, tmp_path): # start() failure modes # --------------------------------------------------------------------------- + class TestStartFailures: def test_called_twice_raises(self, tmp_path, monkeypatch): def fake_popen(argv, **kwargs): @@ -288,9 +291,7 @@ def test_process_exits_during_startup_raises(self, tmp_path, monkeypatch): proc_mock.poll.return_value = 1 # exited immediately proc_mock.stdin = unittest.mock.Mock() - monkeypatch.setattr( - _daemon.subprocess, "Popen", lambda *a, **kw: proc_mock - ) + monkeypatch.setattr(_daemon.subprocess, "Popen", lambda *a, **kw: proc_mock) # _cleanup_tempdir is already a no-op via autouse fixture, but # _force_kill is too — that's fine, we just need the start to raise. monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "0.1") @@ -305,9 +306,7 @@ def test_uds_never_appears_raises_timeout(self, tmp_path, monkeypatch): proc_mock.poll.return_value = None # alive but UDS never arrives proc_mock.stdin = unittest.mock.Mock() - monkeypatch.setattr( - _daemon.subprocess, "Popen", lambda *a, **kw: proc_mock - ) + monkeypatch.setattr(_daemon.subprocess, "Popen", lambda *a, **kw: proc_mock) # Make os.path.exists always False so UDS never appears. monkeypatch.setattr(_daemon.os.path, "exists", lambda p: False) monkeypatch.setenv(_STARTUP_TIMEOUT_ENV_VAR, "0.05") @@ -317,15 +316,17 @@ def test_uds_never_appears_raises_timeout(self, tmp_path, monkeypatch): def test_no_af_unix_raises_os_error(self, tmp_path, monkeypatch): import socket as _socket + monkeypatch.delattr(_socket, "AF_UNIX", raising=False) # Reload the attribute check in _daemon monkeypatch.setattr(_daemon.socket, "AF_UNIX", None, raising=False) # Simpler: just check that hasattr(socket, "AF_UNIX") guards the call. - with unittest.mock.patch.object( - _daemon.socket, "AF_UNIX", None, create=True - ), unittest.mock.patch( - "google.cloud.bigtable.data._accelerator._daemon.hasattr", - return_value=False, + with ( + unittest.mock.patch.object(_daemon.socket, "AF_UNIX", None, create=True), + unittest.mock.patch( + "google.cloud.bigtable.data._accelerator._daemon.hasattr", + return_value=False, + ), ): d = _make_daemon(tmp_path) with pytest.raises(OSError, match="AF_UNIX"): @@ -336,6 +337,7 @@ def test_no_af_unix_raises_os_error(self, tmp_path, monkeypatch): # close() lifecycle # --------------------------------------------------------------------------- + class TestClose: def test_close_before_start_is_noop(self, tmp_path): d = _make_daemon(tmp_path) @@ -351,8 +353,10 @@ def test_close_graceful_via_stdin_eof(self, tmp_path): d = _make_daemon(tmp_path) d._proc = proc_mock - with unittest.mock.patch.object(d, "_wait_for_exit", return_value=True), \ - unittest.mock.patch.object(d, "_cleanup_tempdir"): + with ( + unittest.mock.patch.object(d, "_wait_for_exit", return_value=True), + unittest.mock.patch.object(d, "_cleanup_tempdir"), + ): d.close() proc_mock.stdin.close.assert_called_once() @@ -370,8 +374,10 @@ def test_close_escalates_to_sigterm(self, tmp_path): wait_results = iter([False, True]) # stdin wait fails, sigterm wait succeeds - with unittest.mock.patch.object(d, "_wait_for_exit", side_effect=wait_results), \ - unittest.mock.patch.object(d, "_cleanup_tempdir"): + with ( + unittest.mock.patch.object(d, "_wait_for_exit", side_effect=wait_results), + unittest.mock.patch.object(d, "_cleanup_tempdir"), + ): d.close() proc_mock.terminate.assert_called_once() @@ -389,8 +395,10 @@ def test_close_escalates_to_sigkill(self, tmp_path): # Both stdin and sigterm waits fail → sigkill. wait_results = iter([False, False, True]) - with unittest.mock.patch.object(d, "_wait_for_exit", side_effect=wait_results), \ - unittest.mock.patch.object(d, "_cleanup_tempdir"): + with ( + unittest.mock.patch.object(d, "_wait_for_exit", side_effect=wait_results), + unittest.mock.patch.object(d, "_cleanup_tempdir"), + ): d.close() proc_mock.terminate.assert_called_once() @@ -405,8 +413,10 @@ def test_close_clears_proc_reference(self, tmp_path): d = _make_daemon(tmp_path) d._proc = proc_mock - with unittest.mock.patch.object(d, "_wait_for_exit", return_value=True), \ - unittest.mock.patch.object(d, "_cleanup_tempdir"): + with ( + unittest.mock.patch.object(d, "_wait_for_exit", return_value=True), + unittest.mock.patch.object(d, "_cleanup_tempdir"), + ): d.close() assert d._proc is None @@ -416,6 +426,7 @@ def test_close_clears_proc_reference(self, tmp_path): # _cleanup_tempdir: state clearing # --------------------------------------------------------------------------- + class TestCleanupTempdir: def test_clears_all_state_fields(self, tmp_path): d = _make_daemon(tmp_path) @@ -450,6 +461,7 @@ def test_noop_when_tempdir_is_none(self, tmp_path): # Context manager # --------------------------------------------------------------------------- + class TestContextManager: def test_enter_calls_start_and_returns_self(self, tmp_path, monkeypatch): def fake_popen(argv, **kwargs): @@ -457,8 +469,10 @@ def fake_popen(argv, **kwargs): monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) d = _make_daemon(tmp_path) - with unittest.mock.patch.object(d, "start") as start_mock, \ - unittest.mock.patch.object(d, "close"): + with ( + unittest.mock.patch.object(d, "start") as start_mock, + unittest.mock.patch.object(d, "close"), + ): result = d.__enter__() start_mock.assert_called_once() assert result is d @@ -474,6 +488,7 @@ def test_exit_calls_close(self, tmp_path): # _read_log_tail # --------------------------------------------------------------------------- + class TestReadLogTail: def test_reads_tail_of_log(self, tmp_path): d = _make_daemon(tmp_path) @@ -506,6 +521,7 @@ def test_no_log_path_returns_empty(self, tmp_path): # _wait_for_exit # --------------------------------------------------------------------------- + class TestWaitForExit: def test_returns_true_when_no_proc(self, tmp_path): d = _make_daemon(tmp_path) @@ -513,6 +529,7 @@ def test_returns_true_when_no_proc(self, tmp_path): def test_returns_true_when_proc_exits_in_time(self, tmp_path): import subprocess + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) proc_mock.wait.return_value = 0 # exited d = _make_daemon(tmp_path) @@ -521,6 +538,7 @@ def test_returns_true_when_proc_exits_in_time(self, tmp_path): def test_returns_false_on_timeout(self, tmp_path): import subprocess + proc_mock = unittest.mock.Mock(spec=subprocess.Popen) proc_mock.wait.side_effect = subprocess.TimeoutExpired(cmd="x", timeout=0.1) d = _make_daemon(tmp_path) @@ -532,6 +550,7 @@ def test_returns_false_on_timeout(self, tmp_path): # read_identity: corrupt JSON # --------------------------------------------------------------------------- + class TestReadIdentityCorruptJson: def test_corrupt_json_raises_runtime_error(self, tmp_path): d = _make_daemon(tmp_path) @@ -562,6 +581,7 @@ def test_read_identity_unlinks_even_on_error(self, tmp_path): # CLI flags passed to Popen # --------------------------------------------------------------------------- + class TestCliFlags: def test_cli_flags_appended_to_argv(self, tmp_path, monkeypatch): captured = {} @@ -574,9 +594,12 @@ def fake_popen(argv, **kwargs): d = _make_daemon( tmp_path, cli_flags=[ - "--project", "proj", - "--instance", "inst", - "--app-profile", "prof", + "--project", + "proj", + "--instance", + "inst", + "--app-profile", + "prof", ], ) with pytest.raises(_StopStart): @@ -596,6 +619,7 @@ def fake_popen(argv, **kwargs): # Startup timeout env variable (_resolve_startup_timeout) # --------------------------------------------------------------------------- + class TestResolveStartupTimeout: """Tests for the BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT env-var override. @@ -717,13 +741,18 @@ def tmp_root(self, tmp_path, monkeypatch): def test_removes_stale_dir(self, tmp_root): # Expired deadline + old log + dead spawner PID → should be reaped. - import subprocess, sys + import subprocess + import sys proc = subprocess.Popen([sys.executable, "-c", ""]) dead_pid = proc.pid proc.wait() d = self._make_leftover( - tmp_root, "bt-accel-dead", age_seconds=120, deadline_offset=-70, pid=dead_pid + tmp_root, + "bt-accel-dead", + age_seconds=120, + deadline_offset=-70, + pid=dead_pid, ) _daemon._sweep_stale_tempdirs() assert not d.exists() @@ -754,7 +783,10 @@ def test_ignores_dir_without_log_marker(self, tmp_root): def test_keeps_dir_with_live_pid(self, tmp_root): # Past deadline but spawning process is still alive → must not be reaped. d = self._make_leftover( - tmp_root, "bt-accel-live", age_seconds=120, deadline_offset=-70, + tmp_root, + "bt-accel-live", + age_seconds=120, + deadline_offset=-70, pid=os.getpid(), ) _daemon._sweep_stale_tempdirs() diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py index ce71d84d938c..bd97155a0b7f 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_dispatch.py @@ -25,18 +25,16 @@ from types import SimpleNamespace from unittest import mock +import grpc as _grpc import pytest -from grpc import StatusCode - from google.api_core import exceptions as core_exceptions +from grpc import StatusCode from google.cloud.bigtable.data._accelerator._fallback import ( AcceleratorBreaker, _AcceleratorFallback, ) from google.cloud.bigtable.data._accelerator._routing import is_supported -import grpc as _grpc - from google.cloud.bigtable.data._async.client import ( _AcceleratorIdentityError, _DataApiTargetAsync, @@ -64,6 +62,7 @@ def trailing_metadata(self): # Minimal concrete _DataApiTargetAsync that skips __init__ # --------------------------------------------------------------------------- + class _ConcreteTarget(_DataApiTargetAsync): """Minimal subclass to make the abstract class instantiable.""" @@ -116,6 +115,7 @@ def _bare_target( # _use_accelerator predicate # --------------------------------------------------------------------------- + class TestUseAccelerator: def test_returns_true_when_all_conditions_met(self): t = _bare_target(accel_client=mock.Mock()) @@ -149,6 +149,7 @@ def test_becomes_false_after_trip(self): # _mutate_row_via_accelerator # --------------------------------------------------------------------------- + class TestMutateRowViaAccelerator: def _run(self, coro): return asyncio.get_event_loop().run_until_complete(coro) @@ -164,12 +165,17 @@ def test_raises_fallback_when_client_is_none(self): self._run(t._mutate_row_via_accelerator(req, 30.0)) def test_returns_response_on_success(self): - from google.cloud.bigtable_v2.types.bigtable import MutateRowResponse, MutateRowRequest + from google.cloud.bigtable_v2.types.bigtable import ( + MutateRowRequest, + MutateRowResponse, + ) resp = MutateRowResponse() client_mock = mock.AsyncMock() client_mock.mutate_row.return_value = resp - t = _bare_target(accel_client=client_mock, accel_daemon=mock.Mock(is_running=True)) + t = _bare_target( + accel_client=client_mock, accel_daemon=mock.Mock(is_running=True) + ) req = MutateRowRequest() result = self._run(t._mutate_row_via_accelerator(req, 30.0)) @@ -237,6 +243,7 @@ def test_non_grpc_exception_reraises(self): # _read_row_via_accelerator # --------------------------------------------------------------------------- + class TestReadRowViaAccelerator: def _run(self, coro): return asyncio.get_event_loop().run_until_complete(coro) @@ -310,6 +317,7 @@ def test_non_grpc_exception_reraises(self): # read_row: accelerator dispatch and fallback to native # --------------------------------------------------------------------------- + class TestReadRowAcceleratorDispatch: """Test read_row() routing without spinning up a real client. @@ -329,8 +337,12 @@ def test_routes_to_accelerator_when_enabled(self): async def fake_accel(query, op_timeout, att_timeout): return expected_row - with mock.patch.object(t, "_use_accelerator", return_value=True), \ - mock.patch.object(t, "_read_row_via_accelerator", side_effect=fake_accel) as accel_mock: + with ( + mock.patch.object(t, "_use_accelerator", return_value=True), + mock.patch.object( + t, "_read_row_via_accelerator", side_effect=fake_accel + ) as accel_mock, + ): result = self._run(t.read_row(b"key1")) assert result == expected_row @@ -352,12 +364,16 @@ async def _row_gen(): native_row_merger = mock.Mock() native_row_merger.start_operation.return_value = _row_gen() - with mock.patch.object(t, "_use_accelerator", return_value=True), \ - mock.patch.object(t, "_read_row_via_accelerator", side_effect=fake_accel_fallback), \ - mock.patch( + with ( + mock.patch.object(t, "_use_accelerator", return_value=True), + mock.patch.object( + t, "_read_row_via_accelerator", side_effect=fake_accel_fallback + ), + mock.patch( "google.cloud.bigtable.data._async.client.CrossSync._ReadRowsOperation", return_value=native_row_merger, - ): + ), + ): result = self._run(t.read_row(b"key1")) assert result == expected_row @@ -365,7 +381,9 @@ async def _row_gen(): def test_uses_native_directly_when_accelerator_disabled(self): from google.cloud.bigtable.data.row import Row - t = _bare_target() # _accelerator_client is None → _use_accelerator returns False + t = ( + _bare_target() + ) # _accelerator_client is None → _use_accelerator returns False expected_row = Row(b"key1", cells={}) async def _row_gen(): @@ -374,11 +392,13 @@ async def _row_gen(): native_row_merger = mock.Mock() native_row_merger.start_operation.return_value = _row_gen() - with mock.patch.object(t, "_use_accelerator", return_value=False), \ - mock.patch( + with ( + mock.patch.object(t, "_use_accelerator", return_value=False), + mock.patch( "google.cloud.bigtable.data._async.client.CrossSync._ReadRowsOperation", return_value=native_row_merger, - ): + ), + ): result = self._run(t.read_row(b"key1")) assert result == expected_row @@ -398,11 +418,13 @@ async def _empty_gen(): native_row_merger = mock.Mock() native_row_merger.start_operation.return_value = _empty_gen() - with mock.patch.object(t, "_use_accelerator", return_value=False), \ - mock.patch( + with ( + mock.patch.object(t, "_use_accelerator", return_value=False), + mock.patch( "google.cloud.bigtable.data._async.client.CrossSync._ReadRowsOperation", return_value=native_row_merger, - ): + ), + ): result = self._run(t.read_row(b"missing")) assert result is None @@ -412,6 +434,7 @@ async def _empty_gen(): # mutate_row: accelerator dispatch and fallback to native # --------------------------------------------------------------------------- + class TestMutateRowAcceleratorDispatch: def _run(self, coro): return asyncio.get_event_loop().run_until_complete(coro) @@ -428,8 +451,12 @@ async def fake_accel(request, op_timeout): mutations = [SetCell("cf", "col", b"val", timestamp_micros=1000)] - with mock.patch.object(t, "_use_accelerator", return_value=True), \ - mock.patch.object(t, "_mutate_row_via_accelerator", side_effect=fake_accel) as accel_mock: + with ( + mock.patch.object(t, "_use_accelerator", return_value=True), + mock.patch.object( + t, "_mutate_row_via_accelerator", side_effect=fake_accel + ) as accel_mock, + ): result = self._run(t.mutate_row(b"key1", mutations)) assert result == resp @@ -450,12 +477,16 @@ async def fake_native_mutate(*args, **kwargs): mutations = [SetCell("cf", "col", b"val", timestamp_micros=1000)] - with mock.patch.object(t, "_use_accelerator", return_value=True), \ - mock.patch.object(t, "_mutate_row_via_accelerator", side_effect=fake_accel_fallback), \ - mock.patch( + with ( + mock.patch.object(t, "_use_accelerator", return_value=True), + mock.patch.object( + t, "_mutate_row_via_accelerator", side_effect=fake_accel_fallback + ), + mock.patch( "google.cloud.bigtable.data._async.client.tracked_retry", side_effect=fake_native_mutate, - ): + ), + ): result = self._run(t.mutate_row(b"key1", mutations)) assert result == native_resp @@ -471,11 +502,13 @@ def test_uses_native_directly_when_accelerator_disabled(self): async def fake_native_mutate(*args, **kwargs): return native_resp - with mock.patch.object(t, "_use_accelerator", return_value=False), \ - mock.patch( + with ( + mock.patch.object(t, "_use_accelerator", return_value=False), + mock.patch( "google.cloud.bigtable.data._async.client.tracked_retry", side_effect=fake_native_mutate, - ): + ), + ): result = self._run(t.mutate_row(b"key1", mutations)) assert result == native_resp @@ -490,6 +523,7 @@ def test_raises_value_error_for_empty_mutations(self): # _start_accelerator: startup lifecycle # --------------------------------------------------------------------------- + class TestStartAccelerator: """Tests for _start_accelerator() — the method that spawns the daemon and wires up the AcceleratorClient. These use mocks to avoid real subprocesses.""" @@ -506,13 +540,17 @@ def test_success_sets_daemon_and_client(self): fake_daemon.auth_secret = "secret-token" fake_client = mock.Mock() - with mock.patch( - "google.cloud.bigtable.data._async.client.AcceleratorDaemon", - return_value=fake_daemon, - ), mock.patch( - "google.cloud.bigtable.data._async.client.AcceleratorClientType", - return_value=fake_client, - ), mock.patch.object(t, "_verify_daemon_identity"): + with ( + mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + return_value=fake_daemon, + ), + mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorClientType", + return_value=fake_client, + ), + mock.patch.object(t, "_verify_daemon_identity"), + ): t._start_accelerator() fake_daemon.start.assert_called_once() @@ -529,10 +567,13 @@ def test_daemon_start_failure_calls_close_and_reraises(self): fake_daemon = mock.Mock(spec=AcceleratorDaemon) fake_daemon.start.side_effect = RuntimeError("binary missing") - with mock.patch( - "google.cloud.bigtable.data._async.client.AcceleratorDaemon", - return_value=fake_daemon, - ), pytest.raises(RuntimeError, match="binary missing"): + with ( + mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + return_value=fake_daemon, + ), + pytest.raises(RuntimeError, match="binary missing"), + ): t._start_accelerator() fake_daemon.close.assert_called_once() @@ -554,14 +595,18 @@ def test_identity_verify_failure_closes_daemon_and_propagates(self): fake_daemon.uds_path = "/tmp/fake-sock" fake_daemon.auth_secret = "secret-token" - with mock.patch( - "google.cloud.bigtable.data._async.client.AcceleratorDaemon", - return_value=fake_daemon, - ), mock.patch.object( - t, - "_verify_daemon_identity", - side_effect=_AcceleratorIdentityError("mismatch"), - ), pytest.raises(_AcceleratorIdentityError, match="mismatch"): + with ( + mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + return_value=fake_daemon, + ), + mock.patch.object( + t, + "_verify_daemon_identity", + side_effect=_AcceleratorIdentityError("mismatch"), + ), + pytest.raises(_AcceleratorIdentityError, match="mismatch"), + ): t._start_accelerator() fake_daemon.close.assert_called_once() @@ -587,12 +632,16 @@ def capture_daemon(cli_flags=(), **kwargs): captured["extra_env"] = kwargs.get("extra_env", {}) return fake_daemon - with mock.patch( - "google.cloud.bigtable.data._async.client.AcceleratorDaemon", - side_effect=capture_daemon, - ), mock.patch( - "google.cloud.bigtable.data._async.client.AcceleratorClientType", - ), mock.patch.object(t, "_verify_daemon_identity"): + with ( + mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + side_effect=capture_daemon, + ), + mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorClientType", + ), + mock.patch.object(t, "_verify_daemon_identity"), + ): t._start_accelerator() flags = captured["cli_flags"] @@ -621,12 +670,16 @@ def capture_daemon(cli_flags=(), **kwargs): captured["cli_flags"] = list(cli_flags) return fake_daemon - with mock.patch( - "google.cloud.bigtable.data._async.client.AcceleratorDaemon", - side_effect=capture_daemon, - ), mock.patch( - "google.cloud.bigtable.data._async.client.AcceleratorClientType", - ), mock.patch.object(t, "_verify_daemon_identity"): + with ( + mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + side_effect=capture_daemon, + ), + mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorClientType", + ), + mock.patch.object(t, "_verify_daemon_identity"), + ): t._start_accelerator() flags = captured["cli_flags"] @@ -653,12 +706,16 @@ def capture_daemon(cli_flags=(), **kwargs): captured["cli_flags"] = list(cli_flags) return fake_daemon - with mock.patch( - "google.cloud.bigtable.data._async.client.AcceleratorDaemon", - side_effect=capture_daemon, - ), mock.patch( - "google.cloud.bigtable.data._async.client.AcceleratorClientType", - ), mock.patch.object(t, "_verify_daemon_identity"): + with ( + mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorDaemon", + side_effect=capture_daemon, + ), + mock.patch( + "google.cloud.bigtable.data._async.client.AcceleratorClientType", + ), + mock.patch.object(t, "_verify_daemon_identity"), + ): t._start_accelerator() assert "--app-profile" not in captured["cli_flags"] @@ -668,6 +725,7 @@ def capture_daemon(cli_flags=(), **kwargs): # Event loop fixture so run_until_complete works # --------------------------------------------------------------------------- + @pytest.fixture(autouse=True) def _event_loop(): loop = asyncio.new_event_loop() diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py index 30aca06e0aaa..3ac40c7d17e8 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py @@ -19,12 +19,11 @@ exercise the pure enablement helpers on the async target without spinning up a full client or event loop.""" +import unittest.mock from types import SimpleNamespace import pytest -import unittest.mock - from google.cloud.bigtable.data._async.client import ( _AcceleratorIdentityError, _DataApiTargetAsync, diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_fallback.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_fallback.py index f0856204f239..ad9c187b2520 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_fallback.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_fallback.py @@ -14,13 +14,12 @@ """Unit tests for AcceleratorBreaker, _grpc_code, and handle_accelerator_error in google.cloud.bigtable.data._accelerator._fallback.""" -import pytest from unittest import mock import grpc -from grpc import StatusCode - +import pytest from google.api_core import exceptions as core_exceptions +from grpc import StatusCode from google.cloud.bigtable.data._accelerator._fallback import ( AcceleratorBreaker, @@ -29,11 +28,11 @@ handle_accelerator_error, ) - # --------------------------------------------------------------------------- # AcceleratorBreaker # --------------------------------------------------------------------------- + class TestAcceleratorBreaker: def test_initially_not_bypassed(self): b = AcceleratorBreaker() @@ -59,6 +58,7 @@ def test_bypass_is_sticky(self): def test_thread_safety(self): import threading + b = AcceleratorBreaker() errors = [] @@ -81,6 +81,7 @@ def trip_and_check(): # _grpc_code # --------------------------------------------------------------------------- + class TestGrpcCode: def test_rpc_error_with_code_method(self): # grpc.RpcError doesn't define code() in its class spec, so use plain Mock. @@ -112,6 +113,7 @@ def test_grpc_status_code_wrong_type_returns_none(self): # handle_accelerator_error # --------------------------------------------------------------------------- + class _FakeRpcError(grpc.RpcError): """Minimal grpc.RpcError subclass for tests — passes isinstance checks. From f87850eb8a0f514fa7a88b7f1961cf8edfa1f187 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 1 Sep 2026 18:07:38 +0000 Subject: [PATCH 19/26] feat(bigtable): fall back to native when the accelerator fails health checks fix(bigtable): lock all reads and writes in AcceleratorBreaker Change-Id: I9d1daffbe086788481d6802cec035a3d4d284729 --- .../bigtable/data/_accelerator/_fallback.py | 61 +++- .../bigtable/data/_accelerator/_health.py | 89 +++++ .../data/_async/_accelerator_client.py | 24 ++ .../data/_async/_accelerator_health.py | 232 +++++++++++++ .../cloud/bigtable/data/_async/client.py | 44 ++- .../data/_sync_autogen/_accelerator_client.py | 20 ++ .../data/_sync_autogen/_accelerator_health.py | 184 ++++++++++ .../bigtable/data/_sync_autogen/client.py | 151 ++++---- .../data/_async/test__accelerator_health.py | 323 ++++++++++++++++++ .../_sync_autogen/test__accelerator_health.py | 280 +++++++++++++++ .../unit/data/test_accelerator_health.py | 178 ++++++++++ 11 files changed, 1497 insertions(+), 89 deletions(-) create mode 100644 packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_health.py create mode 100644 packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_health.py create mode 100644 packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_accelerator_health.py create mode 100644 packages/google-cloud-bigtable/tests/unit/data/_async/test__accelerator_health.py create mode 100644 packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__accelerator_health.py create mode 100644 packages/google-cloud-bigtable/tests/unit/data/test_accelerator_health.py diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py index b69ddc304d1e..1eac1742d6a1 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,25 +60,60 @@ 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. """ def __init__(self): - self._tripped = threading.Event() + self._tripped = False + self._degraded = False + self._lock = threading.Lock() 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.""" + with self._lock: + return self._tripped or self._degraded def trip(self) -> None: """Permanently bypass the accelerator (e.g. the daemon process died).""" - self._tripped.set() + with self._lock: + self._tripped = True + + @property + def is_tripped(self) -> bool: + """Whether the accelerator has been permanently abandoned.""" + with self._lock: + return self._tripped + + @property + def is_degraded(self) -> bool: + """Whether the accelerator is currently being skipped for poor health.""" + with self._lock: + return self._degraded + + 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. + """ + with self._lock: + self._degraded = degraded def _grpc_code(exc: BaseException) -> StatusCode | None: 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..af12ccf9b117 --- /dev/null +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_health.py @@ -0,0 +1,232 @@ +# 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 concurrent.futures +import logging + +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 + + +@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 + self._closed = CrossSync.Event() + self._task: CrossSync.Task[None] | None = None + # A dedicated worker rather than the client's shared pool: this task + # runs for the life of the Table, so it must not occupy a slot other + # work is queued behind. + self._sync_executor = ( + concurrent.futures.ThreadPoolExecutor( + max_workers=1, thread_name_prefix="bigtable-accel-health" + ) + if not CrossSync.is_async + else None + ) + + def start(self) -> None: + """Launch the polling task. Idempotent.""" + if self._task is None: + self._task = CrossSync.create_task( + self._run, + sync_executor=self._sync_executor, + task_name="bigtable-accelerator-health", + ) + + @CrossSync.convert + async def _run(self) -> None: + """Poll until closed, or until monitoring turns out to be pointless.""" + while not self._closed.is_set(): + await CrossSync.event_wait(self._closed, 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 alone: an unmonitorable daemon is not a + # degraded one, and treating it as one would disable the + # accelerator against every daemon older than this feature. + _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. Does not clear a degraded flag already set.""" + self._closed.set() + if self._task is not None: + self._task.cancel() + self._task = None + if self._sync_executor is not None: + # Don't wait: a probe in flight holds the worker for up to the probe + # timeout, and closing a Table should not block on that. + self._sync_executor.shutdown(wait=False) 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 64e97ed82e7e..6b75dcd770d6 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 @@ -127,6 +127,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 +151,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, ) @@ -1357,9 +1363,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 @@ -1465,6 +1474,23 @@ 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() + 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. @@ -2348,6 +2374,18 @@ async def close(self): """ Called to close the Table instance and release any resources held by it. """ + self._metrics.close() + if self._register_instance_future: + self._register_instance_future.cancel() + self.client._remove_instance_registration( + self.instance_id, self.app_profile_id, id(self) + ) + # Stop the monitor before the channel it probes over goes away. + if self._accelerator_health is not None: + try: + await self._accelerator_health.close() + 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..ec6c3a2c0848 --- /dev/null +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_accelerator_health.py @@ -0,0 +1,184 @@ +# 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 concurrent.futures +import logging + +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 +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 + + +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 = CrossSync._Sync_Impl.Event() + self._task: CrossSync._Sync_Impl.Task[None] | None = None + self._sync_executor = ( + concurrent.futures.ThreadPoolExecutor( + max_workers=1, thread_name_prefix="bigtable-accel-health" + ) + if not CrossSync._Sync_Impl.is_async + else None + ) + + def start(self) -> None: + """Launch the polling task. Idempotent.""" + if self._task is None: + self._task = CrossSync._Sync_Impl.create_task( + self._run, + sync_executor=self._sync_executor, + task_name="bigtable-accelerator-health", + ) + + def _run(self) -> None: + """Poll until closed, or until monitoring turns out to be pointless.""" + while not self._closed.is_set(): + CrossSync._Sync_Impl.event_wait(self._closed, timeout=self._interval) + if self._closed.is_set(): + return + if self._breaker.is_tripped: + return + healthy = self._probe() + if healthy is None: + _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. Does not clear a degraded flag already set.""" + self._closed.set() + if self._task is not None: + self._task.cancel() + self._task = None + if self._sync_executor is not None: + self._sync_executor.shutdown(wait=False) 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 610fcb7db33a..1ca12c624da0 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 @@ -82,6 +82,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, ) @@ -137,9 +140,8 @@ from google.cloud.bigtable.data.execute_query._sync_autogen.execute_query_iterator import ( ExecuteQueryIterator, ) - _LOGGER = logging.getLogger(__name__) -_PRINCIPAL_RESOLVE_TIMEOUT = 5.0 # seconds; see async client comment +_PRINCIPAL_RESOLVE_TIMEOUT = 5.0 _UNSET: Any = object() @@ -317,7 +319,7 @@ def _init_accelerator_config( client_options, "client_cert_source", None ): self._accelerator_blocked_reason = "client_cert_source (mTLS) cannot be forwarded to the accelerator daemon" - import google.cloud.bigtable # local import avoids circular dependency + import google.cloud.bigtable self._accelerator_flags += [ "--caller-user-agent", @@ -382,7 +384,7 @@ def _with_timeout(*args, **kwargs): except Exception as exc: _LOGGER.debug("Compute Engine credential refresh failed: %s", exc) email = getattr(creds, "service_account_email", None) - principal: str | None = email if (email and email != "default") else None + principal: str | None = email if email and email != "default" else None self._cached_principal = principal return principal @@ -1061,40 +1063,36 @@ def __init__( raise RuntimeError( f"{self.__class__.__name__} must be created within an async event loop context." ) from e - - # Optional in-process accelerator daemon, scoped to this data target's - # (project, instance_id, app_profile_id) tuple. Enabled by default; - # controlled by the client's ``use_accelerator`` option. self._accelerator_daemon: AcceleratorDaemon | None = None self._accelerator_client: AcceleratorClientType | None = None - # Sticky fallback policy: a daemon that cannot serve a routed RPC - # (UNIMPLEMENTED / dead subprocess) is transparently bypassed in favor - # of the native client. See ``_accelerator/_fallback.py``. + self._accelerator_health: AcceleratorHealthMonitorType | None = None self._accelerator_breaker = AcceleratorBreaker() if self.client._use_accelerator is not False: - # None (default) or True: attempt to start. ``True`` is an explicit - # request, so a conflict with the emulator is a hard error. self._maybe_start_accelerator(explicit=self.client._use_accelerator is True) def _maybe_start_accelerator(self, *, explicit: bool) -> None: - """Start the accelerator daemon unless the environment prevents it.""" + """Start the accelerator daemon unless the environment prevents it. + + The accelerator is on by default, so it must never break callers when + it can't run. When the emulator is set, or the daemon fails to start + (e.g. the binary isn't bundled for this platform), fall back to the + native client and warn. The exception is an explicit + ``use_accelerator=True``: a caller who required the accelerator gets a + hard error rather than a silent fallback, both for the emulator (a + genuine misconfiguration) and for a daemon that fails to start.""" if sys.platform == "win32": if explicit: raise RuntimeError( - "use_accelerator=True is not supported on Windows; pass " - "use_accelerator=False to use the native client." + "use_accelerator=True is not supported on Windows; pass use_accelerator=False to use the native client." ) return if self.client._emulator_host is not None: if explicit: raise RuntimeError( - "use_accelerator=True is not supported when " - "BIGTABLE_EMULATOR_HOST is set; unset the emulator, or pass " - "use_accelerator=False to use the emulator." + "use_accelerator=True is not supported when BIGTABLE_EMULATOR_HOST is set; unset the emulator, or pass use_accelerator=False to use the emulator." ) warnings.warn( - "Accelerator disabled because BIGTABLE_EMULATOR_HOST is set; " - "using the native client.", + "Accelerator disabled because BIGTABLE_EMULATOR_HOST is set; using the native client.", RuntimeWarning, stacklevel=2, ) @@ -1111,19 +1109,14 @@ def _maybe_start_accelerator(self, *, explicit: bool) -> None: except _AcceleratorIdentityError as exc: if explicit: raise - warnings.warn( - f"Accelerator disabled: {exc}", - RuntimeWarning, - stacklevel=2, - ) + warnings.warn(f"Accelerator disabled: {exc}", RuntimeWarning, stacklevel=2) self._accelerator_daemon = None self._accelerator_client = None except Exception as exc: if explicit: raise warnings.warn( - "Failed to start the Bigtable accelerator daemon; falling back " - f"to the native client: {exc}", + f"Failed to start the Bigtable accelerator daemon; falling back to the native client: {exc}", RuntimeWarning, stacklevel=2, ) @@ -1131,13 +1124,10 @@ def _maybe_start_accelerator(self, *, explicit: bool) -> None: self._accelerator_client = None def _start_accelerator(self) -> None: - """Spawn the daemon for this Table and connect to its UDS.""" - flags = [ - "--project", - self.client.project, - "--instance", - self.instance_id, - ] + """Spawn the daemon for this Table and connect to its UDS. + + On any failure the daemon is torn down before the error propagates.""" + flags = ["--project", self.client.project, "--instance", self.instance_id] if self.app_profile_id: flags.extend(["--app-profile", self.app_profile_id]) flags.extend(self.client._accelerator_flags) @@ -1154,28 +1144,28 @@ 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() + 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.""" + """Whether this call should be routed through the accelerator daemon. + + False once the fallback breaker has tripped, so a daemon that has proven + unable to serve routed RPCs is skipped without another round trip.""" return ( self._accelerator_client is not None and is_supported(method_name) - and not self._accelerator_breaker.bypass() - ) - - def _create_operation( - self, op_type: OperationType, **kwargs - ) -> ActiveOperationMetric: - table_id = getattr(self, "table_id", None) or getattr( - self, "materialized_view_id", None - ) - return self.client._metrics.create_operation( - op_type, - project_id=self.client.project, - instance_id=self.instance_id, - table_id=table_id, - app_profile_id=self.app_profile_id, - **kwargs, + and (not self._accelerator_breaker.bypass()) ) def _verify_daemon_identity(self, server: AcceleratorDaemon) -> None: @@ -1196,8 +1186,7 @@ def _verify_daemon_identity(self, server: AcceleratorDaemon) -> None: own_principal = self.client._resolve_principal() if own_principal is None or daemon_principal is None: raise _AcceleratorIdentityError( - "Could not verify daemon identity: principal is unknown on " - f"{'this client' if own_principal is None else 'the daemon'} side" + f"Could not verify daemon identity: principal is unknown on {('this client' if own_principal is None else 'the daemon')} side" ) if own_principal != daemon_principal: raise _AcceleratorIdentityError( @@ -1207,14 +1196,28 @@ def _verify_daemon_identity(self, server: AcceleratorDaemon) -> None: own_scopes = _normalize_scopes(self.client._accelerator_scopes) if not own_scopes or not daemon_scopes: raise _AcceleratorIdentityError( - "Could not verify daemon identity: effective scopes are unknown on " - f"{'this client' if not own_scopes else 'the daemon'} side" + f"Could not verify daemon identity: effective scopes are unknown on {('this client' if not own_scopes else 'the daemon')} side" ) if own_scopes != daemon_scopes: raise _AcceleratorIdentityError( f"Accelerator scope mismatch: this client forwarded {sorted(own_scopes)!r} but the daemon resolved {sorted(daemon_scopes)!r}; refusing to route RPCs through the accelerator." ) + def _create_operation( + self, op_type: OperationType, **kwargs + ) -> ActiveOperationMetric: + table_id = getattr(self, "table_id", None) or getattr( + self, "materialized_view_id", None + ) + return self.client._metrics.create_operation( + op_type, + project_id=self.client.project, + instance_id=self.instance_id, + table_id=table_id, + app_profile_id=self.app_profile_id, + **kwargs, + ) + @property @abc.abstractmethod def _request_path(self) -> dict[str, str]: @@ -1329,8 +1332,7 @@ def _read_row_via_accelerator(self, query, operation_timeout, attempt_timeout): Raises ``_AcceleratorFallback`` if the caller should retry on the native client; other gRPC errors are translated to ``google.api_core`` - exceptions and raised. - """ + exceptions and raised.""" client = self._accelerator_client if client is None: raise _AcceleratorFallback() @@ -1349,11 +1351,9 @@ def _read_row_via_accelerator(self, query, operation_timeout, attempt_timeout): results = [a for a in row_merger.merge_rows(chunked_stream)] except Exception as exc: handle_accelerator_error( - exc, - daemon=daemon, - breaker=self._accelerator_breaker, + exc, daemon=daemon, breaker=self._accelerator_breaker ) - raise # unreachable: handle_accelerator_error always raises + raise try: return results[0] except IndexError: @@ -1399,16 +1399,13 @@ def read_row( operation_timeout, attempt_timeout = _get_timeouts( operation_timeout, attempt_timeout, self ) - if self._use_accelerator("read_row"): try: return self._read_row_via_accelerator( query, operation_timeout, attempt_timeout ) except _AcceleratorFallback: - # Daemon can't serve this call; fall through to the native path. pass - retryable_excs = _get_retryable_errors(retryable_errors, self) row_merger = CrossSync._Sync_Impl._ReadRowsOperation( query, @@ -1680,8 +1677,7 @@ def _mutate_row_via_accelerator(self, request, operation_timeout): Raises ``_AcceleratorFallback`` if the caller should retry on the native client; other gRPC errors are translated to ``google.api_core`` - exceptions and raised. - """ + exceptions and raised.""" client = self._accelerator_client if client is None: raise _AcceleratorFallback() @@ -1690,11 +1686,9 @@ def _mutate_row_via_accelerator(self, request, operation_timeout): return client.mutate_row(request, timeout=operation_timeout) except Exception as exc: handle_accelerator_error( - exc, - daemon=daemon, - breaker=self._accelerator_breaker, + exc, daemon=daemon, breaker=self._accelerator_breaker ) - raise # unreachable: handle_accelerator_error always raises + raise def mutate_row( self, @@ -1741,21 +1735,17 @@ def mutate_row( if not mutations: raise ValueError("No mutations provided") mutations_list = mutations if isinstance(mutations, list) else [mutations] - request = MutateRowRequest( row_key=row_key.encode("utf-8") if isinstance(row_key, str) else row_key, mutations=[mutation._to_pb() for mutation in mutations_list], app_profile_id=self.app_profile_id, **self._request_path, ) - if self._use_accelerator("mutate_row"): try: return self._mutate_row_via_accelerator(request, operation_timeout) except _AcceleratorFallback: - # Daemon can't serve this call; fall through to the native path. pass - if all((mutation.is_idempotent() for mutation in mutations_list)): predicate = retries.if_exception_type( *_get_retryable_errors(retryable_errors, self) @@ -1951,6 +1941,17 @@ def read_modify_write_row( def close(self): """Called to close the Table instance and release any resources held by it.""" + self._metrics.close() + if self._register_instance_future: + self._register_instance_future.cancel() + self.client._remove_instance_registration( + self.instance_id, self.app_profile_id, id(self) + ) + if self._accelerator_health is not None: + try: + self._accelerator_health.close() + 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..211d5e4010c9 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test__accelerator_health.py @@ -0,0 +1,323 @@ +# 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. + """ + + def __init__(self, uds_path, *, status=ServingStatus.SERVING, serve_health=True): + self.status = status + 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 + if delay: + time.sleep(delay) + 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: monitor._task.done()) + assert not breaker.bypass() + 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: monitor._task.done()) + 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..f379a5963715 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__accelerator_health.py @@ -0,0 +1,280 @@ +# 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. + """ + + def __init__(self, uds_path, *, status=ServingStatus.SERVING, serve_health=True): + self.status = status + 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) + if delay: + time.sleep(delay) + 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: monitor._task.done()) + assert not breaker.bypass() + 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: monitor._task.done()) + 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..5496f95121c2 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_health.py @@ -0,0 +1,178 @@ +# 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() From e5f088198bc03302a52de840158268d8f9ad42ff Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 13 Oct 2026 19:27:32 +0000 Subject: [PATCH 20/26] fix(bigtable): restore explicit-mode hard-error for identity mismatch, clear stale auth_secret on teardown Change-Id: I02c7ecb2d9ac5aa082bec0dc3325938b621bd210 --- .../google/cloud/bigtable/data/_async/client.py | 3 +++ .../google/cloud/bigtable/data/_sync_autogen/client.py | 3 +++ 2 files changed, 6 insertions(+) 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 6b75dcd770d6..66a92ac22f32 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 @@ -1471,6 +1471,9 @@ def _start_accelerator(self) -> None: server.uds_path, server.auth_secret ) self._accelerator_daemon = server + except _AcceleratorUnverified: + server.close() + raise except BaseException: server.close() raise 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 1ca12c624da0..3f6e7c609c55 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 @@ -1141,6 +1141,9 @@ def _start_accelerator(self) -> None: server.uds_path, server.auth_secret ) self._accelerator_daemon = server + except _AcceleratorUnverified: + server.close() + raise except BaseException: server.close() raise From 7234ccfa5efea3d5b21b48cc84c40b3219bf70c1 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Fri, 16 Oct 2026 19:08:56 +0000 Subject: [PATCH 21/26] fix(bigtable): address accel9 code review findings Change-Id: I6665fec71ca24efe8ec47f95a71de6dda59abe5e --- .../cloud/bigtable/data/_accelerator/_daemon.py | 3 ++- .../cloud/bigtable/data/_async/_accelerator_health.py | 6 ++++-- .../google/cloud/bigtable/data/_async/client.py | 9 +++------ .../cloud/bigtable/data/_sync_autogen/client.py | 11 ++++------- 4 files changed, 13 insertions(+), 16 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index 1caf9dc823d8..7c493c2ecdab 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -172,7 +172,8 @@ def _sweep_stale_tempdirs() -> None: # started the daemon being evaluated. deadline_path = os.path.join(path, _STARTUP_DEADLINE_FILENAME) try: - deadline = float(open(deadline_path).read()) + with open(deadline_path) as _f: + deadline = float(_f.read()) except (OSError, ValueError): # Older daemon without a deadline file: fall back to a simple # age heuristic relative to the log's mtime. 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 index af12ccf9b117..873819b805e0 100644 --- 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 @@ -224,8 +224,10 @@ async def close(self) -> None: """Stop polling. Does not clear a degraded flag already set.""" self._closed.set() if self._task is not None: - self._task.cancel() - self._task = None + task, self._task = self._task, None + task.cancel() + # Retrieve any exception so asyncio doesn't log "Task exception was never retrieved" + task.add_done_callback(lambda t: None if t.cancelled() else t.exception()) if self._sync_executor is not None: # Don't wait: a probe in flight holds the worker for up to the probe # timeout, and closing a Table should not block on that. 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 66a92ac22f32..b9709978ce88 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 @@ -1471,7 +1471,7 @@ def _start_accelerator(self) -> None: server.uds_path, server.auth_secret ) self._accelerator_daemon = server - except _AcceleratorUnverified: + except _AcceleratorIdentityError: server.close() raise except BaseException: @@ -2387,6 +2387,8 @@ async def close(self): 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: @@ -2403,11 +2405,6 @@ async def close(self): _LOGGER.warning("Failed to close accelerator daemon: %s", exc) finally: self._accelerator_daemon = None - if self._register_instance_future: - self._register_instance_future.cancel() - self.client._remove_instance_registration( - self.instance_id, self.app_profile_id, id(self) - ) @CrossSync.convert(sync_name="__enter__") async def __aenter__(self): 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 3f6e7c609c55..e0522378cd72 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 @@ -1141,7 +1141,7 @@ def _start_accelerator(self) -> None: server.uds_path, server.auth_secret ) self._accelerator_daemon = server - except _AcceleratorUnverified: + except _AcceleratorIdentityError: server.close() raise except BaseException: @@ -1953,6 +1953,8 @@ def close(self): 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: @@ -1969,11 +1971,6 @@ def close(self): _LOGGER.warning("Failed to close accelerator daemon: %s", exc) finally: self._accelerator_daemon = None - if self._register_instance_future: - self._register_instance_future.cancel() - self.client._remove_instance_registration( - self.instance_id, self.app_profile_id, id(self) - ) def __enter__(self): """Implement async context manager protocol @@ -1981,7 +1978,7 @@ def __enter__(self): Ensure registration task has time to run, so that grpc channels will be warmed for the specified instance""" if self._register_instance_future: - self._register_instance_future + self._register_instance_future.result() return self def __exit__(self, exc_type, exc_val, exc_tb): From 8fa298bf9bd26e1abb88e4744bf781be14ab26f5 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 20 Oct 2026 14:33:36 +0000 Subject: [PATCH 22/26] fix(bigtable): drop dead identity-error branch, bogus _metrics.close(), and breaker lock Change-Id: I21cf617dbea7aa85dc7b0d093e15c1ea22c7f95c --- .../bigtable/data/_accelerator/_fallback.py | 28 +++++++++---------- .../cloud/bigtable/data/_async/client.py | 18 ++++++------ .../data/_sync_autogen/_accelerator_health.py | 5 ++-- .../bigtable/data/_sync_autogen/client.py | 16 ++++------- 4 files changed, 31 insertions(+), 36 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_fallback.py index 1eac1742d6a1..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 @@ -76,35 +76,33 @@ class AcceleratorBreaker: 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. + 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 = False - self._degraded = False - self._lock = threading.Lock() + self._tripped = threading.Event() + self._degraded = threading.Event() def bypass(self) -> bool: """Whether the accelerator should be skipped for calls made right now.""" - with self._lock: - return self._tripped or self._degraded + return self._tripped.is_set() or self._degraded.is_set() def trip(self) -> None: """Permanently bypass the accelerator (e.g. the daemon process died).""" - with self._lock: - self._tripped = True + self._tripped.set() @property def is_tripped(self) -> bool: """Whether the accelerator has been permanently abandoned.""" - with self._lock: - return self._tripped + return self._tripped.is_set() @property def is_degraded(self) -> bool: """Whether the accelerator is currently being skipped for poor health.""" - with self._lock: - return self._degraded + return self._degraded.is_set() def set_degraded(self, degraded: bool) -> None: """Start or stop bypassing the accelerator for poor health. @@ -112,8 +110,10 @@ def set_degraded(self, degraded: bool) -> None: Reversible, and orthogonal to :meth:`trip`: clearing it does not undo a permanent trip. """ - with self._lock: - self._degraded = degraded + if degraded: + self._degraded.set() + else: + self._degraded.clear() def _grpc_code(exc: BaseException) -> StatusCode | None: 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 b9709978ce88..bb927e8092ba 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 @@ -1471,9 +1471,6 @@ def _start_accelerator(self) -> None: server.uds_path, server.auth_secret ) self._accelerator_daemon = server - except _AcceleratorIdentityError: - server.close() - raise except BaseException: server.close() raise @@ -2377,13 +2374,9 @@ async def close(self): """ Called to close the Table instance and release any resources held by it. """ - self._metrics.close() - if self._register_instance_future: - self._register_instance_future.cancel() - self.client._remove_instance_registration( - self.instance_id, self.app_profile_id, id(self) - ) - # Stop the monitor before the channel it probes over goes away. + # 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() @@ -2405,6 +2398,11 @@ async def close(self): _LOGGER.warning("Failed to close accelerator daemon: %s", exc) finally: self._accelerator_daemon = None + if self._register_instance_future: + self._register_instance_future.cancel() + self.client._remove_instance_registration( + self.instance_id, self.app_profile_id, id(self) + ) @CrossSync.convert(sync_name="__enter__") async def __aenter__(self): 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 index ec6c3a2c0848..19b78198d192 100644 --- 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 @@ -178,7 +178,8 @@ def close(self) -> None: """Stop polling. Does not clear a degraded flag already set.""" self._closed.set() if self._task is not None: - self._task.cancel() - self._task = None + task, self._task = (self._task, None) + task.cancel() + task.add_done_callback(lambda t: None if t.cancelled() else t.exception()) if self._sync_executor is not None: self._sync_executor.shutdown(wait=False) 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 e0522378cd72..94531d09130d 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 @@ -1141,9 +1141,6 @@ def _start_accelerator(self) -> None: server.uds_path, server.auth_secret ) self._accelerator_daemon = server - except _AcceleratorIdentityError: - server.close() - raise except BaseException: server.close() raise @@ -1944,12 +1941,6 @@ def read_modify_write_row( def close(self): """Called to close the Table instance and release any resources held by it.""" - self._metrics.close() - if self._register_instance_future: - self._register_instance_future.cancel() - self.client._remove_instance_registration( - self.instance_id, self.app_profile_id, id(self) - ) if self._accelerator_health is not None: try: self._accelerator_health.close() @@ -1971,6 +1962,11 @@ def close(self): _LOGGER.warning("Failed to close accelerator daemon: %s", exc) finally: self._accelerator_daemon = None + if self._register_instance_future: + self._register_instance_future.cancel() + self.client._remove_instance_registration( + self.instance_id, self.app_profile_id, id(self) + ) def __enter__(self): """Implement async context manager protocol @@ -1978,7 +1974,7 @@ def __enter__(self): Ensure registration task has time to run, so that grpc channels will be warmed for the specified instance""" if self._register_instance_future: - self._register_instance_future.result() + self._register_instance_future return self def __exit__(self, exc_type, exc_val, exc_tb): From bdfb6c880b1a3afc2404083b5051ba54bbb77ad1 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 20 Oct 2026 16:15:54 +0000 Subject: [PATCH 23/26] feat(bigtable): support the accelerator on Windows x86_64 and x86 Change-Id: I7a17dd9919dff6e9d38b0cdfb3c600ffc363900c --- .../bigtable/data/_accelerator/_daemon.py | 62 +++++++- .../data/_async/_accelerator_client.py | 14 +- .../cloud/bigtable/data/_async/client.py | 8 - .../data/_sync_autogen/_accelerator_client.py | 13 +- .../bigtable/data/_sync_autogen/client.py | 7 - .../data/_async/test__accelerator_health.py | 4 +- .../_sync_autogen/test__accelerator_health.py | 2 +- .../unit/data/test_accelerator_daemon.py | 147 ++++++++++++++++-- 8 files changed, 217 insertions(+), 40 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index 7c493c2ecdab..b574a1baacdc 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -30,11 +30,19 @@ import shutil import socket import subprocess +import sys import tempfile import threading import time from typing import Any, Mapping, Sequence +# Windows differs from every other supported platform in two ways that matter +# here, so it gets a named constant rather than scattered ``sys.platform`` +# checks: the daemon binary carries a ``.exe`` suffix, and CPython's socket +# module exposes no ``AF_UNIX`` even though the OS and gRPC's C-core both speak +# it (see ``_socket_is_live``). +_IS_WINDOWS = sys.platform == "win32" + # Environment variable that overrides the daemon startup readiness timeout, in # seconds (float). Lets a deployment tune the cold-start budget without a code # change or rebuild — e.g. raise it on a very constrained/gVisor host, or (when @@ -43,7 +51,10 @@ _STARTUP_TIMEOUT_ENV_VAR = "BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT" # Wheels ship the binary at this path relative to the `_accelerator/` package. +# The Windows wheels (win_amd64, win32) stage it under the ``.exe`` name Go +# gives it, so the suffix is appended there. _DEFAULT_BIN_RELATIVE_PATH = "bin/accelerator" +_WINDOWS_BIN_SUFFIX = ".exe" # The daemon writes the principal it resolved to this file in its tempdir # (alongside the socket) before binding, so the client can verify it matches @@ -91,7 +102,8 @@ def _resolve_binary_path() -> str: """Return the path of the bundled daemon binary, raising if it is absent.""" - bundled = os.path.join(os.path.dirname(__file__), _DEFAULT_BIN_RELATIVE_PATH) + relative = _DEFAULT_BIN_RELATIVE_PATH + (_WINDOWS_BIN_SUFFIX if _IS_WINDOWS else "") + bundled = os.path.join(os.path.dirname(__file__), relative) if not os.path.isfile(bundled): raise FileNotFoundError( "Accelerator binary not found. Install a wheel that bundles the binary." @@ -127,6 +139,23 @@ def _socket_is_live(uds_path: str) -> bool: """Return True if something is currently accepting connections on the UDS.""" if not os.path.exists(uds_path): return False + if not hasattr(socket, "AF_UNIX"): + # Windows. The OS (10 1803+ / Server 2019+) and gRPC's C-core both + # support AF_UNIX, which is why the accelerator runs here at all, but + # CPython's socket module still does not expose the family, so we + # cannot open a probe connection from Python. The gate is + # ``HAVE_SYS_UN_H``: ``PC/pyconfig.h`` leaves it commented out and + # ``Modules/socketmodule.h`` responds with ``#undef AF_UNIX``, so the + # attribute is absent on every current CPython for Windows. Testing + # the attribute rather than the platform keeps this correct for any + # other build that lacks the family, and is what lets the tests + # exercise this branch from a POSIX run. + # + # Fall back to the bind artifact: the daemon's + # ``net.Listen("unix", ...)`` creates this file and puts the socket in + # the listen backlog inside one call, so the path existing means the + # listener is up. + return True with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: probe.settimeout(0.25) try: @@ -154,8 +183,18 @@ def _sweep_stale_tempdirs() -> None: All failures are swallowed — sweeping is opportunistic and must never break startup. + + Skipped entirely on Windows: see the guard below. """ - if not hasattr(socket, "AF_UNIX"): + if _IS_WINDOWS: + # The liveness check below is ``os.kill(pid, 0)``, which is a probe + # only on POSIX. On Windows os.kill routes every signal that is not a + # console control event to TerminateProcess, so that same call would + # *kill* the spawning Python process it is supposed to be asking + # about — someone else's live application. Porting it needs an + # OpenProcess/WaitForSingleObject check instead, which is left for a + # follow-up; until then the safe failure is to leave the tempdirs + # alone. They are four small files under the per-user %TEMP%. return root = tempfile.gettempdir() try: @@ -300,7 +339,12 @@ def start(self) -> None: if self._proc is not None: raise RuntimeError("AcceleratorDaemon.start() called twice") - if not hasattr(socket, "AF_UNIX"): + # The daemon only listens on AF_UNIX, so a platform without it cannot + # run the accelerator at all. Windows is the one case where a missing + # ``socket.AF_UNIX`` says nothing about the platform — the OS and + # gRPC's C-core both support the family, only CPython's socket module + # does not — so it is excluded from the check. + if not hasattr(socket, "AF_UNIX") and not _IS_WINDOWS: raise OSError( "Unix domain sockets (AF_UNIX) are not supported on this platform." ) @@ -310,6 +354,11 @@ def start(self) -> None: # daemon thread so it never blocks startup (sweep probes dead sockets # with a 0.25 s timeout each). threading.Thread(target=_sweep_stale_tempdirs, daemon=True).start() + # mkdtemp is what confines the socket to this user: 0700 on POSIX, and + # on Windows an inherited ACL from the per-user %TEMP% tree. The latter + # is the only access control the socket gets there — Windows has no + # umask and chmod only toggles the read-only bit — so the daemon + # deliberately relies on the directory we pick here. self._tempdir = tempfile.mkdtemp(prefix=_TEMPDIR_PREFIX) self._uds_path = os.path.join(self._tempdir, _SOCKET_FILENAME) # Write the absolute time by which startup must complete. The sweep @@ -405,7 +454,12 @@ def close(self) -> None: except OSError: pass if not self._wait_for_exit(_STDIN_GRACE_SECONDS): - # Step 2: SIGTERM. + # Step 2: SIGTERM. On Windows this is TerminateProcess, so + # steps 2 and 3 are the same hard kill and the daemon's one + # chance to shut down cleanly is the stdin EOF above — + # which is also the signal it uses there to notice the + # parent died, since a ppid watchdog cannot work on + # Windows. try: proc.terminate() except OSError: 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 7b9b94993535..e6953f676664 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 @@ -46,6 +46,18 @@ _READ_ROWS_METHOD = "/google.bigtable.v2.Bigtable/ReadRows" +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: """Thin gRPC client bound to the daemon's UDS. @@ -58,7 +70,7 @@ class _AsyncAcceleratorClient: def __init__(self, uds_path: str, auth_secret: str): self._uds_path = uds_path self._metadata = (("x-accelerator-token", auth_secret),) - self._channel = insecure_channel(f"unix://{uds_path}") + self._channel = insecure_channel(_uds_target(uds_path)) self._mutate_row_stub = self._channel.unary_unary( _MUTATE_ROW_METHOD, request_serializer=MutateRowRequest.serialize, 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 bb927e8092ba..eb45e284bda4 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 @@ -20,7 +20,6 @@ import logging import os import random -import sys import time import warnings from functools import partial @@ -1386,13 +1385,6 @@ def _maybe_start_accelerator(self, *, explicit: bool) -> None: hard error rather than a silent fallback, both for the emulator (a genuine misconfiguration) and for a daemon that fails to start. """ - if sys.platform == "win32": - if explicit: - raise RuntimeError( - "use_accelerator=True is not supported on Windows; pass " - "use_accelerator=False to use the native client." - ) - return if self.client._emulator_host is not None: if explicit: raise RuntimeError( 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 bd0965843620..51905d4b04cd 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 @@ -43,6 +43,17 @@ _READ_ROWS_METHOD = "/google.bigtable.v2.Bigtable/ReadRows" +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}" + + class _AcceleratorClient: """Thin gRPC client bound to the daemon's UDS. @@ -54,7 +65,7 @@ class _AcceleratorClient: def __init__(self, uds_path: str, auth_secret: str): self._uds_path = uds_path self._metadata = (("x-accelerator-token", auth_secret),) - self._channel = insecure_channel(f"unix://{uds_path}") + self._channel = insecure_channel(_uds_target(uds_path)) self._mutate_row_stub = self._channel.unary_unary( _MUTATE_ROW_METHOD, request_serializer=MutateRowRequest.serialize, 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 94531d09130d..2b4cb8418aad 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 @@ -23,7 +23,6 @@ import logging import os import random -import sys import time import warnings from functools import partial @@ -1080,12 +1079,6 @@ def _maybe_start_accelerator(self, *, explicit: bool) -> None: ``use_accelerator=True``: a caller who required the accelerator gets a hard error rather than a silent fallback, both for the emulator (a genuine misconfiguration) and for a daemon that fails to start.""" - if sys.platform == "win32": - if explicit: - raise RuntimeError( - "use_accelerator=True is not supported on Windows; pass use_accelerator=False to use the native client." - ) - return if self.client._emulator_host is not None: if explicit: raise RuntimeError( 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 index 211d5e4010c9..6bb7808025a6 100644 --- 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 @@ -95,7 +95,9 @@ def __init__(self, uds_path, *, status=ServingStatus.SERVING, serve_health=True) self._server.add_generic_rpc_handlers( (grpc.method_handlers_generic_handler(service, {method: handler}),) ) - self._server.add_insecure_port(f"unix://{uds_path}") + # Single-colon form, same as the client — ``unix://`` cannot express a + # Windows path. See ``_accelerator_client._uds_target``. + self._server.add_insecure_port(f"unix:{uds_path}") self._server.start() def _check(self, request, context): 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 index f379a5963715..4d43ffa77187 100644 --- 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 @@ -85,7 +85,7 @@ def __init__(self, uds_path, *, status=ServingStatus.SERVING, serve_health=True) 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.add_insecure_port(f"unix:{uds_path}") self._server.start() def _check(self, request, context): diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index b0bd0f1e575a..12fb53628886 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -201,9 +201,29 @@ def fake_mkdtemp(prefix=""): return str(d) monkeypatch.setattr(_daemon.tempfile, "mkdtemp", fake_mkdtemp) + + # start() also fires the stale-tempdir sweep on a background thread, and + # that reads gettempdir() directly rather than going through mkdtemp. Point + # it somewhere empty: the sweep deletes directories it judges abandoned, and + # a test has no business making that judgement about the real /tmp, where a + # developer's or another test's daemon may be running. + sweep_root = tmp_path / "sweep-root" + os.makedirs(sweep_root, exist_ok=True) + monkeypatch.setattr(_daemon.tempfile, "gettempdir", lambda: str(sweep_root)) yield +@pytest.fixture +def no_af_unix(monkeypatch): + """Make the socket module look the way CPython's does on Windows. + + The family is simply absent there — the OS and gRPC's C-core support + AF_UNIX, the socket module just never exposed it — so every branch keyed on + ``hasattr(socket, "AF_UNIX")`` is reachable from a POSIX test run. + """ + monkeypatch.delattr(_daemon.socket, "AF_UNIX", raising=False) + + # --------------------------------------------------------------------------- # Properties before/after start # --------------------------------------------------------------------------- @@ -257,6 +277,17 @@ def test_missing_binary_raises_file_not_found(self, tmp_path): with pytest.raises(FileNotFoundError, match="Accelerator binary not found"): _resolve_binary_path() + @pytest.mark.parametrize( + "is_windows,expected", + [(False, "bin/accelerator"), (True, "bin/accelerator.exe")], + ) + def test_binary_name_follows_platform(self, monkeypatch, is_windows, expected): + """Windows wheels stage the binary under the .exe name Go gives it.""" + monkeypatch.setattr(_daemon, "_IS_WINDOWS", is_windows) + with unittest.mock.patch("os.path.isfile", return_value=True): + resolved = _daemon._resolve_binary_path() + assert resolved.endswith(expected) + # --------------------------------------------------------------------------- # start() failure modes @@ -314,23 +345,10 @@ def test_uds_never_appears_raises_timeout(self, tmp_path, monkeypatch): with pytest.raises(RuntimeError, match="did not become ready"): d.start() - def test_no_af_unix_raises_os_error(self, tmp_path, monkeypatch): - import socket as _socket - - monkeypatch.delattr(_socket, "AF_UNIX", raising=False) - # Reload the attribute check in _daemon - monkeypatch.setattr(_daemon.socket, "AF_UNIX", None, raising=False) - # Simpler: just check that hasattr(socket, "AF_UNIX") guards the call. - with ( - unittest.mock.patch.object(_daemon.socket, "AF_UNIX", None, create=True), - unittest.mock.patch( - "google.cloud.bigtable.data._accelerator._daemon.hasattr", - return_value=False, - ), - ): - d = _make_daemon(tmp_path) - with pytest.raises(OSError, match="AF_UNIX"): - d.start() + def test_no_af_unix_raises_os_error(self, tmp_path, monkeypatch, no_af_unix): + d = _make_daemon(tmp_path) + with pytest.raises(OSError, match="AF_UNIX"): + d.start() # --------------------------------------------------------------------------- @@ -818,3 +836,98 @@ def boom(*args, **kwargs): with pytest.raises(RuntimeError): daemon.start() assert done.wait(timeout=2.0), "sweep thread did not run within 2 s" + + +# --------------------------------------------------------------------------- +# Windows (win_amd64 / win32) +# --------------------------------------------------------------------------- + + +class TestWindows: + """Windows reaches the daemon over the same AF_UNIX socket as every other + platform. What differs is that CPython's socket module does not expose the + family, so anything the other platforms answer with a probe connection has + to be read off the filesystem instead.""" + + def test_start_is_not_blocked_by_missing_af_unix( + self, tmp_path, monkeypatch, no_af_unix + ): + """The AF_UNIX guard must not fire on Windows: a missing family there + says something about CPython, not about the platform.""" + monkeypatch.setattr(_daemon, "_IS_WINDOWS", True) + # The sweep is not under test here, and leaving it to run would put a + # background thread in Windows mode racing this test's teardown: once + # monkeypatch restores gettempdir the thread reaps against the real + # /tmp, where (AF_UNIX gone) the Windows branch's "the OS will refuse + # to unlink a live daemon's log" assumption does not hold. + monkeypatch.setattr(_daemon, "_sweep_stale_tempdirs", lambda: None) + + def fake_popen(argv, **kwargs): + raise _StopStart() + + monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) + d = _make_daemon(tmp_path) + # Reaching Popen means the guard let us through. + with pytest.raises(_StopStart): + d.start() + + def test_socket_is_live_falls_back_to_path_existence(self, tmp_path, no_af_unix): + sock = tmp_path / _daemon._SOCKET_FILENAME + assert _daemon._socket_is_live(str(sock)) is False + sock.touch() + assert _daemon._socket_is_live(str(sock)) is True + + def test_sweep_is_skipped_on_windows(self, tmp_path, monkeypatch): + """The sweep's liveness check is ``os.kill(pid, 0)``, which terminates + the process on Windows rather than probing it. Until that is ported, + the sweep must not run there at all — leaking a tempdir is the safe + failure, killing the caller's Python process is not.""" + monkeypatch.setattr(_daemon, "_IS_WINDOWS", True) + root = tmp_path / "tmproot" + root.mkdir() + monkeypatch.setattr(_daemon.tempfile, "gettempdir", lambda: str(root)) + d = root / "bt-accel-dead" + d.mkdir() + (d / _daemon._LOG_FILENAME).write_bytes(b"hi") + (d / _daemon._STARTUP_DEADLINE_FILENAME).write_text(repr(time.time() - 70)) + # A PID that is certainly not alive, so POSIX would reap this. + (d / _daemon._PID_FILENAME).write_text("999999999") + + def fail_if_called(*args, **kwargs): + raise AssertionError("os.kill must never run on Windows") + + monkeypatch.setattr(_daemon.os, "kill", fail_if_called) + _daemon._sweep_stale_tempdirs() + assert d.exists() + + def test_uds_target_accepts_a_windows_path(self): + """The channel target must survive a drive-lettered path. + + ``unix://`` opens a URI authority component, so ``unix://C:\\...`` + parses ``C:`` as the authority and gRPC refuses to build the channel at + all — the accelerator would start, verify, and then fail every routed + RPC. Constructing the channel is the assertion; it raises on a bad + target before any connection is attempted. + """ + import grpc + + from google.cloud.bigtable.data._sync_autogen._accelerator_client import ( + _uds_target, + ) + + win_path = r"C:\Users\me\AppData\Local\Temp\bt-accel-abc123\sock" + assert _uds_target(win_path) == "unix:" + win_path + channel = grpc.insecure_channel(_uds_target(win_path)) + try: + stub = channel.unary_unary( + "/google.bigtable.v2.Bigtable/MutateRow", + request_serializer=lambda _: b"", + response_deserializer=lambda payload: payload, + ) + with pytest.raises(grpc.RpcError) as exc_info: + stub(None, timeout=0.5) + finally: + channel.close() + # UNAVAILABLE means the target resolved and nothing was listening, which + # is the expected outcome. A malformed target fails earlier, as UNKNOWN. + assert exc_info.value.code() == grpc.StatusCode.UNAVAILABLE From 84f3b433788498a0c72d8efbbde57b2dd48670dd Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 20 Oct 2026 17:56:17 +0000 Subject: [PATCH 24/26] test(bigtable): align Windows sweep skip condition and shorten health test UDS path --- .../unit/data/_async/test__accelerator_health.py | 15 +++++++++++++++ .../_sync_autogen/test__accelerator_health.py | 9 +++++++++ .../tests/unit/data/test_accelerator_daemon.py | 15 ++++++++------- 3 files changed, 32 insertions(+), 7 deletions(-) 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 index 6bb7808025a6..5b20dc6fe5ce 100644 --- 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 @@ -26,6 +26,8 @@ ``tests/unit/data/test_accelerator_health.py``. """ +import pathlib +import tempfile import threading import time from concurrent import futures @@ -68,6 +70,19 @@ } +# sockaddr_un.sun_path is 108 bytes on Linux and Windows (104 on macOS). +# Pytest's built-in tmp_path nests pytest-of-/pytest-N// under +# %TEMP%, which on Windows and macOS can exceed 107 chars and make +# add_insecure_port fail to bind. Matching the daemon's own +# tempfile.mkdtemp(prefix="bt-accel-") placement keeps the socket path well +# under the limit on every platform. +@pytest.fixture +def tmp_path(): + """Short UDS directory directly under the system temp root.""" + with tempfile.TemporaryDirectory(prefix="bt-accel-") as d: + yield pathlib.Path(d) + + class _FakeDaemonHealthServer: """A real gRPC server standing in for the daemon's health service. 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 index 4d43ffa77187..d5f0f8db75ee 100644 --- 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 @@ -29,6 +29,8 @@ ``tests/unit/data/test_accelerator_health.py``. """ +import pathlib +import tempfile import threading import time from concurrent import futures @@ -58,6 +60,13 @@ } +@pytest.fixture +def tmp_path(): + """Short UDS directory directly under the system temp root.""" + with tempfile.TemporaryDirectory(prefix="bt-accel-") as d: + yield pathlib.Path(d) + + class _FakeDaemonHealthServer: """A real gRPC server standing in for the daemon's health service. diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index 12fb53628886..2a52049ffd17 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -709,8 +709,8 @@ def test_daemon_constructor_invalid_env_var_raises(self, tmp_path, monkeypatch): @pytest.mark.skipif( - not hasattr(_daemon.socket, "AF_UNIX"), - reason="Unix domain sockets (AF_UNIX) are not supported on this platform", + _daemon._IS_WINDOWS or not hasattr(_daemon.socket, "AF_UNIX"), + reason="Stale tempdir sweep is skipped on Windows and platforms without AF_UNIX", ) class TestSweepStaleTempdirs: """_sweep_stale_tempdirs reaps leftover daemon tempdirs, keyed on the @@ -855,11 +855,8 @@ def test_start_is_not_blocked_by_missing_af_unix( """The AF_UNIX guard must not fire on Windows: a missing family there says something about CPython, not about the platform.""" monkeypatch.setattr(_daemon, "_IS_WINDOWS", True) - # The sweep is not under test here, and leaving it to run would put a - # background thread in Windows mode racing this test's teardown: once - # monkeypatch restores gettempdir the thread reaps against the real - # /tmp, where (AF_UNIX gone) the Windows branch's "the OS will refuse - # to unlink a live daemon's log" assumption does not hold. + # The sweep is not under test here; stub it out so start() does not + # spawn a background thread for it. monkeypatch.setattr(_daemon, "_sweep_stale_tempdirs", lambda: None) def fake_popen(argv, **kwargs): @@ -911,12 +908,16 @@ def test_uds_target_accepts_a_windows_path(self): """ import grpc + from google.cloud.bigtable.data._async._accelerator_client import ( + _uds_target as _async_uds_target, + ) from google.cloud.bigtable.data._sync_autogen._accelerator_client import ( _uds_target, ) win_path = r"C:\Users\me\AppData\Local\Temp\bt-accel-abc123\sock" assert _uds_target(win_path) == "unix:" + win_path + assert _async_uds_target(win_path) == _uds_target(win_path) channel = grpc.insecure_channel(_uds_target(win_path)) try: stub = channel.unary_unary( From 32f0f823593744ae6faf31509a617fd61b744cf1 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 20 Oct 2026 17:46:06 +0000 Subject: [PATCH 25/26] feat(bigtable): raise the UDS channel inbound message cap to 256 MB Change-Id: I6adb2982640fe2d6736ea63ed165c9ffd6feacb7 --- .../bigtable/data/_async/_accelerator_client.py | 14 +++++++++++++- .../data/_sync_autogen/_accelerator_client.py | 6 +++++- 2 files changed, 18 insertions(+), 2 deletions(-) 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 e6953f676664..0fc82484e98f 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 @@ -45,6 +45,16 @@ _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``. @@ -70,7 +80,9 @@ class _AsyncAcceleratorClient: def __init__(self, uds_path: str, auth_secret: str): self._uds_path = uds_path self._metadata = (("x-accelerator-token", auth_secret),) - self._channel = insecure_channel(_uds_target(uds_path)) + 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, 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 51905d4b04cd..1d2329bee869 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 @@ -41,6 +41,8 @@ _MUTATE_ROW_METHOD = "/google.bigtable.v2.Bigtable/MutateRow" _READ_ROWS_METHOD = "/google.bigtable.v2.Bigtable/ReadRows" +_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: @@ -65,7 +67,9 @@ class _AcceleratorClient: def __init__(self, uds_path: str, auth_secret: str): self._uds_path = uds_path self._metadata = (("x-accelerator-token", auth_secret),) - self._channel = insecure_channel(_uds_target(uds_path)) + 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, From 3c56e1a8f119cd020b4fb3be15ed99a08ae163af Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 20 Oct 2026 21:48:11 +0000 Subject: [PATCH 26/26] test(bigtable): add accelerator system test harness and concurrency tests --- packages/google-cloud-bigtable/conftest.py | 48 ++ packages/google-cloud-bigtable/pytest.ini | 18 + .../tests/system/data/accelerator/__init__.py | 13 + .../system/data/accelerator/_base_async.py | 109 +++++ .../system/data/accelerator/_base_autogen.py | 91 ++++ .../tests/system/data/accelerator/_harness.py | 425 ++++++++++++++++++ .../tests/system/data/accelerator/conftest.py | 33 ++ .../accelerator/test_concurrency_async.py | 190 ++++++++ .../accelerator/test_concurrency_autogen.py | 122 +++++ 9 files changed, 1049 insertions(+) create mode 100644 packages/google-cloud-bigtable/conftest.py create mode 100644 packages/google-cloud-bigtable/pytest.ini create mode 100644 packages/google-cloud-bigtable/tests/system/data/accelerator/__init__.py create mode 100644 packages/google-cloud-bigtable/tests/system/data/accelerator/_base_async.py create mode 100644 packages/google-cloud-bigtable/tests/system/data/accelerator/_base_autogen.py create mode 100644 packages/google-cloud-bigtable/tests/system/data/accelerator/_harness.py create mode 100644 packages/google-cloud-bigtable/tests/system/data/accelerator/conftest.py create mode 100644 packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_async.py create mode 100644 packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_autogen.py diff --git a/packages/google-cloud-bigtable/conftest.py b/packages/google-cloud-bigtable/conftest.py new file mode 100644 index 000000000000..8a6db20dd159 --- /dev/null +++ b/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 diff --git a/packages/google-cloud-bigtable/pytest.ini b/packages/google-cloud-bigtable/pytest.ini new file mode 100644 index 000000000000..2e5313bfa236 --- /dev/null +++ b/packages/google-cloud-bigtable/pytest.ini @@ -0,0 +1,18 @@ +[pytest] +# --- asyncio event-loop scoping ------------------------------------------- +# The system-test suites (tests/system/**) share session-scoped async +# client/table fixtures across the tests in a class -- see +# tests/system/conftest.py and e.g. tests/system/data/test_system_async.py and +# tests/system/data/accelerator/_base_async.py. A session-scoped async fixture +# and the tests that use it must run on the *same* event loop, otherwise the +# grpc.aio call created in the fixture's loop is awaited from the test's loop and +# raises "got Future <...> attached to a different loop". +# +# pytest-asyncio <=0.21 enforced a single loop via the session-scoped +# ``event_loop`` fixture override in tests/system/conftest.py. pytest-asyncio +# >=0.23 ignores that override and instead reads the loop scope from these two +# options (default: function). Pinning both to "session" preserves the +# single-loop behavior on modern pytest-asyncio; older versions ignore these +# keys (harmless "Unknown config option" warning) and keep using the override. +asyncio_default_fixture_loop_scope = session +asyncio_default_test_loop_scope = session diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/__init__.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/__init__.py new file mode 100644 index 000000000000..58d482ea3866 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/_base_async.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/_base_async.py new file mode 100644 index 000000000000..3d100e11d83e --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/_base_async.py @@ -0,0 +1,109 @@ +# 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. + +import os + +import pytest + +from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable.data.mutations import DeleteAllFromRow, RowMutationEntry + +from .. import SystemTestRunner +from . import _harness + +__CROSS_SYNC_OUTPUT__ = "tests.system.data.accelerator._base_autogen" + + +@CrossSync.convert_class(sync_name="AcceleratorTestBase") +class AcceleratorTestBaseAsync(SystemTestRunner): + @pytest.fixture(scope="session", autouse=True) + def _require_accel_env(self): + _harness.require_real_bigtable_or_skip() + _harness.require_binary_or_skip() + + def _project_id(self): + return os.getenv("GOOGLE_CLOUD_PROJECT") or None + + def _make_client(self, use_accelerator=None, **kwargs): + return CrossSync.DataClient( + project=self._project_id(), use_accelerator=use_accelerator, **kwargs + ) + + def assert_accelerator_active(self, table): + assert table._accelerator_client is not None, ( + "expected the accelerator to be active (use_accelerator=True) but the " + "target fell back to native. Check ADC principal / identity " + "verification and that the bundled daemon binary can start." + ) + + def assert_native(self, table): + assert table._accelerator_client is None, ( + "expected a native target (use_accelerator=False) but an accelerator " + "client was attached." + ) + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def client(self): + async with self._make_client() as client: + yield client + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def accel_client(self): + async with self._make_client(use_accelerator=True) as client: + yield client + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def native_client(self): + async with self._make_client(use_accelerator=False) as client: + yield client + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def accel_table(self, accel_client, instance_id, table_id): + async with accel_client.get_table(instance_id, table_id) as table: + self.assert_accelerator_active(table) + yield table + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def native_table(self, native_client, instance_id, table_id): + async with native_client.get_table(instance_id, table_id) as table: + self.assert_native(table) + yield table + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def default_table(self, client, instance_id, table_id): + async with client.get_table(instance_id, table_id) as table: + yield table + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="function") + async def janitor(self, native_table): + class _Janitor: + def __init__(self): + self.keys = set() + + def track(self, key): + self.keys.add(key) + return key + + j = _Janitor() + yield j + if j.keys: + entries = [RowMutationEntry(key, [DeleteAllFromRow()]) for key in j.keys] + await native_table.bulk_mutate_rows(entries) diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/_base_autogen.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/_base_autogen.py new file mode 100644 index 000000000000..5e81eb347b54 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/_base_autogen.py @@ -0,0 +1,91 @@ +# 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. + +import os +import pytest +from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable.data.mutations import DeleteAllFromRow, RowMutationEntry +from .. import SystemTestRunner +from . import _harness + +class AcceleratorTestBase(SystemTestRunner): + + @pytest.fixture(scope='session', autouse=True) + def _require_accel_env(self): + _harness.require_real_bigtable_or_skip() + _harness.require_binary_or_skip() + + def _project_id(self): + return os.getenv('GOOGLE_CLOUD_PROJECT') or None + + def _make_client(self, use_accelerator=None, **kwargs): + return CrossSync._Sync_Impl.DataClient(project=self._project_id(), use_accelerator=use_accelerator, **kwargs) + + def assert_accelerator_active(self, table): + assert table._accelerator_client is not None, 'expected the accelerator to be active (use_accelerator=True) but the target fell back to native. Check ADC principal / identity verification and that the bundled daemon binary can start.' + + def assert_native(self, table): + assert table._accelerator_client is None, 'expected a native target (use_accelerator=False) but an accelerator client was attached.' + + @pytest.fixture(scope='session') + def client(self): + with self._make_client() as client: + yield client + + @pytest.fixture(scope='session') + def accel_client(self): + with self._make_client(use_accelerator=True) as client: + yield client + + @pytest.fixture(scope='session') + def native_client(self): + with self._make_client(use_accelerator=False) as client: + yield client + + @pytest.fixture(scope='session') + def accel_table(self, accel_client, instance_id, table_id): + with accel_client.get_table(instance_id, table_id) as table: + self.assert_accelerator_active(table) + yield table + + @pytest.fixture(scope='session') + def native_table(self, native_client, instance_id, table_id): + with native_client.get_table(instance_id, table_id) as table: + self.assert_native(table) + yield table + + @pytest.fixture(scope='session') + def default_table(self, client, instance_id, table_id): + with client.get_table(instance_id, table_id) as table: + yield table + + @pytest.fixture(scope='function') + def janitor(self, native_table): + + class _Janitor: + + def __init__(self): + self.keys = set() + + def track(self, key): + self.keys.add(key) + return key + j = _Janitor() + yield j + if j.keys: + entries = [RowMutationEntry(key, [DeleteAllFromRow()]) for key in j.keys] + native_table.bulk_mutate_rows(entries) diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/_harness.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/_harness.py new file mode 100644 index 000000000000..a85f73b872ea --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/_harness.py @@ -0,0 +1,425 @@ +# 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. + +from __future__ import annotations + +import gc +import glob +import os +import stat +import sys +import time +from dataclasses import dataclass +from typing import Any, Callable, Iterable, Sequence + +from google.cloud.bigtable.data import ( + DeleteAllFromFamily, + DeleteAllFromRow, + DeleteRangeFromColumn, + SetCell, +) +from google.cloud.bigtable.data._accelerator import _daemon +from google.cloud.bigtable.data.row import Row + +# Millisecond timestamp granularity (Bigtable default). +MS = 1000 + +# Column families created by the shared system-test conftest. +TEST_FAMILY = "test-family" +TEST_FAMILY_2 = "test-family-2" + +# Env var used by tests / runners to override the binary location. +BIN_ENV_VAR = "BIGTABLE_ACCELERATOR_BIN" + +Expected = tuple # (family: str, qualifier: bytes, timestamp_micros: int, value: bytes) + + +class ExpectedState: + def __init__(self) -> None: + self._rows: dict[bytes, dict[tuple[str, bytes], dict[int, bytes]]] = {} + + def set_cell( + self, row_key: bytes, family: str, qualifier: bytes, value: bytes, ts: int + ) -> SetCell: + if ts % MS != 0: + raise ValueError( + "timestamps must be millisecond-granular (multiple of 1000)" + ) + col = self._rows.setdefault(row_key, {}).setdefault((family, qualifier), {}) + col[ts] = value + return SetCell(family, qualifier, value, timestamp_micros=ts) + + def delete_range_from_column( + self, + row_key: bytes, + family: str, + qualifier: bytes, + start: int | None = None, + end: int | None = None, + ) -> DeleteRangeFromColumn: + col = self._rows.get(row_key, {}).get((family, qualifier)) + if col is not None: + lo = 0 if start is None else start + for ts in list(col): + if ts >= lo and (end is None or ts < end): + del col[ts] + if not col: + self._rows[row_key].pop((family, qualifier), None) + return DeleteRangeFromColumn(family, qualifier, start, end) + + def delete_from_family(self, row_key: bytes, family: str) -> DeleteAllFromFamily: + row = self._rows.get(row_key) + if row is not None: + for key in [k for k in row if k[0] == family]: + del row[key] + return DeleteAllFromFamily(family) + + def delete_from_row(self, row_key: bytes) -> DeleteAllFromRow: + self._rows.pop(row_key, None) + return DeleteAllFromRow() + + def expected_cells(self, row_key: bytes) -> list[Expected]: + row = self._rows.get(row_key, {}) + out: list[Expected] = [] + for family, qualifier in sorted(row, key=lambda k: (k[0], k[1])): + for ts in sorted(row[(family, qualifier)], reverse=True): + out.append((family, qualifier, ts, row[(family, qualifier)][ts])) + return out + + def row_is_empty(self, row_key: bytes) -> bool: + return not self._rows.get(row_key) + + def keys(self) -> Iterable[bytes]: + return list(self._rows) + + +def normalize_row(row: Row | None) -> list[Expected]: + if row is None: + return [] + return [ + (cell.family, cell.qualifier, cell.timestamp_micros, cell.value) + for cell in row.cells + ] + + +def _fmt(cells: Sequence[Expected]) -> str: + return ( + "\n".join(f" {fam}:{qual!r}@{ts} = {val!r}" for (fam, qual, ts, val) in cells) + or " " + ) + + +def assert_rows_equivalent( + label_a: str, cells_a: Sequence[Expected], label_b: str, cells_b: Sequence[Expected] +) -> None: + if list(cells_a) != list(cells_b): + raise AssertionError( + f"row mismatch between {label_a} and {label_b}:\n" + f"{label_a}:\n{_fmt(cells_a)}\n{label_b}:\n{_fmt(cells_b)}" + ) + + +def resolve_binary() -> str | None: + override = os.environ.get(BIN_ENV_VAR) + if override: + return override if os.path.isfile(override) else None + try: + return _daemon._resolve_binary_path() + except FileNotFoundError: + return None + + +def require_binary_or_skip() -> str: + import pytest + + path = resolve_binary() + if path is None: + pytest.skip( + "No accelerator daemon binary available " + f"(set {BIN_ENV_VAR} or install a wheel that bundles it)." + ) + return path + + +def require_real_bigtable_or_skip() -> None: + import pytest + from google.cloud.environment_vars import BIGTABLE_EMULATOR + + if os.environ.get(BIGTABLE_EMULATOR): + pytest.skip("accelerator is not supported against the emulator") + if not (os.environ.get("GOOGLE_CLOUD_PROJECT") or os.environ.get("PROJECT_ID")): + pytest.skip("no GOOGLE_CLOUD_PROJECT set for live accelerator tests") + + +def _write_script( + directory: str, + name: str, + body: str, + *, + executable: bool = True, + monkeypatch: Any = None, +) -> str: + path = os.path.join(directory, name) + with open(path, "w") as f: + f.write(body) + if executable: + mode = os.stat(path).st_mode + os.chmod(path, mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) + if monkeypatch is not None: + monkeypatch.setenv(BIN_ENV_VAR, path) + monkeypatch.setattr(_daemon, "_resolve_binary_path", lambda: path) + return path + + +def nonexecutable_binary(directory: str, monkeypatch: Any = None) -> str: + return _write_script( + directory, + "not-exec", + "#!/bin/sh\nexit 0\n", + executable=False, + monkeypatch=monkeypatch, + ) + + +def immediately_exiting_binary( + directory: str, exit_code: int = 1, monkeypatch: Any = None +) -> str: + return _write_script( + directory, + "exits-now", + f"#!/bin/sh\nprintf 'boom\\n' 1>&2\nexit {exit_code}\n", + monkeypatch=monkeypatch, + ) + + +broken_binary = immediately_exiting_binary + + +def never_binds_binary(directory: str, monkeypatch: Any = None) -> str: + return _write_script( + directory, + "never-binds", + "#!/bin/sh\nexec sleep 3600\n", + monkeypatch=monkeypatch, + ) + + +hanging_binary = never_binds_binary + + +def slow_bind_binary( + directory: str, delay_seconds: float, monkeypatch: Any = None +) -> str: + body = f"""#!{sys.executable} +import os, socket, sys, time, signal + +def uds_path(argv): + for i, a in enumerate(argv): + if a == "--uds-path": + return argv[i + 1] + raise SystemExit("no --uds-path") + +path = uds_path(sys.argv) +try: + sys.stdin.readline() +except Exception: + pass +time.sleep({delay_seconds!r}) +try: + os.unlink(path) +except OSError: + pass +srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) +srv.bind(path) +srv.listen(8) +signal.signal(signal.SIGTERM, lambda *a: os._exit(0)) +try: + while True: + line = sys.stdin.readline() + if line == "": + break +except Exception: + pass +os._exit(0) +""" + return _write_script(directory, "slow-bind", body, monkeypatch=monkeypatch) + + +ACCEL_TEMPDIR_GLOB = os.path.join("/tmp", "bt-accel-*") + + +def accel_tempdirs() -> set[str]: + return set(glob.glob(ACCEL_TEMPDIR_GLOB)) + + +@dataclass +class LeakSnapshot: + child_pids: set[int] + num_fds: int + tempdirs: set[str] + + +class ProcessIntrospector: + def __init__(self) -> None: + import psutil + + self._psutil = psutil + self._proc = psutil.Process(os.getpid()) + + def snapshot(self) -> LeakSnapshot: + children = set() + for child in self._proc.children(recursive=True): + try: + children.add(child.pid) + except self._psutil.Error: + pass + try: + num_fds = self._proc.num_fds() + except (self._psutil.Error, AttributeError): + num_fds = -1 + return LeakSnapshot( + child_pids=children, num_fds=num_fds, tempdirs=accel_tempdirs() + ) + + def assert_no_leaks( + self, + before: LeakSnapshot, + *, + fd_slack: int = 8, + label: str = "", + settle_timeout: float = 5.0, + ) -> None: + after = self.snapshot() + prefix = f"[{label}] " if label else "" + leaked_children = after.child_pids - before.child_pids + leaked_dirs = after.tempdirs - before.tempdirs + assert not leaked_children, ( + f"{prefix}leaked {len(leaked_children)} daemon subprocess(es): " + f"{sorted(leaked_children)}" + ) + assert not leaked_dirs, f"{prefix}leaked accel tempdirs: {sorted(leaked_dirs)}" + if before.num_fds < 0 or after.num_fds < 0: + return + gc.collect() + min_fds = after.num_fds + deadline = time.monotonic() + settle_timeout + while min_fds > before.num_fds + fd_slack and time.monotonic() < deadline: + time.sleep(0.1) + cur = self.snapshot().num_fds + if cur >= 0: + min_fds = min(min_fds, cur) + assert min_fds <= before.num_fds + fd_slack, ( + f"{prefix}fd count grew from {before.num_fds} to {min_fds} " + f"(slack {fd_slack}) and did not drain within {settle_timeout:g}s" + ) + + +def call_with_timeout(fn: Callable[[], Any], timeout: float) -> Any: + import threading + + box: dict[str, Any] = {} + + def _run() -> None: + try: + box["result"] = fn() + except BaseException as exc: # noqa: BLE001 - re-raised on caller thread + box["error"] = exc + + thread = threading.Thread(target=_run, daemon=True) + thread.start() + thread.join(timeout) + if thread.is_alive(): + raise TimeoutError(f"call did not return within {timeout}s") + if "error" in box: + raise box["error"] + return box.get("result") + + +def force_kill_daemon(daemon: Any) -> None: + proc = getattr(daemon, "_proc", None) + if proc is None: + return + try: + proc.kill() + except Exception: + pass + try: + proc.wait(timeout=2.0) + except Exception: + pass + + +def daemon_pid(table: Any) -> int | None: + daemon = getattr(table, "_accelerator_daemon", None) + if daemon is None: + return None + try: + return daemon.pid + except RuntimeError: + return None + + +def pid_alive(pid: int) -> bool: + import psutil + + try: + proc = psutil.Process(pid) + return proc.is_running() and proc.status() != psutil.STATUS_ZOMBIE + except psutil.Error: + return False + + +class RandomOps: + def __init__(self, seed: int, key_prefix: bytes = b"accel-"): + import random + + self._rand = random.Random(seed) + self._keyspace = [key_prefix + f"{i:08d}".encode() for i in range(256)] + self._families = [TEST_FAMILY, TEST_FAMILY_2] + self._qualifiers = [f"q{i}".encode() for i in range(self._rand.randint(2, 16))] + + def key(self) -> bytes: + return self._rand.choice(self._keyspace) + + def value(self, max_len: int = 64) -> bytes: + n = self._rand.randint(0, max_len) + return bytes(self._rand.getrandbits(8) for _ in range(n)) + + def ms_timestamp(self) -> int: + return self._rand.randint(1, 2_000_000) * MS + + def family(self) -> str: + return self._rand.choice(self._families) + + def qualifier(self) -> bytes: + return self._rand.choice(self._qualifiers) + + def build_mutation(self, expected_state: ExpectedState) -> tuple[bytes, object]: + roll = self._rand.random() + key = self.key() + if roll < 0.70: + return key, expected_state.set_cell( + key, self.family(), self.qualifier(), self.value(), self.ms_timestamp() + ) + if roll < 0.82: + a = self._rand.choice([None, self.ms_timestamp()]) + b = self._rand.choice([None, self.ms_timestamp()]) + if a is not None and b is not None and a > b: + a, b = b, a + return key, expected_state.delete_range_from_column( + key, self.family(), self.qualifier(), a, b + ) + if roll < 0.92: + return key, expected_state.delete_from_family(key, self.family()) + return key, expected_state.delete_from_row(key) diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/conftest.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/conftest.py new file mode 100644 index 000000000000..83b1deb27f73 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/conftest.py @@ -0,0 +1,33 @@ +# 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. + +import tempfile + +import pytest + + +def pytest_configure(config): + config.addinivalue_line( + "markers", + "accelerator: pre-release integration tests for the Bigtable accelerator", + ) + config.addinivalue_line( + "markers", "slow: longer-running accelerator tests (leak loops, startup races)" + ) + + +@pytest.fixture +def tmp_bin_dir(): + with tempfile.TemporaryDirectory(prefix="accel-test-bin-") as d: + yield d diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_async.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_async.py new file mode 100644 index 000000000000..ffbf8d7ecb8d --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_async.py @@ -0,0 +1,190 @@ +# 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. + +import concurrent.futures +import functools +import multiprocessing +import queue +import time +import uuid + +from google.cloud.bigtable.data._cross_sync import CrossSync + +from . import _harness + +if CrossSync.is_async: + from ._base_async import AcceleratorTestBaseAsync as AcceleratorTestBase +else: + from ._base_autogen import AcceleratorTestBase + +__CROSS_SYNC_OUTPUT__ = "tests.system.data.accelerator.test_concurrency_autogen" + + +def _run_worker(table, worker_id, ops_per_worker): + prefix = f"conc-{worker_id}-{uuid.uuid4().hex[:8]}-".encode() + ops = _harness.RandomOps(worker_id, key_prefix=prefix) + expected_state = _harness.ExpectedState() + touched = set() + for _ in range(ops_per_worker): + key, mutation = ops.build_mutation(expected_state) + table.mutate_row(key, mutation) + touched.add(key) + for key in touched: + row = table.read_row(key) + _harness.assert_rows_equivalent( + f"worker-{worker_id}", + _harness.normalize_row(row), + "model", + expected_state.expected_cells(key), + ) + return sorted(touched) + + +def _fork_worker( + project, + instance_id, + table_id, + base_worker_id, + workers_per_process, + ops_per_worker, + result_queue, +): + from google.cloud.bigtable.data import BigtableDataClient + + try: + with BigtableDataClient(project=project, use_accelerator=True) as client: + with client.get_table(instance_id, table_id) as table: + if table._accelerator_client is None: + result_queue.put((False, [], "accelerator fell back to native")) + return + with concurrent.futures.ThreadPoolExecutor( + max_workers=workers_per_process + ) as executor: + futures = [ + executor.submit( + _run_worker, table, base_worker_id + w, ops_per_worker + ) + for w in range(workers_per_process) + ] + touched = [] + for future in futures: + touched.extend(future.result()) + result_queue.put((True, touched, None)) + except BaseException as exc: # noqa: BLE001 - reported to the parent process + result_queue.put((False, [], repr(exc))) + + +_process_workload = _fork_worker + + +@CrossSync.convert_class(sync_name="TestConcurrency") +class TestConcurrencyAsync(AcceleratorTestBase): + NUM_WORKERS = 16 + OPS_PER_WORKER = 50 + NUM_PROCESSES = 4 + WORKERS_PER_PROCESS = 8 + + @CrossSync.convert + async def _worker(self, table, worker_id): + prefix = f"conc-{worker_id}-{uuid.uuid4().hex[:8]}-".encode() + ops = _harness.RandomOps(worker_id, key_prefix=prefix) + expected_state = _harness.ExpectedState() + touched = set() + for _ in range(self.OPS_PER_WORKER): + key, mutation = ops.build_mutation(expected_state) + await table.mutate_row(key, mutation) + touched.add(key) + for key in touched: + row = await table.read_row(key) + _harness.assert_rows_equivalent( + f"worker-{worker_id}", + _harness.normalize_row(row), + "model", + expected_state.expected_cells(key), + ) + return sorted(touched) + + @CrossSync.pytest + async def test_concurrent_workers_stay_consistent(self, accel_table, janitor): + executor = ( + concurrent.futures.ThreadPoolExecutor(max_workers=self.NUM_WORKERS) + if not CrossSync.is_async + else None + ) + try: + partials = [ + functools.partial(self._worker, accel_table, i) + for i in range(self.NUM_WORKERS) + ] + results = await CrossSync.gather_partials( + partials, return_exceptions=True, sync_executor=executor + ) + finally: + if executor is not None: + executor.shutdown(wait=True) + + for i, result in enumerate(results): + assert not isinstance(result, BaseException), ( + f"worker {i} failed: {result!r}" + ) + for key in result: + janitor.track(key) + + @CrossSync.pytest + async def test_multiprocess_concurrent_clients_stay_consistent( + self, instance_id, table_id, janitor + ): + project = self._project_id() + ctx = multiprocessing.get_context("spawn") + ctx = multiprocessing.get_context("fork") + result_queue = ctx.Queue() + procs = [] + for p in range(self.NUM_PROCESSES): + proc = ctx.Process( + target=_fork_worker, + args=( + project, + instance_id, + table_id, + p * self.WORKERS_PER_PROCESS, + self.WORKERS_PER_PROCESS, + self.OPS_PER_WORKER, + result_queue, + ), + ) + proc.start() + procs.append(proc) + results = [] + deadline = time.monotonic() + 300.0 + try: + while len(results) < self.NUM_PROCESSES and time.monotonic() < deadline: + try: + results.append(result_queue.get(timeout=0.2)) + except queue.Empty: + if not any(p.is_alive() for p in procs) and result_queue.empty(): + break + finally: + for proc in procs: + proc.join(timeout=30) + if proc.is_alive(): + proc.kill() + proc.join() + + assert len(results) == self.NUM_PROCESSES, ( + f"expected {self.NUM_PROCESSES} results, got {len(results)}" + ) + for ok, touched, error in results: + assert ok, f"a worker process failed: {error}" + for key in touched: + janitor.track(key) diff --git a/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_autogen.py b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_autogen.py new file mode 100644 index 000000000000..beee9591b701 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/accelerator/test_concurrency_autogen.py @@ -0,0 +1,122 @@ +# 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. + +import concurrent.futures +import functools +import multiprocessing +import queue +import time +import uuid +from google.cloud.bigtable.data._cross_sync import CrossSync +from . import _harness +from ._base_autogen import AcceleratorTestBase + +def _run_worker(table, worker_id, ops_per_worker): + prefix = f'conc-{worker_id}-{uuid.uuid4().hex[:8]}-'.encode() + ops = _harness.RandomOps(worker_id, key_prefix=prefix) + expected_state = _harness.ExpectedState() + touched = set() + for _ in range(ops_per_worker): + key, mutation = ops.build_mutation(expected_state) + table.mutate_row(key, mutation) + touched.add(key) + for key in touched: + row = table.read_row(key) + _harness.assert_rows_equivalent(f'worker-{worker_id}', _harness.normalize_row(row), 'model', expected_state.expected_cells(key)) + return sorted(touched) + +def _fork_worker(project, instance_id, table_id, base_worker_id, workers_per_process, ops_per_worker, result_queue): + from google.cloud.bigtable.data import BigtableDataClient + try: + with BigtableDataClient(project=project, use_accelerator=True) as client: + with client.get_table(instance_id, table_id) as table: + if table._accelerator_client is None: + result_queue.put((False, [], 'accelerator fell back to native')) + return + with concurrent.futures.ThreadPoolExecutor(max_workers=workers_per_process) as executor: + futures = [executor.submit(_run_worker, table, base_worker_id + w, ops_per_worker) for w in range(workers_per_process)] + touched = [] + for future in futures: + touched.extend(future.result()) + result_queue.put((True, touched, None)) + except BaseException as exc: + result_queue.put((False, [], repr(exc))) +_process_workload = _fork_worker + +class TestConcurrency(AcceleratorTestBase): + NUM_WORKERS = 16 + OPS_PER_WORKER = 50 + NUM_PROCESSES = 4 + WORKERS_PER_PROCESS = 8 + + def _worker(self, table, worker_id): + prefix = f'conc-{worker_id}-{uuid.uuid4().hex[:8]}-'.encode() + ops = _harness.RandomOps(worker_id, key_prefix=prefix) + expected_state = _harness.ExpectedState() + touched = set() + for _ in range(self.OPS_PER_WORKER): + key, mutation = ops.build_mutation(expected_state) + table.mutate_row(key, mutation) + touched.add(key) + for key in touched: + row = table.read_row(key) + _harness.assert_rows_equivalent(f'worker-{worker_id}', _harness.normalize_row(row), 'model', expected_state.expected_cells(key)) + return sorted(touched) + + def test_concurrent_workers_stay_consistent(self, accel_table, janitor): + executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.NUM_WORKERS) if not CrossSync._Sync_Impl.is_async else None + try: + partials = [functools.partial(self._worker, accel_table, i) for i in range(self.NUM_WORKERS)] + results = CrossSync._Sync_Impl.gather_partials(partials, return_exceptions=True, sync_executor=executor) + finally: + if executor is not None: + executor.shutdown(wait=True) + for i, result in enumerate(results): + assert not isinstance(result, BaseException), f'worker {i} failed: {result!r}' + for key in result: + janitor.track(key) + + def test_multiprocess_concurrent_clients_stay_consistent(self, instance_id, table_id, janitor): + project = self._project_id() + ctx = multiprocessing.get_context('spawn') + ctx = multiprocessing.get_context('fork') + result_queue = ctx.Queue() + procs = [] + for p in range(self.NUM_PROCESSES): + proc = ctx.Process(target=_fork_worker, args=(project, instance_id, table_id, p * self.WORKERS_PER_PROCESS, self.WORKERS_PER_PROCESS, self.OPS_PER_WORKER, result_queue)) + proc.start() + procs.append(proc) + results = [] + deadline = time.monotonic() + 300.0 + try: + while len(results) < self.NUM_PROCESSES and time.monotonic() < deadline: + try: + results.append(result_queue.get(timeout=0.2)) + except queue.Empty: + if not any((p.is_alive() for p in procs)) and result_queue.empty(): + break + finally: + for proc in procs: + proc.join(timeout=30) + if proc.is_alive(): + proc.kill() + proc.join() + assert len(results) == self.NUM_PROCESSES, f'expected {self.NUM_PROCESSES} results, got {len(results)}' + for ok, touched, error in results: + assert ok, f'a worker process failed: {error}' + for key in touched: + janitor.track(key)