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..8f9a6b0ce813 --- /dev/null +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -0,0 +1,312 @@ +# 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 socket +import subprocess +import tempfile +import time +from typing import Sequence + +# 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() -> 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( + "Accelerator binary not found. 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] = (), + *, + 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``). + startup_timeout: seconds ``start()`` waits for the daemon's UDS to + become connectable before raising. + + Raises: + FileNotFoundError: the bundled binary is not present in this wheel. + """ + self._binary_path = _resolve_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") + 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 + # 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. + 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, + 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. + if log_file is not None: + 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. + 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. + try: + proc.kill() + except OSError: + pass + 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.kill() + 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 + self._log_path = None