From 26ed3b9d1f814811f9d0ecfd9bd862130ac285b8 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 15 Sep 2026 20:26:28 +0000 Subject: [PATCH 1/4] feat: add accelerator daemon process wrapper Change-Id: I576365629b0f737edc86815208351759d548b734 --- .../bigtable/data/_accelerator/__init__.py | 15 + .../bigtable/data/_accelerator/_daemon.py | 332 ++++++++++++++++++ 2 files changed, 347 insertions(+) create mode 100644 packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/__init__.py create mode 100644 packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/__init__.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/__init__.py new file mode 100644 index 000000000000..914739165cf6 --- /dev/null +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/__init__.py @@ -0,0 +1,15 @@ +# 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. +# +"""Accelerator daemon subprocess and routing support.""" 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 new file mode 100644 index 000000000000..c3a8d269225b --- /dev/null +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -0,0 +1,332 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Subprocess lifecycle wrapper for the Go accelerator daemon binary. + +The daemon binary embeds an in-process Go Bigtable client and exposes the +``google.bigtable.v2.Bigtable`` service over a Unix domain socket. This module +owns spawning the binary, waiting for the UDS to become connectable, and +tearing the process down. It does NOT speak gRPC; that's the +``_AcceleratorClient`` companion's job. +""" + +from __future__ import annotations + +import os +import shutil +import signal +import socket +import subprocess +import tempfile +import time +from typing import Sequence + +# Environment variable that overrides the bundled binary location. Primarily +# for development against a locally-built daemon, and for tests pointing at a +# fake binary. +_BIN_ENV_VAR = "BIGTABLE_ACCELERATOR_BIN" + +# Wheels ship the binary at this path relative to the `_accelerator/` package. +_DEFAULT_BIN_RELATIVE_PATH = "bin/accelerator" + +# How long to wait for the daemon to start listening on its UDS before giving +# up at startup. +_DEFAULT_STARTUP_TIMEOUT = 10.0 + +# Sequence: close stdin, wait this long; SIGTERM, wait again; SIGKILL. +_STDIN_GRACE_SECONDS = 2.0 +_SIGTERM_GRACE_SECONDS = 2.0 +_SIGKILL_GRACE_SECONDS = 2.0 + + +def _resolve_binary_path(explicit_path: str | None = None) -> str: + """Resolve the daemon binary path, validating that it is a regular file. + + Precedence: an explicit ``binary_path`` argument, then the + ``BIGTABLE_ACCELERATOR_BIN`` env var, then the binary bundled in the wheel. + An explicit path or env override that does not point at a regular file is a + hard error (a caller who named a path meant it); a missing bundled binary + reports how to supply one. + """ + if explicit_path is not None: + if not os.path.isfile(explicit_path): + raise FileNotFoundError( + f"binary_path={explicit_path!r} does not point at a regular file" + ) + return explicit_path + override = os.environ.get(_BIN_ENV_VAR) + if override: + if not os.path.isfile(override): + raise FileNotFoundError( + f"{_BIN_ENV_VAR}={override!r} does not point at a regular file" + ) + return override + bundled = os.path.join(os.path.dirname(__file__), _DEFAULT_BIN_RELATIVE_PATH) + if not os.path.isfile(bundled): + raise FileNotFoundError( + "No accelerator binary found. Set the " + f"{_BIN_ENV_VAR} env var to a daemon binary path, or install a " + "wheel that bundles the binary." + ) + return bundled + + +class AcceleratorDaemon: + """Manages the Go accelerator daemon subprocess. + + The Python class is named for the thing it runs (the daemon hosts the + actual gRPC server). Lifecycle: + + 1. ``__init__`` validates and resolves the binary, picks the UDS path. + 2. ``start()`` spawns the subprocess and blocks until the UDS is + connectable, or raises if the process dies first. + 3. ``close()`` closes stdin (the daemon shuts down on EOF), then escalates + to SIGTERM and SIGKILL if it doesn't exit promptly. Cleans up the temp + directory holding the socket. + + This class is intentionally sync-only: ``subprocess.Popen`` works + identically for async and sync callers, and spawn/close happen once per + client lifetime — there's nothing to await. + """ + + def __init__( + self, + cli_flags: Sequence[str] = (), + *, + binary_path: str | None = None, + startup_timeout: float = _DEFAULT_STARTUP_TIMEOUT, + ): + """Resolve the binary and pick the UDS path (does not spawn anything). + + Args: + cli_flags: extra arguments appended after ``--uds-path`` when + spawning the daemon (e.g. ``--project``/``--instance``). + binary_path: explicit path to the daemon binary. When omitted, the + path is resolved from the ``BIGTABLE_ACCELERATOR_BIN`` env var + and then the binary bundled in the wheel. + startup_timeout: seconds ``start()`` waits for the daemon's UDS to + become connectable before raising. + + Raises: + FileNotFoundError: no binary could be resolved, or an explicit + ``binary_path``/env override does not point at a regular file. + """ + self._binary_path = _resolve_binary_path(binary_path) + self._cli_flags = list(cli_flags) + self._startup_timeout = startup_timeout + self._tempdir: str | None = None + self._uds_path: str | None = None + self._log_path: str | None = None + self._proc: subprocess.Popen[bytes] | None = None + + def __enter__(self) -> "AcceleratorDaemon": + """Start the daemon on ``with`` entry and return it.""" + self.start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + """Tear the daemon down on ``with`` exit.""" + self.close() + + @property + def uds_path(self) -> str: + if self._uds_path is None: + raise RuntimeError("AcceleratorDaemon has not been started") + return self._uds_path + + @property + def log_path(self) -> str: + if self._log_path is None: + raise RuntimeError("AcceleratorDaemon has not been started") + return self._log_path + + @property + def pid(self) -> int: + if self._proc is None: + raise RuntimeError("AcceleratorDaemon has not been started") + return self._proc.pid + + @property + def is_running(self) -> bool: + return self._proc is not None and self._proc.poll() is None + + def start(self) -> None: + """Spawn the daemon and wait for the UDS to become connectable. + + Raises: + RuntimeError: called twice, the daemon failed to spawn, exited + during startup, or did not become ready within + ``startup_timeout``. On any of these the child is killed and the + tempdir removed before the error propagates. + """ + if self._proc is not None: + raise RuntimeError("AcceleratorDaemon.start() called twice") + self._tempdir = tempfile.mkdtemp(prefix="bt-accel-") + self._uds_path = os.path.join(self._tempdir, "sock") + # Redirect the daemon's stdout/stderr to a log file rather than + # subprocess.PIPE. Nothing drains those pipes for the daemon's + # lifetime, so a PIPE's fixed OS buffer would eventually fill and + # block (deadlock) the daemon on its next write. A regular file has no + # such limit. The log lives alongside the socket in the daemon's + # 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") + # The log file handle only needs to live long enough for Popen to dup + # it into the child, so it stays local to start() rather than being an + # attribute. Startup failures read the tail back from the path. + log_file = open(self._log_path, "wb") + argv = [self._binary_path, "--uds-path", self._uds_path, *self._cli_flags] + try: + self._proc = subprocess.Popen( + argv, + stdin=subprocess.PIPE, + stdout=log_file, + stderr=subprocess.STDOUT, + close_fds=True, + ) + except OSError as exc: + self._cleanup_tempdir() + raise RuntimeError( + f"Failed to spawn accelerator daemon at {self._binary_path}: {exc}" + ) from exc + finally: + # Whether or not the spawn succeeded, the parent no longer needs its + # copy of the log fd: on success the child holds its own dup, and on + # failure there is nothing to keep open. + try: + log_file.close() + except OSError: + pass + try: + self._wait_until_ready(self._startup_timeout) + except BaseException: + self._force_kill() + self._cleanup_tempdir() + raise + + def close(self) -> None: + """Tear down the daemon and clean up the UDS tempdir.""" + proc = self._proc + if proc is None: + return + try: + if proc.poll() is None: + # Step 1: close stdin → daemon's stdin-EOF watchdog triggers + # graceful shutdown. + if proc.stdin is not None: + try: + proc.stdin.close() + except OSError: + pass + if not self._wait_for_exit(_STDIN_GRACE_SECONDS): + # Step 2: SIGTERM. + proc.terminate() + if not self._wait_for_exit(_SIGTERM_GRACE_SECONDS): + # Step 3: SIGKILL. Bounded wait so teardown can't hang + # forever if the process is stuck unreapable. + proc.kill() + self._wait_for_exit(_SIGKILL_GRACE_SECONDS) + finally: + self._proc = None + self._cleanup_tempdir() + + def _wait_until_ready(self, timeout: float) -> None: + """Poll until the UDS accepts a connection, the child dies, or timeout. + + Raises RuntimeError if the daemon exits during startup or does not + become connectable within ``timeout`` seconds. + """ + if self._proc is None or self._uds_path is None: + raise RuntimeError("AcceleratorDaemon._wait_until_ready() before spawn") + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + exit_code = self._proc.poll() + if exit_code is not None: + log_tail = self._read_log_tail() + raise RuntimeError( + "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 + time.sleep(0.05) + log_tail = self._read_log_tail() + raise RuntimeError( + "Accelerator daemon did not become ready within " + f"{timeout}s. log: {log_tail!r}" + ) + + def _read_log_tail(self, max_bytes: int = 4096) -> str: + """Best-effort read of the tail of the daemon's log file. + + Used only to enrich startup-failure error messages with whatever the + daemon wrote to stdout/stderr (both are redirected to the log file). + The log is a regular file, so this is a plain bounded read with no risk + of blocking on an alive-but-silent daemon. Returns "" if the log is + unavailable. + """ + if self._log_path is None: + return "" + try: + with open(self._log_path, "rb") as fh: + try: + fh.seek(-max_bytes, os.SEEK_END) + except OSError: + fh.seek(0) + # Bound the read too, so a log smaller than max_bytes (seek + # failed, fell back to seek(0)) is still capped. + data = fh.read(max_bytes) + except OSError: + return "" + return data.decode("utf-8", errors="replace") + + def _wait_for_exit(self, timeout: float) -> bool: + """Wait up to ``timeout`` seconds for the child to exit. + + Returns True if it has exited (or there is no child), False on timeout. + """ + if self._proc is None: + return True + try: + self._proc.wait(timeout=timeout) + return True + except subprocess.TimeoutExpired: + return False + + def _force_kill(self) -> None: + """SIGKILL the child and reap it; a no-op if it is already gone.""" + if self._proc is None or self._proc.poll() is not None: + return + try: + self._proc.send_signal(signal.SIGKILL) + except (OSError, ProcessLookupError): + pass + try: + self._proc.wait(timeout=1.0) + except subprocess.TimeoutExpired: + pass + + def _cleanup_tempdir(self) -> None: + """Remove the tempdir holding the socket and log; clear derived paths.""" + if self._tempdir is not None and os.path.isdir(self._tempdir): + shutil.rmtree(self._tempdir, ignore_errors=True) + self._tempdir = None + self._uds_path = None From dd73e1e1e8afce4461df8675173f714186cc650a Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 16 Sep 2026 01:48:24 +0000 Subject: [PATCH 2/4] fix: address review comments on accelerator daemon wrapper Change-Id: I7c218be39c2f62a87e8ca8e924717d6091d97ed7 --- .../cloud/bigtable/data/_accelerator/_daemon.py | 12 ++++++++++-- 1 file changed, 10 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 c3a8d269225b..b784da55b814 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,7 +25,6 @@ import os import shutil -import signal import socket import subprocess import tempfile @@ -197,6 +196,10 @@ def start(self) -> None: close_fds=True, ) except OSError as exc: + try: + log_file.close() + except OSError: + pass self._cleanup_tempdir() raise RuntimeError( f"Failed to spawn accelerator daemon at {self._binary_path}: {exc}" @@ -260,6 +263,10 @@ def _wait_until_ready(self, timeout: float) -> None: f"(exit code {exit_code}). log: {log_tail!r}" ) if os.path.exists(self._uds_path): + if not hasattr(socket, "AF_UNIX"): + raise OSError( + "Unix domain sockets (AF_UNIX) are not supported on this platform." + ) with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: probe.settimeout(0.25) try: @@ -316,7 +323,7 @@ def _force_kill(self) -> None: if self._proc is None or self._proc.poll() is not None: return try: - self._proc.send_signal(signal.SIGKILL) + self._proc.kill() except (OSError, ProcessLookupError): pass try: @@ -330,3 +337,4 @@ def _cleanup_tempdir(self) -> None: shutil.rmtree(self._tempdir, ignore_errors=True) self._tempdir = None self._uds_path = None + self._log_path = None From 398c73fafcde0eb5a5d8b6ad84993c4e51f228a0 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Fri, 18 Sep 2026 20:00:45 +0000 Subject: [PATCH 3/4] fix: address second round of review comments on accelerator daemon wrapper Change-Id: Ia8b08c13324612b478e47782db6042eecf4fcac7 --- .../bigtable/data/_accelerator/_daemon.py | 34 +++++++++++-------- 1 file changed, 19 insertions(+), 15 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 b784da55b814..7d92eaebb10c 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 @@ -171,6 +171,10 @@ def start(self) -> None: """ if self._proc is not None: raise RuntimeError("AcceleratorDaemon.start() called twice") + if not hasattr(socket, "AF_UNIX"): + raise OSError( + "Unix domain sockets (AF_UNIX) are not supported on this platform." + ) self._tempdir = tempfile.mkdtemp(prefix="bt-accel-") self._uds_path = os.path.join(self._tempdir, "sock") # Redirect the daemon's stdout/stderr to a log file rather than @@ -185,9 +189,10 @@ def start(self) -> None: # The log file handle only needs to live long enough for Popen to dup # it into the child, so it stays local to start() rather than being an # attribute. Startup failures read the tail back from the path. - log_file = open(self._log_path, "wb") argv = [self._binary_path, "--uds-path", self._uds_path, *self._cli_flags] + log_file = None try: + log_file = open(self._log_path, "wb") self._proc = subprocess.Popen( argv, stdin=subprocess.PIPE, @@ -196,10 +201,6 @@ def start(self) -> None: close_fds=True, ) except OSError as exc: - try: - log_file.close() - except OSError: - pass self._cleanup_tempdir() raise RuntimeError( f"Failed to spawn accelerator daemon at {self._binary_path}: {exc}" @@ -208,10 +209,11 @@ def start(self) -> None: # Whether or not the spawn succeeded, the parent no longer needs its # copy of the log fd: on success the child holds its own dup, and on # failure there is nothing to keep open. - try: - log_file.close() - except OSError: - pass + if log_file is not None: + try: + log_file.close() + except OSError: + pass try: self._wait_until_ready(self._startup_timeout) except BaseException: @@ -235,11 +237,17 @@ def close(self) -> None: pass if not self._wait_for_exit(_STDIN_GRACE_SECONDS): # Step 2: SIGTERM. - proc.terminate() + try: + proc.terminate() + except OSError: + pass if not self._wait_for_exit(_SIGTERM_GRACE_SECONDS): # Step 3: SIGKILL. Bounded wait so teardown can't hang # forever if the process is stuck unreapable. - proc.kill() + try: + proc.kill() + except OSError: + pass self._wait_for_exit(_SIGKILL_GRACE_SECONDS) finally: self._proc = None @@ -263,10 +271,6 @@ def _wait_until_ready(self, timeout: float) -> None: f"(exit code {exit_code}). log: {log_tail!r}" ) if os.path.exists(self._uds_path): - if not hasattr(socket, "AF_UNIX"): - raise OSError( - "Unix domain sockets (AF_UNIX) are not supported on this platform." - ) with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: probe.settimeout(0.25) try: From 16d683f45c579f4f505d732033b87dc2a79e59c6 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Fri, 18 Sep 2026 20:24:23 +0000 Subject: [PATCH 4/4] fix: remove binary_path override and env var; only use bundled binary Change-Id: I9e3cbbc360a39955459269526bd6efe417bb764f --- .../bigtable/data/_accelerator/_daemon.py | 42 +++---------------- 1 file changed, 5 insertions(+), 37 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 7d92eaebb10c..8f9a6b0ce813 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,11 +31,6 @@ import time from typing import Sequence -# Environment variable that overrides the bundled binary location. Primarily -# for development against a locally-built daemon, and for tests pointing at a -# fake binary. -_BIN_ENV_VAR = "BIGTABLE_ACCELERATOR_BIN" - # Wheels ship the binary at this path relative to the `_accelerator/` package. _DEFAULT_BIN_RELATIVE_PATH = "bin/accelerator" @@ -49,34 +44,12 @@ _SIGKILL_GRACE_SECONDS = 2.0 -def _resolve_binary_path(explicit_path: str | None = None) -> str: - """Resolve the daemon binary path, validating that it is a regular file. - - Precedence: an explicit ``binary_path`` argument, then the - ``BIGTABLE_ACCELERATOR_BIN`` env var, then the binary bundled in the wheel. - An explicit path or env override that does not point at a regular file is a - hard error (a caller who named a path meant it); a missing bundled binary - reports how to supply one. - """ - if explicit_path is not None: - if not os.path.isfile(explicit_path): - raise FileNotFoundError( - f"binary_path={explicit_path!r} does not point at a regular file" - ) - return explicit_path - override = os.environ.get(_BIN_ENV_VAR) - if override: - if not os.path.isfile(override): - raise FileNotFoundError( - f"{_BIN_ENV_VAR}={override!r} does not point at a regular file" - ) - return override +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) if not os.path.isfile(bundled): raise FileNotFoundError( - "No accelerator binary found. Set the " - f"{_BIN_ENV_VAR} env var to a daemon binary path, or install a " - "wheel that bundles the binary." + "Accelerator binary not found. Install a wheel that bundles the binary." ) return bundled @@ -103,7 +76,6 @@ def __init__( self, cli_flags: Sequence[str] = (), *, - binary_path: str | None = None, startup_timeout: float = _DEFAULT_STARTUP_TIMEOUT, ): """Resolve the binary and pick the UDS path (does not spawn anything). @@ -111,17 +83,13 @@ def __init__( Args: cli_flags: extra arguments appended after ``--uds-path`` when spawning the daemon (e.g. ``--project``/``--instance``). - binary_path: explicit path to the daemon binary. When omitted, the - path is resolved from the ``BIGTABLE_ACCELERATOR_BIN`` env var - and then the binary bundled in the wheel. startup_timeout: seconds ``start()`` waits for the daemon's UDS to become connectable before raising. Raises: - FileNotFoundError: no binary could be resolved, or an explicit - ``binary_path``/env override does not point at a regular file. + FileNotFoundError: the bundled binary is not present in this wheel. """ - self._binary_path = _resolve_binary_path(binary_path) + self._binary_path = _resolve_binary_path() self._cli_flags = list(cli_flags) self._startup_timeout = startup_timeout self._tempdir: str | None = None