fix: stop compute log capture deadlocking a command at exit

Makefile

@@ -179,7 +179,8 @@ E2E_DUCKDB_CLI_PATHS := \
 	tests/e2e/src/sqlbuild/cli/commands/main/snowflake \
 	tests/e2e/src/sqlbuild/cli/commands/main/sqlite \
 	tests/e2e/src/sqlbuild/cli/commands/main/sqlserver \
-	tests/e2e/src/sqlbuild/cli/commands/main/test
+	tests/e2e/src/sqlbuild/cli/commands/main/test \
+	tests/e2e/src/sqlbuild/cli/commands/shared
 
 E2E_DUCKDB_INSPECTION_PATHS := \
 	tests/e2e/src/sqlbuild/cli/commands/main/clone \

crates/sqlbuild-python/src/bindings/_helpers/compute_logs/storage.rs

@@ -90,6 +90,7 @@ fn exact_count(value: &Bound<'_, PyAny>, fallback: i64) -> i64 {
 }
 
 impl LocalFilesystemComputeLogStorage {
+    /// Callers wait for this lock holding the GIL, so never run Python while holding it.
     fn locked(&self) -> PyResult<MutexGuard<'_, StorageState>> {
         self.state
             .lock()
@@ -116,6 +117,7 @@ impl LocalFilesystemComputeLogStorage {
     ) -> PyResult<()> {
         let mut state = self.locked()?;
         let Some(writers) = state.writers.get_mut(invocation_id) else {
+            drop(state);
             return Err(not_active(py));
         };
         writers[stream].write_all(data)?;

tests/e2e/src/sqlbuild/cli/commands/shared/_test_types.py

@@ -0,0 +1,24 @@
+"""Test case types for the shared E2E command helpers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class HangDumpTestCase:
+    """A Python command that never finishes and what its hang report must show."""
+
+    description: str
+    script: str
+    expected_fragments: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class FinishedCommandTestCase:
+    """A Python command that finishes and the result the helper must return for it."""
+
+    description: str
+    script: str
+    input_text: str | None
+    expected_result: tuple[int, str, str]

tests/e2e/src/sqlbuild/cli/commands/shared/helpers.py

@@ -6,6 +6,8 @@
 import os
 import pty
 import re
+import shlex
+import signal
 import struct
 import subprocess
 import termios
@@ -22,6 +24,10 @@
 SOURCE_LOADER_STRATEGIES_DIR: Path = (
     REPO_ROOT / "tests" / "e2e" / "fixtures" / "source_loader_strategies"
 )
+# Far above any healthy command, so only a hang trips it: one test fails instead of the CI job.
+SQB_COMMAND_TIMEOUT_SECONDS: float = 600.0
+_HANG_DUMP_GRACE_SECONDS: float = 30.0
+_HANG_OUTPUT_TAIL_CHARACTERS: int = 20_000
 
 
 def prepare_waffle_shop(tmp_path: Path) -> Path:
@@ -208,18 +214,78 @@ def run_sqb(
 
     process_env: dict[str, str] = dict(os.environ)
     process_env.update(env or {})
+    process_env.setdefault("PYTHONFAULTHANDLER", "1")
 
-    return subprocess.run(
-        ["uv", "run", "sqb", "--project-dir", str(project_dir), *command],
+    return run_with_hang_dump(
+        argv=("uv", "run", "sqb", "--project-dir", str(project_dir), *command),
         cwd=working_dir or REPO_ROOT,
-        capture_output=True,
-        text=True,
-        input=input_text,
         env=process_env,
-        check=False,
+        input_text=input_text,
+        timeout_seconds=SQB_COMMAND_TIMEOUT_SECONDS,
+    )
+
+
+def run_with_hang_dump(
+    *,
+    argv: tuple[str, ...],
+    cwd: Path,
+    env: Mapping[str, str],
+    input_text: str | None,
+    timeout_seconds: float,
+) -> subprocess.CompletedProcess[str]:
+    """Run `argv` to completion, or fail with every Python thread's stack once it hangs.
+
+    Standard input carries `input_text`, or ends at once without it. The command runs in its own
+    process group. Past `timeout_seconds` the whole group receives
+    SIGABRT, so Python processes started with `PYTHONFAULTHANDLER` print all thread stacks to
+    stderr before they exit, and the group is then killed.
+    """
+
+    process: subprocess.Popen[str] = subprocess.Popen(
+        argv,
+        cwd=cwd,
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        env=dict(env),
+        process_group=0,
+    )
+    try:
+        stdout, stderr = process.communicate(input=input_text, timeout=timeout_seconds)
+    except subprocess.TimeoutExpired:
+        stdout, stderr = _dump_and_kill_group(process=process)
+        raise AssertionError(
+            f"{shlex.join(argv)} did not finish within {timeout_seconds:.0f}s; "
+            "its process group was aborted to dump thread stacks.\n"
+            f"--- stdout (last {_HANG_OUTPUT_TAIL_CHARACTERS} characters) ---\n"
+            f"{stdout[-_HANG_OUTPUT_TAIL_CHARACTERS:]}\n"
+            f"--- stderr (last {_HANG_OUTPUT_TAIL_CHARACTERS} characters) ---\n"
+            f"{stderr[-_HANG_OUTPUT_TAIL_CHARACTERS:]}"
+        ) from None
+    return subprocess.CompletedProcess(
+        args=list(argv), returncode=process.returncode, stdout=stdout, stderr=stderr
     )
 
 
+def _dump_and_kill_group(*, process: subprocess.Popen[str]) -> tuple[str, str]:
+    _signal_group(process_group=process.pid, signal_number=signal.SIGABRT)
+    try:
+        stdout, stderr = process.communicate(timeout=_HANG_DUMP_GRACE_SECONDS)
+    except subprocess.TimeoutExpired:
+        _signal_group(process_group=process.pid, signal_number=signal.SIGKILL)
+        stdout, stderr = process.communicate()
+    _signal_group(process_group=process.pid, signal_number=signal.SIGKILL)
+    return stdout or "", stderr or ""
+
+
+def _signal_group(*, process_group: int, signal_number: signal.Signals) -> None:
+    try:
+        os.killpg(process_group, signal_number)
+    except ProcessLookupError:
+        return
+
+
 def run_sqb_with_pty(
     *,
     command: tuple[str, ...],

tests/e2e/src/sqlbuild/cli/commands/shared/test_run_with_hang_dump.py

@@ -0,0 +1,95 @@
+"""A hung command fails its own test with every Python thread's stack instead of hanging the run."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import pytest
+
+from tests.e2e.src.sqlbuild.cli.commands.shared._test_types import (
+    FinishedCommandTestCase,
+    HangDumpTestCase,
+)
+from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import run_with_hang_dump
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        HangDumpTestCase(
+            description="main thread joins a worker blocked forever",
+            script=(
+                "import threading\n"
+                "def blocked_worker_wait():\n"
+                "    threading.Event().wait()\n"
+                "def joining_main_thread():\n"
+                "    worker = threading.Thread(target=blocked_worker_wait)\n"
+                "    worker.start()\n"
+                "    print('started', flush=True)\n"
+                "    worker.join()\n"
+                "joining_main_thread()\n"
+            ),
+            expected_fragments=("started", "blocked_worker_wait", "joining_main_thread"),
+        ),
+    ],
+    ids=lambda case: case.description,
+)
+def test_given_hung_python_command_when_running_with_hang_dump_then_fails_with_thread_stacks(
+    test_case: HangDumpTestCase, tmp_path: Path
+) -> None:
+    started: float = time.monotonic()
+    with pytest.raises(AssertionError) as raised:
+        run_with_hang_dump(
+            argv=(sys.executable, "-c", test_case.script),
+            cwd=tmp_path,
+            env={**os.environ, "PYTHONFAULTHANDLER": "1"},
+            input_text=None,
+            timeout_seconds=2.0,
+        )
+
+    message: str = str(raised.value)
+    assert "did not finish within 2s" in message
+    assert [fragment in message for fragment in test_case.expected_fragments] == [True] * len(
+        test_case.expected_fragments
+    ), message
+    assert time.monotonic() - started < 30.0
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        FinishedCommandTestCase(
+            description="input reaches the command and its exit code and output return",
+            script="import sys; sys.stdout.write(sys.stdin.read().upper()); sys.exit(3)",
+            input_text="orders\n",
+            expected_result=(3, "ORDERS\n", ""),
+        ),
+        FinishedCommandTestCase(
+            description="without input the command reads end of input",
+            script="import sys; print(repr(sys.stdin.read()))",
+            input_text=None,
+            expected_result=(0, "''\n", ""),
+        ),
+    ],
+    ids=lambda case: case.description,
+)
+def test_given_finishing_command_when_running_with_hang_dump_then_returns_its_result(
+    test_case: FinishedCommandTestCase, tmp_path: Path
+) -> None:
+    completed: subprocess.CompletedProcess[str] = run_with_hang_dump(
+        argv=(sys.executable, "-c", test_case.script),
+        cwd=tmp_path,
+        env=dict(os.environ),
+        input_text=test_case.input_text,
+        timeout_seconds=60.0,
+    )
+
+    assert (completed.returncode, completed.stdout, completed.stderr) == test_case.expected_result
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-vv"])

