dedupe helpers in buildbarn tests and support restarting

src/python/pants/engine/internals/buildbarn_integration_tests/buildbarn_remote_cache_integration_test.py

@@ -3,40 +3,24 @@
 
 from __future__ import annotations
 
-import os
-import shutil
-import subprocess
 import time
 from collections.abc import Callable
 
 import pytest
 
 from pants.engine.fs import CreateDigest, Digest, DigestContents, Directory
-from pants.engine.internals.buildbarn_integration_tests.stack import LocalBuildbarnStack
+from pants.engine.internals.buildbarn_integration_tests.stack import (
+    CACHE_SPECULATION_DELAY_MILLIS,
+    LocalBuildbarnStack,
+    should_skip_for_missing_docker,
+)
 from pants.engine.process import Process, ProcessResult
 from pants.engine.rules import QueryRule
 from pants.testutil.rule_runner import RuleRunner
 from pants.util.logging import LogLevel
 
-
-def _docker_available() -> bool:
-    docker = shutil.which("docker")
-    if docker is None:
-        return False
-    result = subprocess.run(
-        [docker, "version", "--format", "{{.Server.Version}}"],
-        capture_output=True,
-        text=True,
-    )
-    return result.returncode == 0
-
-
-def _should_skip_for_missing_docker() -> bool:
-    return "CI" not in os.environ and not _docker_available()
-
-
 pytestmark = pytest.mark.skipif(
-    _should_skip_for_missing_docker(), reason="Docker is required for Buildbarn tests"
+    should_skip_for_missing_docker(), reason="Docker is required for Buildbarn tests"
 )
 
 
@@ -94,6 +78,7 @@ def _working_directory_process(rule_runner: RuleRunner) -> Process:
         working_directory="workdir",
         output_directories=[""],
         level=LogLevel.INFO,
+        remote_cache_speculation_delay_millis=CACHE_SPECULATION_DELAY_MILLIS,
     )
 
 
@@ -111,6 +96,7 @@ def test_buildbarn_remote_cache_roundtrips_outputs(subtests) -> None:
                 output_files=["file.txt"],
                 output_directories=["out"],
                 level=LogLevel.INFO,
