From 80021ec204b82e44a6d8d84bba122db089c24a13 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 21 Oct 2026 20:24:05 +0000 Subject: [PATCH 1/2] 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/2] 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(