tests/integration/src/sqlbuild/runtime/compute_logs/_test_types.py

@@ -0,0 +1,14 @@
+"""Test case types for compute log storage integration tests."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class StorageLockTestCase:
+    """A Python program making concurrent storage calls from two threads, and what it prints."""
+
+    description: str
+    script: str
+    expected_stdout: str

tests/integration/src/sqlbuild/runtime/compute_logs/test_storage_lock_without_python.py

@@ -0,0 +1,130 @@
+"""Compute log storage never runs Python while holding its lock, so it cannot deadlock on the GIL."""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+from pathlib import Path
+from textwrap import dedent
+
+import pytest
+
+from tests.integration.src.sqlbuild.runtime.compute_logs._test_types import (
+    StorageLockTestCase,
+)
+
+_DEADLOCK_TIMEOUT_SECONDS: int = 30
+
+# A late write from a background thread reaches a capture that was already marked complete.
+# Raising its rejection imports the storage exceptions, here forced to load afresh, and the
+# import's file system calls release the GIL. The finder below holds that import open until the main thread is
+# pruning, which takes the storage lock while holding the GIL.
+_LATE_APPEND_DURING_PRUNE: str = dedent(
+    """
+    import dataclasses
+    import sys
+    import threading
+    import time
+    from datetime import UTC, datetime
+    from pathlib import Path
+
+    from sqlbuild.compute_logs import (
+        COMPUTE_LOG_FORMAT_VERSION,
+        CaptureMetadata,
+        ComputeLogStream,
+        FinalCaptureMetadata,
+        LocalFilesystemComputeLogStorage,
+    )
+
+    EXCEPTIONS_MODULE = "sqlbuild.runtime.compute_logs.exceptions"
+    importing = threading.Event()
+
+
+    class SlowExceptionsImport:
+        def find_spec(self, name, path=None, target=None):
+            if name == EXCEPTIONS_MODULE and not importing.is_set():
+                importing.set()
+                time.sleep(0.5)
+            return None
+
+
+    project_dir = Path(sys.argv[1])
+    started_at = datetime.now(UTC)
+    metadata = CaptureMetadata(
+        format_version=COMPUTE_LOG_FORMAT_VERSION,
+        invocation_id="orders0000000000000000000000000a",
+        command="compile",
+        project_dir=str(project_dir),
+        started_at=started_at,
+        capture_date=started_at.date().isoformat(),
+    )
+    storage = LocalFilesystemComputeLogStorage(project_dir=project_dir)
+    storage.start_capture(metadata)
+    storage.mark_complete(
+        FinalCaptureMetadata(
+            **dataclasses.asdict(metadata),
+            completed_at=datetime.now(UTC),
+            exit_code=1,
+            stdout_bytes=0,
+            stderr_bytes=0,
+            diagnostics_bytes=0,
+        )
+    )
+    rejected = []
+    del sys.modules[EXCEPTIONS_MODULE]
+    sys.meta_path.insert(0, SlowExceptionsImport())
+
+
+    def late_append():
+        try:
+            storage.append(
+                invocation_id=metadata.invocation_id,
+                stream=ComputeLogStream.STDERR,
+                data=b"late customer export notice\\n",
+            )
+        except Exception as error:
+            rejected.append(type(error).__name__)
+
+
+    writer = threading.Thread(target=late_append)
+    writer.start()
+    importing.wait()
+    storage.prune()
+    writer.join()
+    print(rejected)
+    """
+)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    (
+        StorageLockTestCase(
+            description="late append rejected while the main thread prunes",
+            script=_LATE_APPEND_DURING_PRUNE,
+            expected_stdout="['CaptureStateError']\n",
+        ),
+    ),
+    ids=lambda case: case.description,
+)
+def test_given_rejected_append_releasing_gil_when_pruning_then_both_threads_finish(
+    test_case: StorageLockTestCase, tmp_path: Path
+) -> None:
+    try:
+        completed: subprocess.CompletedProcess[str] = subprocess.run(
+            [sys.executable, "-c", test_case.script, str(tmp_path)],
+            capture_output=True,
+            text=True,
+            timeout=_DEADLOCK_TIMEOUT_SECONDS,
+            check=False,
+        )
+    except subprocess.TimeoutExpired:
+        pytest.fail(f"storage calls deadlocked for {_DEADLOCK_TIMEOUT_SECONDS}s")
+
+    assert (completed.returncode, completed.stdout) == (0, test_case.expected_stdout), (
+        completed.stderr
+    )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-vv"])