+                remote_cache_speculation_delay_millis=CACHE_SPECULATION_DELAY_MILLIS,
             ),
             {
                 "file.txt": b"file-output\n",
@@ -128,6 +114,7 @@ def test_buildbarn_remote_cache_roundtrips_outputs(subtests) -> None:
                 description="Create root output directory contents",
                 output_directories=["."],
                 level=LogLevel.INFO,
+                remote_cache_speculation_delay_millis=CACHE_SPECULATION_DELAY_MILLIS,
             ),
             {
                 "root.txt": b"root\n",
@@ -145,6 +132,7 @@ def test_buildbarn_remote_cache_roundtrips_outputs(subtests) -> None:
                 description="Create empty root output directory contents",
                 output_directories=[""],
                 level=LogLevel.INFO,
+                remote_cache_speculation_delay_millis=CACHE_SPECULATION_DELAY_MILLIS,
             ),
             {
                 "root.txt": b"empty-root\n",

src/python/pants/engine/internals/buildbarn_integration_tests/buildbarn_remote_execution_integration_test.py

@@ -4,9 +4,6 @@
 from __future__ import annotations
 
 import itertools
-import os
-import shutil
-import subprocess
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
 from dataclasses import dataclass
@@ -18,6 +15,7 @@
 from pants.engine.internals.buildbarn_integration_tests.stack import (
     LocalBuildbarnStack,
     RemoteExecutionBuildbarn,
+    should_skip_for_missing_docker,
 )
 from pants.engine.internals.engine_testutil import WorkunitTracker
 from pants.engine.process import Process, ProcessResult
@@ -39,24 +37,8 @@ class RemoteExecutionRun:
     remote_command_digest: FileDigest
 
 
-def _docker_available() -> bool:
-    docker = shutil.which("docker")
-    if docker is None:
-        return False
-    result = subprocess.run(
-        [docker, "version", "--format", "{{.Server.Version}}"],
-        capture_output=True,
-        text=True,
-    )
-    return result.returncode == 0
-
-
-def _should_skip_for_missing_docker() -> bool:
-    return "CI" not in os.environ and not _docker_available()
-
-
 pytestmark = pytest.mark.skipif(
-    _should_skip_for_missing_docker(), reason="Docker is required for Buildbarn tests"
+    should_skip_for_missing_docker(), reason="Docker is required for Buildbarn tests"
 )
 
 # With the default of 0, Pants starts a remote execution before the cache lookup can answer,

src/python/pants/engine/internals/buildbarn_integration_tests/buildbarn_test.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import socket
 import subprocess
 from collections.abc import Sequence
 from pathlib import Path
@@ -66,6 +67,67 @@ def test_write_runtime_overlay_uses_instance_name(tmp_path: Path) -> None:
     assert "ghcr.io/example/executor:tag@sha256:" in overlay
 
 
+def _configured_cache_stack(
+    tmp_path: Path, run_command: buildbarn.RunCommand
+) -> buildbarn.LocalBuildbarnStack:
+    stack = buildbarn.LocalBuildbarnStack(temp_dir=tmp_path, run_command=run_command)
+    stack._project_name = "pants-buildbarn-test"
+    stack._runtime_root = tmp_path
+    stack._env_file = tmp_path / "compose.env"
+    stack._compose_file = tmp_path / "docker-compose.yaml"
+    stack._compose_profiles = ("cache",)
+    stack._active_services = ("cache",)
+    return stack
+
+
+def test_stop_cache_service_stops_compose_service(tmp_path: Path) -> None:
+    commands: list[tuple[tuple[str, ...], bool]] = []
+
+    def run_command(args: Sequence[str], check: bool) -> subprocess.CompletedProcess[str]:
+        commands.append((tuple(args), check))
+        return completed_process(args)
+
+    stack = _configured_cache_stack(tmp_path, run_command)
+    stack.stop_cache_service()
+
+    assert [(command[-2:], check) for command, check in commands] == [(("stop", "cache"), True)]
+
+
+def test_start_cache_service_rediscovers_ephemeral_port(tmp_path: Path) -> None:
+    commands: list[tuple[tuple[str, ...], bool]] = []
+
+    # The readiness check makes a real TCP connection, so a real listener stands in for the
+    # restarted service.
+    with socket.create_server(("127.0.0.1", 0)) as listener:
+        new_port = listener.getsockname()[1]
+
+        def run_command(args: Sequence[str], check: bool) -> subprocess.CompletedProcess[str]:
+            commands.append((tuple(args), check))
+            if "port" in args:
+                return completed_process(args, stdout=f"0.0.0.0:{new_port}\n")
+            if "ps" in args:
+                return completed_process(args, stdout="")
+            return completed_process(args)
+
+        stack = _configured_cache_stack(tmp_path, run_command)
+        stale = buildbarn.CacheOnlyBuildbarn(
+            address="grpc://127.0.0.1:1",
+            instance_name="fuse",
+            grpc_port=1,
+            project_name="pants-buildbarn-test",
+            temp_dir=tmp_path,
+            logs_path=tmp_path / "logs" / "compose.log",
+        )
+        stack._launched = stale
+
+        relaunched = stack.start_cache_service()
+
+    assert relaunched.grpc_port == new_port
+    assert relaunched.address == f"grpc://127.0.0.1:{new_port}"
+    assert stack._launched is relaunched
+    assert (("start", "cache"), True) in [(command[-2:], check) for command, check in commands]
+
+
 def test_write_compose_logs_uses_compose_command(tmp_path: Path) -> None:
     commands: list[tuple[tuple[str, ...], bool]] = []
 

src/python/pants/engine/internals/buildbarn_integration_tests/stack.py

@@ -13,14 +13,19 @@
 import uuid
 from collections.abc import Callable, Sequence
 from contextlib import AbstractContextManager
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from importlib.resources import files
 from pathlib import Path
 
 DEFAULT_INSTANCE_NAME = "fuse"
 DEFAULT_GRPC_PORT = 8980
 DEFAULT_READINESS_TIMEOUT_SECONDS = 30.0
 
+# For Processes whose test asserts a remote cache hit: with `Process`'s default of 0, Pants can
+# complete the process (locally, or via remote execution) before the cache lookup answers, and
+# the cached run then counts as uncached.
+CACHE_SPECULATION_DELAY_MILLIS = 10_000
+
 _CONFIG_ROOT = files(__package__).joinpath("config")
 _COMPOSE_FILE_NAME = "docker-compose.yaml"
 _IMAGE_CONFIG_FILE_NAME = "docker-compose.images.yaml"
@@ -39,6 +44,26 @@
 RunCommand = Callable[[Sequence[str], bool], subprocess.CompletedProcess[str]]
 
 
+def docker_available(docker_binary: str = "docker") -> bool:
+    docker = shutil.which(docker_binary)
+    if docker is None:
+        return False
+    result = subprocess.run(
+        [docker, "version", "--format", "{{.Server.Version}}"],
+        capture_output=True,
+        text=True,
+    )
+    return result.returncode == 0
+
+
+def should_skip_for_missing_docker() -> bool:
+    """Whether a Buildbarn-stack test should be skipped rather than fail for lack of Docker.
+
+    In CI, Docker is expected to be present, so a missing/broken Docker fails loudly there.
+    """
+    return "CI" not in os.environ and not docker_available()
+
+
 class FetchError(ValueError):
     pass
 
@@ -230,6 +255,37 @@ def launch_remote_execution(self) -> RemoteExecutionBuildbarn:
         self._launched = launched
         return launched
 
+    def stop_cache_service(self) -> None:
+        """Stop the cache-only Buildbarn service, e.g. to manipulate its storage on disk."""
+        self._run_compose(["stop", "cache"], check=True)
+
+    def start_cache_service(self) -> CacheOnlyBuildbarn:
+        """Start the cache-only Buildbarn service again after `stop_cache_service`.
+
+        Returns an updated `CacheOnlyBuildbarn`: Docker assigns a fresh ephemeral host port on
+        every container start, so the address from before the stop is stale.
+        """
+        launched = self._launched
+        if not isinstance(launched, CacheOnlyBuildbarn):
+            raise FetchError("start_cache_service requires a launched cache-only stack")
+        self._run_compose(["start", "cache"], check=True)
+        grpc_port = _discover_compose_host_port(
+            service="cache",
+            container_port=DEFAULT_GRPC_PORT,
+            timeout_seconds=self.readiness_timeout_seconds,
+            stack=self,
+            required_services=self._active_services,
+        )
+        _wait_for_tcp_readiness(
+            port=grpc_port,
+            timeout_seconds=self.readiness_timeout_seconds,
+            stack=self,
+            required_services=self._active_services,
+        )
+        relaunched = replace(launched, address=f"grpc://127.0.0.1:{grpc_port}", grpc_port=grpc_port)
+        self._launched = relaunched
+        return relaunched
+
     def teardown(self, *, remove_temp_dir: bool = True) -> None:
         if self._project_name is not None and self._runtime_root is not None:
             _write_compose_logs(self)