From 80021ec204b82e44a6d8d84bba122db089c24a13 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 21 Oct 2026 20:24:05 +0000 Subject: [PATCH 1/3] feat(bigtable): sweep stale accelerator daemon tempdirs at startup Change-Id: Icd8495d9c7f1bdcf39ae984b0037f813ef4c4011 --- .../bigtable/data/_accelerator/_daemon.py | 135 ++++++++++-- .../bigtable/data/_accelerator/_fallback.py | 3 +- .../cloud/bigtable/data/_async/client.py | 14 +- .../unit/data/test_accelerator_credentials.py | 2 +- .../unit/data/test_accelerator_daemon.py | 203 +++++++++++++++--- .../unit/data/test_accelerator_dispatch.py | 188 ++++++++++------ .../unit/data/test_accelerator_enablement.py | 3 +- .../unit/data/test_accelerator_fallback.py | 10 +- 8 files changed, 440 insertions(+), 118 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 25888209395a..4e9fef9cc112 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 math import os @@ -31,10 +32,10 @@ import socket import subprocess import tempfile +import threading 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 @@ -50,6 +51,25 @@ # 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" +# 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" +# 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 # sandbox) spend ~10s in Go runtime init + ADC token fetch from the metadata @@ -64,6 +84,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.""" @@ -103,6 +128,79 @@ 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 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. + """ + if not hasattr(socket, "AF_UNIX"): + return + 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: + # 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 + 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) + + class AcceleratorDaemon: """Manages the Go accelerator daemon subprocess. @@ -139,7 +237,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 @@ -203,12 +303,27 @@ 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. 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) + # 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)) + 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 # lifetime, so a PIPE's fixed OS buffer would eventually fill and @@ -217,7 +332,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. @@ -368,14 +483,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/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 8ef6cb138316..8360e1a60ff2 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/tests/unit/data/test_accelerator_credentials.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_credentials.py index a50f8c81ea51..ee4f755b02d9 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 5293f100ccec..379d09b65613 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 @@ -207,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) @@ -246,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 @@ -259,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): @@ -287,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") @@ -304,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") @@ -316,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"): @@ -335,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) @@ -350,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() @@ -369,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() @@ -388,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() @@ -404,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 @@ -415,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) @@ -449,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): @@ -456,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 @@ -473,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) @@ -505,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) @@ -512,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) @@ -520,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) @@ -531,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) @@ -561,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 = {} @@ -573,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): @@ -595,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. @@ -672,3 +697,133 @@ 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) + + +@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.""" + + @staticmethod + 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() + 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)) + 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 + 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): + # Expired deadline + old log + dead spawner PID → should be reaped. + 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, + ) + _daemon._sweep_stale_tempdirs() + assert not d.exists() + + def test_keeps_young_dir(self, tmp_root): + # 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" + d.mkdir() + (d / "something-else").write_bytes(b"x") + _daemon._sweep_stale_tempdirs() + assert d.exists() + + 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, + pid=os.getpid(), + ) + _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() + + def test_start_invokes_sweep(self, tmp_path, monkeypatch): + 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") + + monkeypatch.setattr(_daemon.subprocess, "Popen", boom) + daemon = _make_daemon(tmp_path) + with pytest.raises(RuntimeError): + daemon.start() + assert done.wait(timeout=2.0), "sweep thread did not run within 2 s" 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 d6db1582aa3d2dca9aa3799bb6f6d94eece58310 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Thu, 22 Oct 2026 02:36:22 +0000 Subject: [PATCH 2/3] fix(bigtable): bound the tempdir sweep, clean up on metadata write failure Change-Id: Ie5c45b75e0785a284a0d241415cd5b8492e14e9c --- .../bigtable/data/_accelerator/_daemon.py | 53 ++++++++++++++----- .../unit/data/test_accelerator_daemon.py | 41 ++++++++++++++ 2 files changed, 81 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 4e9fef9cc112..5e2f6c896061 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 @@ -84,11 +84,17 @@ _SIGTERM_GRACE_SECONDS = 2.0 _SIGKILL_GRACE_SECONDS = 2.0 -# A freshly-created tempdir may not have bound its socket yet. Never sweep one +# A freshly-created tempdir has not written its daemon.pid 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 +# Upper bound on how long a single sweep may spend walking candidates. The +# sweep runs on a fire-and-forget thread, so nothing is waiting on it; this +# only stops a pathological temp directory (thousands of leftovers, a slow +# network filesystem) from keeping the thread alive indefinitely. +_SWEEP_BUDGET_SECONDS = 5.0 + def _resolve_binary_path() -> str: """Return the path of the bundled daemon binary, raising if it is absent.""" @@ -158,7 +164,9 @@ def _sweep_stale_tempdirs() -> None: the process exits for any reason, including SIGKILL. All failures are swallowed — sweeping is opportunistic and must never break - startup. + startup. The walk also gives up after ``_SWEEP_BUDGET_SECONDS`` so the + thread cannot linger on a pathological temp directory; whatever it did not + reach is simply left for the next start(). """ if not hasattr(socket, "AF_UNIX"): return @@ -168,7 +176,10 @@ def _sweep_stale_tempdirs() -> None: except OSError: return now = time.time() + budget_ends = now + _SWEEP_BUDGET_SECONDS for log_path in logs: + if time.time() >= budget_ends: + return path = os.path.dirname(log_path) try: # Each tempdir records its own startup deadline so this sweep is @@ -177,7 +188,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. @@ -303,7 +315,6 @@ 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." @@ -311,18 +322,34 @@ def start(self) -> None: # Reap tempdirs leaked by daemons from prior runs that crashed or were # 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). + # daemon thread: the walk is only stat()s and signal-0 probes, but it + # scales with the number of leftovers, and start() (which the async + # client calls synchronously, on the event loop thread) must not wait + # on it. The sweep bounds its own runtime; see _SWEEP_BUDGET_SECONDS. 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)) - with open(os.path.join(self._tempdir, _PID_FILENAME), "w") as _f: - _f.write(str(os.getpid())) + # Record the metadata the sweep needs to judge this tempdir later: the + # absolute time by which startup must complete (read per-directory, so + # the sweep is right regardless of what + # BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT the *sweeping* process has) and + # the spawning PID. A failure here means the tempdir is unwritable, so + # the log file and daemon would fail next anyway: clean up and report + # it the same way a failed spawn does, rather than leaving behind a + # directory that has no daemon.log and so would never be swept. + try: + 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())) + except OSError as exc: + tempdir = self._tempdir # _cleanup_tempdir() clears the attribute. + self._cleanup_tempdir() + raise RuntimeError( + f"Failed to write accelerator daemon metadata to {tempdir}: {exc}" + ) from exc # 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 379d09b65613..6e6917a08427 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 @@ -284,6 +284,27 @@ def bad_popen(argv, **kwargs): with pytest.raises(RuntimeError, match="Failed to spawn"): d.start() + def test_metadata_write_failure_cleans_up_and_raises(self, tmp_path, monkeypatch): + # An unwritable tempdir fails start() the same way a failed spawn does. + # The tempdir must not survive: it has no daemon.log yet, and the sweep + # keys off that marker, so a leftover here would never be reaped. + monkeypatch.setattr( + AcceleratorDaemon, "_cleanup_tempdir", _real_cleanup_tempdir + ) + real_open = open + + def bad_open(path, *args, **kwargs): + if os.path.basename(str(path)) == _daemon._PID_FILENAME: + raise OSError("no space left on device") + return real_open(path, *args, **kwargs) + + monkeypatch.setattr(_daemon, "open", bad_open, raising=False) + d = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="Failed to write accelerator daemon"): + d.start() + assert d._tempdir is None + assert not (tmp_path / "accel-tmp").exists() + def test_process_exits_during_startup_raises(self, tmp_path, monkeypatch): import subprocess @@ -801,6 +822,26 @@ def test_keeps_dir_with_live_pid(self, tmp_root): _daemon._sweep_stale_tempdirs() assert d.exists() + def test_budget_exhausted_leaves_dirs_for_next_sweep(self, tmp_root, monkeypatch): + # Same inputs as test_removes_stale_dir, but the time budget is already + # spent: the walk bails out rather than running long on a huge /tmp. + 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, + ) + monkeypatch.setattr(_daemon, "_SWEEP_BUDGET_SECONDS", -1.0) + _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( From 2fbb38431877b40e20decd7ca4a1ec491c18acda Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Fri, 23 Oct 2026 14:16:37 +0000 Subject: [PATCH 3/3] fix(bigtable): only sweep a tempdir when its spawner is provably gone Change-Id: I2dcd04d457e90e628ebd17d69987174a58489379 --- .../bigtable/data/_accelerator/_daemon.py | 272 ++++++++---- .../unit/data/test_accelerator_daemon.py | 419 +++++++++++++++--- 2 files changed, 554 insertions(+), 137 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 5e2f6c896061..598164d8a335 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,10 +64,12 @@ # 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). +# Identity of the Python process that spawned the daemon, written at start() +# as ``" "``. The sweep signals this PID to decide whether +# the spawner is still alive: if it is, the daemon's stdin is open and the +# daemon is (or is about to be) serving; if it is gone, the daemon's stdin is +# closed and it has exited or is exiting. See ``_sweep_stale_tempdirs`` for +# why the namespace is recorded next to the number. _PID_FILENAME = "daemon.pid" # How long to wait for the daemon to start listening on its UDS before giving @@ -84,11 +86,23 @@ _SIGTERM_GRACE_SECONDS = 2.0 _SIGKILL_GRACE_SECONDS = 2.0 -# A freshly-created tempdir has not written its daemon.pid 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. +# Grace period past a tempdir's startup deadline before the sweep will judge +# it. A daemon that has not bound by its own deadline has failed, but give it +# this much slack anyway 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 +# Written in place of the PID namespace on platforms that do not have one to +# report (no /proc, e.g. macOS). Comparing "none" against "none" is a match -- +# where namespaces do not exist, PIDs cannot be ambiguous between them. +_NO_PID_NAMESPACE = "-" + +# How long a single connect() probe of a daemon's socket may take. A UDS +# connect against a listening socket completes in the kernel without the +# daemon doing anything, so this only ever elapses on a badly overloaded +# host — where the honest answer is "I cannot tell", not "it is dead". +_SOCKET_PROBE_TIMEOUT_SECONDS = 0.25 + # Upper bound on how long a single sweep may spend walking candidates. The # sweep runs on a fire-and-forget thread, so nothing is waiting on it; this # only stops a pathological temp directory (thousands of leftovers, a slow @@ -135,11 +149,17 @@ def _resolve_startup_timeout() -> float: def _socket_is_live(uds_path: str) -> bool: - """Return True if something is currently accepting connections on the UDS.""" + """Return True if something is currently accepting connections on the UDS. + + Used by ``_wait_until_ready`` to poll for the daemon coming up, so every + failure simply answers "not yet" and the caller tries again. Note that + this is a *measurement*, which is why the sweep does not use it to decide + whether a daemon has died -- see ``_sweep_stale_tempdirs``. + """ if not os.path.exists(uds_path): return False with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: - probe.settimeout(0.25) + probe.settimeout(_SOCKET_PROBE_TIMEOUT_SECONDS) try: probe.connect(uds_path) return True @@ -147,70 +167,130 @@ def _socket_is_live(uds_path: str) -> bool: return False +def _pid_namespace() -> str: + """Identify the PID namespace this process's PIDs are numbered in. + + A PID is only meaningful inside the namespace that issued it, but the + tempdir holding it can be shared across namespaces -- two containers with + the same ``/tmp`` volume, or a host ``/tmp`` bind-mounted into one. The + sweep records this alongside the PID so it can tell whether a number it + reads is one it is entitled to interpret. + + Returns ``_NO_PID_NAMESPACE`` where there is nothing to report, which + compares equal to itself: on a platform with no PID namespaces, PIDs + cannot be ambiguous between them. + """ + try: + st = os.stat("/proc/self/ns/pid") + except OSError: + return _NO_PID_NAMESPACE + return f"{st.st_dev}:{st.st_ino}" + + +def _write_metadata_file(path: str, contents: str) -> None: + """Write ``contents`` to ``path`` so a reader never sees a partial file. + + The sweep reads these files out of *other* processes' tempdirs, possibly + while they are still being written. A half-written pid file would be + unreadable or, worse, parse as some other process's PID, and a truncated + deadline could parse as a plausible but wrong time. Writing to a sibling + temp name and renaming into place makes the swap atomic: a reader sees + either no file at all or the complete contents. A leftover ``.tmp`` from + a process killed mid-write goes away with the tempdir. + """ + tmp_path = path + ".tmp" + with open(tmp_path, "w") as _f: + _f.write(contents) + os.replace(tmp_path, path) + + 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 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. The walk also gives up after ``_SWEEP_BUDGET_SECONDS`` so the - thread cannot linger on a pathological temp directory; whatever it did not - reach is simply left for the next start(). + A clean ``close()`` removes its own tempdir; a crash or ``SIGKILL`` leaves + the socket and log behind. We tidy those up on the next start. + + Best-effort in both directions, and that is the whole design: + + * It never breaks startup. Every failure is swallowed, the walk abandons + whatever it has not reached once ``_SWEEP_BUDGET_SECONDS`` is up, and it + runs on a thread nobody waits on. + * It never deletes without proof. The only door to ``rmtree`` is + ``ProcessLookupError`` from ``os.kill(pid, 0)`` -- the kernel stating no + such process exists. Everything else (unreadable file, truncated write, + ``EPERM``, a PID from another namespace) means "could not establish + death", and skips. + + The asymmetry is deliberate: a missed leftover costs a few KB of /tmp + until the next start() looks again, while deleting a live daemon's + directory destroys the socket it is serving on with no way back. So two + leaks are accepted on purpose -- a recycled PID reads as alive and pins + its directory forever, and a pid file truncated mid-write is skipped + forever. Both are cheap; neither is worth a heuristic that can delete a + live daemon. + + Mechanically: candidates are ``bt-accel-*`` directories containing our + ``daemon.log`` marker, so an unrelated directory sharing the prefix is + never touched. The spawning Python process is the liveness signal -- + while it lives the daemon's stdin is open and the daemon is serving; when + it dies the daemon sees EOF and exits. Each tempdir also records the PID + namespace that numbered it (see :func:`_pid_namespace`), since a shared + ``/tmp`` can cross namespaces and a foreign PID would otherwise give a + confident wrong answer; and its own startup deadline, which covers the + window before the pid file exists and keeps the sweep correct whatever + ``BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT`` the sweeping process has. """ if not hasattr(socket, "AF_UNIX"): return - root = tempfile.gettempdir() - try: - logs = glob.glob(os.path.join(root, _TEMPDIR_PREFIX + "*", _LOG_FILENAME)) - except OSError: - return + # Start the clock before the glob: on the pathological temp directory this + # budget exists for, enumerating it is itself the slow part. now = time.time() budget_ends = now + _SWEEP_BUDGET_SECONDS - for log_path in logs: - if time.time() >= budget_ends: - return - path = os.path.dirname(log_path) - try: - # 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: - 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. - deadline = os.path.getmtime(log_path) + _STALE_SWEEP_MIN_AGE_SECONDS - if now < deadline + _STALE_SWEEP_MIN_AGE_SECONDS: - continue - pid_path = os.path.join(path, _PID_FILENAME) + own_namespace = _pid_namespace() + root = tempfile.gettempdir() + try: + logs = glob.iglob(os.path.join(root, _TEMPDIR_PREFIX + "*", _LOG_FILENAME)) + for log_path in logs: + if time.time() >= budget_ends: + return + path = os.path.dirname(log_path) 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. + deadline_path = os.path.join(path, _STARTUP_DEADLINE_FILENAME) + try: + with open(deadline_path) as _f: + deadline = float(_f.read()) + except (OSError, ValueError): + # Older tempdir without a deadline file, or one we could + # not read: fall back to a simple age heuristic relative + # to the log's mtime. This gate only decides *when* to + # look; the pid check below still has to prove death. + deadline = os.path.getmtime(log_path) + _STALE_SWEEP_MIN_AGE_SECONDS + if now < deadline + _STALE_SWEEP_MIN_AGE_SECONDS: + continue + pid_path = os.path.join(path, _PID_FILENAME) + try: + with open(pid_path) as _f: + fields = _f.read().split() + pid = int(fields[0]) + if fields[1] != own_namespace: + # Someone else's numbering. IndexError lands in the + # same place: a tempdir written before the namespace + # was recorded cannot be checked either. + continue + os.kill(pid, 0) + continue # spawner alive → daemon alive → skip + except ProcessLookupError: + pass # ESRCH: spawner is dead → stdin closed → daemon gone + except (OSError, ValueError, IndexError): + # Missing, unreadable, empty or truncated pid file, or + # EPERM from os.kill (the PID exists, it just is not ours + # to signal). None of these is proof of death. + continue + except OSError: continue - except OSError: - continue - shutil.rmtree(path, ignore_errors=True) + shutil.rmtree(path, ignore_errors=True) + except OSError: + return class AcceleratorDaemon: @@ -326,30 +406,15 @@ def start(self) -> None: # scales with the number of leftovers, and start() (which the async # client calls synchronously, on the event loop thread) must not wait # on it. The sweep bounds its own runtime; see _SWEEP_BUDGET_SECONDS. - threading.Thread(target=_sweep_stale_tempdirs, daemon=True).start() + # Even spawning the thread is best-effort: a process already at its + # thread ceiling should still get a working client, not an exception + # from the one step of startup that exists purely to tidy up /tmp. + try: + threading.Thread(target=_sweep_stale_tempdirs, daemon=True).start() + except RuntimeError: + pass self._tempdir = tempfile.mkdtemp(prefix=_TEMPDIR_PREFIX) self._uds_path = os.path.join(self._tempdir, _SOCKET_FILENAME) - # Record the metadata the sweep needs to judge this tempdir later: the - # absolute time by which startup must complete (read per-directory, so - # the sweep is right regardless of what - # BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT the *sweeping* process has) and - # the spawning PID. A failure here means the tempdir is unwritable, so - # the log file and daemon would fail next anyway: clean up and report - # it the same way a failed spawn does, rather than leaving behind a - # directory that has no daemon.log and so would never be swept. - try: - 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())) - except OSError as exc: - tempdir = self._tempdir # _cleanup_tempdir() clears the attribute. - self._cleanup_tempdir() - raise RuntimeError( - f"Failed to write accelerator daemon metadata to {tempdir}: {exc}" - ) from exc # Redirect the daemon's stdout/stderr to a log file rather than # subprocess.PIPE. Nothing drains those pipes for the daemon's @@ -363,13 +428,46 @@ 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. - argv = [self._binary_path, "--uds-path", self._uds_path, *self._cli_flags] log_file = None + try: + # Create the log before anything else can fail. It is the marker + # the sweep keys on, so a crash from here on still leaves a + # directory that a later sweep can recognise and reap. + log_file = open(self._log_path, "wb") + # Record the absolute time by which startup must complete. The + # sweep reads this per-directory, so it judges correctly however + # the *sweeping* process has BIGTABLE_ACCELERATOR_STARTUP_TIMEOUT + # set. Written atomically, since a sweep in another process may + # read it at any moment. + _write_metadata_file( + os.path.join(self._tempdir, _STARTUP_DEADLINE_FILENAME), + repr(time.time() + self._startup_timeout), + ) + # Record who owns this tempdir, as " ". The + # namespace qualifies the number: a sweep running where that PID + # would mean something else entirely can see that and decline to + # judge, rather than confidently signalling an unrelated process. + _write_metadata_file( + os.path.join(self._tempdir, _PID_FILENAME), + f"{os.getpid()} {_pid_namespace()}", + ) + except OSError as exc: + if log_file is not None: + try: + log_file.close() + except OSError: + pass + tempdir = self._tempdir # _cleanup_tempdir() clears the attribute. + self._cleanup_tempdir() + raise RuntimeError( + f"Failed to write accelerator daemon metadata to {tempdir}: {exc}" + ) from exc + + argv = [self._binary_path, "--uds-path", self._uds_path, *self._cli_flags] env = None if self._extra_env: env = {**os.environ, **self._extra_env} try: - log_file = open(self._log_path, "wb") self._proc = subprocess.Popen( argv, stdin=subprocess.PIPE, 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 6e6917a08427..8713d01a7931 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,7 +14,14 @@ import json import os +import pathlib import re +import shutil +import socket +import subprocess +import sys +import tempfile +import threading import time import unittest.mock @@ -44,6 +51,17 @@ def _make_daemon(tmp_path, **kwargs): # directly as _real_cleanup_tempdir(daemon_instance). _real_cleanup_tempdir = AcceleratorDaemon._cleanup_tempdir +# Likewise for tempfile.mkdtemp/gettempdir, which the autouse fixture +# redirects into pytest's tmp_path. The sweep tests need real, short /tmp +# paths: they bind Unix sockets, whose paths are capped at ~107 bytes, and +# pytest's per-test directory names alone eat most of that budget. +_real_mkdtemp = tempfile.mkdtemp +_real_gettempdir = tempfile.gettempdir + +# Distinguishes "caller did not ask for a namespace stamp" (use our own, the +# ordinary case) from "caller asked for no stamp at all" (None). +_UNSET = object() + class TestExtraEnv: def test_default_env_none(self, tmp_path, monkeypatch): @@ -201,6 +219,12 @@ def fake_mkdtemp(prefix=""): return str(d) monkeypatch.setattr(_daemon.tempfile, "mkdtemp", fake_mkdtemp) + + # start() kicks off _sweep_stale_tempdirs on a background thread, and the + # sweep reads gettempdir() rather than going through mkdtemp. Without this + # the suite would glob — and rmtree — real bt-accel-* directories in the + # developer's or CI machine's /tmp, including live ones. + monkeypatch.setattr(_daemon.tempfile, "gettempdir", lambda: str(tmp_path)) yield @@ -284,17 +308,30 @@ def bad_popen(argv, **kwargs): with pytest.raises(RuntimeError, match="Failed to spawn"): d.start() - def test_metadata_write_failure_cleans_up_and_raises(self, tmp_path, monkeypatch): - # An unwritable tempdir fails start() the same way a failed spawn does. - # The tempdir must not survive: it has no daemon.log yet, and the sweep - # keys off that marker, so a leftover here would never be reaped. + @pytest.mark.parametrize( + "doomed", + [ + _daemon._LOG_FILENAME, + _daemon._STARTUP_DEADLINE_FILENAME, + _daemon._PID_FILENAME, + ], + ) + def test_metadata_write_failure_cleans_up_and_raises( + self, tmp_path, monkeypatch, doomed + ): + # An unwritable tempdir fails start() the same way a failed spawn + # does, whichever of the files start() lays down hits the error. The + # tempdir must not survive a failure before daemon.log exists, since + # the sweep keys off that marker and would never reap it. monkeypatch.setattr( AcceleratorDaemon, "_cleanup_tempdir", _real_cleanup_tempdir ) real_open = open def bad_open(path, *args, **kwargs): - if os.path.basename(str(path)) == _daemon._PID_FILENAME: + # The deadline file is written to a ".tmp" sibling and renamed, so + # match on the prefix rather than the final name. + if os.path.basename(str(path)).startswith(doomed): raise OSError("no space left on device") return real_open(path, *args, **kwargs) @@ -726,7 +763,24 @@ def test_daemon_constructor_invalid_env_var_raises(self, tmp_path, monkeypatch): ) class TestSweepStaleTempdirs: """_sweep_stale_tempdirs reaps leftover daemon tempdirs, keyed on the - daemon.log marker, without touching live or too-young ones.""" + daemon.log marker, and deletes only on proof that the spawner is gone.""" + + @pytest.fixture + def tmp_root(self, monkeypatch): + # Point the sweep's tempdir root at an isolated directory of our own. + # dir= is explicit because mkdtemp would otherwise resolve its default + # through the gettempdir() the autouse fixture has already patched. + root = pathlib.Path(_real_mkdtemp(prefix="bt-sweep-", dir=_real_gettempdir())) + monkeypatch.setattr(_daemon.tempfile, "gettempdir", lambda: str(root)) + yield root + shutil.rmtree(root, ignore_errors=True) + + @staticmethod + def _dead_pid(): + """A PID that is guaranteed not to resolve to a running process.""" + proc = subprocess.Popen([sys.executable, "-c", ""]) + proc.wait() + return proc.pid @staticmethod def _make_leftover( @@ -737,14 +791,19 @@ def _make_leftover( age_seconds=0.0, deadline_offset=None, pid=None, + namespace=_UNSET, ): """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". + pid: if given, write a daemon.pid naming this PID. Pass a dead PID to + simulate a crashed spawner, or os.getpid() for "alive". + namespace: what to stamp next to the PID. Defaults to this process's + own, so the sweep accepts it; pass something else to simulate a + tempdir written in a different PID namespace, or None to omit the + stamp entirely (a tempdir predating it). """ d = root / name d.mkdir() @@ -758,44 +817,176 @@ def _make_leftover( 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)) + if namespace is _UNSET: + namespace = _daemon._pid_namespace() + stamp = str(pid) if namespace is None else f"{pid} {namespace}" + (d / _daemon._PID_FILENAME).write_text(stamp) 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 + # -- the one path that deletes ------------------------------------------- def test_removes_stale_dir(self, tmp_root): - # Expired deadline + old log + dead spawner PID → should be reaped. - import subprocess - import sys - - proc = subprocess.Popen([sys.executable, "-c", ""]) - dead_pid = proc.pid - proc.wait() + # Expired deadline, old log, and the kernel says the spawner's PID + # does not resolve to any process. That ESRCH is the only thing that + # may lead to a delete. d = self._make_leftover( tmp_root, "bt-accel-dead", age_seconds=120, deadline_offset=-70, - pid=dead_pid, + pid=self._dead_pid(), + ) + _daemon._sweep_stale_tempdirs() + assert not d.exists() + + def test_removes_stale_dir_without_deadline_file(self, tmp_root): + # A tempdir predating the deadline file still gets judged, via the + # log-mtime fallback -- but only because the pid check proves death. + d = self._make_leftover( + tmp_root, + "bt-accel-nodl", + age_seconds=120, + deadline_offset=None, + pid=self._dead_pid(), ) _daemon._sweep_stale_tempdirs() assert not d.exists() + # -- everything else keeps ------------------------------------------------ + + def test_keeps_dir_with_live_spawner(self, tmp_root): + # Well past every age gate, but the spawning process is alive, so its + # daemon is alive and the directory is not a leftover. + d = self._make_leftover( + tmp_root, + "bt-accel-live", + age_seconds=86400, + deadline_offset=-86400, + pid=os.getpid(), + ) + _daemon._sweep_stale_tempdirs() + assert d.exists() + + def test_keeps_live_dir_from_an_older_client(self, tmp_root): + # A tempdir written by a previously released client has a daemon.log + # and no metadata at all. There is no owner to probe, so there is no + # proof of death, so it must survive -- its daemon may well be + # serving. Mixed versions share one /tmp whenever two virtualenvs (or + # a rolling restart) run side by side. + d = self._make_leftover( + tmp_root, "bt-accel-legacy", age_seconds=86400, deadline_offset=None + ) + _daemon._sweep_stale_tempdirs() + assert d.exists() + + def test_keeps_dir_from_another_pid_namespace(self, tmp_root): + # Containers sharing a /tmp volume have separate PID numbering, so a + # PID written over there means something else (or nothing) here. + # Signalling it would give a confident, repeatable wrong answer, so + # the sweep declines to judge. The dead PID makes the point: without + # the namespace check this directory would be deleted. + d = self._make_leftover( + tmp_root, + "bt-accel-otherns", + age_seconds=86400, + deadline_offset=-86400, + pid=self._dead_pid(), + namespace="4:4026531999", + ) + _daemon._sweep_stale_tempdirs() + assert d.exists() + + def test_keeps_dir_with_unstamped_pid_file(self, tmp_root): + # A pid file written before the namespace was recorded cannot be + # checked either, for the same reason. + d = self._make_leftover( + tmp_root, + "bt-accel-nostamp", + age_seconds=86400, + deadline_offset=-86400, + pid=self._dead_pid(), + namespace=None, + ) + _daemon._sweep_stale_tempdirs() + assert d.exists() + + @pytest.mark.parametrize("contents", ["", " ", "not-a-pid", "garbage stamp"]) + def test_keeps_dir_with_unusable_pid_file(self, tmp_root, contents): + # A crash mid-write can leave the pid file empty or truncated. That is + # a deliberate, permanent leak: unreadable is not proof of death. + d = self._make_leftover( + tmp_root, "bt-accel-badpid", age_seconds=86400, deadline_offset=-86400 + ) + (d / _daemon._PID_FILENAME).write_text(contents) + _daemon._sweep_stale_tempdirs() + assert d.exists() + + def test_keeps_dir_when_pid_file_is_unreadable(self, tmp_root, monkeypatch): + # An OSError other than "missing" tells us nothing about the owner. + d = self._make_leftover( + tmp_root, + "bt-accel-unreadable", + age_seconds=86400, + deadline_offset=-86400, + pid=self._dead_pid(), + ) + real_open = open + + def bad_open(path, *args, **kwargs): + if os.path.basename(str(path)) == _daemon._PID_FILENAME: + raise PermissionError("EACCES") + return real_open(path, *args, **kwargs) + + monkeypatch.setattr(_daemon, "open", bad_open, raising=False) + _daemon._sweep_stale_tempdirs() + assert d.exists() + + def test_keeps_dir_when_pid_is_not_ours_to_signal(self, tmp_root, monkeypatch): + # EPERM from os.kill means the PID exists under another user: alive. + d = self._make_leftover( + tmp_root, + "bt-accel-eperm", + age_seconds=86400, + deadline_offset=-86400, + pid=os.getpid(), + ) + + def eperm(pid, sig): + raise PermissionError("EPERM") + + monkeypatch.setattr(_daemon.os, "kill", eperm) + _daemon._sweep_stale_tempdirs() + assert d.exists() + + def test_keeps_dir_whose_pid_was_recycled(self, tmp_root): + # A recycled PID now belonging to some unrelated long-lived process + # reads as alive, so the directory is pinned forever. That leak is the + # accepted cost of never deleting a live daemon; detecting reuse means + # a heuristic, and a wrong "recycled" verdict deletes a live daemon. + d = self._make_leftover( + tmp_root, + "bt-accel-recycled", + age_seconds=86400, + deadline_offset=-86400, + pid=os.getpid(), + ) + pid_file = d / _daemon._PID_FILENAME + old = time.time() - 86400 + os.utime(pid_file, (old, old)) + _daemon._sweep_stale_tempdirs() + assert d.exists() + def test_keeps_young_dir(self, tmp_root): # No deadline file: fallback uses log mtime + buffer; log is fresh. - d = self._make_leftover(tmp_root, "bt-accel-young", age_seconds=0) + d = self._make_leftover( + tmp_root, "bt-accel-young", age_seconds=0, pid=self._dead_pid() + ) _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. + # Log is old (a long-timeout override) but the daemon is still inside + # its own startup budget, when the pid file may not be written yet. d = self._make_leftover( tmp_root, "bt-accel-slow", age_seconds=120, deadline_offset=+200 ) @@ -810,49 +1001,45 @@ def test_ignores_dir_without_log_marker(self, tmp_root): _daemon._sweep_stale_tempdirs() assert d.exists() - def test_keeps_dir_with_live_pid(self, tmp_root): - # Past deadline but spawning process is still alive → must not be reaped. + def test_budget_exhausted_leaves_dirs_for_next_sweep(self, tmp_root, monkeypatch): + # Same inputs as test_removes_stale_dir, but the time budget is already + # spent: the walk bails out rather than running long on a huge /tmp. d = self._make_leftover( tmp_root, - "bt-accel-live", + "bt-accel-dead", age_seconds=120, deadline_offset=-70, - pid=os.getpid(), + pid=self._dead_pid(), ) + monkeypatch.setattr(_daemon, "_SWEEP_BUDGET_SECONDS", -1.0) _daemon._sweep_stale_tempdirs() assert d.exists() - def test_budget_exhausted_leaves_dirs_for_next_sweep(self, tmp_root, monkeypatch): - # Same inputs as test_removes_stale_dir, but the time budget is already - # spent: the walk bails out rather than running long on a huge /tmp. - import subprocess - import sys + def test_glob_failure_is_swallowed(self, tmp_root, monkeypatch): + # Enumeration can fail outright (a vanished or unreadable /tmp). The + # sweep is opportunistic, so that must not propagate to start(). + def boom(*args, **kwargs): + raise OSError("EIO") - proc = subprocess.Popen([sys.executable, "-c", ""]) - dead_pid = proc.pid - proc.wait() + monkeypatch.setattr(_daemon.glob, "iglob", boom) + _daemon._sweep_stale_tempdirs() + + def test_no_af_unix_is_a_noop(self, tmp_root, monkeypatch): + # Without AF_UNIX the accelerator cannot have run here at all. d = self._make_leftover( tmp_root, "bt-accel-dead", age_seconds=120, deadline_offset=-70, - pid=dead_pid, + pid=self._dead_pid(), ) - monkeypatch.setattr(_daemon, "_SWEEP_BUDGET_SECONDS", -1.0) + monkeypatch.delattr(_daemon.socket, "AF_UNIX") _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() + # -- wiring into start() -------------------------------------------------- def test_start_invokes_sweep(self, tmp_path, monkeypatch): - import threading - done = threading.Event() def _fake_sweep(): @@ -868,3 +1055,135 @@ def boom(*args, **kwargs): with pytest.raises(RuntimeError): daemon.start() assert done.wait(timeout=2.0), "sweep thread did not run within 2 s" + + def test_start_survives_sweep_thread_failure(self, tmp_path, monkeypatch): + # At the process thread ceiling, Thread.start() raises. Tidying /tmp + # is the least important thing start() does; it must not be the thing + # that costs the caller a working client. + class NoThreads: + def __init__(self, *args, **kwargs): + pass + + def start(self): + raise RuntimeError("can't start new thread") + + monkeypatch.setattr(_daemon.threading, "Thread", NoThreads) + popen_called = [] + + def boom(*args, **kwargs): + popen_called.append(True) + raise OSError("refusing to spawn in test") + + monkeypatch.setattr(_daemon.subprocess, "Popen", boom) + daemon = _make_daemon(tmp_path) + # Reaches Popen and fails there, rather than at Thread.start(). + with pytest.raises(RuntimeError, match="Failed to spawn"): + daemon.start() + assert popen_called + + def test_start_stamps_pid_and_namespace(self, tmp_path, monkeypatch): + # The file start() writes must be the one the sweep knows how to read. + def boom(*args, **kwargs): + raise OSError("refusing to spawn in test") + + monkeypatch.setattr(_daemon.subprocess, "Popen", boom) + monkeypatch.setattr(AcceleratorDaemon, "_cleanup_tempdir", lambda self: None) + daemon = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="Failed to spawn"): + daemon.start() + written = (pathlib.Path(daemon._tempdir) / _daemon._PID_FILENAME).read_text() + assert written.split() == [str(os.getpid()), _daemon._pid_namespace()] + + +class TestPidNamespace: + """_pid_namespace identifies the namespace a process's PIDs belong to.""" + + @pytest.mark.skipif( + not os.path.exists("/proc/self/ns/pid"), reason="requires Linux namespaces" + ) + def test_reports_own_namespace(self): + ns = _daemon._pid_namespace() + assert ns != _daemon._NO_PID_NAMESPACE + # Stable for the life of the process, and shared with our children. + assert ns == _daemon._pid_namespace() + child = subprocess.run( + [ + sys.executable, + "-c", + "import os;st=os.stat('/proc/self/ns/pid');" + "print(f'{st.st_dev}:{st.st_ino}')", + ], + capture_output=True, + text=True, + ) + assert child.stdout.strip() == ns + + def test_falls_back_when_proc_is_absent(self, monkeypatch): + # macOS and friends have no /proc, and also no PID namespaces, so the + # sentinel compares equal to itself and the check is a no-op there. + def no_proc(path, *args, **kwargs): + raise FileNotFoundError(path) + + monkeypatch.setattr(_daemon.os, "stat", no_proc) + assert _daemon._pid_namespace() == _daemon._NO_PID_NAMESPACE + + +@pytest.mark.skipif(not hasattr(_daemon.socket, "AF_UNIX"), reason="requires AF_UNIX") +class TestSocketIsLive: + """_socket_is_live backs the _wait_until_ready poll: is the daemon up?""" + + def test_true_while_listening(self, tmp_path): + path = str(tmp_path / "sock") + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: + sock.bind(path) + sock.listen(1) + assert _daemon._socket_is_live(path) is True + + def test_false_when_path_is_absent(self, tmp_path): + assert _daemon._socket_is_live(str(tmp_path / "nope")) is False + + def test_false_for_socket_file_whose_owner_exited(self, tmp_path): + # The socket file outlives the process that bound it; nothing is + # accepting, so the daemon is not up yet (or any more). + path = str(tmp_path / "sock") + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + sock.bind(path) + sock.listen(1) + sock.close() + assert os.path.exists(path) + assert _daemon._socket_is_live(path) is False + + +class TestWriteMetadataFile: + """_write_metadata_file swaps contents in atomically.""" + + def test_writes_contents(self, tmp_path): + target = tmp_path / "daemon.pid" + _daemon._write_metadata_file(str(target), "1234") + assert target.read_text() == "1234" + # The staging file must not be left behind on success. + assert not (tmp_path / "daemon.pid.tmp").exists() + + def test_target_is_absent_until_contents_are_complete(self, tmp_path): + # A reader racing the write sees no file rather than a truncated one: + # everything lands in the ".tmp" sibling and is complete and closed by + # the time the rename publishes it under the real name. + target = tmp_path / "daemon.pid" + observed = {} + real_replace = _daemon.os.replace + + def spy_replace(src, dst): + observed["target_existed"] = target.exists() + observed["staged"] = pathlib.Path(src).read_text() + return real_replace(src, dst) + + with unittest.mock.patch.object(_daemon.os, "replace", spy_replace): + _daemon._write_metadata_file(str(target), "1234") + assert observed == {"target_existed": False, "staged": "1234"} + assert target.read_text() == "1234" + + def test_overwrites_existing_file(self, tmp_path): + target = tmp_path / "daemon.pid" + target.write_text("stale-and-longer") + _daemon._write_metadata_file(str(target), "7") + assert target.read_text() == "7"