From e2dd5f53547511306e9dbb2acc4c47b9973cb811 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:57:53 -0400 Subject: [PATCH 001/161] fix(executor): bound registry artifact cache with leased LRU eviction The executor cache at /tmp/tracecat/registry-cache grew without bound (~22 GB per replica) until kubelet evicted the pod for ephemeral-storage pressure. Each mounted squashfs artifact also pinned one loop device forever; exhausting the container's loop nodes flipped a sticky process-wide flag that downgraded every later artifact to full unsquashfs extraction. - Add a lease API so action subprocess execution pins its artifacts against eviction for the full run - Evict least-recently-used idle entries at materialize time under entry-count and byte budgets (new config vars, defaults 8 / 10 GiB) - Unmount before deleting so eviction releases loop devices; skip busy or unmountable entries instead of forcing - Sweep orphaned scratch, stale mount dirs, and over-budget entries at startup - Only treat a mount failure as a capability probe if no mount ever succeeded; otherwise reclaim an idle mount and retry once ENG-1568 --- tests/unit/test_action_runner.py | 67 +++ tests/unit/test_registry_artifacts.py | 510 +++++++++++++++++++++ tracecat/config.py | 16 + tracecat/executor/action_runner.py | 83 ++-- tracecat/executor/registry_artifacts.py | 584 +++++++++++++++++++++++- 5 files changed, 1200 insertions(+), 60 deletions(-) diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index 8b06ca69dc..00981d0131 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -22,6 +22,7 @@ from tracecat.dsl.schemas import ActionStatement, RunActionInput, RunContext from tracecat.executor import action_runner from tracecat.executor.action_runner import ActionRunner +from tracecat.executor.registry_artifacts import compute_registry_artifact_cache_key from tracecat.executor.schemas import ( ActionImplementation, ExecutorActionErrorInfo, @@ -573,3 +574,69 @@ async def test_execute_action_masks_stderr_on_subprocess_crash( assert "temp_token" not in result.message assert "temp_secret" not in result.message assert "***" in result.message + + @pytest.mark.anyio + async def test_execute_action_holds_registry_lease_for_whole_subprocess( + self, + temp_cache_dir: Path, + mock_run_action_input: RunActionInput, + mock_role: Role, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """The artifact stays pinned until the action subprocess has exited.""" + runner = ActionRunner(cache_dir=temp_cache_dir) + artifact_uri = "s3://bucket/execute.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + entry_dir = temp_cache_dir / f"tarball-{cache_key}" + entry_dir.mkdir() + + monkeypatch.setattr( + action_runner.config, "TRACECAT__EXECUTOR_SANDBOX_ENABLED", False + ) + + success_response = orjson.dumps({"success": True, "result": {"data": "test"}}) + refcounts: list[int] = [] + registry_paths: list[str] = [] + + resolved_context = ResolvedContext( + action_impl=ActionImplementation( + type="udf", + action_name="core.table.search_rows", + module="tracecat_registry.core.table", + name="search_rows", + ), + evaluated_args={"table": "customers"}, + workspace_id=str(mock_role.workspace_id), + workflow_id=str(mock_run_action_input.run_context.wf_id), + run_id=str(mock_run_action_input.run_context.wf_run_id), + executor_token="test-executor-token", + secret_projection=_empty_secret_projection(), + ) + + async def create_subprocess_exec_side_effect(*args, **kwargs): # noqa: ARG001 + refcounts.append(runner.registry_artifacts._refcount(cache_key)) + env = kwargs.get("env") + assert isinstance(env, dict) + registry_paths.append(env["PYTHONPATH"]) + + mock_proc = AsyncMock() + mock_proc.returncode = 0 + mock_proc.communicate = AsyncMock(return_value=(success_response, b"")) + return mock_proc + + with patch( + "asyncio.create_subprocess_exec", + side_effect=create_subprocess_exec_side_effect, + ): + result = await runner.execute_action( + input=mock_run_action_input, + role=mock_role, + resolved_context=resolved_context, + artifact_uris=[artifact_uri], + timeout=10.0, + ) + + assert result == {"data": "test"} + assert refcounts == [1] + assert registry_paths[0].startswith(str(entry_dir)) + assert runner.registry_artifacts._refcount(cache_key) == 0 diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 6954ef57ae..671cca1eee 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import os import tarfile import tempfile from pathlib import Path @@ -23,6 +24,37 @@ ) from tracecat.registry.artifact_keys import parse_s3_uri +MAX_ENTRIES_CONFIG = ( + "tracecat.executor.registry_artifacts.config" + ".TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES" +) +MAX_BYTES_CONFIG = ( + "tracecat.executor.registry_artifacts.config" + ".TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES" +) +SQUASHFS_ENABLED_CONFIG = ( + "tracecat.executor.registry_artifacts.config" + ".TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED" +) + + +def _write_tarball_entry(cache_dir: Path, cache_key: str) -> Path: + """Create a materialized tarball cache entry on disk.""" + target_dir = cache_dir / f"tarball-{cache_key}" + target_dir.mkdir(parents=True) + (target_dir / "module.py").write_text("VALUE = 1") + return target_dir + + +def _write_image_entry( + cache_dir: Path, cache_key: str, *, size: int, mtime: float +) -> Path: + """Create a downloaded SquashFS image cache entry with a fixed mtime.""" + image_path = cache_dir / f"squashfs-{cache_key}.squashfs" + image_path.write_bytes(b"x" * size) + os.utime(image_path, (mtime, mtime)) + return image_path + @pytest.fixture def temp_cache_dir(): @@ -683,3 +715,481 @@ async def test_lock_for_different_keys(self, temp_cache_dir): lock2 = await cache._lock_for("key2") assert lock1 is not lock2 + + +class TestRegistryArtifactCacheLease: + """Tests for lease-based pinning of registry artifact cache entries.""" + + @pytest.mark.anyio + async def test_lease_refcounts_and_touches_image_mtime(self, temp_cache_dir): + """A lease pins its entry and refreshes the restart-safe LRU timestamp.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/leased.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + target_dir = _write_tarball_entry(temp_cache_dir, cache_key) + image_path = _write_image_entry(temp_cache_dir, cache_key, size=16, mtime=100.0) + + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [target_dir] + assert cache._refcount(cache_key) == 1 + assert image_path.stat().st_mtime > 100.0 + + assert cache._refcount(cache_key) == 0 + + @pytest.mark.anyio + async def test_lease_releases_refcount_when_materialization_fails( + self, temp_cache_dir + ): + """A failed materialization must not leak a permanent pin.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/broken.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def mock_download(self, ctx, path): + raise RuntimeError("download failed") + + with patch.object(TarballArtifact, "download", mock_download): + with pytest.raises(RuntimeError, match="download failed"): + async with cache.lease([artifact_uri]): + pass + + assert cache._refcount(cache_key) == 0 + + @pytest.mark.anyio + async def test_lease_without_uris_returns_base_pythonpath_dir(self, temp_cache_dir): + """No artifact URIs still yields the base PYTHONPATH directory.""" + cache = RegistryArtifactCache(temp_cache_dir) + + async with cache.lease(None) as registry_paths: + assert registry_paths == [temp_cache_dir / "base"] + assert registry_paths[0].is_dir() + + assert cache._leases == {} + + @pytest.mark.anyio + async def test_lease_preserves_uri_order(self, temp_cache_dir): + """Multiple artifacts keep their deterministic PYTHONPATH order.""" + cache = RegistryArtifactCache(temp_cache_dir) + uris = ["s3://bucket/first.tar.gz", "s3://bucket/second.tar.gz"] + expected = [ + _write_tarball_entry( + temp_cache_dir, compute_registry_artifact_cache_key(uri) + ) + for uri in uris + ] + + async with cache.lease(uris) as registry_paths: + assert registry_paths == expected + + @pytest.mark.anyio + async def test_builtin_artifact_is_exempt_from_cache_accounting( + self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch + ): + """The bundled builtin registry is never a cache entry.""" + version = "1.2.3" + site_packages = temp_cache_dir / "venv" / "site-packages" + package_dir = site_packages / "tracecat_registry" + package_dir.mkdir(parents=True) + package_file = package_dir / "__init__.py" + package_file.write_text("") + + monkeypatch.setattr(tracecat_registry, "__version__", version) + monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) + monkeypatch.setattr( + "tracecat.executor.registry_artifacts.sysconfig.get_path", + lambda name: str(site_packages) if name == "purelib" else None, + ) + + cache = RegistryArtifactCache(temp_cache_dir) + + with patch.object( + cache, + "_enforce_cache_budget", + new_callable=AsyncMock, + ) as enforce_cache_budget: + async with cache.lease([bundled_builtin_registry_uri(version)]) as paths: + assert paths == [site_packages.resolve()] + assert cache._leases == {} + + enforce_cache_budget.assert_not_awaited() + + +class TestRegistryArtifactCacheEviction: + """Tests for bounded eviction of registry artifact cache entries.""" + + @pytest.mark.anyio + async def test_leased_entry_survives_eviction_of_idle_entry(self, temp_cache_dir): + """Eviction must never remove an entry a live action is importing from.""" + cache = RegistryArtifactCache(temp_cache_dir) + leased_uri = "s3://bucket/leased.tar.gz" + idle_uri = "s3://bucket/idle.tar.gz" + new_uri = "s3://bucket/new.tar.gz" + leased_dir = _write_tarball_entry( + temp_cache_dir, compute_registry_artifact_cache_key(leased_uri) + ) + idle_dir = _write_tarball_entry( + temp_cache_dir, compute_registry_artifact_cache_key(idle_uri) + ) + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch(MAX_ENTRIES_CONFIG, 2), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + async with cache.lease([leased_uri]): + await cache.materialize( + compute_registry_artifact_cache_key(new_uri), new_uri + ) + + assert leased_dir.is_dir() + assert not idle_dir.exists() + + @pytest.mark.anyio + async def test_enforce_budget_evicts_least_recently_used_until_under_max_bytes( + self, temp_cache_dir + ): + """Size eviction stops as soon as the cache is back within budget.""" + cache = RegistryArtifactCache(temp_cache_dir) + oldest = _write_image_entry(temp_cache_dir, "oldest", size=4096, mtime=100.0) + older = _write_image_entry(temp_cache_dir, "older", size=4096, mtime=200.0) + newest = _write_image_entry(temp_cache_dir, "newest", size=4096, mtime=300.0) + + with patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 9000): + await cache._enforce_cache_budget(protected_key="pending") + + assert not oldest.exists() + assert older.exists() + assert newest.exists() + + @pytest.mark.anyio + async def test_enforce_budget_counts_the_pending_entry(self, temp_cache_dir): + """The entry about to be materialized counts against the entry budget.""" + cache = RegistryArtifactCache(temp_cache_dir) + existing = _write_image_entry(temp_cache_dir, "existing", size=16, mtime=100.0) + + with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): + await cache._enforce_cache_budget(protected_key="pending") + + assert not existing.exists() + + @pytest.mark.anyio + async def test_enforce_budget_never_evicts_the_protected_key(self, temp_cache_dir): + """The key being materialized is exempt even when it is the LRU entry.""" + cache = RegistryArtifactCache(temp_cache_dir) + protected = _write_image_entry( + temp_cache_dir, "protected", size=4096, mtime=100.0 + ) + other = _write_image_entry(temp_cache_dir, "other", size=4096, mtime=300.0) + + with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): + await cache._enforce_cache_budget(protected_key="protected") + + assert protected.exists() + assert not other.exists() + + @pytest.mark.anyio + async def test_enforce_budget_proceeds_over_budget_when_everything_is_leased( + self, temp_cache_dir + ): + """An over-budget cache must degrade, not fail the action.""" + cache = RegistryArtifactCache(temp_cache_dir) + leased = _write_image_entry(temp_cache_dir, "leased", size=4096, mtime=100.0) + cache._acquire_lease("leased") + + with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): + await cache._enforce_cache_budget(protected_key="pending") + + assert leased.exists() + + @pytest.mark.anyio + async def test_eviction_unmounts_before_deleting_the_image(self, temp_cache_dir): + """Unlinking a mounted image would strand an open-file zombie.""" + cache = RegistryArtifactCache(temp_cache_dir) + paths = cache._paths_for("mounted") + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + mounted = {paths.squashfs_mount_dir} + image_present_at_umount: list[bool] = [] + + process = AsyncMock() + process.communicate.return_value = (b"", b"") + process.returncode = 0 + + async def mock_umount(*args, **kwargs): + image_present_at_umount.append(paths.squashfs_image_path.exists()) + mounted.discard(paths.squashfs_mount_dir) + return process + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/umount", + ), + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + side_effect=mock_umount, + ) as create_subprocess_exec, + ): + evicted = await cache._evict_entry("mounted") + + assert evicted is True + assert image_present_at_umount == [True] + assert not paths.squashfs_image_path.exists() + assert not paths.squashfs_mount_dir.exists() + create_subprocess_exec.assert_called_once() + assert create_subprocess_exec.call_args.args == ( + "/sbin/umount", + str(paths.squashfs_mount_dir), + ) + + @pytest.mark.anyio + async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): + """A failed unmount skips the key instead of forcing a lazy detach.""" + cache = RegistryArtifactCache(temp_cache_dir) + stuck = cache._paths_for("stuck") + stuck.squashfs_image_path.write_bytes(b"squashfs") + stuck.squashfs_mount_dir.mkdir() + os.utime(stuck.squashfs_image_path, (100.0, 100.0)) + idle = _write_image_entry(temp_cache_dir, "idle", size=16, mtime=300.0) + mounted = {stuck.squashfs_mount_dir} + + process = AsyncMock() + process.communicate.return_value = (b"", b"target is busy") + process.returncode = 32 + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/umount", + ), + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ), + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + await cache._enforce_cache_budget(protected_key="pending") + + assert stuck.squashfs_image_path.exists() + assert stuck.squashfs_mount_dir.exists() + assert not idle.exists() + + @pytest.mark.anyio + async def test_eviction_drops_in_memory_bookkeeping(self, temp_cache_dir): + """Evicting a key releases its lock and lease records.""" + cache = RegistryArtifactCache(temp_cache_dir) + _write_tarball_entry(temp_cache_dir, "bookkeeping") + cache._acquire_lease("bookkeeping") + cache._release_lease("bookkeeping") + await cache._lock_for("bookkeeping") + + assert await cache._evict_entry("bookkeeping") is True + assert "bookkeeping" not in cache._locks + assert "bookkeeping" not in cache._leases + + @pytest.mark.anyio + async def test_eviction_skips_busy_key(self, temp_cache_dir): + """A key another task is materializing is never evicted underneath it.""" + cache = RegistryArtifactCache(temp_cache_dir) + target_dir = _write_tarball_entry(temp_cache_dir, "busy") + lock = await cache._lock_for("busy") + + async with lock: + assert await cache._evict_entry("busy") is False + + assert target_dir.is_dir() + + +class TestRegistryArtifactCacheStartupSweep: + """Tests for the startup sweep that reclaims state from a dead process.""" + + def test_sweep_tolerates_missing_cache_dir(self, temp_cache_dir): + """A cache directory that does not exist yet is a no-op.""" + cache_dir = temp_cache_dir / "missing" + + cache = RegistryArtifactCache(cache_dir) + + assert cache.cache_dir == cache_dir + assert not cache_dir.exists() + + def test_sweep_removes_orphaned_scratch_and_stale_mount_dirs(self, temp_cache_dir): + """Interrupted materializations and dead mount dirs are reclaimed.""" + orphaned = temp_cache_dir / "abc123.999999.4321.squashfs" + orphaned.write_bytes(b"partial") + orphaned_dir = temp_cache_dir / "abc123.999999.4321.tmp" + orphaned_dir.mkdir() + own = temp_cache_dir / f"abc123.{os.getpid()}.4321.tar.gz" + own.write_bytes(b"in flight") + stale_mount_dir = temp_cache_dir / "squashfs-abc123" + stale_mount_dir.mkdir() + entry_dir = _write_tarball_entry(temp_cache_dir, "abc123") + + RegistryArtifactCache(temp_cache_dir) + + assert not orphaned.exists() + assert not orphaned_dir.exists() + assert own.exists() + assert not stale_mount_dir.exists() + assert entry_dir.is_dir() + + def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): + """A live mountpoint belongs to a running process and must survive.""" + mount_dir = temp_cache_dir / "squashfs-abc123" + mount_dir.mkdir() + + with patch.object(Path, "is_mount", lambda self: self == mount_dir): + RegistryArtifactCache(temp_cache_dir) + + assert mount_dir.is_dir() + + def test_sweep_trims_to_budget_using_image_mtimes(self, temp_cache_dir): + """Startup LRU order comes from image mtimes, which survive a restart.""" + oldest = _write_image_entry(temp_cache_dir, "oldest", size=64, mtime=100.0) + older = _write_image_entry(temp_cache_dir, "older", size=64, mtime=200.0) + newest = _write_image_entry(temp_cache_dir, "newest", size=64, mtime=300.0) + + with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): + RegistryArtifactCache(temp_cache_dir) + + assert not oldest.exists() + assert not older.exists() + assert newest.exists() + + +class TestSquashfsMountCapability: + """Tests for process-wide SquashFS mount capability tracking.""" + + @pytest.mark.anyio + async def test_first_mount_failure_disables_squashfs_process_wide( + self, temp_cache_dir + ): + """With no prior success the failure is a capability probe.""" + cache = RegistryArtifactCache(temp_cache_dir) + + async def mock_mount(self, ctx, image_path): + raise RuntimeError("operation not permitted") + + async def mock_extract(self, ctx, image_path): + target_dir = ctx.paths.squashfs_extract_dir + target_dir.mkdir(parents=True, exist_ok=True) + return target_dir + + with ( + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + patch.object(SquashfsArtifact, "extract", mock_extract), + patch.object( + cache, + "_release_mounted_slot", + new_callable=AsyncMock, + return_value=False, + ) as release_mounted_slot, + ): + await cache.materialize( + "probe-key", "s3://bucket/path/site-packages.squashfs" + ) + + assert cache._squashfs_mount_state.disabled is True + release_mounted_slot.assert_not_awaited() + + @pytest.mark.anyio + async def test_mount_failure_after_success_reclaims_loop_device_and_retries( + self, temp_cache_dir + ): + """Loop-device exhaustion evicts an idle mount instead of going sticky.""" + cache = RegistryArtifactCache(temp_cache_dir) + cache._squashfs_mount_state.mounted_once = True + idle = cache._paths_for("idle") + idle.squashfs_image_path.write_bytes(b"squashfs") + idle.squashfs_mount_dir.mkdir() + mounted = {idle.squashfs_mount_dir} + attempts: list[str] = [] + + async def mock_mount(self, ctx, image_path): + attempts.append(ctx.cache_key) + if mounted: + raise RuntimeError("failed to setup loop device") + target_dir = ctx.paths.squashfs_mount_dir + target_dir.mkdir(parents=True, exist_ok=True) + (target_dir / "module.py").write_text("VALUE = 1") + mounted.add(target_dir) + return target_dir + + umount_process = AsyncMock() + umount_process.communicate.return_value = (b"", b"") + umount_process.returncode = 0 + + async def mock_umount(*args, **kwargs): + mounted.discard(idle.squashfs_mount_dir) + return umount_process + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/umount", + ), + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + side_effect=mock_umount, + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + ): + result = await cache.materialize( + "new", "s3://bucket/path/site-packages.squashfs" + ) + + assert result == [cache._paths_for("new").squashfs_mount_dir] + assert (result[0] / "module.py").read_text() == "VALUE = 1" + assert attempts == ["new", "new"] + assert cache._squashfs_mount_state.disabled is False + assert not idle.squashfs_image_path.exists() + assert not idle.squashfs_mount_dir.exists() + + @pytest.mark.anyio + async def test_mount_failure_without_reclaimable_slot_falls_back_to_extraction( + self, temp_cache_dir + ): + """Only this artifact degrades when no idle mount can be reclaimed.""" + cache = RegistryArtifactCache(temp_cache_dir) + cache._squashfs_mount_state.mounted_once = True + + async def mock_mount(self, ctx, image_path): + raise RuntimeError("failed to setup loop device") + + async def mock_extract(self, ctx, image_path): + target_dir = ctx.paths.squashfs_extract_dir + target_dir.mkdir(parents=True, exist_ok=True) + (target_dir / "module.py").write_text("VALUE = 1") + return target_dir + + with ( + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + patch.object(SquashfsArtifact, "extract", mock_extract), + ): + result = await cache.materialize( + "no-slot", "s3://bucket/path/site-packages.squashfs" + ) + + assert result[0].name.startswith("unsquashfs-") + assert cache._squashfs_mount_state.disabled is False diff --git a/tracecat/config.py b/tracecat/config.py index 8a9773b9de..45bb35e548 100644 --- a/tracecat/config.py +++ b/tracecat/config.py @@ -190,6 +190,22 @@ class RLSMode(StrEnum): ) """Prefer SquashFS registry artifacts when sidecars and mount support are available.""" +TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES = int( + os.environ.get("TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES") or 8 +) +"""Maximum number of registry artifacts kept in the executor-local cache. + +Each mounted SquashFS artifact pins one loop device, so this also bounds loop +device usage. Set to 0 to disable entry-count eviction.""" + +TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES = int( + os.environ.get("TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES") or 10 * 1024**3 +) +"""Maximum on-disk size of the executor-local registry artifact cache, in bytes. + +Mounted artifacts only account for their backing image file. Set to 0 to +disable size-based eviction.""" + TRACECAT__AGENT_SKILL_CACHE_DIR = os.environ.get( "TRACECAT__AGENT_SKILL_CACHE_DIR", "/tmp/tracecat/agent-skill-cache" ) diff --git a/tracecat/executor/action_runner.py b/tracecat/executor/action_runner.py index 79b7052abd..2ce64e4fbf 100644 --- a/tracecat/executor/action_runner.py +++ b/tracecat/executor/action_runner.py @@ -158,23 +158,14 @@ async def ensure_registry_environment(self, artifact_uri: str | None) -> list[Pa async def resolve_registry_paths( self, artifact_uris: list[str] | None = None ) -> list[Path]: - """Materialize registry artifacts and return importable Python paths.""" - registry_paths: list[Path] = [] - if artifact_uris: - for artifact_uri in artifact_uris: - registry_paths.extend( - await self.ensure_registry_environment(artifact_uri) - ) - logger.info( - "Using registry artifact environments", - count=len(registry_paths), - ) - return registry_paths + """Materialize registry artifacts and return importable Python paths. - base_dir = self.cache_dir / "base" - base_dir.mkdir(parents=True, exist_ok=True) - logger.info("No registry artifact URIs provided, using base PYTHONPATH") - return [base_dir] + The artifacts are only pinned for the duration of this call. Callers that + execute a subprocess against the returned paths should hold + ``registry_artifacts.lease`` for the whole execution instead. + """ + async with self.registry_artifacts.lease(artifact_uris) as registry_paths: + return registry_paths async def execute_action( self, @@ -207,38 +198,38 @@ async def execute_action( timeout = timeout or config.TRACECAT__EXECUTOR_CLIENT_TIMEOUT # Materialize each registry artifact, collect paths in deterministic order. - registry_paths = await self.resolve_registry_paths(artifact_uris) - - # Check if sandbox execution is enabled and available - # force_sandbox=True overrides config (used by ephemeral backend) - use_sandbox = force_sandbox or ( - config.TRACECAT__EXECUTOR_SANDBOX_ENABLED and _is_sandbox_available() - ) - logger.debug( - "Using sandbox execution", - use_sandbox=use_sandbox, - force_sandbox=force_sandbox, - ) - - secret_projection = resolved_context.secret_projection - if secret_projection is None: - secret_projection = await project_secret_env( - secrets=resolved_context.secrets, - role=role, - run_context=input.run_context, + # The lease is held for the whole subprocess execution so cache eviction + # cannot delete a directory the subprocess is still importing from. + async with self.registry_artifacts.lease(artifact_uris) as registry_paths: + # Check if sandbox execution is enabled and available + # force_sandbox=True overrides config (used by ephemeral backend) + use_sandbox = force_sandbox or ( + config.TRACECAT__EXECUTOR_SANDBOX_ENABLED and _is_sandbox_available() ) - - if use_sandbox: - return await self._execute_sandboxed( - input=input, - role=role, - registry_paths=registry_paths, - secret_projection=secret_projection, - env_vars=env_vars, - timeout=timeout, - resolved_context=resolved_context, + logger.debug( + "Using sandbox execution", + use_sandbox=use_sandbox, + force_sandbox=force_sandbox, ) - else: + + secret_projection = resolved_context.secret_projection + if secret_projection is None: + secret_projection = await project_secret_env( + secrets=resolved_context.secrets, + role=role, + run_context=input.run_context, + ) + + if use_sandbox: + return await self._execute_sandboxed( + input=input, + role=role, + registry_paths=registry_paths, + secret_projection=secret_projection, + env_vars=env_vars, + timeout=timeout, + resolved_context=resolved_context, + ) return await self._execute_direct( input=input, role=role, diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 2a82a28490..08d6110bf2 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -5,11 +5,14 @@ import asyncio import hashlib import os +import re import shutil import sysconfig import tarfile import time from abc import ABC, abstractmethod +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable +from contextlib import asynccontextmanager from dataclasses import dataclass from enum import StrEnum from pathlib import Path @@ -43,6 +46,20 @@ class RegistryArtifactFormat(StrEnum): BUNDLED_BUILTIN_REGISTRY_URI_PREFIX = f"tracecat-builtin://{DEFAULT_REGISTRY_ORIGIN}/" """Pseudo-URI for the builtin registry already installed in the executor image.""" +BASE_PYTHONPATH_DIR_NAME = "base" +"""Cache subdirectory used as the PYTHONPATH entry when no artifact is requested.""" + +CACHE_ENTRY_PREFIXES = ("squashfs-", "unsquashfs-", "tarball-") +"""On-disk name prefixes owned by a registry artifact cache entry.""" + +TEMP_ARTIFACT_PATTERN = re.compile( + r"^[^.]+\.(?P\d+)\.\d+\.(?:squashfs|unsquashfs|tar\.gz|tmp)$" +) +"""Matches in-flight materialization scratch names produced by ``_temp_path``.""" + +type MountSlotReleaser = Callable[[str], Awaitable[bool]] +"""Evicts one idle mounted artifact, excluding the given cache key.""" + @dataclass(frozen=True, slots=True) class RegistryArtifactPaths: @@ -59,6 +76,24 @@ class SquashfsMountState: """Shared process-local SquashFS mount state.""" disabled: bool = False + mounted_once: bool = False + + +@dataclass(slots=True) +class RegistryArtifactLease: + """In-process lease bookkeeping for one registry artifact cache key.""" + + refcount: int = 0 + last_used: float = 0.0 + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactCacheEntry: + """Measured on-disk footprint and recency for one registry artifact key.""" + + cache_key: str + size_bytes: int + last_used: float @dataclass(slots=True) @@ -69,6 +104,7 @@ class RegistryArtifactMaterializationContext: cache_dir: Path paths: RegistryArtifactPaths squashfs_mount_state: SquashfsMountState + mount_slot_releaser: MountSlotReleaser | None = None def can_mount_squashfs(self) -> bool: return ( @@ -80,6 +116,20 @@ def can_mount_squashfs(self) -> bool: def disable_squashfs_mount(self) -> None: self.squashfs_mount_state.disabled = True + def record_squashfs_mount(self) -> None: + """Record that this process has successfully mounted a SquashFS image.""" + self.squashfs_mount_state.mounted_once = True + + def has_mounted_squashfs(self) -> bool: + """Return whether any SquashFS mount has ever succeeded in this process.""" + return self.squashfs_mount_state.mounted_once + + async def release_mounted_slot(self) -> bool: + """Evict one idle mounted artifact to free a loop device.""" + if self.mount_slot_releaser is None: + return False + return await self.mount_slot_releaser(self.cache_key) + @dataclass(frozen=True, slots=True) class RegistryArtifact(ABC): @@ -172,20 +222,75 @@ async def materialize( ) -> list[Path]: image_path = ctx.paths.squashfs_image_path if ctx.can_mount_squashfs(): - try: - return [await self.mount(ctx, image_path)] - except Exception as e: - ctx.disable_squashfs_mount() - logger.warning( - "Failed to mount SquashFS registry artifact, trying extraction", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - error=str(e), - ) + if (mount_dir := await self._try_mount(ctx, image_path)) is not None: + return [mount_dir] return [await self.extract(ctx, image_path)] + async def _try_mount( + self, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> Path | None: + """Mount the image, retrying once after reclaiming a loop device. + + The first mount failure in a process is treated as a capability probe and + disables mounting process-wide. Once any mount has succeeded, later + failures are attributed to exhausted loop devices instead: one idle + mounted artifact is evicted and the mount is retried once, so a single + failure never downgrades the whole process to extraction. + + Args: + ctx: Materialization context for the artifact being mounted. + image_path: Local path of the SquashFS image. + + Returns: + The mount directory, or None if the caller should extract instead. + """ + try: + return await self.mount(ctx, image_path) + except Exception as e: + mount_error = e + + if not ctx.has_mounted_squashfs(): + ctx.disable_squashfs_mount() + logger.warning( + "Failed to mount SquashFS registry artifact, trying extraction", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + error=str(mount_error), + ) + return None + + logger.warning( + "Failed to mount SquashFS registry artifact, reclaiming an idle mount", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + error=str(mount_error), + ) + if not await ctx.release_mounted_slot(): + logger.warning( + "No idle SquashFS mount to reclaim, trying extraction", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + ) + return None + + try: + return await self.mount(ctx, image_path) + except Exception as e: + logger.warning( + "SquashFS mount retry failed, trying extraction", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + error=str(e), + ) + return None + async def download( self, ctx: RegistryArtifactMaterializationContext, @@ -215,6 +320,7 @@ async def mount( ) -> Path: target_dir = ctx.paths.squashfs_mount_dir if target_dir.is_mount(): + ctx.record_squashfs_mount() return target_dir ctx.cache_dir.mkdir(parents=True, exist_ok=True) @@ -231,6 +337,7 @@ async def mount( mount_start = time.monotonic() await self._mount_image(image_path, target_dir) + ctx.record_squashfs_mount() mount_elapsed = (time.monotonic() - mount_start) * 1000 total_elapsed = (time.monotonic() - start_time) * 1000 @@ -547,6 +654,68 @@ def _artifact_format(artifact_uri: str) -> RegistryArtifactFormat: return RegistryArtifactFormat.TAR_GZ +def _is_cache_entry_uri(artifact_uri: str) -> bool: + """Return whether an artifact URI materializes into an evictable cache entry. + + The bundled builtin registry is served from the executor image and never + writes into the cache directory, so it is exempt from eviction accounting. + """ + return _bundled_builtin_registry_version(artifact_uri) is None + + +def _cache_key_from_entry_name(name: str) -> str | None: + """Return the cache key owning a cache directory entry name, if any.""" + for prefix in CACHE_ENTRY_PREFIXES: + if name.startswith(prefix): + cache_key = name.removeprefix(prefix).removesuffix(".squashfs") + return cache_key or None + return None + + +def _directory_footprint(directory: Path) -> tuple[int, float]: + """Return the total file size and local creation time of a cache directory. + + Directory ``mtime`` values inside extracted artifacts come from the artifact + build, so ``ctime`` (updated when the staging directory is renamed into + place) is used as the local recency signal instead. + + Args: + directory: Cache directory to measure. + + Returns: + Total byte size of contained files and the directory's ctime, or + ``(0, 0.0)`` when the directory is missing. + """ + if not directory.is_dir(): + return 0, 0.0 + + try: + created_at = directory.stat().st_ctime + except OSError: + created_at = 0.0 + + total_bytes = 0 + for root, _dirs, files in os.walk(directory): + for file_name in files: + try: + total_bytes += os.lstat(os.path.join(root, file_name)).st_size + except OSError: + continue + return total_bytes, created_at + + +def _delete_entry_paths(paths: RegistryArtifactPaths) -> None: + """Delete every on-disk path owned by a registry artifact cache key. + + The caller must unmount ``squashfs_mount_dir`` first: unlinking the image + file behind a live mount leaves an open-file zombie pinning a loop device. + """ + shutil.rmtree(paths.squashfs_extract_dir, ignore_errors=True) + shutil.rmtree(paths.tarball_target_dir, ignore_errors=True) + paths.squashfs_image_path.unlink(missing_ok=True) + shutil.rmtree(paths.squashfs_mount_dir, ignore_errors=True) + + class RegistryArtifactCache: """Materializes registry artifacts into executor-local Python paths.""" @@ -555,13 +724,56 @@ def __init__(self, cache_dir: Path): self._locks: dict[str, asyncio.Lock] = {} self._locks_lock = asyncio.Lock() self._squashfs_mount_state = SquashfsMountState() + self._leases: dict[str, RegistryArtifactLease] = {} + self._sweep_startup_state() + + @asynccontextmanager + async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Path]]: + """Materialize registry artifacts and pin them for the life of the context. + + Leased cache entries are never evicted, so callers may keep importing + from the returned paths until the context exits. + + Args: + artifact_uris: Registry artifact URIs in deterministic PYTHONPATH + order, or None to use the base PYTHONPATH directory. + + Yields: + Importable Python paths for the requested artifacts. + """ + if not artifact_uris: + logger.info("No registry artifact URIs provided, using base PYTHONPATH") + yield [self._base_pythonpath_dir()] + return + + leased_keys: list[str] = [] + try: + registry_paths: list[Path] = [] + for artifact_uri in artifact_uris: + cache_key = compute_registry_artifact_cache_key(artifact_uri) + if _is_cache_entry_uri(artifact_uri): + self._acquire_lease(cache_key) + leased_keys.append(cache_key) + registry_paths.extend(await self.materialize(cache_key, artifact_uri)) + logger.info( + "Using registry artifact environments", + count=len(registry_paths), + ) + yield registry_paths + finally: + for cache_key in leased_keys: + self._release_lease(cache_key) async def ensure_environment(self, artifact_uri: str | None) -> list[Path]: - """Materialize an optional registry artifact and return PYTHONPATH entries.""" + """Materialize an optional registry artifact and return PYTHONPATH entries. + + The artifact is only pinned for the duration of this call. Callers that + execute against the returned paths should hold a ``lease`` instead. + """ if not artifact_uri: return [] - cache_key = compute_registry_artifact_cache_key(artifact_uri) - return await self.materialize(cache_key, artifact_uri) + async with self.lease([artifact_uri]) as registry_paths: + return registry_paths async def materialize(self, cache_key: str, artifact_uri: str) -> list[Path]: """Materialize a registry artifact as local importable directories.""" @@ -571,6 +783,11 @@ async def materialize(self, cache_key: str, artifact_uri: str) -> list[Path]: if cached_paths := self._first_cached_path(candidates, ctx): return cached_paths + if _is_cache_entry_uri(artifact_uri): + # Make room before downloading or expanding a new entry. This runs + # outside the per-key lock so eviction never nests key locks. + await self._enforce_cache_budget(protected_key=cache_key) + lock = await self._lock_for(cache_key) async with lock: candidates = await self._artifact_candidates(ctx, artifact_uri) @@ -615,8 +832,50 @@ def _context_for(self, cache_key: str) -> RegistryArtifactMaterializationContext cache_dir=self.cache_dir, paths=self._paths_for(cache_key), squashfs_mount_state=self._squashfs_mount_state, + mount_slot_releaser=self._release_mounted_slot, ) + def _base_pythonpath_dir(self) -> Path: + """Return the base PYTHONPATH directory used when no artifact is requested.""" + base_dir = self.cache_dir / BASE_PYTHONPATH_DIR_NAME + base_dir.mkdir(parents=True, exist_ok=True) + return base_dir + + def _acquire_lease(self, cache_key: str) -> None: + """Pin a cache entry against eviction and mark it as recently used.""" + lease = self._leases.setdefault(cache_key, RegistryArtifactLease()) + lease.refcount += 1 + lease.last_used = time.time() + self._touch_image(cache_key) + + def _release_lease(self, cache_key: str) -> None: + """Release one pin on a cache entry.""" + lease = self._leases.get(cache_key) + if lease is None: + return + lease.refcount = max(0, lease.refcount - 1) + lease.last_used = time.time() + + def _refcount(self, cache_key: str) -> int: + """Return the number of live leases on a cache entry.""" + lease = self._leases.get(cache_key) + return 0 if lease is None else lease.refcount + + def _touch_image(self, cache_key: str) -> None: + """Best-effort refresh of an artifact image mtime for restart-safe LRU. + + Only the downloaded image file has a locally meaningful mtime; mount and + extraction directory timestamps come from the artifact build. + """ + image_path = self._paths_for(cache_key).squashfs_image_path + try: + os.utime(image_path) + except OSError: + logger.debug( + "Could not refresh registry artifact image mtime", + cache_key=cache_key, + ) + def _paths_for(self, cache_key: str) -> RegistryArtifactPaths: """Return local cache paths for a registry artifact key.""" return RegistryArtifactPaths( @@ -732,3 +991,300 @@ async def _sidecar_exists( def _can_try_squashfs(self) -> bool: """Return whether this process should prefer SquashFS artifacts.""" return config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED + + async def _enforce_cache_budget(self, *, protected_key: str) -> None: + """Evict least-recently-used idle entries until the cache fits its budget. + + Args: + protected_key: Cache key about to be materialized. It is counted + against the budget but never evicted. + """ + max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES + max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES + if max_entries <= 0 and max_bytes <= 0: + return + + entries = await asyncio.to_thread(self._scan_cache_entries) + # The protected key is not on disk yet when it is a fresh entry. + pending_entries = 0 if protected_key in entries else 1 + total_bytes = sum(entry.size_bytes for entry in entries.values()) + skipped: set[str] = set() + + while (max_entries > 0 and len(entries) + pending_entries > max_entries) or ( + max_bytes > 0 and total_bytes > max_bytes + ): + candidate = self._least_recently_used( + entries.values(), + excluded=skipped | {protected_key}, + ) + if candidate is None: + logger.warning( + "Registry artifact cache is over budget but every entry is in use", + cache_dir=str(self.cache_dir), + entries=len(entries) + pending_entries, + max_entries=max_entries, + total_bytes=total_bytes, + max_bytes=max_bytes, + ) + return + + if await self._evict_entry(candidate.cache_key): + del entries[candidate.cache_key] + total_bytes -= candidate.size_bytes + else: + skipped.add(candidate.cache_key) + + async def _release_mounted_slot(self, protected_key: str) -> bool: + """Evict one idle mounted artifact so its loop device can be reused. + + Args: + protected_key: Cache key that must not be evicted. + + Returns: + Whether a mounted artifact was unmounted and removed. + """ + entries = await asyncio.to_thread(self._scan_cache_entries) + mounted = [ + entry + for entry in entries.values() + if entry.cache_key != protected_key + and self._paths_for(entry.cache_key).squashfs_mount_dir.is_mount() + ] + skipped: set[str] = set() + while ( + candidate := self._least_recently_used(mounted, excluded=skipped) + ) is not None: + if await self._evict_entry(candidate.cache_key): + return True + skipped.add(candidate.cache_key) + return False + + def _least_recently_used( + self, + entries: Iterable[RegistryArtifactCacheEntry], + *, + excluded: set[str], + ) -> RegistryArtifactCacheEntry | None: + """Return the least recently used idle entry eligible for eviction.""" + eligible = [ + entry + for entry in entries + if entry.cache_key not in excluded and self._refcount(entry.cache_key) == 0 + ] + if not eligible: + return None + return min(eligible, key=self._recency) + + def _recency(self, entry: RegistryArtifactCacheEntry) -> float: + """Return the most recent known use time for a cache entry.""" + lease = self._leases.get(entry.cache_key) + if lease is None: + return entry.last_used + return max(entry.last_used, lease.last_used) + + async def _evict_entry(self, cache_key: str) -> bool: + """Remove one cache entry from disk, unmounting it first. + + The entry is skipped rather than forced when it is leased, busy, or + cannot be unmounted: deleting the image file behind a live mount would + leave an open-file zombie holding the loop device. + + Args: + cache_key: Cache key to evict. + + Returns: + Whether the entry was removed. + """ + lock = await self._lock_for(cache_key) + if lock.locked(): + logger.debug( + "Skipping eviction of busy registry artifact", + cache_key=cache_key, + ) + return False + + async with lock: + if self._refcount(cache_key) > 0: + return False + + paths = self._paths_for(cache_key) + if paths.squashfs_mount_dir.is_mount() and not await self._unmount( + paths.squashfs_mount_dir + ): + logger.warning( + "Failed to unmount registry artifact, skipping eviction", + cache_key=cache_key, + mount_dir=str(paths.squashfs_mount_dir), + ) + return False + + _delete_entry_paths(paths) + logger.info("Evicted registry artifact from cache", cache_key=cache_key) + + await self._forget(cache_key, lock) + return True + + async def _forget(self, cache_key: str, lock: asyncio.Lock) -> None: + """Drop in-memory bookkeeping for an evicted cache key.""" + async with self._locks_lock: + self._leases.pop(cache_key, None) + if self._locks.get(cache_key) is lock and not lock.locked(): + del self._locks[cache_key] + + async def _unmount(self, mount_dir: Path) -> bool: + """Unmount a SquashFS artifact directory, releasing its loop device.""" + umount = shutil.which("umount") + if umount is None: + return False + + proc = await asyncio.create_subprocess_exec( + umount, + str(mount_dir), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await proc.communicate() + if proc.returncode == 0 or not mount_dir.is_mount(): + return True + + logger.warning( + "umount command failed", + mount_dir=str(mount_dir), + output=(stderr or stdout).decode(errors="replace").strip(), + ) + return False + + def _scan_cache_entries(self) -> dict[str, RegistryArtifactCacheEntry]: + """Measure every registry artifact entry currently on disk.""" + return { + cache_key: self._measure_entry(cache_key) + for cache_key in self._discover_cache_keys() + } + + def _discover_cache_keys(self) -> set[str]: + """Return the cache keys with at least one path in the cache directory.""" + try: + names = os.listdir(self.cache_dir) + except OSError: + return set() + + cache_keys: set[str] = set() + for name in names: + if (cache_key := _cache_key_from_entry_name(name)) is not None: + cache_keys.add(cache_key) + return cache_keys + + def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: + """Measure the on-disk footprint and recency of one cache entry. + + The mount directory is excluded because a mounted view only costs the + image file that backs it. + """ + paths = self._paths_for(cache_key) + size_bytes = 0 + image_mtime = 0.0 + created_at = 0.0 + + if paths.squashfs_image_path.is_file(): + image_stat = paths.squashfs_image_path.stat() + size_bytes += image_stat.st_size + image_mtime = image_stat.st_mtime + + for directory in (paths.squashfs_extract_dir, paths.tarball_target_dir): + directory_bytes, directory_created_at = _directory_footprint(directory) + size_bytes += directory_bytes + created_at = max(created_at, directory_created_at) + + # Image mtimes are refreshed on lease; directory ctimes are only a + # fallback for entries that have no locally downloaded image. + last_used = image_mtime or created_at + + return RegistryArtifactCacheEntry( + cache_key=cache_key, + size_bytes=size_bytes, + last_used=last_used, + ) + + def _sweep_startup_state(self) -> None: + """Reclaim orphaned cache state left behind by a previous process. + + Mounts never survive a container restart, so any non-mountpoint mount + directory is stale. Scratch paths from interrupted materializations are + removed, and the cache is trimmed to budget using image mtimes as LRU + order. A missing or empty cache directory is a no-op. + """ + if not self.cache_dir.is_dir(): + return + + try: + self._remove_orphaned_temp_paths() + self._remove_stale_mount_dirs() + self._trim_startup_cache() + except OSError as e: + logger.warning( + "Failed to sweep registry artifact cache", + cache_dir=str(self.cache_dir), + error=str(e), + ) + + def _remove_orphaned_temp_paths(self) -> None: + """Delete materialization scratch paths owned by dead processes.""" + current_pid = os.getpid() + for name in os.listdir(self.cache_dir): + match = TEMP_ARTIFACT_PATTERN.match(name) + if match is None or int(match.group("pid")) == current_pid: + continue + path = self.cache_dir / name + if path.is_dir(): + shutil.rmtree(path, ignore_errors=True) + else: + path.unlink(missing_ok=True) + logger.info("Removed orphaned registry artifact scratch path", path=name) + + def _remove_stale_mount_dirs(self) -> None: + """Remove empty mount directories left over from a previous process.""" + for cache_key in self._discover_cache_keys(): + mount_dir = self._paths_for(cache_key).squashfs_mount_dir + if not mount_dir.is_dir() or mount_dir.is_mount(): + continue + try: + mount_dir.rmdir() + except OSError: + continue + logger.debug( + "Removed stale registry artifact mount directory", + cache_key=cache_key, + ) + + def _trim_startup_cache(self) -> None: + """Trim the cache to budget before any artifact is leased.""" + max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES + max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES + if max_entries <= 0 and max_bytes <= 0: + return + + entries = self._scan_cache_entries() + total_bytes = sum(entry.size_bytes for entry in entries.values()) + # Mounted entries belong to a live process sharing this cache directory. + candidates = sorted( + ( + entry + for entry in entries.values() + if not self._paths_for(entry.cache_key).squashfs_mount_dir.is_mount() + ), + key=lambda entry: entry.last_used, + ) + + for entry in candidates: + if (max_entries <= 0 or len(entries) <= max_entries) and ( + max_bytes <= 0 or total_bytes <= max_bytes + ): + return + _delete_entry_paths(self._paths_for(entry.cache_key)) + del entries[entry.cache_key] + total_bytes -= entry.size_bytes + logger.info( + "Evicted stale registry artifact during startup sweep", + cache_key=entry.cache_key, + size_bytes=entry.size_bytes, + ) From a01aa327899605e93407864c2e382b3f21a71e55 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:18:08 -0400 Subject: [PATCH 002/161] fix(executor): make cache leases atomic with eviction, converge budgets Hardens the eviction change after independent review: - Serialize lease admission with eviction on the per-key lock: the refcount increment and cached-path check happen under the same lock eviction holds, so a lease can never be handed a path an in-flight eviction is about to delete. Locks now live for the process lifetime. - Hold the lease across core.script.run_python sandbox execution and TestBackend execution; delete the resolve_registry_paths / ensure_registry_environment / ensure_environment wrappers so lease is the only way to obtain cache paths. - Converge budgets once entries go idle: a dirty flag set on new materialization triggers re-enforcement with real on-disk sizes on lease release; steady-state cache hits never rescan the disk. - Scope mount policy to genuine mount-command failures via SquashfsMountCommandError; transient download errors no longer disable mounting process-wide or evict an unrelated idle mount. - Offload runtime entry deletion to a thread; fix an is_file/stat TOCTOU in cache scanning. - Add a privileged-docker integration test exercising the real mount -> evict -> loop-device-release -> remount lifecycle and the startup sweep. ENG-1568 --- ...registry_artifact_cache_mount_lifecycle.py | 375 ++++++++++++++++++ tests/integration/test_syncv2_execv2_e2e.py | 23 +- .../executor/test_run_python_sdk_context.py | 86 +++- .../test_test_backend_no_registry_action.py | 86 ++++ tests/unit/test_action_runner.py | 12 +- tests/unit/test_registry_artifacts.py | 264 +++++++++++- tracecat/executor/action_runner.py | 25 -- tracecat/executor/backends/base.py | 103 +++-- tracecat/executor/backends/test.py | 61 ++- tracecat/executor/registry_artifacts.py | 215 ++++++++-- 10 files changed, 1085 insertions(+), 165 deletions(-) create mode 100644 tests/integration/test_registry_artifact_cache_mount_lifecycle.py diff --git a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py new file mode 100644 index 0000000000..1abeeb0a1c --- /dev/null +++ b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py @@ -0,0 +1,375 @@ +"""Dockerized lifecycle test for executor registry artifact SquashFS mounts. + +The executor caches registry environments as SquashFS images and mounts them, +which consumes one loop device per mounted artifact. Unit tests stub the mount +and umount commands, so this test drives the real thing: it runs inside the +privileged executor image, builds tiny SquashFS images with ``mksquashfs``, and +asserts that materialization mounts them, that eviction unmounts them and +releases their loop devices, and that the startup sweep reclaims stale state. + +Run it with the ``integration`` marker; it is skipped when Docker is unavailable. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any + +import pytest + +_MOUNT_LIFECYCLE_CHILD_ENV = "TRACECAT__REGISTRY_MOUNT_LIFECYCLE_CHILD" +_MOUNT_LIFECYCLE_RESULT = "TRACECAT_REGISTRY_MOUNT_LIFECYCLE_RESULT:" +_MOUNT_LIFECYCLE_FLAG = "--run-registry-mount-lifecycle" + + +def _run_mount_lifecycle_in_docker_or_skip() -> dict[str, Any]: + """Run the mount lifecycle child inside the privileged executor image. + + Returns: + The JSON payload emitted by the in-container child run. + """ + if os.environ.get(_MOUNT_LIFECYCLE_CHILD_ENV) == "1": + pytest.skip("already inside registry mount lifecycle Docker child") + if shutil.which("docker") is None: + pytest.skip("Docker CLI unavailable for registry mount lifecycle") + if ( + subprocess.run( + ["docker", "info"], + capture_output=True, + text=True, + timeout=10, + check=False, + ).returncode + != 0 + ): + pytest.skip("Docker daemon unavailable for registry mount lifecycle") + + repo_root = Path(__file__).resolve().parents[2] + + compose_env = os.environ.copy() + compose_env.setdefault( + "TRACECAT__LOCAL_REPOSITORY_PATH", + str(repo_root / "packages"), + ) + compose_env.setdefault("PUBLIC_APP_PORT", "80") + compose_env.setdefault("BASE_DOMAIN", ":80") + compose_env.setdefault("ADDRESS", "0.0.0.0") + compose_env["LOG_LEVEL"] = "INFO" + compose_env["TRACECAT__APP_ENV"] = "development" + compose_env["TRACECAT__SERVICE_KEY"] = "test-service-key" + compose_env["TRACECAT__LOCAL_REPOSITORY_ENABLED"] = "false" + compose_env["TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED"] = "true" + compose_env[_MOUNT_LIFECYCLE_CHILD_ENV] = "1" + compose_env["PYTHONDONTWRITEBYTECODE"] = "1" + + override_path = Path( + tempfile.mkstemp(prefix="tracecat-registry-mount-lifecycle-", suffix=".yml")[1] + ) + override_path.write_text( + "\n".join( + [ + "services:", + " executor:", + " build:", + " target: test", + # Mounting SquashFS images needs both the capability and a uid + # that holds it, so the child runs as root in a privileged + # container. Nothing else in this test touches the host. + " privileged: true", + ' user: "0:0"', + " security_opt:", + " - seccomp:unconfined", + " - systempaths=unconfined", + " environment:", + f" - {_MOUNT_LIFECYCLE_CHILD_ENV}", + " - TRACECAT__APP_ENV", + " - TRACECAT__SERVICE_KEY", + " - TRACECAT__LOCAL_REPOSITORY_ENABLED", + " - TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", + " - PYTHONDONTWRITEBYTECODE", + "", + ] + ) + ) + try: + result = subprocess.run( + [ + "docker", + "compose", + "-f", + str(repo_root / "docker-compose.dev.yml"), + "-f", + str(override_path), + "run", + "--rm", + "--no-deps", + "--build", + "-T", + "--entrypoint", + "/app/.venv/bin/python", + "executor", + "-m", + "tests.integration.test_registry_artifact_cache_mount_lifecycle", + _MOUNT_LIFECYCLE_FLAG, + ], + cwd=repo_root, + env=compose_env, + capture_output=True, + text=True, + timeout=900, + check=False, + ) + finally: + override_path.unlink(missing_ok=True) + + if result.returncode != 0: + pytest.fail( + "Dockerized registry mount lifecycle failed." + f"\n\nstdout:\n{result.stdout}\n\nstderr:\n{result.stderr}" + ) + + output = f"{result.stdout}\n{result.stderr}" + for line in output.splitlines(): + if line.startswith(_MOUNT_LIFECYCLE_RESULT): + return json.loads(line.removeprefix(_MOUNT_LIFECYCLE_RESULT)) + + pytest.fail( + "Dockerized registry mount lifecycle did not emit result sentinel." + f"\n\nstdout:\n{result.stdout}\n\nstderr:\n{result.stderr}" + ) + + +@pytest.mark.integration +def test_registry_artifact_cache_mount_lifecycle() -> None: + """Real mounts, evictions, and loop devices behave as the cache assumes.""" + payload = _run_mount_lifecycle_in_docker_or_skip() + + if skipped := payload.get("skipped"): + pytest.skip(f"SquashFS mounts unsupported in this container: {skipped}") + + # Every materialized artifact is mounted and holds its own loop device. + assert payload["mounted_targets"] == 3 + assert payload["mounted_loop_devices"] == 3 + assert payload["module_readable_through_mount"] is True + + # Eviction unmounts, frees the loop device, and deletes the entry. + assert payload["evicted"] is True + assert payload["evicted_target_unmounted"] is True + assert payload["evicted_loop_device_released"] is True + assert payload["evicted_paths_removed"] is True + + # Re-materializing the evicted key mounts again: no sticky disable flag. + assert payload["remounted"] is True + assert payload["squashfs_disabled_after_eviction"] is False + + # Releasing a lease converges an over-budget cache, unmounting as it goes. + assert payload["converged_entry_unmounted"] is True + assert payload["converged_loop_device_released"] is True + assert payload["converged_entries_remaining"] == 3 + + # The startup sweep trims to budget and removes stale mount directories. + assert payload["startup_sweep_trimmed"] is True + assert payload["startup_sweep_removed_stale_mount_dir"] is True + + +def _squashfs_mounts(cache_dir: Path) -> dict[str, str]: + """Map mount target to backing device for SquashFS mounts under cache_dir. + + Args: + cache_dir: Registry artifact cache directory. + + Returns: + Mount target path to backing device (a loop device when mounted). + """ + mounts: dict[str, str] = {} + for line in Path("/proc/mounts").read_text().splitlines(): + fields = line.split() + if len(fields) < 3 or fields[2] != "squashfs": + continue + if fields[1].startswith(f"{cache_dir}/"): + mounts[fields[1]] = fields[0] + return mounts + + +def _build_squashfs_image(source_dir: Path, image_path: Path, module_name: str) -> None: + """Build a tiny SquashFS image containing a single Python module. + + Args: + source_dir: Scratch directory holding the module. + image_path: Destination image path inside the cache directory. + module_name: Module file name to place in the image. + """ + source_dir.mkdir(parents=True, exist_ok=True) + (source_dir / module_name).write_text("VALUE = 1\n") + image_path.unlink(missing_ok=True) + subprocess.run( + ["mksquashfs", str(source_dir), str(image_path), "-noappend", "-quiet"], + check=True, + capture_output=True, + ) + + +async def _run_mount_lifecycle_child() -> None: + """Exercise the real mount, eviction, and sweep lifecycle in a container.""" + from tracecat import config + from tracecat.executor.registry_artifacts import ( + RegistryArtifactCache, + compute_registry_artifact_cache_key, + ) + + config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED = True + config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES = 8 + config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES = 0 + + payload: dict[str, Any] = {} + with tempfile.TemporaryDirectory(prefix="registry-mount-lifecycle-") as tmp: + root = Path(tmp) + cache_dir = root / "registry-cache" + cache_dir.mkdir() + cache = RegistryArtifactCache(cache_dir) + + uris = [ + f"s3://bucket/lifecycle/{index}/site-packages.squashfs" + for index in range(3) + ] + keys = [compute_registry_artifact_cache_key(uri) for uri in uris] + for index, key in enumerate(keys): + _build_squashfs_image( + root / f"source-{index}", + cache._paths_for(key).squashfs_image_path, + f"module_{index}.py", + ) + + # (a) Materialize every artifact through the real cache. + module_readable = True + for index, uri in enumerate(uris): + async with cache.lease([uri]) as registry_paths: + mount_dir = cache._paths_for(keys[index]).squashfs_mount_dir + if registry_paths != [mount_dir]: + if index == 0 and not mount_dir.is_mount(): + # Kernels without SquashFS or loop support fall back to + # extraction; that is an environment limit, not a bug. + payload["skipped"] = ( + f"materialized {registry_paths} instead of a mount" + ) + print( + f"{_MOUNT_LIFECYCLE_RESULT}" + f"{json.dumps(payload, sort_keys=True)}" + ) + return + raise AssertionError( + f"Expected mount for {uri}, got {registry_paths}" + ) + module_readable = ( + module_readable + and (mount_dir / f"module_{index}.py").read_text() == "VALUE = 1\n" + ) + + mounts = _squashfs_mounts(cache_dir) + payload["mounted_targets"] = len(mounts) + payload["mounted_loop_devices"] = len( + {device for device in mounts.values() if device.startswith("/dev/loop")} + ) + payload["module_readable_through_mount"] = module_readable + + # (b) Evict one entry: it must unmount, free its loop device, and vanish. + evicted_paths = cache._paths_for(keys[0]) + evicted_device = mounts[str(evicted_paths.squashfs_mount_dir)] + payload["evicted"] = await cache._evict_entry(keys[0]) + mounts_after_eviction = _squashfs_mounts(cache_dir) + payload["evicted_target_unmounted"] = ( + str(evicted_paths.squashfs_mount_dir) not in mounts_after_eviction + ) + payload["evicted_loop_device_released"] = ( + evicted_device not in mounts_after_eviction.values() + ) + payload["evicted_paths_removed"] = not ( + evicted_paths.squashfs_image_path.exists() + or evicted_paths.squashfs_mount_dir.exists() + ) + + # (c) The evicted key must mount again: eviction is not a capability probe. + _build_squashfs_image( + root / "source-0", + evicted_paths.squashfs_image_path, + "module_0.py", + ) + async with cache.lease([uris[0]]) as registry_paths: + payload["remounted"] = registry_paths == [evicted_paths.squashfs_mount_dir] + payload["squashfs_disabled_after_eviction"] = ( + cache._squashfs_mount_state.disabled + ) + + # (d) A released lease converges an over-budget cache, unmounting as it + # goes. The fourth entry is materialized under the old budget, so only + # the release-time check can bring the cache back within the new one. + fourth_uri = "s3://bucket/lifecycle/3/site-packages.squashfs" + fourth_key = compute_registry_artifact_cache_key(fourth_uri) + _build_squashfs_image( + root / "source-3", + cache._paths_for(fourth_key).squashfs_image_path, + "module_3.py", + ) + # keys[2] is the least recently used idle entry: keys[0] was re-leased + # above and keys[1] is leased again below. + converged_paths = cache._paths_for(keys[2]) + async with cache.lease([fourth_uri]): + mounts_before_converge = _squashfs_mounts(cache_dir) + converged_device = mounts_before_converge[ + str(converged_paths.squashfs_mount_dir) + ] + config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES = 3 + async with cache.lease([uris[1]]): + pass + mounts_after_converge = _squashfs_mounts(cache_dir) + payload["converged_entry_unmounted"] = ( + str(converged_paths.squashfs_mount_dir) not in mounts_after_converge + ) + payload["converged_loop_device_released"] = ( + converged_device not in mounts_after_converge.values() + ) + payload["converged_entries_remaining"] = len(cache._discover_cache_keys()) + + # (e) The startup sweep trims to budget and drops stale mount directories. + sweep_dir = root / "sweep-cache" + sweep_dir.mkdir() + sweep_keys = ["aaaa1111", "bbbb2222"] + for index, sweep_key in enumerate(sweep_keys): + image_path = sweep_dir / f"squashfs-{sweep_key}.squashfs" + image_path.write_bytes(b"x" * 4096) + os.utime(image_path, (100.0 + index, 100.0 + index)) + stale_mount_dir = sweep_dir / f"squashfs-{sweep_keys[0]}" + stale_mount_dir.mkdir() + + config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES = 1 + RegistryArtifactCache(sweep_dir) + payload["startup_sweep_trimmed"] = ( + not (sweep_dir / f"squashfs-{sweep_keys[0]}.squashfs").exists() + and (sweep_dir / f"squashfs-{sweep_keys[1]}.squashfs").exists() + ) + payload["startup_sweep_removed_stale_mount_dir"] = not stale_mount_dir.exists() + + # Leave no mounts behind for the container teardown. + for cache_key in sorted(cache._discover_cache_keys()): + await cache._evict_entry(cache_key) + + print(f"{_MOUNT_LIFECYCLE_RESULT}{json.dumps(payload, sort_keys=True)}") + + +if __name__ == "__main__": + if os.environ.get(_MOUNT_LIFECYCLE_CHILD_ENV) == "1" and sys.argv[1:] == [ + _MOUNT_LIFECYCLE_FLAG + ]: + asyncio.run(_run_mount_lifecycle_child()) + else: + raise SystemExit( + "Usage: python -m tests.integration." + f"test_registry_artifact_cache_mount_lifecycle {_MOUNT_LIFECYCLE_FLAG}" + ) diff --git a/tests/integration/test_syncv2_execv2_e2e.py b/tests/integration/test_syncv2_execv2_e2e.py index 1d1c4339ea..28634de432 100644 --- a/tests/integration/test_syncv2_execv2_e2e.py +++ b/tests/integration/test_syncv2_execv2_e2e.py @@ -865,18 +865,17 @@ async def test_action_runner_downloads_and_extracts_tarball( runner = ActionRunner(cache_dir=temp_cache_dir) # Download and materialize the registry artifact - extracted_paths = await runner.ensure_registry_environment(tarball_uri) - - # Verify extraction - assert len(extracted_paths) == 1 - extracted_path = extracted_paths[0] - assert extracted_path.exists() - assert extracted_path.is_dir() - - # Should have Python packages extracted - # The tarball contains site-packages content - files = list(extracted_path.rglob("*")) - assert len(files) > 0 + async with runner.registry_artifacts.lease([tarball_uri]) as extracted_paths: + # Verify extraction + assert len(extracted_paths) == 1 + extracted_path = extracted_paths[0] + assert extracted_path.exists() + assert extracted_path.is_dir() + + # Should have Python packages extracted + # The tarball contains site-packages content + files = list(extracted_path.rglob("*")) + assert len(files) > 0 @pytest.mark.anyio async def test_execute_action_with_ephemeral_backend( diff --git a/tests/unit/executor/test_run_python_sdk_context.py b/tests/unit/executor/test_run_python_sdk_context.py index caf22119cf..9adb24b45d 100644 --- a/tests/unit/executor/test_run_python_sdk_context.py +++ b/tests/unit/executor/test_run_python_sdk_context.py @@ -10,9 +10,11 @@ import sysconfig import tempfile import uuid +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from datetime import UTC, datetime from pathlib import Path -from typing import TYPE_CHECKING, Any, assert_type +from typing import TYPE_CHECKING, Any, Never, assert_type import pytest import tracecat_registry @@ -465,16 +467,37 @@ async def _run_sandbox_registry_ctx_smoke( ) -class _FakeRunPythonRegistryPathRunner: +class _FakeRunPythonRegistryArtifacts: + """Stand-in for the executor registry artifact cache.""" + def __init__(self, paths: list[Path]) -> None: self.paths = paths self.artifact_uris: list[str] | None = None + self.leased = False - async def resolve_registry_paths( + @asynccontextmanager + async def lease( self, artifact_uris: list[str] | None = None - ) -> list[Path]: + ) -> AsyncIterator[list[Path]]: self.artifact_uris = artifact_uris - return self.paths + self.leased = True + try: + yield self.paths + finally: + self.leased = False + + +class _FakeRunPythonRegistryPathRunner: + def __init__(self, paths: list[Path]) -> None: + self.registry_artifacts = _FakeRunPythonRegistryArtifacts(paths) + + @property + def paths(self) -> list[Path]: + return self.registry_artifacts.paths + + @property + def artifact_uris(self) -> list[str] | None: + return self.registry_artifacts.artifact_uris async def _run_backend_registry_ctx_smoke( @@ -1169,6 +1192,48 @@ async def run_python(self, **kwargs: Any) -> dict[str, bool]: ] +@pytest.mark.anyio +async def test_run_python_backend_holds_registry_lease_for_whole_sandbox_run( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Registry artifacts stay pinned until the run_python sandbox has finished.""" + artifact_path = tmp_path / "registry-artifact" + artifact_path.mkdir() + fake_runner = _FakeRunPythonRegistryPathRunner([artifact_path]) + leased_during_run: list[bool] = [] + + async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: + return ["s3://tracecat-registry/test/site-packages.tar.gz"] + + class FakeSandboxService: + async def run_python(self, **kwargs: Any) -> dict[str, bool]: + del kwargs + leased_during_run.append(fake_runner.registry_artifacts.leased) + return {"ok": True} + + monkeypatch.setattr( + "tracecat.executor.backends.base.SandboxService", + FakeSandboxService, + ) + monkeypatch.setattr( + "tracecat.executor.backends.base.get_action_runner", + lambda: fake_runner, + ) + backend = DirectBackend() + monkeypatch.setattr(backend, "_get_artifact_uris", _get_artifact_uris) + + result = await backend.execute( + input=_make_run_python_input(), + role=_make_role(), + resolved_context=_make_run_python_context(), + ) + + assert result.type == "success" + assert leased_during_run == [True] + assert fake_runner.registry_artifacts.leased is False + + @pytest.mark.anyio async def test_run_python_backend_fails_without_registry_artifacts( monkeypatch: pytest.MonkeyPatch, @@ -1221,14 +1286,15 @@ async def run_python(self, **kwargs: Any) -> dict[str, bool]: captured.update(kwargs) return {"ok": True} - class FailingRegistryPathRunner: - async def resolve_registry_paths( - self, artifact_uris: list[str] | None = None - ) -> list[Path]: + class FailingRegistryArtifacts: + def lease(self, artifact_uris: list[str] | None = None) -> Never: raise AssertionError( - f"local repository mode should not resolve {artifact_uris=}" + f"local repository mode should not lease {artifact_uris=}" ) + class FailingRegistryPathRunner: + registry_artifacts = FailingRegistryArtifacts() + monkeypatch.setattr(executor_backend_module, "SandboxService", FakeSandboxService) monkeypatch.setattr( executor_backend_module, diff --git a/tests/unit/executor/test_test_backend_no_registry_action.py b/tests/unit/executor/test_test_backend_no_registry_action.py index 4a23926667..faf96571f4 100644 --- a/tests/unit/executor/test_test_backend_no_registry_action.py +++ b/tests/unit/executor/test_test_backend_no_registry_action.py @@ -8,8 +8,12 @@ from __future__ import annotations +import sys import uuid +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from datetime import UTC, datetime +from pathlib import Path import pytest from tracecat_registry import secrets as registry_secrets @@ -242,6 +246,88 @@ async def test_execute_udf_missing_module_fails( finally: await backend.shutdown() + @pytest.mark.anyio + async def test_execute_holds_artifact_leases_for_the_whole_execution( + self, + test_role: Role, + test_resolved_context: ResolvedContext, + test_run_action_input: RunActionInput, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + """Leases span execution, and one bad artifact does not drop the others.""" + good_path = tmp_path / "good-artifact" + good_path.mkdir() + broken_uri = "s3://bucket/broken.tar.gz" + good_uri = "s3://bucket/good.tar.gz" + + class FakeRegistryArtifacts: + def __init__(self) -> None: + self.active = 0 + + @asynccontextmanager + async def lease( + self, artifact_uris: list[str] | None = None + ) -> AsyncIterator[list[Path]]: + if artifact_uris == [broken_uri]: + raise RuntimeError("artifact unavailable") + self.active += 1 + try: + yield [good_path] + finally: + self.active -= 1 + + class FakeActionRunner: + def __init__(self) -> None: + self.registry_artifacts = FakeRegistryArtifacts() + + fake_runner = FakeActionRunner() + observed: list[tuple[int, bool]] = [] + + async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: + return [broken_uri, good_uri] + + backend = TestBackend() + await backend.start() + + try: + monkeypatch.setattr( + "tracecat.executor.backends.test.config" + ".TRACECAT__LOCAL_REPOSITORY_ENABLED", + False, + ) + monkeypatch.setattr( + "tracecat.executor.backends.test.get_action_runner", + lambda: fake_runner, + ) + monkeypatch.setattr(backend, "_get_artifact_uris", _get_artifact_uris) + monkeypatch.setattr( + backend, + "_load_udf_callable", + lambda _action_impl: ( + lambda **_kwargs: observed.append( + ( + fake_runner.registry_artifacts.active, + str(good_path) in sys.path, + ) + ) + ), + ) + + result = await backend.execute( + input=test_run_action_input, + role=test_role, + resolved_context=test_resolved_context, + timeout=30.0, + ) + + assert result.type == "success" + assert observed == [(1, True)] + assert fake_runner.registry_artifacts.active == 0 + assert str(good_path) not in sys.path + finally: + await backend.shutdown() + @pytest.mark.anyio async def test_execute_udf_reuses_cached_secret_projection( self, diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index 00981d0131..14e3f6959a 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -117,15 +117,15 @@ async def communicate( return communication @pytest.mark.anyio - async def test_ensure_registry_environment_no_tarball(self, temp_cache_dir): - """Test that an empty list is returned when no tarball URI provided.""" + async def test_lease_without_artifacts_yields_base_pythonpath(self, temp_cache_dir): + """Test that the base PYTHONPATH directory is used without artifacts.""" runner = ActionRunner(cache_dir=temp_cache_dir) - result = await runner.ensure_registry_environment(None) - assert result == [] + async with runner.registry_artifacts.lease(None) as registry_paths: + assert registry_paths == [temp_cache_dir / "base"] - result = await runner.ensure_registry_environment("") - assert result == [] + async with runner.registry_artifacts.lease([]) as registry_paths: + assert registry_paths == [temp_cache_dir / "base"] @pytest.mark.anyio async def test_execute_action_timeout( diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 671cca1eee..253b724e2c 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -18,6 +18,7 @@ RegistryArtifactCache, RegistryArtifactFormat, SquashfsArtifact, + SquashfsMountCommandError, TarballArtifact, bundled_builtin_registry_uri, compute_registry_artifact_cache_key, @@ -154,7 +155,7 @@ async def mock_download_file_to_path( assert output_path.read_bytes() == b"squashfs" @pytest.mark.anyio - async def test_ensure_environment_uses_bundled_current_builtin( + async def test_lease_uses_bundled_current_builtin( self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch ): """In-tree builtin registry returns only the installed site-packages.""" @@ -173,12 +174,11 @@ async def test_ensure_environment_uses_bundled_current_builtin( ) cache = RegistryArtifactCache(temp_cache_dir) - result = await cache.ensure_environment(bundled_builtin_registry_uri(version)) - - assert result == [site_packages.resolve()] + async with cache.lease([bundled_builtin_registry_uri(version)]) as result: + assert result == [site_packages.resolve()] @pytest.mark.anyio - async def test_ensure_environment_exposes_editable_builtin_parent( + async def test_lease_exposes_editable_builtin_parent( self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch ): """Editable builtin registry exposes the package wrapper + site-packages.""" @@ -201,12 +201,11 @@ async def test_ensure_environment_exposes_editable_builtin_parent( ) cache = RegistryArtifactCache(temp_cache_dir) - result = await cache.ensure_environment(bundled_builtin_registry_uri(version)) - - assert result == [source_root.resolve(), site_packages.resolve()] + async with cache.lease([bundled_builtin_registry_uri(version)]) as result: + assert result == [source_root.resolve(), site_packages.resolve()] @pytest.mark.anyio - async def test_ensure_environment_rejects_stale_bundled_builtin( + async def test_lease_rejects_stale_bundled_builtin( self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch ): """Bundled pseudo-URIs must match this executor's installed package.""" @@ -214,7 +213,8 @@ async def test_ensure_environment_rejects_stale_bundled_builtin( cache = RegistryArtifactCache(temp_cache_dir) with pytest.raises(RuntimeError, match="does not match installed version"): - await cache.ensure_environment(bundled_builtin_registry_uri("1.2.4")) + async with cache.lease([bundled_builtin_registry_uri("1.2.4")]): + pass @pytest.mark.anyio async def test_download_artifact_normalizes_missing_objects_to_http_404( @@ -509,7 +509,7 @@ async def test_materialize_extracts_squashfs_when_mount_fails(self, temp_cache_d cache = RegistryArtifactCache(temp_cache_dir) async def mock_mount(self, ctx, image_path): - raise RuntimeError("operation not permitted") + raise SquashfsMountCommandError("operation not permitted") async def mock_extract(self, ctx, image_path): target_dir = ctx.paths.squashfs_extract_dir @@ -601,7 +601,7 @@ async def mock_tarball_download(self, ctx, path): tar.add(source / "module.py", arcname="module.py") async def mock_mount(self, ctx, image_path): - raise RuntimeError("operation not permitted") + raise SquashfsMountCommandError("operation not permitted") async def mock_extract(self, ctx, image_path): raise RuntimeError("unsquashfs unavailable") @@ -781,6 +781,76 @@ async def test_lease_preserves_uri_order(self, temp_cache_dir): async with cache.lease(uris) as registry_paths: assert registry_paths == expected + @pytest.mark.anyio + async def test_lease_is_never_admitted_across_an_in_flight_eviction( + self, temp_cache_dir + ): + """A lease must not return a mount an in-flight eviction is deleting.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/site-packages.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + paths = cache._paths_for(cache_key) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + mounted = {paths.squashfs_mount_dir} + umount_started = asyncio.Event() + finish_umount = asyncio.Event() + remounts: list[str] = [] + + umount_process = AsyncMock() + umount_process.communicate.return_value = (b"", b"") + umount_process.returncode = 0 + + async def mock_umount(*args, **kwargs): + umount_started.set() + await finish_umount.wait() + mounted.discard(paths.squashfs_mount_dir) + return umount_process + + async def mock_mount(self, ctx, image_path): + remounts.append(ctx.cache_key) + target_dir = ctx.paths.squashfs_mount_dir + target_dir.mkdir(parents=True, exist_ok=True) + (target_dir / "module.py").write_text("VALUE = 1") + mounted.add(target_dir) + return target_dir + + leased_paths: list[Path] = [] + leased_path_exists: list[bool] = [] + + async def take_lease() -> None: + async with cache.lease([artifact_uri]) as registry_paths: + leased_paths.extend(registry_paths) + leased_path_exists.append(registry_paths[0].is_dir()) + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/umount", + ), + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + side_effect=mock_umount, + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + ): + eviction = asyncio.create_task(cache._evict_entry(cache_key)) + await umount_started.wait() + lease = asyncio.create_task(take_lease()) + # Let the lease block on the per-key lock the eviction holds. + await asyncio.sleep(0) + finish_umount.set() + evicted, _ = await asyncio.gather(eviction, lease) + + assert evicted is True + # The lease waited for the eviction and re-materialized the entry. + assert remounts == [cache_key] + assert leased_paths == [paths.squashfs_mount_dir] + assert leased_path_exists == [True] + assert (paths.squashfs_mount_dir / "module.py").read_text() == "VALUE = 1" + @pytest.mark.anyio async def test_builtin_artifact_is_exempt_from_cache_accounting( self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch @@ -850,6 +920,78 @@ async def mock_extract(self, tarball_path, target_dir): assert leased_dir.is_dir() assert not idle_dir.exists() + @pytest.mark.anyio + async def test_releasing_a_lease_converges_the_cache_to_budget( + self, temp_cache_dir + ): + """Enforcement before materialization cannot see the new entry's size.""" + cache = RegistryArtifactCache(temp_cache_dir) + idle = _write_image_entry(temp_cache_dir, "idle", size=4096, mtime=100.0) + new_uri = "s3://bucket/new.tar.gz" + new_key = compute_registry_artifact_cache_key(new_uri) + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_bytes(b"x" * 4096) + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 6000), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + async with cache.lease([new_uri]) as registry_paths: + # Both entries fit only because the new one is still leased. + assert registry_paths == [temp_cache_dir / f"tarball-{new_key}"] + assert idle.exists() + + assert not idle.exists() + assert (temp_cache_dir / f"tarball-{new_key}").is_dir() + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( + self, temp_cache_dir + ): + """Steady-state cache hits must not pay for a full cache scan.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/cached.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + _write_tarball_entry(temp_cache_dir, cache_key) + cache._budget_dirty = False + + with patch.object( + cache, + "_scan_cache_entries", + side_effect=AssertionError("cache hits must not scan the cache dir"), + ): + async with cache.lease([artifact_uri]): + pass + + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_release_keeps_retrying_while_the_cache_stays_over_budget( + self, temp_cache_dir + ): + """A cache that cannot shrink yet must stay marked for re-enforcement.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/pinned.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + _write_image_entry(temp_cache_dir, cache_key, size=4096, mtime=100.0) + _write_tarball_entry(temp_cache_dir, cache_key) + cache._budget_dirty = True + # A second holder keeps the entry pinned past the inner lease. + cache._acquire_lease(cache_key) + + with patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 1): + async with cache.lease([artifact_uri]): + pass + + assert cache._budget_dirty is True + @pytest.mark.anyio async def test_enforce_budget_evicts_least_recently_used_until_under_max_bytes( self, temp_cache_dir @@ -985,16 +1127,18 @@ async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): assert not idle.exists() @pytest.mark.anyio - async def test_eviction_drops_in_memory_bookkeeping(self, temp_cache_dir): - """Evicting a key releases its lock and lease records.""" + async def test_eviction_drops_lease_records_but_keeps_the_key_lock( + self, temp_cache_dir + ): + """Locks are stable for the process lifetime; lease records are not.""" cache = RegistryArtifactCache(temp_cache_dir) _write_tarball_entry(temp_cache_dir, "bookkeeping") cache._acquire_lease("bookkeeping") cache._release_lease("bookkeeping") - await cache._lock_for("bookkeeping") + lock = await cache._lock_for("bookkeeping") assert await cache._evict_entry("bookkeeping") is True - assert "bookkeeping" not in cache._locks + assert cache._locks["bookkeeping"] is lock assert "bookkeeping" not in cache._leases @pytest.mark.anyio @@ -1077,7 +1221,7 @@ async def test_first_mount_failure_disables_squashfs_process_wide( cache = RegistryArtifactCache(temp_cache_dir) async def mock_mount(self, ctx, image_path): - raise RuntimeError("operation not permitted") + raise SquashfsMountCommandError("operation not permitted") async def mock_extract(self, ctx, image_path): target_dir = ctx.paths.squashfs_extract_dir @@ -1122,7 +1266,7 @@ async def test_mount_failure_after_success_reclaims_loop_device_and_retries( async def mock_mount(self, ctx, image_path): attempts.append(ctx.cache_key) if mounted: - raise RuntimeError("failed to setup loop device") + raise SquashfsMountCommandError("failed to setup loop device") target_dir = ctx.paths.squashfs_mount_dir target_dir.mkdir(parents=True, exist_ok=True) (target_dir / "module.py").write_text("VALUE = 1") @@ -1161,6 +1305,88 @@ async def mock_umount(*args, **kwargs): assert not idle.squashfs_image_path.exists() assert not idle.squashfs_mount_dir.exists() + @pytest.mark.anyio + async def test_download_failure_does_not_disable_squashfs_process_wide( + self, temp_cache_dir + ): + """A transient download error is not a missing mount capability.""" + cache = RegistryArtifactCache(temp_cache_dir) + tarball_dir = temp_cache_dir / "gzip-fallback" + tarball_dir.mkdir() + + async def mock_download(self, ctx, image_path): + raise RuntimeError("connection reset by peer") + + with ( + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "download", mock_download), + patch.object( + TarballArtifact, + "materialize", + new_callable=AsyncMock, + return_value=[tarball_dir], + ) as tarball_materialize, + patch.object( + cache, + "_release_mounted_slot", + new_callable=AsyncMock, + return_value=False, + ) as release_mounted_slot, + ): + result = await cache.materialize( + "download-failure", "s3://bucket/path/site-packages.squashfs" + ) + + assert result == [tarball_dir] + assert cache._squashfs_mount_state.disabled is False + tarball_materialize.assert_awaited_once() + release_mounted_slot.assert_not_awaited() + + @pytest.mark.anyio + async def test_download_failure_after_a_mount_success_does_not_reclaim_a_slot( + self, temp_cache_dir + ): + """A transient download error must not evict an unrelated idle mount.""" + cache = RegistryArtifactCache(temp_cache_dir) + cache._squashfs_mount_state.mounted_once = True + tarball_dir = temp_cache_dir / "gzip-fallback" + tarball_dir.mkdir() + + async def mock_download(self, ctx, image_path): + raise RuntimeError("connection reset by peer") + + with ( + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "download", mock_download), + patch.object( + TarballArtifact, + "materialize", + new_callable=AsyncMock, + return_value=[tarball_dir], + ), + patch.object( + cache, + "_release_mounted_slot", + new_callable=AsyncMock, + return_value=True, + ) as release_mounted_slot, + ): + result = await cache.materialize( + "download-failure", "s3://bucket/path/site-packages.squashfs" + ) + + assert result == [tarball_dir] + assert cache._squashfs_mount_state.disabled is False + release_mounted_slot.assert_not_awaited() + @pytest.mark.anyio async def test_mount_failure_without_reclaimable_slot_falls_back_to_extraction( self, temp_cache_dir @@ -1170,7 +1396,7 @@ async def test_mount_failure_without_reclaimable_slot_falls_back_to_extraction( cache._squashfs_mount_state.mounted_once = True async def mock_mount(self, ctx, image_path): - raise RuntimeError("failed to setup loop device") + raise SquashfsMountCommandError("failed to setup loop device") async def mock_extract(self, ctx, image_path): target_dir = ctx.paths.squashfs_extract_dir diff --git a/tracecat/executor/action_runner.py b/tracecat/executor/action_runner.py index 2ce64e4fbf..438ab64307 100644 --- a/tracecat/executor/action_runner.py +++ b/tracecat/executor/action_runner.py @@ -142,31 +142,6 @@ def __init__(self, cache_dir: Path | None = None): self.registry_artifacts = RegistryArtifactCache(self.cache_dir) logger.info("ActionRunner initialized", cache_dir=str(self.cache_dir)) - async def ensure_registry_environment(self, artifact_uri: str | None) -> list[Path]: - """Ensure the registry environment is set up and return PYTHONPATH entries. - - This is the public API for pool workers to get the paths to add to PYTHONPATH. - - Args: - artifact_uri: S3 URI to the registry execution artifact. - - Returns: - Paths to add to PYTHONPATH (empty if no artifact is available). - """ - return await self.registry_artifacts.ensure_environment(artifact_uri) - - async def resolve_registry_paths( - self, artifact_uris: list[str] | None = None - ) -> list[Path]: - """Materialize registry artifacts and return importable Python paths. - - The artifacts are only pinned for the duration of this call. Callers that - execute a subprocess against the returned paths should hold - ``registry_artifacts.lease`` for the whole execution instead. - """ - async with self.registry_artifacts.lease(artifact_uris) as registry_paths: - return registry_paths - async def execute_action( self, input: RunActionInput, diff --git a/tracecat/executor/backends/base.py b/tracecat/executor/backends/base.py index effa97a0f9..b4b3790231 100644 --- a/tracecat/executor/backends/base.py +++ b/tracecat/executor/backends/base.py @@ -15,7 +15,7 @@ import os from abc import ABC, abstractmethod from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from tracecat import config from tracecat.dsl.enums import PlatformAction @@ -132,10 +132,56 @@ async def _execute_run_python( resolved_context, user_env_vars=args.get("env_vars"), ) - registry_paths = await self._resolve_run_python_registry_paths(input, role) - if isinstance(registry_paths, ExecutorActionErrorInfo): - return ExecutorResultFailure(error=registry_paths) + artifact_uris = await self._get_artifact_uris(input, role) + if not artifact_uris: + # Local-repository mode imports straight from the repo checkout, so + # there is no cache entry to lease. + local_registry_paths = self._resolve_run_python_local_registry_paths() + if isinstance(local_registry_paths, ExecutorActionErrorInfo): + return ExecutorResultFailure(error=local_registry_paths) + return await self._run_python_in_sandbox( + script=script, + args=args, + env_vars=env_vars, + registry_paths=local_registry_paths, + resolved_context=resolved_context, + ) + + # The lease is held for the whole sandbox run so cache eviction cannot + # delete a directory the script is still importing from. + registry_artifacts = get_action_runner().registry_artifacts + async with registry_artifacts.lease(artifact_uris) as registry_paths: + return await self._run_python_in_sandbox( + script=script, + args=args, + env_vars=env_vars, + registry_paths=registry_paths, + resolved_context=resolved_context, + ) + + async def _run_python_in_sandbox( + self, + *, + script: str, + args: dict[str, Any], + env_vars: dict[str, str], + registry_paths: list[Path], + resolved_context: ResolvedContext, + ) -> ExecutorResult: + """Run a validated run_python script against resolved registry paths. + + Args: + script: Validated run_python script source. + args: Evaluated run_python action arguments. + env_vars: Environment variables for the sandboxed script. + registry_paths: Import roots for the sandboxed script. + resolved_context: Pre-resolved execution context. + + Returns: + ExecutorResultSuccess on success, ExecutorResultFailure on a + sandbox, validation, or dependency-install failure. + """ service = SandboxService() try: result = await service.run_python( @@ -165,36 +211,31 @@ async def _execute_run_python( ) return ExecutorResultFailure(error=error_info) - async def _resolve_run_python_registry_paths( + def _resolve_run_python_local_registry_paths( self, - input: RunActionInput, - role: Role, ) -> list[Path] | ExecutorActionErrorInfo: - """Resolve registry artifact paths for run_python SDK imports.""" - artifact_uris = await self._get_artifact_uris(input, role) - if not artifact_uris: - if config.TRACECAT__LOCAL_REPOSITORY_ENABLED: - if local_registry_paths := self._get_run_python_local_registry_paths(): - return local_registry_paths - message = ( - "No local registry paths available for run_python execution. " - "Check TRACECAT__BUILTIN_REGISTRY_SOURCE_PATH, " - "TRACECAT__LOCAL_REPOSITORY_CONTAINER_PATH, and PYTHONUSERBASE." - ) - else: - message = ( - "No registry artifacts available for run_python execution. " - "Check that the registry is synced and the registry_lock is valid." - ) - - return ExecutorActionErrorInfo( - action_name=PlatformAction.RUN_PYTHON, - type="RegistryError", - message=message, - filename="base.py", - function="_execute_run_python", + """Resolve run_python import roots when no registry artifact is available.""" + if config.TRACECAT__LOCAL_REPOSITORY_ENABLED: + if local_registry_paths := self._get_run_python_local_registry_paths(): + return local_registry_paths + message = ( + "No local registry paths available for run_python execution. " + "Check TRACECAT__BUILTIN_REGISTRY_SOURCE_PATH, " + "TRACECAT__LOCAL_REPOSITORY_CONTAINER_PATH, and PYTHONUSERBASE." ) - return await get_action_runner().resolve_registry_paths(artifact_uris) + else: + message = ( + "No registry artifacts available for run_python execution. " + "Check that the registry is synced and the registry_lock is valid." + ) + + return ExecutorActionErrorInfo( + action_name=PlatformAction.RUN_PYTHON, + type="RegistryError", + message=message, + filename="base.py", + function="_execute_run_python", + ) def _get_run_python_local_registry_paths(self) -> list[Path]: """Return local registry import roots for run_python local-repository mode.""" diff --git a/tracecat/executor/backends/test.py b/tracecat/executor/backends/test.py index abfbee50b0..1f7cae6952 100644 --- a/tracecat/executor/backends/test.py +++ b/tracecat/executor/backends/test.py @@ -20,7 +20,7 @@ import asyncio import sys import threading -from contextlib import contextmanager +from contextlib import AsyncExitStack, contextmanager from typing import TYPE_CHECKING, Any from tracecat_registry import secrets as registry_secrets @@ -121,20 +121,24 @@ async def _execute( task_ref=input.task.ref, ) - artifact_paths = await self._ensure_registry_artifacts(input, role) - if artifact_paths: - logger.debug( - "Adding artifact paths to sys.path for test execution", - paths=artifact_paths, - ) - try: - with _temporary_sys_path(artifact_paths): - result = await asyncio.wait_for( - self._execute_with_context(input, role, resolved_context), - timeout=timeout, + # The leases are held for the whole in-process execution so cache + # eviction cannot delete a directory sys.path still points at. + async with AsyncExitStack() as leases: + artifact_paths = await self._lease_registry_artifacts( + leases, input, role ) - return ExecutorResultSuccess(result=result) + if artifact_paths: + logger.debug( + "Adding artifact paths to sys.path for test execution", + paths=artifact_paths, + ) + with _temporary_sys_path(artifact_paths): + result = await asyncio.wait_for( + self._execute_with_context(input, role, resolved_context), + timeout=timeout, + ) + return ExecutorResultSuccess(result=result) except TimeoutError: logger.error( "Test backend execution timed out", @@ -247,10 +251,26 @@ def _load_udf_callable(self, action_impl: ActionImplementation): ) return load_udf_impl(udf_impl) - async def _ensure_registry_artifacts( - self, input: RunActionInput, role: Role + async def _lease_registry_artifacts( + self, + leases: AsyncExitStack, + input: RunActionInput, + role: Role, ) -> list[str]: - """Materialize registry artifacts, returning paths for sys.path.""" + """Lease registry artifacts, returning paths for sys.path. + + Each artifact is leased independently so one unavailable artifact only + drops its own paths: the remaining artifacts still load, matching the + best-effort behaviour tests rely on. + + Args: + leases: Exit stack that owns the leases for the caller's execution. + input: Action input used to resolve artifact URIs. + role: Role used to resolve artifact URIs. + + Returns: + Importable paths for every artifact that could be materialized. + """ if config.TRACECAT__LOCAL_REPOSITORY_ENABLED: return [] @@ -259,14 +279,13 @@ async def _ensure_registry_artifacts( logger.debug("No artifact URIs found, using empty paths") return [] - runner = get_action_runner() + registry_artifacts = get_action_runner().registry_artifacts extracted_paths: list[str] = [] for artifact_uri in artifact_uris: try: - extracted_paths.extend( - str(p) - for p in await runner.ensure_registry_environment(artifact_uri) + artifact_paths = await leases.enter_async_context( + registry_artifacts.lease([artifact_uri]) ) except Exception as e: logger.warning( @@ -274,6 +293,8 @@ async def _ensure_registry_artifacts( artifact_uri=artifact_uri, error=str(e), ) + continue + extracted_paths.extend(str(path) for path in artifact_paths) logger.debug( "Materialized registry artifacts for test execution", diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 08d6110bf2..a75abc075f 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -61,6 +61,15 @@ class RegistryArtifactFormat(StrEnum): """Evicts one idle mounted artifact, excluding the given cache key.""" +class SquashfsMountCommandError(RuntimeError): + """The ``mount`` command itself failed for a SquashFS registry artifact. + + Only this error drives SquashFS mount policy. Download, mkdir, and other + preparation failures must not be mistaken for a missing mount capability or + for loop-device exhaustion. + """ + + @dataclass(frozen=True, slots=True) class RegistryArtifactPaths: """Executor-local cache paths for one registry artifact key.""" @@ -234,22 +243,29 @@ async def _try_mount( ) -> Path | None: """Mount the image, retrying once after reclaiming a loop device. - The first mount failure in a process is treated as a capability probe and - disables mounting process-wide. Once any mount has succeeded, later - failures are attributed to exhausted loop devices instead: one idle + The first mount-command failure in a process is treated as a capability + probe and disables mounting process-wide. Once any mount has succeeded, + later failures are attributed to exhausted loop devices instead: one idle mounted artifact is evicted and the mount is retried once, so a single failure never downgrades the whole process to extraction. + Only ``SquashfsMountCommandError`` drives this policy. Download and + preparation errors propagate to the caller so a transient S3 failure + never disables mounting or evicts an unrelated idle mount. + Args: ctx: Materialization context for the artifact being mounted. image_path: Local path of the SquashFS image. Returns: The mount directory, or None if the caller should extract instead. + + Raises: + Exception: Any non-mount failure raised while preparing the image. """ try: return await self.mount(ctx, image_path) - except Exception as e: + except SquashfsMountCommandError as e: mount_error = e if not ctx.has_mounted_squashfs(): @@ -281,7 +297,7 @@ async def _try_mount( try: return await self.mount(ctx, image_path) - except Exception as e: + except SquashfsMountCommandError as e: logger.warning( "SquashFS mount retry failed, trying extraction", cache_key=ctx.cache_key, @@ -318,6 +334,19 @@ async def mount( ctx: RegistryArtifactMaterializationContext, image_path: Path, ) -> Path: + """Download the image if needed and mount it read-only. + + Args: + ctx: Materialization context for the artifact being mounted. + image_path: Local path of the SquashFS image. + + Returns: + The mount directory. + + Raises: + SquashfsMountCommandError: The ``mount`` command failed. + Exception: The image could not be downloaded or prepared. + """ target_dir = ctx.paths.squashfs_mount_dir if target_dir.is_mount(): ctx.record_squashfs_mount() @@ -408,7 +437,15 @@ async def extract( return target_dir async def _mount_image(self, image_path: Path, target_dir: Path) -> None: - """Mount a SquashFS image read-only at target_dir.""" + """Mount a SquashFS image read-only at target_dir. + + Args: + image_path: Local path of the SquashFS image. + target_dir: Existing directory to mount the image at. + + Raises: + SquashfsMountCommandError: The ``mount`` command failed. + """ if target_dir.is_mount(): return @@ -429,7 +466,7 @@ async def _mount_image(self, image_path: Path, target_dir: Path) -> None: return output = (stderr or stdout).decode(errors="replace").strip() - raise RuntimeError(output or "mount command failed") + raise SquashfsMountCommandError(output or "mount command failed") async def _extract_image(self, image_path: Path, target_dir: Path) -> None: """Extract a SquashFS image to target_dir using unsquashfs.""" @@ -721,10 +758,17 @@ class RegistryArtifactCache: def __init__(self, cache_dir: Path): self.cache_dir = cache_dir + # Per-key locks live for the process lifetime: eviction and lease + # admission must serialize on the same object for a given key, so a + # lock is never dropped and re-created underneath a waiter. self._locks: dict[str, asyncio.Lock] = {} self._locks_lock = asyncio.Lock() self._squashfs_mount_state = SquashfsMountState() self._leases: dict[str, RegistryArtifactLease] = {} + # Whether the on-disk cache may exceed its budget. Set when a new entry + # is materialized and cleared once enforcement measures a cache that + # fits, so steady-state cache hits never pay for a disk scan. + self._budget_dirty = True self._sweep_startup_state() @asynccontextmanager @@ -751,9 +795,17 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat registry_paths: list[Path] = [] for artifact_uri in artifact_uris: cache_key = compute_registry_artifact_cache_key(artifact_uri) - if _is_cache_entry_uri(artifact_uri): - self._acquire_lease(cache_key) - leased_keys.append(cache_key) + if not _is_cache_entry_uri(artifact_uri): + registry_paths.extend( + await self.materialize(cache_key, artifact_uri) + ) + continue + + cached_paths = await self._admit_lease(cache_key, artifact_uri) + leased_keys.append(cache_key) + if cached_paths is not None: + registry_paths.extend(cached_paths) + continue registry_paths.extend(await self.materialize(cache_key, artifact_uri)) logger.info( "Using registry artifact environments", @@ -763,17 +815,41 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat finally: for cache_key in leased_keys: self._release_lease(cache_key) + if leased_keys: + await self._converge_cache_budget() + + async def _admit_lease( + self, cache_key: str, artifact_uri: str + ) -> list[Path] | None: + """Pin a cache entry and return its already-materialized paths, if any. + + The refcount increment and the cached-path check run under the same + per-key lock that eviction holds. An in-flight eviction therefore always + completes before a lease is admitted, and once the refcount is raised no + eviction can delete the entry, so a lease can never be handed a path + that is about to disappear. + + Args: + cache_key: Cache key to pin. + artifact_uri: Registry artifact URI backing the cache key. - async def ensure_environment(self, artifact_uri: str | None) -> list[Path]: - """Materialize an optional registry artifact and return PYTHONPATH entries. + Returns: + Importable paths when the entry is already materialized, else None. - The artifact is only pinned for the duration of this call. Callers that - execute against the returned paths should hold a ``lease`` instead. + Raises: + Exception: Any failure while resolving artifact candidates. The + lease is released before the error propagates. """ - if not artifact_uri: - return [] - async with self.lease([artifact_uri]) as registry_paths: - return registry_paths + lock = await self._lock_for(cache_key) + async with lock: + self._acquire_lease(cache_key) + try: + ctx = self._context_for(cache_key) + candidates = await self._artifact_candidates(ctx, artifact_uri) + except Exception: + self._release_lease(cache_key) + raise + return self._first_cached_path(candidates, ctx) async def materialize(self, cache_key: str, artifact_uri: str) -> list[Path]: """Materialize a registry artifact as local importable directories.""" @@ -804,7 +880,13 @@ async def materialize(self, cache_key: str, artifact_uri: str) -> list[Path]: candidate=index + 1, candidates=len(candidates), ) - return await artifact.materialize(ctx) + registry_paths = await artifact.materialize(ctx) + if _is_cache_entry_uri(artifact.uri): + # A new entry landed on disk after the budget was + # measured, so the cache must be re-checked once the + # entry goes idle. + self._budget_dirty = True + return registry_paths except Exception as e: if index == len(candidates) - 1: raise @@ -842,7 +924,11 @@ def _base_pythonpath_dir(self) -> Path: return base_dir def _acquire_lease(self, cache_key: str) -> None: - """Pin a cache entry against eviction and mark it as recently used.""" + """Pin a cache entry against eviction and mark it as recently used. + + Callers must hold the per-key lock so the increment is ordered against + in-flight eviction of the same key. + """ lease = self._leases.setdefault(cache_key, RegistryArtifactLease()) lease.refcount += 1 lease.last_used = time.time() @@ -992,22 +1078,50 @@ def _can_try_squashfs(self) -> bool: """Return whether this process should prefer SquashFS artifacts.""" return config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED - async def _enforce_cache_budget(self, *, protected_key: str) -> None: + async def _converge_cache_budget(self) -> None: + """Bring an idle cache back under budget after a lease is released. + + Materialization enforces the budget before a new entry exists, so the + cache can legitimately sit over budget while that entry is leased. This + runs on release, when the real on-disk size is known and the entry is + evictable. The scan is skipped entirely unless a new entry has landed + since the last successful enforcement. + """ + if not self._budget_dirty: + return + + try: + self._budget_dirty = not await self._enforce_cache_budget() + except OSError as e: + logger.warning( + "Failed to converge registry artifact cache to budget", + cache_dir=str(self.cache_dir), + error=str(e), + ) + + async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bool: """Evict least-recently-used idle entries until the cache fits its budget. Args: protected_key: Cache key about to be materialized. It is counted - against the budget but never evicted. + against the budget but never evicted. None when enforcing + against the entries already on disk. + + Returns: + Whether the cache is within budget once eviction has finished. """ max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES if max_entries <= 0 and max_bytes <= 0: - return + return True entries = await asyncio.to_thread(self._scan_cache_entries) # The protected key is not on disk yet when it is a fresh entry. - pending_entries = 0 if protected_key in entries else 1 + pending_entries = ( + 1 if protected_key is not None and protected_key not in entries else 0 + ) total_bytes = sum(entry.size_bytes for entry in entries.values()) + protected = set() if protected_key is None else {protected_key} skipped: set[str] = set() while (max_entries > 0 and len(entries) + pending_entries > max_entries) or ( @@ -1015,7 +1129,7 @@ async def _enforce_cache_budget(self, *, protected_key: str) -> None: ): candidate = self._least_recently_used( entries.values(), - excluded=skipped | {protected_key}, + excluded=skipped | protected, ) if candidate is None: logger.warning( @@ -1026,7 +1140,7 @@ async def _enforce_cache_budget(self, *, protected_key: str) -> None: total_bytes=total_bytes, max_bytes=max_bytes, ) - return + return False if await self._evict_entry(candidate.cache_key): del entries[candidate.cache_key] @@ -1034,6 +1148,8 @@ async def _enforce_cache_budget(self, *, protected_key: str) -> None: else: skipped.add(candidate.cache_key) + return True + async def _release_mounted_slot(self, protected_key: str) -> bool: """Evict one idle mounted artifact so its loop device can be reused. @@ -1089,6 +1205,10 @@ async def _evict_entry(self, cache_key: str) -> bool: cannot be unmounted: deleting the image file behind a live mount would leave an open-file zombie holding the loop device. + The whole sequence - refcount check, unmount, delete, and dropping the + lease record - happens under the per-key lock that lease admission also + takes, so no lease can be admitted across the unmount await point. + Args: cache_key: Cache key to evict. @@ -1118,19 +1238,12 @@ async def _evict_entry(self, cache_key: str) -> bool: ) return False - _delete_entry_paths(paths) + await asyncio.to_thread(_delete_entry_paths, paths) + self._leases.pop(cache_key, None) logger.info("Evicted registry artifact from cache", cache_key=cache_key) - await self._forget(cache_key, lock) return True - async def _forget(self, cache_key: str, lock: asyncio.Lock) -> None: - """Drop in-memory bookkeeping for an evicted cache key.""" - async with self._locks_lock: - self._leases.pop(cache_key, None) - if self._locks.get(cache_key) is lock and not lock.locked(): - del self._locks[cache_key] - async def _unmount(self, mount_dir: Path) -> bool: """Unmount a SquashFS artifact directory, releasing its loop device.""" umount = shutil.which("umount") @@ -1178,15 +1291,19 @@ def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: """Measure the on-disk footprint and recency of one cache entry. The mount directory is excluded because a mounted view only costs the - image file that backs it. + image file that backs it. The image is measured with a single ``stat`` + so a concurrent eviction deleting it cannot fail the scan. """ paths = self._paths_for(cache_key) size_bytes = 0 image_mtime = 0.0 created_at = 0.0 - if paths.squashfs_image_path.is_file(): + try: image_stat = paths.squashfs_image_path.stat() + except OSError: + pass + else: size_bytes += image_stat.st_size image_mtime = image_stat.st_mtime @@ -1212,8 +1329,12 @@ def _sweep_startup_state(self) -> None: directory is stale. Scratch paths from interrupted materializations are removed, and the cache is trimmed to budget using image mtimes as LRU order. A missing or empty cache directory is a no-op. + + The sweep is deliberately synchronous: it runs once during construction, + before the process serves any action. """ if not self.cache_dir.is_dir(): + self._budget_dirty = False return try: @@ -1257,10 +1378,15 @@ def _remove_stale_mount_dirs(self) -> None: ) def _trim_startup_cache(self) -> None: - """Trim the cache to budget before any artifact is leased.""" + """Trim the cache to budget before any artifact is leased. + + Clears the budget-dirty flag when the cache ends up within budget, so a + healthy cache never rescans until a new entry is materialized. + """ max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES if max_entries <= 0 and max_bytes <= 0: + self._budget_dirty = False return entries = self._scan_cache_entries() @@ -1275,11 +1401,14 @@ def _trim_startup_cache(self) -> None: key=lambda entry: entry.last_used, ) - for entry in candidates: - if (max_entries <= 0 or len(entries) <= max_entries) and ( + def within_budget() -> bool: + return (max_entries <= 0 or len(entries) <= max_entries) and ( max_bytes <= 0 or total_bytes <= max_bytes - ): - return + ) + + for entry in candidates: + if within_budget(): + break _delete_entry_paths(self._paths_for(entry.cache_key)) del entries[entry.cache_key] total_bytes -= entry.size_bytes @@ -1288,3 +1417,5 @@ def _trim_startup_cache(self) -> None: cache_key=entry.cache_key, size_bytes=entry.size_bytes, ) + + self._budget_dirty = not within_budget() From 55680b6797827025f7a6a61c5dbd9993f4fc94fd Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 27 Jul 2026 18:40:36 -0400 Subject: [PATCH 003/161] fix(executor): preserve concurrent dirty flag, sweep same-PID scratch - Convergence now clears the budget-dirty flag before its awaited scan and loops, so a materialization that lands mid-scan re-arms the flag instead of being overwritten; over-budget or failed scans restore the flag and break to avoid spinning. - The startup sweep removes all materialization scratch paths: it runs before this process can have in-flight work, and a restarted container commonly reuses PID 1, so the same-PID guard leaked invisible partial downloads and extractions forever. ENG-1568 --- tests/unit/test_registry_artifacts.py | 53 ++++++++++++++++++++++++- tracecat/executor/registry_artifacts.py | 45 ++++++++++++++------- 2 files changed, 82 insertions(+), 16 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 253b724e2c..617fd13269 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -992,6 +992,57 @@ async def test_release_keeps_retrying_while_the_cache_stays_over_budget( assert cache._budget_dirty is True + @pytest.mark.anyio + async def test_convergence_rescans_after_concurrent_materialization( + self, temp_cache_dir + ): + """A materialization during a budget scan must schedule another scan.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/concurrent.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + scan_started = asyncio.Event() + materialized = asyncio.Event() + convergence_scans = 0 + + async def mock_enforce_cache_budget( + *, protected_key: str | None = None + ) -> bool: + nonlocal convergence_scans + if protected_key is not None: + return True + + convergence_scans += 1 + if convergence_scans == 1: + scan_started.set() + await materialized.wait() + return True + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 1") + + cache._budget_dirty = True + with ( + patch.object( + cache, + "_enforce_cache_budget", + side_effect=mock_enforce_cache_budget, + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + convergence = asyncio.create_task(cache._converge_cache_budget()) + await scan_started.wait() + registry_paths = await cache.materialize(cache_key, artifact_uri) + materialized.set() + await convergence + + assert registry_paths == [temp_cache_dir / f"tarball-{cache_key}"] + assert convergence_scans == 2 + assert cache._budget_dirty is False + @pytest.mark.anyio async def test_enforce_budget_evicts_least_recently_used_until_under_max_bytes( self, temp_cache_dir @@ -1182,7 +1233,7 @@ def test_sweep_removes_orphaned_scratch_and_stale_mount_dirs(self, temp_cache_di assert not orphaned.exists() assert not orphaned_dir.exists() - assert own.exists() + assert not own.exists() assert not stale_mount_dir.exists() assert entry_dir.is_dir() diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index a75abc075f..c9f8659ab9 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -53,7 +53,7 @@ class RegistryArtifactFormat(StrEnum): """On-disk name prefixes owned by a registry artifact cache entry.""" TEMP_ARTIFACT_PATTERN = re.compile( - r"^[^.]+\.(?P\d+)\.\d+\.(?:squashfs|unsquashfs|tar\.gz|tmp)$" + r"^[^.]+\.\d+\.\d+\.(?:squashfs|unsquashfs|tar\.gz|tmp)$" ) """Matches in-flight materialization scratch names produced by ``_temp_path``.""" @@ -1086,18 +1086,29 @@ async def _converge_cache_budget(self) -> None: runs on release, when the real on-disk size is known and the entry is evictable. The scan is skipped entirely unless a new entry has landed since the last successful enforcement. + + Each successful pass consumes the dirty signal before its awaited scan. + A follow-up pass therefore occurs only when a concurrent materialization + sets the flag again. Without new materializations the loop terminates, + while an over-budget or failed scan restores the flag and breaks so it + cannot spin while entries remain leased. """ - if not self._budget_dirty: - return + while self._budget_dirty: + self._budget_dirty = False + try: + within_budget = await self._enforce_cache_budget() + except OSError as e: + logger.warning( + "Failed to converge registry artifact cache to budget", + cache_dir=str(self.cache_dir), + error=str(e), + ) + self._budget_dirty = True + break - try: - self._budget_dirty = not await self._enforce_cache_budget() - except OSError as e: - logger.warning( - "Failed to converge registry artifact cache to budget", - cache_dir=str(self.cache_dir), - error=str(e), - ) + if not within_budget: + self._budget_dirty = True + break async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bool: """Evict least-recently-used idle entries until the cache fits its budget. @@ -1349,11 +1360,15 @@ def _sweep_startup_state(self) -> None: ) def _remove_orphaned_temp_paths(self) -> None: - """Delete materialization scratch paths owned by dead processes.""" - current_pid = os.getpid() + """Delete every materialization scratch path during startup. + + The sweep runs during cache construction, before this process can start + a materialization in the cache directory. Every matching path is + therefore interrupted scratch from an earlier process and is safe to + remove even when the operating system reused that process's PID. + """ for name in os.listdir(self.cache_dir): - match = TEMP_ARTIFACT_PATTERN.match(name) - if match is None or int(match.group("pid")) == current_pid: + if TEMP_ARTIFACT_PATTERN.match(name) is None: continue path = self.cache_dir / name if path.is_dir(): From 4c6153172c48eee7a60d00cf20715e168a4a6fe6 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:36:05 -0400 Subject: [PATCH 004/161] fix(executor): cancellation-safe eviction, serialize budget passes - Eviction now synchronously renames entry paths to startup-sweepable scratch names under the per-key lock, then deletes them in a worker thread with no lock held. Cancelling the deletion await can no longer release the lock while the thread is mid-delete; the worst case is doomed scratch that the next startup sweep removes. - Cache-key discovery explicitly rejects scratch-pattern names so doomed paths are invisible to budget accounting. - A single budget lock serializes the scan/select/evict pass so concurrent cache misses and lease releases cannot double-evict from the same over-budget snapshot. Lock order stays budget -> per-key; the mount slot-release path keeps out of the budget lock since it runs under a per-key lock. ENG-1568 --- tests/unit/test_registry_artifacts.py | 170 ++++++++++++++++++++++++ tracecat/executor/registry_artifacts.py | 141 ++++++++++++++------ 2 files changed, 272 insertions(+), 39 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 617fd13269..772cc3acf1 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -6,6 +6,7 @@ import os import tarfile import tempfile +import threading from pathlib import Path from unittest.mock import AsyncMock, patch @@ -15,11 +16,14 @@ from tracecat.executor.registry_artifacts import ( SQUASHFS_MOUNT_OPTIONS, + TEMP_ARTIFACT_PATTERN, RegistryArtifactCache, RegistryArtifactFormat, + RegistryArtifactPaths, SquashfsArtifact, SquashfsMountCommandError, TarballArtifact, + _delete_entry_paths, bundled_builtin_registry_uri, compute_registry_artifact_cache_key, ) @@ -1060,6 +1064,81 @@ async def test_enforce_budget_evicts_least_recently_used_until_under_max_bytes( assert older.exists() assert newest.exists() + @pytest.mark.anyio + async def test_concurrent_budget_passes_only_evict_once(self, temp_cache_dir): + """A waiting budget pass must re-scan after the active pass evicts.""" + cache = RegistryArtifactCache(temp_cache_dir) + oldest = _write_image_entry(temp_cache_dir, "oldest", size=16, mtime=100.0) + retained = _write_image_entry(temp_cache_dir, "retained", size=16, mtime=200.0) + original_scan = cache._scan_cache_entries + first_scan_started = threading.Event() + second_scan_started = threading.Event() + release_first_scan = threading.Event() + release_second_scan = threading.Event() + scan_count_lock = threading.Lock() + scan_count = 0 + + def controlled_scan(): + nonlocal scan_count + entries = original_scan() + with scan_count_lock: + scan_index = scan_count + scan_count += 1 + if scan_index == 0: + first_scan_started.set() + release_first_scan.wait(timeout=5) + elif scan_index == 1: + second_scan_started.set() + release_second_scan.wait(timeout=5) + return entries + + eviction_started = asyncio.Event() + finish_eviction = asyncio.Event() + extra_eviction_finished = asyncio.Event() + evicted_keys: list[str] = [] + + async def controlled_evict(cache_key: str) -> bool: + if cache_key == "oldest": + if eviction_started.is_set(): + return False + eviction_started.set() + await finish_eviction.wait() + oldest.unlink(missing_ok=True) + else: + retained.unlink(missing_ok=True) + extra_eviction_finished.set() + evicted_keys.append(cache_key) + return True + + with ( + patch.object(cache, "_scan_cache_entries", side_effect=controlled_scan), + patch.object(cache, "_evict_entry", side_effect=controlled_evict), + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + first_pass = asyncio.create_task(cache._enforce_cache_budget()) + assert await asyncio.to_thread(first_scan_started.wait, 1) + second_pass = asyncio.create_task(cache._enforce_cache_budget()) + second_scan_overlapped = await asyncio.to_thread( + second_scan_started.wait, 0.5 + ) + + release_first_scan.set() + await asyncio.wait_for(eviction_started.wait(), timeout=1) + release_second_scan.set() + try: + await asyncio.wait_for(extra_eviction_finished.wait(), timeout=0.2) + except TimeoutError: + pass + finish_eviction.set() + await asyncio.gather(first_pass, second_pass) + + assert second_scan_overlapped is False + assert scan_count == 2 + assert evicted_keys == ["oldest"] + assert not oldest.exists() + assert retained.exists() + @pytest.mark.anyio async def test_enforce_budget_counts_the_pending_entry(self, temp_cache_dir): """The entry about to be materialized counts against the entry budget.""" @@ -1142,6 +1221,97 @@ async def mock_umount(*args, **kwargs): str(paths.squashfs_mount_dir), ) + @pytest.mark.anyio + async def test_cancelled_background_deletion_leaves_a_clean_miss( + self, temp_cache_dir + ): + """Cancellation cannot expose live paths that deletion still owns.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/cancelled-eviction.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + original_target = _write_tarball_entry(temp_cache_dir, cache_key) + delete_started = threading.Event() + finish_delete = threading.Event() + delete_finished = threading.Event() + doomed: list[RegistryArtifactPaths] = [] + + def blocked_delete(paths: RegistryArtifactPaths) -> None: + doomed.append(paths) + delete_started.set() + finish_delete.wait(timeout=5) + _delete_entry_paths(paths) + delete_finished.set() + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch( + "tracecat.executor.registry_artifacts._delete_entry_paths", + side_effect=blocked_delete, + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + eviction = asyncio.create_task(cache._evict_entry(cache_key)) + assert await asyncio.to_thread(delete_started.wait, 1) + eviction.cancel() + with pytest.raises(asyncio.CancelledError): + await eviction + + try: + assert not original_target.exists() + assert doomed[0].tarball_target_dir.is_dir() + assert cache._discover_cache_keys() == set() + + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [original_target] + assert original_target.is_dir() + assert doomed[0].tarball_target_dir.is_dir() + assert (original_target / "module.py").read_text() == "VALUE = 2" + finally: + finish_delete.set() + + assert await asyncio.to_thread(delete_finished.wait, 1) + assert original_target.is_dir() + assert not doomed[0].tarball_target_dir.exists() + + @pytest.mark.anyio + async def test_doomed_eviction_names_are_startup_scratch(self, temp_cache_dir): + """Every renamed entry root is invisible and reclaimed on startup.""" + cache = RegistryArtifactCache(temp_cache_dir) + # This key makes doomed names look like live image paths unless cache + # discovery rejects the shared scratch pattern first. + cache_key = "squashfs-doomed" + paths = cache._paths_for(cache_key) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + paths.squashfs_extract_dir.mkdir() + paths.tarball_target_dir.mkdir() + + with patch( + "tracecat.executor.registry_artifacts._delete_entry_paths" + ) as delete_entry_paths: + assert await cache._evict_entry(cache_key) is True + + doomed_paths = delete_entry_paths.call_args.args[0] + renamed = ( + doomed_paths.squashfs_image_path, + doomed_paths.squashfs_mount_dir, + doomed_paths.squashfs_extract_dir, + doomed_paths.tarball_target_dir, + ) + assert all(path.exists() for path in renamed) + assert all(TEMP_ARTIFACT_PATTERN.fullmatch(path.name) for path in renamed) + assert cache._discover_cache_keys() == set() + + cache._sweep_startup_state() + + assert not any(path.exists() for path in renamed) + @pytest.mark.anyio async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): """A failed unmount skips the key instead of forcing a lazy detach.""" diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index c9f8659ab9..5a3fc97d4e 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -55,7 +55,7 @@ class RegistryArtifactFormat(StrEnum): TEMP_ARTIFACT_PATTERN = re.compile( r"^[^.]+\.\d+\.\d+\.(?:squashfs|unsquashfs|tar\.gz|tmp)$" ) -"""Matches in-flight materialization scratch names produced by ``_temp_path``.""" +"""Matches materialization scratch and doomed eviction paths.""" type MountSlotReleaser = Callable[[str], Awaitable[bool]] """Evicts one idle mounted artifact, excluding the given cache key.""" @@ -702,6 +702,8 @@ def _is_cache_entry_uri(artifact_uri: str) -> bool: def _cache_key_from_entry_name(name: str) -> str | None: """Return the cache key owning a cache directory entry name, if any.""" + if TEMP_ARTIFACT_PATTERN.fullmatch(name) is not None: + return None for prefix in CACHE_ENTRY_PREFIXES: if name.startswith(prefix): cache_key = name.removeprefix(prefix).removesuffix(".squashfs") @@ -753,6 +755,48 @@ def _delete_entry_paths(paths: RegistryArtifactPaths) -> None: shutil.rmtree(paths.squashfs_mount_dir, ignore_errors=True) +def _rename_entry_paths( + paths: RegistryArtifactPaths, + *, + cache_dir: Path, + cache_key: str, +) -> RegistryArtifactPaths: + """Synchronously rename live entry paths to unique startup-sweep scratch.""" + unique_id = time.time_ns() + while True: + doomed_paths = RegistryArtifactPaths( + squashfs_image_path=cache_dir + / f"{cache_key}.{os.getpid()}.{unique_id}.squashfs", + squashfs_mount_dir=cache_dir / f"{cache_key}.{os.getpid()}.{unique_id}.tmp", + squashfs_extract_dir=cache_dir + / f"{cache_key}.{os.getpid()}.{unique_id}.unsquashfs", + tarball_target_dir=cache_dir + / f"{cache_key}.{os.getpid()}.{unique_id}.tar.gz", + ) + if not any( + path.exists() + for path in ( + doomed_paths.squashfs_image_path, + doomed_paths.squashfs_mount_dir, + doomed_paths.squashfs_extract_dir, + doomed_paths.tarball_target_dir, + ) + ): + break + unique_id += 1 + + for source, target in ( + (paths.squashfs_extract_dir, doomed_paths.squashfs_extract_dir), + (paths.tarball_target_dir, doomed_paths.tarball_target_dir), + (paths.squashfs_image_path, doomed_paths.squashfs_image_path), + (paths.squashfs_mount_dir, doomed_paths.squashfs_mount_dir), + ): + if source.exists(): + source.rename(target) + + return doomed_paths + + class RegistryArtifactCache: """Materializes registry artifacts into executor-local Python paths.""" @@ -763,6 +807,7 @@ def __init__(self, cache_dir: Path): # lock is never dropped and re-created underneath a waiter. self._locks: dict[str, asyncio.Lock] = {} self._locks_lock = asyncio.Lock() + self._budget_lock = asyncio.Lock() self._squashfs_mount_state = SquashfsMountState() self._leases: dict[str, RegistryArtifactLease] = {} # Whether the on-disk cache may exceed its budget. Set when a new entry @@ -1113,6 +1158,10 @@ async def _converge_cache_budget(self) -> None: async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bool: """Evict least-recently-used idle entries until the cache fits its budget. + The budget lock serializes the complete scan/select/evict pass. It is + always acquired before any candidate's per-key lock, and callers must + invoke enforcement without holding a per-key lock. + Args: protected_key: Cache key about to be materialized. It is counted against the budget but never evicted. None when enforcing @@ -1121,49 +1170,55 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo Returns: Whether the cache is within budget once eviction has finished. """ - max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES - max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES - if max_entries <= 0 and max_bytes <= 0: - return True - - entries = await asyncio.to_thread(self._scan_cache_entries) - # The protected key is not on disk yet when it is a fresh entry. - pending_entries = ( - 1 if protected_key is not None and protected_key not in entries else 0 - ) - total_bytes = sum(entry.size_bytes for entry in entries.values()) - protected = set() if protected_key is None else {protected_key} - skipped: set[str] = set() + async with self._budget_lock: + max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES + max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES + if max_entries <= 0 and max_bytes <= 0: + return True - while (max_entries > 0 and len(entries) + pending_entries > max_entries) or ( - max_bytes > 0 and total_bytes > max_bytes - ): - candidate = self._least_recently_used( - entries.values(), - excluded=skipped | protected, + entries = await asyncio.to_thread(self._scan_cache_entries) + # The protected key is not on disk yet when it is a fresh entry. + pending_entries = ( + 1 if protected_key is not None and protected_key not in entries else 0 ) - if candidate is None: - logger.warning( - "Registry artifact cache is over budget but every entry is in use", - cache_dir=str(self.cache_dir), - entries=len(entries) + pending_entries, - max_entries=max_entries, - total_bytes=total_bytes, - max_bytes=max_bytes, + total_bytes = sum(entry.size_bytes for entry in entries.values()) + protected = set() if protected_key is None else {protected_key} + skipped: set[str] = set() + + while ( + max_entries > 0 and len(entries) + pending_entries > max_entries + ) or (max_bytes > 0 and total_bytes > max_bytes): + candidate = self._least_recently_used( + entries.values(), + excluded=skipped | protected, ) - return False + if candidate is None: + logger.warning( + "Registry artifact cache is over budget but every entry is in use", + cache_dir=str(self.cache_dir), + entries=len(entries) + pending_entries, + max_entries=max_entries, + total_bytes=total_bytes, + max_bytes=max_bytes, + ) + return False - if await self._evict_entry(candidate.cache_key): - del entries[candidate.cache_key] - total_bytes -= candidate.size_bytes - else: - skipped.add(candidate.cache_key) + if await self._evict_entry(candidate.cache_key): + del entries[candidate.cache_key] + total_bytes -= candidate.size_bytes + else: + skipped.add(candidate.cache_key) - return True + return True async def _release_mounted_slot(self, protected_key: str) -> bool: """Evict one idle mounted artifact so its loop device can be reused. + This path deliberately does not take the budget lock: ``_try_mount`` may + call it while holding ``protected_key``'s per-key lock. It excludes that + key and only tries candidate per-key locks, so it cannot invert the + budget-lock-to-per-key-lock ordering used by budget enforcement. + Args: protected_key: Cache key that must not be evicted. @@ -1216,9 +1271,12 @@ async def _evict_entry(self, cache_key: str) -> bool: cannot be unmounted: deleting the image file behind a live mount would leave an open-file zombie holding the loop device. - The whole sequence - refcount check, unmount, delete, and dropping the - lease record - happens under the per-key lock that lease admission also - takes, so no lease can be admitted across the unmount await point. + After unmounting, live paths are synchronously renamed under the per-key + lock to scratch names ignored by cache discovery and budget accounting. + The lease record is then dropped and the lock released before physical + deletion runs in a worker thread. If that await is cancelled, the live + key remains a clean cache miss and the next startup sweep removes any + doomed scratch left behind. Args: cache_key: Cache key to evict. @@ -1249,10 +1307,15 @@ async def _evict_entry(self, cache_key: str) -> bool: ) return False - await asyncio.to_thread(_delete_entry_paths, paths) + doomed_paths = _rename_entry_paths( + paths, + cache_dir=self.cache_dir, + cache_key=cache_key, + ) self._leases.pop(cache_key, None) logger.info("Evicted registry artifact from cache", cache_key=cache_key) + await asyncio.to_thread(_delete_entry_paths, doomed_paths) return True async def _unmount(self, mount_dir: Path) -> bool: From 6c6312d4eb2b9f96fe0727c8d5f433da92bfdc06 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:28:49 -0400 Subject: [PATCH 005/161] fix(executor): shield pool-worker tarball paths from runtime eviction Warm pool workers (TRACECAT__EXECUTOR_BACKEND=pool) are long-lived subprocesses that add tarball-* cache directories to their own sys.path and import from them lazily, invisibly to in-process lease refcounts. Runtime eviction and mount-slot reclamation now skip entries with a tarball directory when the pool backend is configured, preserving the pre-eviction status quo for pool deployments. The startup sweep still trims tarball entries: it runs before this process spawns any workers, and dead processes' workers died with them. ENG-1568 --- tests/unit/test_registry_artifacts.py | 132 ++++++++++++++++++++++-- tracecat/executor/registry_artifacts.py | 17 ++- 2 files changed, 142 insertions(+), 7 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 772cc3acf1..b8e75ea4cb 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -27,6 +27,7 @@ bundled_builtin_registry_uri, compute_registry_artifact_cache_key, ) +from tracecat.executor.schemas import ExecutorBackendType from tracecat.registry.artifact_keys import parse_s3_uri MAX_ENTRIES_CONFIG = ( @@ -41,6 +42,9 @@ "tracecat.executor.registry_artifacts.config" ".TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED" ) +BACKEND_CONFIG = ( + "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_BACKEND" +) def _write_tarball_entry(cache_dir: Path, cache_key: str) -> Path: @@ -1047,23 +1051,133 @@ async def mock_extract(self, tarball_path, target_dir): assert convergence_scans == 2 assert cache._budget_dirty is False + @pytest.mark.parametrize( + "oldest_has_tarball", + [False, True], + ids=["squashfs-only", "tarball-bearing"], + ) @pytest.mark.anyio async def test_enforce_budget_evicts_least_recently_used_until_under_max_bytes( - self, temp_cache_dir + self, temp_cache_dir, oldest_has_tarball: bool ): - """Size eviction stops as soon as the cache is back within budget.""" + """Direct-backend size eviction stops once the cache is within budget.""" cache = RegistryArtifactCache(temp_cache_dir) oldest = _write_image_entry(temp_cache_dir, "oldest", size=4096, mtime=100.0) + if oldest_has_tarball: + _write_tarball_entry(temp_cache_dir, "oldest") older = _write_image_entry(temp_cache_dir, "older", size=4096, mtime=200.0) newest = _write_image_entry(temp_cache_dir, "newest", size=4096, mtime=300.0) - with patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 9000): - await cache._enforce_cache_budget(protected_key="pending") + with ( + patch(BACKEND_CONFIG, ExecutorBackendType.DIRECT.value), + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 9000), + ): + within_budget = await cache._enforce_cache_budget(protected_key="pending") + assert within_budget is True assert not oldest.exists() + assert not cache._paths_for("oldest").tarball_target_dir.exists() assert older.exists() assert newest.exists() + @pytest.mark.anyio + async def test_pool_backend_skips_tarball_lru_and_evicts_next_entry( + self, temp_cache_dir + ): + """Warm-worker tarball paths are ineligible runtime victims.""" + cache = RegistryArtifactCache(temp_cache_dir) + pool_visible_image = _write_image_entry( + temp_cache_dir, "pool-visible", size=16, mtime=100.0 + ) + pool_visible_tarball = _write_tarball_entry(temp_cache_dir, "pool-visible") + next_lru = _write_image_entry(temp_cache_dir, "next-lru", size=16, mtime=200.0) + newest = _write_image_entry(temp_cache_dir, "newest", size=16, mtime=300.0) + + with ( + patch(BACKEND_CONFIG, ExecutorBackendType.POOL.value), + patch(MAX_ENTRIES_CONFIG, 2), + patch(MAX_BYTES_CONFIG, 0), + ): + within_budget = await cache._enforce_cache_budget() + + assert within_budget is True + assert pool_visible_image.exists() + assert pool_visible_tarball.is_dir() + assert not next_lru.exists() + assert newest.exists() + + @pytest.mark.anyio + async def test_pool_backend_all_tarballs_remain_dirty_when_over_budget( + self, temp_cache_dir + ): + """An all-tarball pool cache warns and retries convergence later.""" + cache = RegistryArtifactCache(temp_cache_dir) + first_image = _write_image_entry(temp_cache_dir, "first", size=16, mtime=100.0) + first_tarball = _write_tarball_entry(temp_cache_dir, "first") + second_image = _write_image_entry( + temp_cache_dir, "second", size=16, mtime=200.0 + ) + second_tarball = _write_tarball_entry(temp_cache_dir, "second") + cache._budget_dirty = True + + with ( + patch(BACKEND_CONFIG, ExecutorBackendType.POOL.value), + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch("tracecat.executor.registry_artifacts.logger.warning") as warning, + ): + await cache._converge_cache_budget() + + assert first_image.exists() + assert first_tarball.is_dir() + assert second_image.exists() + assert second_tarball.is_dir() + assert cache._budget_dirty is True + warning.assert_called_once_with( + "Registry artifact cache is over budget but every entry is in use", + cache_dir=str(temp_cache_dir), + entries=2, + max_entries=1, + total_bytes=50, + max_bytes=0, + ) + + @pytest.mark.anyio + async def test_pool_backend_release_mounted_slot_skips_tarball_entry( + self, temp_cache_dir + ): + """Loop-device recovery must preserve paths visible to warm workers.""" + cache = RegistryArtifactCache(temp_cache_dir) + pool_visible = cache._paths_for("pool-visible") + pool_visible.squashfs_image_path.write_bytes(b"squashfs") + os.utime(pool_visible.squashfs_image_path, (100.0, 100.0)) + pool_visible.squashfs_mount_dir.mkdir() + _write_tarball_entry(temp_cache_dir, "pool-visible") + eligible = cache._paths_for("eligible") + eligible.squashfs_image_path.write_bytes(b"squashfs") + os.utime(eligible.squashfs_image_path, (200.0, 200.0)) + eligible.squashfs_mount_dir.mkdir() + mounted = { + pool_visible.squashfs_mount_dir, + eligible.squashfs_mount_dir, + } + + with ( + patch(BACKEND_CONFIG, ExecutorBackendType.POOL.value), + patch.object(Path, "is_mount", lambda self: self in mounted), + patch.object( + cache, + "_evict_entry", + new_callable=AsyncMock, + return_value=True, + ) as evict_entry, + ): + released = await cache._release_mounted_slot("protected") + + assert released is True + evict_entry.assert_awaited_once_with("eligible") + @pytest.mark.anyio async def test_concurrent_budget_passes_only_evict_once(self, temp_cache_dir): """A waiting budget pass must re-scan after the active pass evicts.""" @@ -1418,15 +1532,21 @@ def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): assert mount_dir.is_dir() def test_sweep_trims_to_budget_using_image_mtimes(self, temp_cache_dir): - """Startup LRU order comes from image mtimes, which survive a restart.""" + """Pool startup still trims tarballs using persistent image mtimes.""" oldest = _write_image_entry(temp_cache_dir, "oldest", size=64, mtime=100.0) + oldest_tarball = _write_tarball_entry(temp_cache_dir, "oldest") older = _write_image_entry(temp_cache_dir, "older", size=64, mtime=200.0) newest = _write_image_entry(temp_cache_dir, "newest", size=64, mtime=300.0) - with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): + with ( + patch(BACKEND_CONFIG, ExecutorBackendType.POOL.value), + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): RegistryArtifactCache(temp_cache_dir) assert not oldest.exists() + assert not oldest_tarball.exists() assert not older.exists() assert newest.exists() diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 5a3fc97d4e..e791d3eef4 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -21,6 +21,7 @@ import tracecat_registry from tracecat import config +from tracecat.executor.schemas import ExecutorBackendType from tracecat.logger import logger from tracecat.registry.artifact_keys import parse_s3_uri from tracecat.registry.constants import DEFAULT_REGISTRY_ORIGIN @@ -1241,6 +1242,18 @@ async def _release_mounted_slot(self, protected_key: str) -> bool: skipped.add(candidate.cache_key) return False + def _is_pool_worker_visible(self, cache_key: str) -> bool: + """Return whether a warm pool worker may retain this tarball path. + + This runtime guard intentionally does not apply to the startup sweep, + which runs before this process spawns pool workers and may trim tarball + entries left behind by dead processes. + """ + return ( + config.TRACECAT__EXECUTOR_BACKEND == ExecutorBackendType.POOL + and self._paths_for(cache_key).tarball_target_dir.exists() + ) + def _least_recently_used( self, entries: Iterable[RegistryArtifactCacheEntry], @@ -1251,7 +1264,9 @@ def _least_recently_used( eligible = [ entry for entry in entries - if entry.cache_key not in excluded and self._refcount(entry.cache_key) == 0 + if entry.cache_key not in excluded + and self._refcount(entry.cache_key) == 0 + and not self._is_pool_worker_visible(entry.cache_key) ] if not eligible: return None From 66cce97571eaaf57b79511cb7bf9d8f5ce40f388 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:59:42 -0400 Subject: [PATCH 006/161] fix(executor): plug lease cancellation leak, resolve pool backend - Release the admission lease on BaseException so CancelledError during the sidecar lookup cannot pin an entry forever. - Base the pool-worker tarball guard on resolve_backend_type() so TRACECAT__EXECUTOR_BACKEND=auto resolving to pool is recognized. - Apply the same guard to the startup trim: pool workers can spawn before the cache is lazily constructed and inherit tarball paths, so the sweep may no longer assume it runs before workers exist. Scratch and stale mount dirs are still removed unconditionally. ENG-1568 --- tests/unit/test_registry_artifacts.py | 83 ++++++++++++++++++++++--- tracecat/executor/registry_artifacts.py | 37 +++++++---- 2 files changed, 99 insertions(+), 21 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index b8e75ea4cb..d0e2a2c5c0 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -45,6 +45,7 @@ BACKEND_CONFIG = ( "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_BACKEND" ) +RESOLVE_BACKEND = "tracecat.executor.registry_artifacts.resolve_backend_type" def _write_tarball_entry(cache_dir: Path, cache_key: str) -> Path: @@ -763,6 +764,40 @@ async def mock_download(self, ctx, path): assert cache._refcount(cache_key) == 0 + @pytest.mark.anyio + async def test_cancelled_lease_admission_releases_refcount(self, temp_cache_dir): + """Cancellation during candidate lookup must not leak a permanent pin.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/site-packages.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + target_dir = _write_tarball_entry(temp_cache_dir, cache_key) + lookup_started = asyncio.Event() + finish_lookup = asyncio.Event() + + async def blocked_sidecar_lookup(**kwargs): + lookup_started.set() + await finish_lookup.wait() + return False + + async def take_lease() -> None: + async with cache.lease([artifact_uri]): + pass + + with ( + patch(SQUASHFS_ENABLED_CONFIG, True), + patch.object(cache, "_sidecar_exists", blocked_sidecar_lookup), + ): + acquisition = asyncio.create_task(take_lease()) + await lookup_started.wait() + assert cache._refcount(cache_key) == 1 + acquisition.cancel() + with pytest.raises(asyncio.CancelledError): + await acquisition + + assert cache._refcount(cache_key) == 0 + assert await cache._evict_entry(cache_key) is True + assert not target_dir.exists() + @pytest.mark.anyio async def test_lease_without_uris_returns_base_pythonpath_dir(self, temp_cache_dir): """No artifact URIs still yields the base PYTHONPATH directory.""" @@ -1082,10 +1117,10 @@ async def test_enforce_budget_evicts_least_recently_used_until_under_max_bytes( assert newest.exists() @pytest.mark.anyio - async def test_pool_backend_skips_tarball_lru_and_evicts_next_entry( + async def test_auto_pool_backend_skips_tarball_lru_and_evicts_next_entry( self, temp_cache_dir ): - """Warm-worker tarball paths are ineligible runtime victims.""" + """Auto-resolved pool workers protect tarballs from runtime eviction.""" cache = RegistryArtifactCache(temp_cache_dir) pool_visible_image = _write_image_entry( temp_cache_dir, "pool-visible", size=16, mtime=100.0 @@ -1095,7 +1130,8 @@ async def test_pool_backend_skips_tarball_lru_and_evicts_next_entry( newest = _write_image_entry(temp_cache_dir, "newest", size=16, mtime=300.0) with ( - patch(BACKEND_CONFIG, ExecutorBackendType.POOL.value), + patch(BACKEND_CONFIG, ExecutorBackendType.AUTO.value), + patch(RESOLVE_BACKEND, return_value=ExecutorBackendType.POOL), patch(MAX_ENTRIES_CONFIG, 2), patch(MAX_BYTES_CONFIG, 0), ): @@ -1107,6 +1143,29 @@ async def test_pool_backend_skips_tarball_lru_and_evicts_next_entry( assert not next_lru.exists() assert newest.exists() + @pytest.mark.anyio + async def test_auto_non_pool_backend_evicts_tarball_lru(self, temp_cache_dir): + """Auto-resolved non-pool backends may evict tarball entries normally.""" + cache = RegistryArtifactCache(temp_cache_dir) + oldest_image = _write_image_entry( + temp_cache_dir, "oldest", size=16, mtime=100.0 + ) + oldest_tarball = _write_tarball_entry(temp_cache_dir, "oldest") + newest = _write_image_entry(temp_cache_dir, "newest", size=16, mtime=200.0) + + with ( + patch(BACKEND_CONFIG, ExecutorBackendType.AUTO.value), + patch(RESOLVE_BACKEND, return_value=ExecutorBackendType.DIRECT), + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + within_budget = await cache._enforce_cache_budget() + + assert within_budget is True + assert not oldest_image.exists() + assert not oldest_tarball.exists() + assert newest.exists() + @pytest.mark.anyio async def test_pool_backend_all_tarballs_remain_dirty_when_over_budget( self, temp_cache_dir @@ -1531,24 +1590,28 @@ def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): assert mount_dir.is_dir() - def test_sweep_trims_to_budget_using_image_mtimes(self, temp_cache_dir): - """Pool startup still trims tarballs using persistent image mtimes.""" + def test_auto_pool_sweep_protects_tarballs_and_trims_other_entries( + self, temp_cache_dir + ): + """Startup trimming preserves paths inherited by auto-resolved workers.""" oldest = _write_image_entry(temp_cache_dir, "oldest", size=64, mtime=100.0) oldest_tarball = _write_tarball_entry(temp_cache_dir, "oldest") older = _write_image_entry(temp_cache_dir, "older", size=64, mtime=200.0) newest = _write_image_entry(temp_cache_dir, "newest", size=64, mtime=300.0) with ( - patch(BACKEND_CONFIG, ExecutorBackendType.POOL.value), - patch(MAX_ENTRIES_CONFIG, 1), + patch(BACKEND_CONFIG, ExecutorBackendType.AUTO.value), + patch(RESOLVE_BACKEND, return_value=ExecutorBackendType.POOL), + patch(MAX_ENTRIES_CONFIG, 2), patch(MAX_BYTES_CONFIG, 0), ): - RegistryArtifactCache(temp_cache_dir) + cache = RegistryArtifactCache(temp_cache_dir) - assert not oldest.exists() - assert not oldest_tarball.exists() + assert oldest.exists() + assert oldest_tarball.is_dir() assert not older.exists() assert newest.exists() + assert cache._budget_dirty is False class TestSquashfsMountCapability: diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index e791d3eef4..68355aecf2 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -21,7 +21,7 @@ import tracecat_registry from tracecat import config -from tracecat.executor.schemas import ExecutorBackendType +from tracecat.executor.schemas import ExecutorBackendType, resolve_backend_type from tracecat.logger import logger from tracecat.registry.artifact_keys import parse_s3_uri from tracecat.registry.constants import DEFAULT_REGISTRY_ORIGIN @@ -883,8 +883,8 @@ async def _admit_lease( Importable paths when the entry is already materialized, else None. Raises: - Exception: Any failure while resolving artifact candidates. The - lease is released before the error propagates. + BaseException: Any failure or cancellation while resolving artifact + candidates. The lease is released before it propagates. """ lock = await self._lock_for(cache_key) async with lock: @@ -892,7 +892,8 @@ async def _admit_lease( try: ctx = self._context_for(cache_key) candidates = await self._artifact_candidates(ctx, artifact_uri) - except Exception: + except BaseException: + # Cancellation is a BaseException and must not leak the pin. self._release_lease(cache_key) raise return self._first_cached_path(candidates, ctx) @@ -1242,15 +1243,20 @@ async def _release_mounted_slot(self, protected_key: str) -> bool: skipped.add(candidate.cache_key) return False - def _is_pool_worker_visible(self, cache_key: str) -> bool: + def _is_pool_worker_visible( + self, + cache_key: str, + backend_type: ExecutorBackendType, + ) -> bool: """Return whether a warm pool worker may retain this tarball path. - This runtime guard intentionally does not apply to the startup sweep, - which runs before this process spawns pool workers and may trim tarball - entries left behind by dead processes. + Pool workers may start before the cache is constructed, then inherit + tarball paths that remain invisible to in-process lease refcounts. Both + runtime eviction and startup trimming must therefore protect these paths + whenever the resolved backend is the pool. """ return ( - config.TRACECAT__EXECUTOR_BACKEND == ExecutorBackendType.POOL + backend_type == ExecutorBackendType.POOL and self._paths_for(cache_key).tarball_target_dir.exists() ) @@ -1261,12 +1267,13 @@ def _least_recently_used( excluded: set[str], ) -> RegistryArtifactCacheEntry | None: """Return the least recently used idle entry eligible for eviction.""" + backend_type = resolve_backend_type() eligible = [ entry for entry in entries if entry.cache_key not in excluded and self._refcount(entry.cache_key) == 0 - and not self._is_pool_worker_visible(entry.cache_key) + and not self._is_pool_worker_visible(entry.cache_key, backend_type) ] if not eligible: return None @@ -1417,7 +1424,11 @@ def _sweep_startup_state(self) -> None: Mounts never survive a container restart, so any non-mountpoint mount directory is stale. Scratch paths from interrupted materializations are removed, and the cache is trimmed to budget using image mtimes as LRU - order. A missing or empty cache directory is a no-op. + order. Scratch and stale empty mount directories are never worker import + paths, so they are removed unconditionally. Tarball-bearing entries are + protected for the pool backend because cache construction may happen + after warm workers have already inherited those paths. A missing or + empty cache directory is a no-op. The sweep is deliberately synchronous: it runs once during construction, before the process serves any action. @@ -1473,6 +1484,8 @@ def _remove_stale_mount_dirs(self) -> None: def _trim_startup_cache(self) -> None: """Trim the cache to budget before any artifact is leased. + Tarball-bearing entries are ineligible when the resolved backend is the + pool because existing workers may already import from those paths. Clears the budget-dirty flag when the cache ends up within budget, so a healthy cache never rescans until a new entry is materialized. """ @@ -1484,12 +1497,14 @@ def _trim_startup_cache(self) -> None: entries = self._scan_cache_entries() total_bytes = sum(entry.size_bytes for entry in entries.values()) + backend_type = resolve_backend_type() # Mounted entries belong to a live process sharing this cache directory. candidates = sorted( ( entry for entry in entries.values() if not self._paths_for(entry.cache_key).squashfs_mount_dir.is_mount() + and not self._is_pool_worker_visible(entry.cache_key, backend_type) ), key=lambda entry: entry.last_used, ) From 1ba54bb2fc28deb5a02892a2653a74576c85238d Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:29:27 -0400 Subject: [PATCH 007/161] fix(executor): kill and reap mount/umount subprocesses on cancellation Cancelling the await on proc.communicate() abandons a still-running mount or umount process while unwinding releases the per-key lock; a new lease could then be handed a mounted path the orphaned umount was about to remove. Both subprocess awaits now kill and reap the process before propagating cancellation, so the lock covers the complete (u)mount lifecycle: either the syscall never happened and the entry stays consistent, or it completed and the entry is a plain cache miss on the next admission. ENG-1568 --- tests/unit/test_registry_artifacts.py | 102 ++++++++++++++++++++++++ tracecat/executor/registry_artifacts.py | 31 ++++++- 2 files changed, 130 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index d0e2a2c5c0..22821b8190 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -66,6 +66,31 @@ def _write_image_entry( return image_path +class _BlockingSubprocess: + """Fake subprocess that blocks in communicate until it is cancelled.""" + + def __init__(self) -> None: + self.communicate_started = asyncio.Event() + self.cleanup_calls: list[str] = [] + self.returncode: int | None = None + + async def communicate(self) -> tuple[bytes, bytes]: + """Block until the task awaiting subprocess completion is cancelled.""" + self.communicate_started.set() + await asyncio.Event().wait() + return b"", b"" + + def kill(self) -> None: + """Record that the subprocess was killed.""" + self.cleanup_calls.append("kill") + self.returncode = -9 + + async def wait(self) -> int: + """Record that the killed subprocess was reaped.""" + self.cleanup_calls.append("wait") + return -9 + + @pytest.fixture def temp_cache_dir(): """Create a temporary cache directory.""" @@ -512,6 +537,40 @@ async def test_mount_squashfs_uses_hardened_read_only_options( stderr=asyncio.subprocess.PIPE, ) + @pytest.mark.anyio + async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): + """Cancellation cannot leave an orphan mount process after lock release.""" + cache_key = "cancelled-mount" + cache = RegistryArtifactCache(temp_cache_dir) + ctx = cache._context_for(cache_key) + artifact = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key=cache_key, + ) + image_path = ctx.paths.squashfs_image_path + target_dir = ctx.paths.squashfs_mount_dir + image_path.write_bytes(b"squashfs") + target_dir.mkdir() + process = _BlockingSubprocess() + + with patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ): + mounting = asyncio.create_task( + artifact._mount_image(image_path, target_dir) + ) + await process.communicate_started.wait() + mounting.cancel() + + with pytest.raises(asyncio.CancelledError): + await mounting + + assert process.cleanup_calls == ["kill", "wait"] + assert target_dir.is_dir() + assert not target_dir.is_mount() + @pytest.mark.anyio async def test_materialize_extracts_squashfs_when_mount_fails(self, temp_cache_dir): """Test that SquashFS mount failures fall back to unsquashfs extraction.""" @@ -1394,6 +1453,49 @@ async def mock_umount(*args, **kwargs): str(paths.squashfs_mount_dir), ) + @pytest.mark.anyio + async def test_cancelled_unmount_kills_and_reaps_before_releasing_key_lock( + self, temp_cache_dir + ): + """Cancellation leaves a consistent entry for the next admission.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/cancelled-unmount.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + paths = cache._paths_for(cache_key) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") + mounted = {paths.squashfs_mount_dir} + process = _BlockingSubprocess() + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/umount", + ), + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ), + ): + eviction = asyncio.create_task(cache._evict_entry(cache_key)) + await process.communicate_started.wait() + eviction.cancel() + + with pytest.raises(asyncio.CancelledError): + await eviction + + assert process.cleanup_calls == ["kill", "wait"] + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [paths.squashfs_mount_dir] + assert registry_paths[0].is_dir() + assert (registry_paths[0] / "module.py").read_text() == "VALUE = 1" + + assert paths.squashfs_image_path.is_file() + assert paths.squashfs_mount_dir.is_dir() + @pytest.mark.anyio async def test_cancelled_background_deletion_leaves_a_clean_miss( self, temp_cache_dir diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 68355aecf2..ee36f1210d 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import contextlib import hashlib import os import re @@ -440,6 +441,11 @@ async def extract( async def _mount_image(self, image_path: Path, target_dir: Path) -> None: """Mount a SquashFS image read-only at target_dir. + Cancellation kills and reaps the mount subprocess before propagating, + so the caller's per-key lock covers the complete mount lifecycle. The + target therefore remains an unmounted cache miss if mount never took + effect, or is already mounted and reusable by the next admission. + Args: image_path: Local path of the SquashFS image. target_dir: Existing directory to mount the image at. @@ -461,7 +467,13 @@ async def _mount_image(self, image_path: Path, target_dir: Path) -> None: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - stdout, stderr = await proc.communicate() + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + proc.kill() + await proc.wait() + raise if proc.returncode == 0 or target_dir.is_mount(): return @@ -1341,7 +1353,14 @@ async def _evict_entry(self, cache_key: str) -> bool: return True async def _unmount(self, mount_dir: Path) -> bool: - """Unmount a SquashFS artifact directory, releasing its loop device.""" + """Unmount a SquashFS artifact directory, releasing its loop device. + + Cancellation kills and reaps the umount subprocess before propagating, + so the caller's per-key lock covers the complete unmount lifecycle. If + umount never took effect, the mounted entry stays consistent and can be + reused; if it already took effect, the missing extraction directory + makes the entry a plain cache miss on the next admission. + """ umount = shutil.which("umount") if umount is None: return False @@ -1352,7 +1371,13 @@ async def _unmount(self, mount_dir: Path) -> bool: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - stdout, stderr = await proc.communicate() + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + proc.kill() + await proc.wait() + raise if proc.returncode == 0 or not mount_dir.is_mount(): return True From 5354ca405895a3eeb1c9166b867df0ee68db9f0a Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:59:39 -0400 Subject: [PATCH 008/161] fix(executor): move pool experimental warning off eviction hot path --- tracecat/executor/backends/__init__.py | 7 +++++++ tracecat/executor/schemas.py | 6 ------ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/tracecat/executor/backends/__init__.py b/tracecat/executor/backends/__init__.py index 158e4ee6bb..7d26e7fefd 100644 --- a/tracecat/executor/backends/__init__.py +++ b/tracecat/executor/backends/__init__.py @@ -46,6 +46,13 @@ def _create_backend(backend_type: ExecutorBackendType) -> ExecutorBackend: case ExecutorBackendType.POOL: from tracecat.executor.backends.pool import PoolBackend + # Warn here rather than in resolve_backend_type(), which is also + # called on hot paths such as cache eviction passes. + logger.warning( + "The 'pool' executor backend is experimental and not production " + "ready: its registry cache is exempt from eviction and can grow " + "without bound. Use 'ephemeral' for production nsjail isolation.", + ) return PoolBackend() case ExecutorBackendType.EPHEMERAL: from tracecat.executor.backends.ephemeral import EphemeralBackend diff --git a/tracecat/executor/schemas.py b/tracecat/executor/schemas.py index eaf0307964..35f9cc9040 100644 --- a/tracecat/executor/schemas.py +++ b/tracecat/executor/schemas.py @@ -96,12 +96,6 @@ def resolve_backend_type() -> ExecutorBackendType: "Auto-selecting 'direct' backend (nsjail not available)", ) backend_type = ExecutorBackendType.DIRECT - elif backend_type == ExecutorBackendType.POOL: - logger.warning( - "The 'pool' executor backend is experimental and not production " - "ready: its registry cache is exempt from eviction and can grow " - "without bound. Use 'ephemeral' for production nsjail isolation.", - ) return backend_type From 72b18b49f16965d3357a6a83c0ae0c269e16a164 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:12:36 -0400 Subject: [PATCH 009/161] fix(executor): kill and reap action subprocesses on cancellation Cancelling the await on proc.communicate() left the direct and nsjail action children running while unwinding released the registry-path lease, so a concurrent budget pass could evict site-packages out from under a live importing process. Both sites now kill and reap the child before propagating cancellation, mirroring the mount/umount fix. ENG-1568 --- tests/unit/test_action_runner.py | 62 +++++++++++++++++++++- tests/unit/test_executor_sandbox_nsjail.py | 56 +++++++++++++++++++ tracecat/executor/action_runner.py | 6 ++- tracecat/sandbox/executor.py | 9 ++++ 4 files changed, 131 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index 14e3f6959a..d24f79064f 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -93,13 +93,20 @@ def mock_process_group_communication( self, monkeypatch: pytest.MonkeyPatch ) -> AsyncMock: """Keep subprocess unit tests focused on ActionRunner behavior.""" + real_communication = action_runner.communicate_process_group async def communicate( - process: asyncio.subprocess.Process, + process: asyncio.subprocess.Process | AsyncMock, *, input: bytes | None = None, # noqa: A002 timeout: float | None = None, ) -> tuple[bytes, bytes]: + if isinstance(process, asyncio.subprocess.Process): + return await real_communication( + process, + input=input, + timeout=timeout, + ) stdout, stderr = await asyncio.wait_for( process.communicate(input=input), timeout=timeout, @@ -175,6 +182,59 @@ async def slow_communicate(input=None): # noqa: A002 assert result.type == "TimeoutError" mock_process_group_communication.assert_awaited_once() + @pytest.mark.anyio + async def test_cancelled_direct_action_kills_and_reaps_subprocess( + self, temp_cache_dir, mock_run_action_input, mock_role + ) -> None: + """Cancellation propagates only after the direct child is reaped.""" + runner = ActionRunner(cache_dir=temp_cache_dir) + base_dir = temp_cache_dir / "base" + base_dir.mkdir() + real_create_subprocess_exec = asyncio.create_subprocess_exec + process_started = asyncio.Event() + process: asyncio.subprocess.Process | None = None + + async def capture_subprocess(*args, **kwargs): + nonlocal process + process = await real_create_subprocess_exec(*args, **kwargs) + process_started.set() + return process + + with ( + patch.object( + action_runner, + "_direct_subprocess_command", + return_value=["/bin/sleep", "30"], + ), + patch( + "tracecat.executor.action_runner.asyncio.create_subprocess_exec", + side_effect=capture_subprocess, + ), + ): + execution = asyncio.create_task( + runner._execute_direct( + input=mock_run_action_input, + role=mock_role, + registry_paths=[base_dir], + secret_projection=_empty_secret_projection(), + timeout=60.0, + ) + ) + try: + await process_started.wait() + await asyncio.sleep(0) + execution.cancel() + + with pytest.raises(asyncio.CancelledError): + await execution + + assert process is not None + assert process.returncode is not None + finally: + if process is not None and process.returncode is None: + process.kill() + await process.wait() + @pytest.mark.anyio async def test_execute_action_subprocess_crash( self, temp_cache_dir, mock_run_action_input, mock_role diff --git a/tests/unit/test_executor_sandbox_nsjail.py b/tests/unit/test_executor_sandbox_nsjail.py index 22e91f96e7..6966dcd6d5 100644 --- a/tests/unit/test_executor_sandbox_nsjail.py +++ b/tests/unit/test_executor_sandbox_nsjail.py @@ -46,6 +46,7 @@ from tracecat.executor.secret_preprocessors import SecretEnvProjection from tracecat.identifiers.workflow import WorkflowUUID from tracecat.registry.lock.types import RegistryLock +from tracecat.sandbox.executor import ActionSandboxConfig, NsjailExecutor _DOCKER_CHILD_ENV = "TRACECAT__EXECUTOR_ACTION_SMOKE_DOCKER_CHILD" _SKIP_SENTINEL = "TRACE_CAT_EXECUTOR_ACTION_SMOKE_SKIP:" @@ -671,6 +672,61 @@ async def _run_current_builtin_smoke_case( } +@pytest.mark.anyio +async def test_cancelled_nsjail_action_kills_and_reaps_subprocess( + tmp_path: Path, +) -> None: + """Cancellation propagates only after the nsjail child is reaped.""" + job_dir = tmp_path / "job" + job_dir.mkdir() + rootfs_dir = tmp_path / "rootfs" + rootfs_dir.mkdir() + runner = NsjailExecutor( + nsjail_path=str(tmp_path / "nsjail"), + rootfs_path=str(rootfs_dir), + cache_dir=str(tmp_path / "cache"), + ) + sandbox_config = ActionSandboxConfig( + registry_paths=[], + tracecat_app_dir=tmp_path, + timeout_seconds=30, + ) + real_create_subprocess_exec = asyncio.create_subprocess_exec + process_started = asyncio.Event() + process: asyncio.subprocess.Process | None = None + + async def capture_subprocess(*args, **kwargs): + nonlocal process + process = await real_create_subprocess_exec( + "/bin/sleep", + "30", + **kwargs, + ) + process_started.set() + return process + + with patch( + "tracecat.sandbox.executor.asyncio.create_subprocess_exec", + side_effect=capture_subprocess, + ): + execution = asyncio.create_task(runner.execute_action(job_dir, sandbox_config)) + try: + await process_started.wait() + await asyncio.sleep(0) + execution.cancel() + + with pytest.raises(asyncio.CancelledError): + await execution + + assert process is not None + assert process.returncode is not None + assert not (job_dir / "nsjail.cfg").exists() + finally: + if process is not None and process.returncode is None: + process.kill() + await process.wait() + + @pytest.mark.parametrize( "smoke_case", [ diff --git a/tracecat/executor/action_runner.py b/tracecat/executor/action_runner.py index 438ab64307..63b9c3606d 100644 --- a/tracecat/executor/action_runner.py +++ b/tracecat/executor/action_runner.py @@ -369,7 +369,11 @@ async def _execute_direct( timeout: float | None = None, resolved_context: ResolvedContext | None = None, ) -> ExecutionResult: - """Execute an action in a direct subprocess (no sandbox).""" + """Execute an action in a direct subprocess (no sandbox). + + Every exit kills lingering descendants before the registry-path lease + protecting their imports can unwind. + """ timeout = timeout or config.TRACECAT__EXECUTOR_CLIENT_TIMEOUT # Prepare input JSON for subprocess diff --git a/tracecat/sandbox/executor.py b/tracecat/sandbox/executor.py index 5ad8b8c4d1..d7872ff565 100644 --- a/tracecat/sandbox/executor.py +++ b/tracecat/sandbox/executor.py @@ -1,6 +1,7 @@ """nsjail executor for sandboxed Python execution.""" import asyncio +import contextlib import json import os import re @@ -889,6 +890,14 @@ async def execute_action( timeout=timeout, ) + except asyncio.CancelledError: + # The registry-path lease is released as cancellation unwinds, so + # the importing child must be dead and reaped before propagation. + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + raise + except TimeoutError as e: process.kill() await process.wait() From 8012347e23086769d1904a920480f47f4c4a48dd Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:13:38 -0400 Subject: [PATCH 010/161] fix(executor): serialize first squashfs mount capability probe The process-wide mount state was mutated under per-key locks, so two concurrent first mounts could interleave a failure and a success into disabled=True with mounted_once=True, disabling mounts forever despite proven capability. The first probe now runs under a shared probe lock with state rechecks: disablement can only happen while holding the probe with no successful mount, making the contradictory state unreachable. Reclaim and retry stay outside the probe lock. ENG-1568 --- tests/unit/test_registry_artifacts.py | 129 ++++++++++++++++++++++++ tracecat/executor/registry_artifacts.py | 40 +++++--- 2 files changed, 157 insertions(+), 12 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 22821b8190..161ce2ef4c 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1719,6 +1719,135 @@ def test_auto_pool_sweep_protects_tarballs_and_trims_other_entries( class TestSquashfsMountCapability: """Tests for process-wide SquashFS mount capability tracking.""" + @pytest.mark.anyio + async def test_concurrent_first_mount_failure_serializes_capability_probe( + self, temp_cache_dir + ) -> None: + """A failed first probe disables a waiter without racing its mount.""" + cache = RegistryArtifactCache(temp_cache_dir) + first_ctx = cache._context_for("first-probe") + second_ctx = cache._context_for("second-probe") + assert first_ctx.squashfs_mount_state is second_ctx.squashfs_mount_state + first_ctx.paths.squashfs_image_path.write_bytes(b"squashfs") + second_ctx.paths.squashfs_image_path.write_bytes(b"squashfs") + first_artifact = SquashfsArtifact( + uri="s3://bucket/path/first.squashfs", + cache_key=first_ctx.cache_key, + ) + second_artifact = SquashfsArtifact( + uri="s3://bucket/path/second.squashfs", + cache_key=second_ctx.cache_key, + ) + first_mount_started = asyncio.Event() + release_first_mount = asyncio.Event() + mount_attempts: list[Path] = [] + + async def mock_mount_image(self, image_path, target_dir): + mount_attempts.append(target_dir) + if target_dir == first_ctx.paths.squashfs_mount_dir: + first_mount_started.set() + await release_first_mount.wait() + raise SquashfsMountCommandError("operation not permitted") + + with patch.object(SquashfsArtifact, "_mount_image", mock_mount_image): + first_mount = asyncio.create_task( + first_artifact._try_mount( + first_ctx, + first_ctx.paths.squashfs_image_path, + ) + ) + await first_mount_started.wait() + second_mount = asyncio.create_task( + second_artifact._try_mount( + second_ctx, + second_ctx.paths.squashfs_image_path, + ) + ) + await asyncio.sleep(0) + + assert first_ctx.squashfs_mount_state.probe_lock.locked() + assert not second_mount.done() + assert mount_attempts == [first_ctx.paths.squashfs_mount_dir] + + release_first_mount.set() + first_result, second_result = await asyncio.gather( + first_mount, + second_mount, + ) + + state = first_ctx.squashfs_mount_state + assert first_result is None + assert second_result is None + assert state.disabled is True + assert state.mounted_once is False + assert not (state.disabled and state.mounted_once) + assert mount_attempts == [first_ctx.paths.squashfs_mount_dir] + + @pytest.mark.anyio + async def test_concurrent_probe_waiter_mounts_after_first_success( + self, temp_cache_dir + ) -> None: + """A waiter mounts after the successful probe releases serialization.""" + cache = RegistryArtifactCache(temp_cache_dir) + first_ctx = cache._context_for("first-success") + second_ctx = cache._context_for("second-success") + assert first_ctx.squashfs_mount_state is second_ctx.squashfs_mount_state + first_ctx.paths.squashfs_image_path.write_bytes(b"squashfs") + second_ctx.paths.squashfs_image_path.write_bytes(b"squashfs") + first_artifact = SquashfsArtifact( + uri="s3://bucket/path/first.squashfs", + cache_key=first_ctx.cache_key, + ) + second_artifact = SquashfsArtifact( + uri="s3://bucket/path/second.squashfs", + cache_key=second_ctx.cache_key, + ) + first_mount_started = asyncio.Event() + release_first_mount = asyncio.Event() + mount_attempts: list[Path] = [] + + async def mock_mount_image(self, image_path, target_dir): + mount_attempts.append(target_dir) + if target_dir == first_ctx.paths.squashfs_mount_dir: + first_mount_started.set() + await release_first_mount.wait() + + with patch.object(SquashfsArtifact, "_mount_image", mock_mount_image): + first_mount = asyncio.create_task( + first_artifact._try_mount( + first_ctx, + first_ctx.paths.squashfs_image_path, + ) + ) + await first_mount_started.wait() + second_mount = asyncio.create_task( + second_artifact._try_mount( + second_ctx, + second_ctx.paths.squashfs_image_path, + ) + ) + await asyncio.sleep(0) + + assert first_ctx.squashfs_mount_state.probe_lock.locked() + assert not second_mount.done() + assert mount_attempts == [first_ctx.paths.squashfs_mount_dir] + + release_first_mount.set() + first_result, second_result = await asyncio.gather( + first_mount, + second_mount, + ) + + state = first_ctx.squashfs_mount_state + assert first_result == first_ctx.paths.squashfs_mount_dir + assert second_result == second_ctx.paths.squashfs_mount_dir + assert state.disabled is False + assert state.mounted_once is True + assert mount_attempts == [ + first_ctx.paths.squashfs_mount_dir, + second_ctx.paths.squashfs_mount_dir, + ] + @pytest.mark.anyio async def test_first_mount_failure_disables_squashfs_process_wide( self, temp_cache_dir diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index ee36f1210d..3b3acaa029 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -14,7 +14,7 @@ from abc import ABC, abstractmethod from collections.abc import AsyncIterator, Awaitable, Callable, Iterable from contextlib import asynccontextmanager -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import StrEnum from pathlib import Path @@ -88,6 +88,7 @@ class SquashfsMountState: disabled: bool = False mounted_once: bool = False + probe_lock: asyncio.Lock = field(default_factory=asyncio.Lock) @dataclass(slots=True) @@ -125,6 +126,7 @@ def can_mount_squashfs(self) -> bool: ) def disable_squashfs_mount(self) -> None: + """Disable mounting after the serialized first capability probe fails.""" self.squashfs_mount_state.disabled = True def record_squashfs_mount(self) -> None: @@ -251,6 +253,10 @@ async def _try_mount( mounted artifact is evicted and the mount is retried once, so a single failure never downgrades the whole process to extraction. + The probe lock deliberately covers the first mount's download and mount + command. It nests inside a per-key materialization lock and acquires no + other lock; reclaim and retry remain outside it. + Only ``SquashfsMountCommandError`` drives this policy. Download and preparation errors propagate to the caller so a transient S3 failure never disables mounting or evicts an unrelated idle mount. @@ -265,22 +271,32 @@ async def _try_mount( Raises: Exception: Any non-mount failure raised while preparing the image. """ + if not ctx.has_mounted_squashfs(): + async with ctx.squashfs_mount_state.probe_lock: + if ctx.squashfs_mount_state.disabled: + return None + if not ctx.has_mounted_squashfs(): + try: + return await self.mount(ctx, image_path) + except SquashfsMountCommandError as e: + # This is the only disable path: the probe lock is held + # and no mount has succeeded, so disabled and mounted_once + # cannot both become true. + ctx.disable_squashfs_mount() + logger.warning( + "Failed to mount SquashFS registry artifact, trying extraction", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + error=str(e), + ) + return None + try: return await self.mount(ctx, image_path) except SquashfsMountCommandError as e: mount_error = e - if not ctx.has_mounted_squashfs(): - ctx.disable_squashfs_mount() - logger.warning( - "Failed to mount SquashFS registry artifact, trying extraction", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - error=str(mount_error), - ) - return None - logger.warning( "Failed to mount SquashFS registry artifact, reclaiming an idle mount", cache_key=ctx.cache_key, From 04e5595c6bffb6792723926ef9f3543cf917e6b3 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:14:25 -0400 Subject: [PATCH 011/161] fix(executor): run cache startup sweep off the worker event loop The registry artifact cache was constructed lazily inside the first action, running its synchronous startup sweep (a full-cache stat scan plus rmtree deletes) on the event loop and stalling heartbeats and concurrent activities on incident-sized caches. The worker now constructs the cache via asyncio.to_thread at startup, before the backend spawns workers or any activity can run; lazy construction remains only as a fallback. ENG-1568 --- tracecat/executor/registry_artifacts.py | 4 ++-- tracecat/executor/worker.py | 5 +++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 3b3acaa029..e2d72cd810 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -1471,8 +1471,8 @@ def _sweep_startup_state(self) -> None: after warm workers have already inherited those paths. A missing or empty cache directory is a no-op. - The sweep is deliberately synchronous: it runs once during construction, - before the process serves any action. + Construction runs eagerly in a worker-startup thread before activities + can run; lazy in-activity construction remains a fallback. """ if not self.cache_dir.is_dir(): self._budget_dirty = False diff --git a/tracecat/executor/worker.py b/tracecat/executor/worker.py index f7beb69c89..68924e73a4 100644 --- a/tracecat/executor/worker.py +++ b/tracecat/executor/worker.py @@ -56,6 +56,7 @@ from tracecat import config from tracecat.dsl.client import get_temporal_client from tracecat.executor.action_gateway.server import ActionGateway + from tracecat.executor.action_runner import get_action_runner from tracecat.executor.activities import ExecutorActivities from tracecat.executor.backends import ( initialize_executor_backend, @@ -136,6 +137,10 @@ async def main(shutdown_event: asyncio.Event | None = None) -> None: # socket path is available in their immutable process environment. await action_gateway.start() + # Construct the registry artifact cache and run its synchronous startup + # sweep in a thread before the backend spawns workers or activities run. + await asyncio.to_thread(get_action_runner) + # Initialize the executor backend before accepting tasks await initialize_executor_backend() From 20770ae352eaf0ceb5ac3f064c5d57cb48de98a7 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:50:45 -0400 Subject: [PATCH 012/161] fix(sandbox): kill and reap run-python subprocesses on cancellation run_python holds a registry-artifact lease for the whole sandbox run, but the nsjail and PID-namespace executors killed their children only on TimeoutError. Cancellation left the child alive while the lease unwound, so a concurrent budget pass could delete an extraction the child still lazily imports from. All five communicate() sites now kill and reap the child before propagating cancellation. ENG-1568 --- tests/unit/test_executor_sandbox_nsjail.py | 51 ++++++++++++++++++++++ tracecat/sandbox/executor.py | 14 ++++++ tracecat/sandbox/unsafe_pid_executor.py | 11 +++++ 3 files changed, 76 insertions(+) diff --git a/tests/unit/test_executor_sandbox_nsjail.py b/tests/unit/test_executor_sandbox_nsjail.py index 6966dcd6d5..52778e4cdf 100644 --- a/tests/unit/test_executor_sandbox_nsjail.py +++ b/tests/unit/test_executor_sandbox_nsjail.py @@ -47,6 +47,7 @@ from tracecat.identifiers.workflow import WorkflowUUID from tracecat.registry.lock.types import RegistryLock from tracecat.sandbox.executor import ActionSandboxConfig, NsjailExecutor +from tracecat.sandbox.types import SandboxConfig _DOCKER_CHILD_ENV = "TRACECAT__EXECUTOR_ACTION_SMOKE_DOCKER_CHILD" _SKIP_SENTINEL = "TRACE_CAT_EXECUTOR_ACTION_SMOKE_SKIP:" @@ -672,6 +673,56 @@ async def _run_current_builtin_smoke_case( } +@pytest.mark.anyio +async def test_cancelled_nsjail_execute_kills_and_reaps_subprocess( + tmp_path: Path, +) -> None: + """Cancellation propagates only after the nsjail child is reaped.""" + job_dir = tmp_path / "job" + job_dir.mkdir() + rootfs_dir = tmp_path / "rootfs" + rootfs_dir.mkdir() + runner = NsjailExecutor( + nsjail_path=str(tmp_path / "nsjail"), + rootfs_path=str(rootfs_dir), + cache_dir=str(tmp_path / "cache"), + ) + real_create_subprocess_exec = asyncio.create_subprocess_exec + process_started = asyncio.Event() + process: asyncio.subprocess.Process | None = None + + async def capture_subprocess(*args, **kwargs): + nonlocal process + process = await real_create_subprocess_exec( + "/bin/sleep", + "30", + **kwargs, + ) + process_started.set() + return process + + with patch( + "tracecat.sandbox.executor.asyncio.create_subprocess_exec", + side_effect=capture_subprocess, + ): + execution = asyncio.create_task(runner.execute(job_dir, SandboxConfig())) + try: + await process_started.wait() + await asyncio.sleep(0) + execution.cancel() + + with pytest.raises(asyncio.CancelledError): + await execution + + assert process is not None + assert process.returncode is not None + assert not (job_dir / "nsjail.cfg").exists() + finally: + if process is not None and process.returncode is None: + process.kill() + await process.wait() + + @pytest.mark.anyio async def test_cancelled_nsjail_action_kills_and_reaps_subprocess( tmp_path: Path, diff --git a/tracecat/sandbox/executor.py b/tracecat/sandbox/executor.py index d7872ff565..af861a9988 100644 --- a/tracecat/sandbox/executor.py +++ b/tracecat/sandbox/executor.py @@ -477,6 +477,14 @@ async def execute( timeout=timeout, ) + except asyncio.CancelledError: + # The registry-path lease is released as cancellation unwinds, so + # the importing child must be dead and reaped before propagation. + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + raise + except TimeoutError as e: # Kill the process if it times out process.kill() @@ -622,6 +630,12 @@ async def execute_install( timeout=timeout, ) + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + raise + except TimeoutError as e: process.kill() await process.wait() diff --git a/tracecat/sandbox/unsafe_pid_executor.py b/tracecat/sandbox/unsafe_pid_executor.py index ccd5cda937..7c8ebdc400 100644 --- a/tracecat/sandbox/unsafe_pid_executor.py +++ b/tracecat/sandbox/unsafe_pid_executor.py @@ -6,6 +6,7 @@ """ import asyncio +import contextlib import hashlib import json import logging @@ -333,6 +334,11 @@ async def _create_venv(self, venv_path: Path) -> None: ) try: _, stderr = await asyncio.wait_for(process.communicate(), timeout=60) + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + raise except TimeoutError as e: process.kill() await process.wait() @@ -378,6 +384,11 @@ async def _install_packages( process.communicate(), timeout=timeout_seconds, ) + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + raise except TimeoutError as e: process.kill() await process.wait() From 6cfcb3f17dfa765dda78d0a3def8096e635b12e2 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:51:17 -0400 Subject: [PATCH 013/161] fix(executor): re-arm budget dirty flag on cancelled convergence _converge_cache_budget consumes the dirty signal before its awaited enforcement pass, but only OSError re-armed it. Cancellation mid-scan propagated with the flag consumed, so later cache hits skipped enforcement and an over-budget cache stayed over budget until the next materialization. Any non-OSError exit now restores the flag first. ENG-1568 --- tests/unit/test_registry_artifacts.py | 32 +++++++++++++++++++++++++ tracecat/executor/registry_artifacts.py | 7 +++++- 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 161ce2ef4c..386deb9425 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1145,6 +1145,38 @@ async def mock_extract(self, tarball_path, target_dir): assert convergence_scans == 2 assert cache._budget_dirty is False + @pytest.mark.anyio + async def test_cancelled_convergence_rearms_budget_dirty( + self, temp_cache_dir + ) -> None: + """Cancellation must preserve the need for a later budget pass.""" + cache = RegistryArtifactCache(temp_cache_dir) + enforcement_started = asyncio.Event() + finish_enforcement = asyncio.Event() + + async def mock_enforce_cache_budget( + *, protected_key: str | None = None + ) -> bool: + del protected_key + enforcement_started.set() + await finish_enforcement.wait() + return True + + cache._budget_dirty = True + with patch.object( + cache, + "_enforce_cache_budget", + side_effect=mock_enforce_cache_budget, + ): + convergence = asyncio.create_task(cache._converge_cache_budget()) + await enforcement_started.wait() + convergence.cancel() + + with pytest.raises(asyncio.CancelledError): + await convergence + + assert cache._budget_dirty is True + @pytest.mark.parametrize( "oldest_has_tarball", [False, True], diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index e2d72cd810..3c875b0abd 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -1166,7 +1166,8 @@ async def _converge_cache_budget(self) -> None: A follow-up pass therefore occurs only when a concurrent materialization sets the flag again. Without new materializations the loop terminates, while an over-budget or failed scan restores the flag and breaks so it - cannot spin while entries remain leased. + cannot spin while entries remain leased. Cancellation also restores the + consumed flag before propagating. """ while self._budget_dirty: self._budget_dirty = False @@ -1180,6 +1181,10 @@ async def _converge_cache_budget(self) -> None: ) self._budget_dirty = True break + except BaseException: + # Cancellation must re-arm the consumed dirty signal before propagating. + self._budget_dirty = True + raise if not within_budget: self._budget_dirty = True From 51d8623d32da35b79724df991e39915555c66e85 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:03:02 -0400 Subject: [PATCH 014/161] refactor(executor): once-guard the registry cache startup sweep The off-loop, before-first-lease sweep invariant lived in a worker.py startup convention, and the lazy in-activity construction fallback would have blocked the event loop again. Construction is now cheap and ensure_swept() enforces the invariant at the lease and materialization boundaries: exactly once, in a thread, retried on failure, regardless of who constructs the cache first. Worker startup keeps a warm-up call so the sweep still lands before the backend spawns workers. ENG-1568 --- ...registry_artifact_cache_mount_lifecycle.py | 3 +- tests/unit/test_registry_artifacts.py | 88 +++++++++++++++++-- tracecat/executor/registry_artifacts.py | 36 ++++++-- tracecat/executor/worker.py | 6 +- 4 files changed, 116 insertions(+), 17 deletions(-) diff --git a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py index 1abeeb0a1c..6a40456bcf 100644 --- a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py +++ b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py @@ -349,7 +349,8 @@ async def _run_mount_lifecycle_child() -> None: stale_mount_dir.mkdir() config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES = 1 - RegistryArtifactCache(sweep_dir) + sweep_cache = RegistryArtifactCache(sweep_dir) + await sweep_cache.ensure_swept() payload["startup_sweep_trimmed"] = ( not (sweep_dir / f"squashfs-{sweep_keys[0]}.squashfs").exists() and (sweep_dir / f"squashfs-{sweep_keys[1]}.squashfs").exists() diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 386deb9425..cf4d7ecd09 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1059,6 +1059,7 @@ async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( ): """Steady-state cache hits must not pay for a full cache scan.""" cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() artifact_uri = "s3://bucket/cached.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) _write_tarball_entry(temp_cache_dir, cache_key) @@ -1080,6 +1081,7 @@ async def test_release_keeps_retrying_while_the_cache_stays_over_budget( ): """A cache that cannot shrink yet must stay marked for re-enforcement.""" cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() artifact_uri = "s3://bucket/pinned.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) _write_image_entry(temp_cache_dir, cache_key, size=4096, mtime=100.0) @@ -1534,6 +1536,7 @@ async def test_cancelled_background_deletion_leaves_a_clean_miss( ): """Cancellation cannot expose live paths that deletion still owns.""" cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() artifact_uri = "s3://bucket/cancelled-eviction.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) original_target = _write_tarball_entry(temp_cache_dir, cache_key) @@ -1685,16 +1688,21 @@ async def test_eviction_skips_busy_key(self, temp_cache_dir): class TestRegistryArtifactCacheStartupSweep: """Tests for the startup sweep that reclaims state from a dead process.""" - def test_sweep_tolerates_missing_cache_dir(self, temp_cache_dir): + @pytest.mark.anyio + async def test_sweep_tolerates_missing_cache_dir(self, temp_cache_dir): """A cache directory that does not exist yet is a no-op.""" cache_dir = temp_cache_dir / "missing" cache = RegistryArtifactCache(cache_dir) + await cache.ensure_swept() assert cache.cache_dir == cache_dir assert not cache_dir.exists() - def test_sweep_removes_orphaned_scratch_and_stale_mount_dirs(self, temp_cache_dir): + @pytest.mark.anyio + async def test_sweep_removes_orphaned_scratch_and_stale_mount_dirs( + self, temp_cache_dir + ): """Interrupted materializations and dead mount dirs are reclaimed.""" orphaned = temp_cache_dir / "abc123.999999.4321.squashfs" orphaned.write_bytes(b"partial") @@ -1706,7 +1714,8 @@ def test_sweep_removes_orphaned_scratch_and_stale_mount_dirs(self, temp_cache_di stale_mount_dir.mkdir() entry_dir = _write_tarball_entry(temp_cache_dir, "abc123") - RegistryArtifactCache(temp_cache_dir) + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() assert not orphaned.exists() assert not orphaned_dir.exists() @@ -1714,17 +1723,20 @@ def test_sweep_removes_orphaned_scratch_and_stale_mount_dirs(self, temp_cache_di assert not stale_mount_dir.exists() assert entry_dir.is_dir() - def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): + @pytest.mark.anyio + async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): """A live mountpoint belongs to a running process and must survive.""" mount_dir = temp_cache_dir / "squashfs-abc123" mount_dir.mkdir() with patch.object(Path, "is_mount", lambda self: self == mount_dir): - RegistryArtifactCache(temp_cache_dir) + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() assert mount_dir.is_dir() - def test_auto_pool_sweep_protects_tarballs_and_trims_other_entries( + @pytest.mark.anyio + async def test_auto_pool_sweep_protects_tarballs_and_trims_other_entries( self, temp_cache_dir ): """Startup trimming preserves paths inherited by auto-resolved workers.""" @@ -1740,6 +1752,7 @@ def test_auto_pool_sweep_protects_tarballs_and_trims_other_entries( patch(MAX_BYTES_CONFIG, 0), ): cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() assert oldest.exists() assert oldest_tarball.is_dir() @@ -1747,6 +1760,69 @@ def test_auto_pool_sweep_protects_tarballs_and_trims_other_entries( assert newest.exists() assert cache._budget_dirty is False + @pytest.mark.anyio + async def test_ensure_swept_runs_once(self, temp_cache_dir): + """Repeated startup-sweep requests invoke the sweep once.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with patch.object( + cache, + "_sweep_startup_state", + wraps=cache._sweep_startup_state, + ) as sweep: + await cache.ensure_swept() + await cache.ensure_swept() + + assert sweep.call_count == 1 + + @pytest.mark.anyio + async def test_concurrent_ensure_swept_runs_once(self, temp_cache_dir): + """Concurrent startup-sweep requests invoke the sweep once.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with patch.object( + cache, + "_sweep_startup_state", + wraps=cache._sweep_startup_state, + ) as sweep: + await asyncio.gather(*(cache.ensure_swept() for _ in range(4))) + + assert sweep.call_count == 1 + + @pytest.mark.anyio + async def test_lease_triggers_startup_sweep(self, temp_cache_dir): + """Lease admission reclaims startup scratch before yielding paths.""" + orphaned_dir = temp_cache_dir / "abc123.999999.4321.tmp" + orphaned_dir.mkdir() + assert TEMP_ARTIFACT_PATTERN.match(orphaned_dir.name) is not None + cache = RegistryArtifactCache(temp_cache_dir) + + async with cache.lease(None): + assert not orphaned_dir.exists() + + @pytest.mark.anyio + async def test_failed_ensure_swept_retries(self, temp_cache_dir): + """A failed startup sweep is retried by the next caller.""" + orphaned_dir = temp_cache_dir / "abc123.999999.4321.tmp" + orphaned_dir.mkdir() + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch.object( + cache, + "_sweep_startup_state", + side_effect=OSError("simulated sweep failure"), + ), + pytest.raises(OSError), + ): + await cache.ensure_swept() + + assert orphaned_dir.is_dir() + + await cache.ensure_swept() + + assert not orphaned_dir.exists() + class TestSquashfsMountCapability: """Tests for process-wide SquashFS mount capability tracking.""" diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 3c875b0abd..caaa6c223d 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -837,13 +837,31 @@ def __init__(self, cache_dir: Path): self._locks: dict[str, asyncio.Lock] = {} self._locks_lock = asyncio.Lock() self._budget_lock = asyncio.Lock() + # Guard the off-loop startup sweep independently from cache operations. + self._swept: bool = False + self._sweep_lock = asyncio.Lock() self._squashfs_mount_state = SquashfsMountState() self._leases: dict[str, RegistryArtifactLease] = {} # Whether the on-disk cache may exceed its budget. Set when a new entry # is materialized and cleared once enforcement measures a cache that # fits, so steady-state cache hits never pay for a disk scan. self._budget_dirty = True - self._sweep_startup_state() + + async def ensure_swept(self) -> None: + """Run the startup sweep exactly once, off the event loop. + + Idempotent and safe under concurrency: the first caller performs the + sweep in a thread; concurrent callers wait; later callers return + immediately. Runs before the first lease or materialization so the + sweep never observes in-flight cache entries. + """ + if self._swept: + return + async with self._sweep_lock: + if self._swept: + return + await asyncio.to_thread(self._sweep_startup_state) + self._swept = True @asynccontextmanager async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Path]]: @@ -859,6 +877,8 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat Yields: Importable Python paths for the requested artifacts. """ + await self.ensure_swept() + if not artifact_uris: logger.info("No registry artifact URIs provided, using base PYTHONPATH") yield [self._base_pythonpath_dir()] @@ -928,6 +948,8 @@ async def _admit_lease( async def materialize(self, cache_key: str, artifact_uri: str) -> list[Path]: """Materialize a registry artifact as local importable directories.""" + await self.ensure_swept() + ctx = self._context_for(cache_key) candidates = await self._artifact_candidates(ctx, artifact_uri) @@ -1476,8 +1498,8 @@ def _sweep_startup_state(self) -> None: after warm workers have already inherited those paths. A missing or empty cache directory is a no-op. - Construction runs eagerly in a worker-startup thread before activities - can run; lazy in-activity construction remains a fallback. + The worker warms this sweep before activities can run; lazy first-use + sweeping remains a safe fallback. """ if not self.cache_dir.is_dir(): self._budget_dirty = False @@ -1497,10 +1519,10 @@ def _sweep_startup_state(self) -> None: def _remove_orphaned_temp_paths(self) -> None: """Delete every materialization scratch path during startup. - The sweep runs during cache construction, before this process can start - a materialization in the cache directory. Every matching path is - therefore interrupted scratch from an earlier process and is safe to - remove even when the operating system reused that process's PID. + The sweep runs before the first lease or materialization in this process. + Every matching path is therefore interrupted scratch from an earlier + process and is safe to remove even when the operating system reused that + process's PID. """ for name in os.listdir(self.cache_dir): if TEMP_ARTIFACT_PATTERN.match(name) is None: diff --git a/tracecat/executor/worker.py b/tracecat/executor/worker.py index 68924e73a4..7ec431230d 100644 --- a/tracecat/executor/worker.py +++ b/tracecat/executor/worker.py @@ -137,9 +137,9 @@ async def main(shutdown_event: asyncio.Event | None = None) -> None: # socket path is available in their immutable process environment. await action_gateway.start() - # Construct the registry artifact cache and run its synchronous startup - # sweep in a thread before the backend spawns workers or activities run. - await asyncio.to_thread(get_action_runner) + # Warm the registry artifact cache sweep before the backend spawns + # workers or activities run; cache construction itself is cheap. + await get_action_runner().registry_artifacts.ensure_swept() # Initialize the executor backend before accepting tasks await initialize_executor_backend() From f0a8ac974b9b0fa3677e2b2e7b41f65b6f0e7b93 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:15:24 -0400 Subject: [PATCH 015/161] fix(executor): pre-arm budget dirty flag before materialization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A SquashFS candidate downloads its canonical image before mounting or extracting, so a failed or cancelled materialization can leave the image on disk with the dirty flag never set — after a clean startup sweep, lease release then skips convergence and orphaned images hold the cache over budget indefinitely. The flag is now set before every cache-entry materialization attempt, covering all failure paths. ENG-1568 --- tests/unit/test_registry_artifacts.py | 30 +++++++++++++++++++++++++ tracecat/executor/registry_artifacts.py | 8 +++---- 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index cf4d7ecd09..f28d6d7264 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1075,6 +1075,36 @@ async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( assert cache._budget_dirty is False + @pytest.mark.anyio + async def test_failed_materialization_rearms_budget_dirty(self, temp_cache_dir): + """A failed materialization may leave a canonical image to evict.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + assert cache._budget_dirty is False + + artifact_uri = "s3://bucket/path/site-packages.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) + + async def mock_materialize(self, ctx): + ctx.paths.squashfs_image_path.write_bytes(b"orphaned image") + raise RuntimeError("mount failed") + + with ( + patch.object( + cache, + "_artifact_candidates", + new_callable=AsyncMock, + return_value=[artifact], + ), + patch.object(SquashfsArtifact, "materialize", mock_materialize), + ): + with pytest.raises(RuntimeError, match="mount failed"): + await cache.materialize(cache_key, artifact_uri) + + assert cache._paths_for(cache_key).squashfs_image_path.is_file() + assert cache._budget_dirty is True + @pytest.mark.anyio async def test_release_keeps_retrying_while_the_cache_stays_over_budget( self, temp_cache_dir diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index caaa6c223d..440b7d30cd 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -977,12 +977,12 @@ async def materialize(self, cache_key: str, artifact_uri: str) -> list[Path]: candidate=index + 1, candidates=len(candidates), ) - registry_paths = await artifact.materialize(ctx) if _is_cache_entry_uri(artifact.uri): - # A new entry landed on disk after the budget was - # measured, so the cache must be re-checked once the - # entry goes idle. + # Any attempt may deposit the canonical image before + # failing or being cancelled, so the budget must be + # re-checked once the entry goes idle. self._budget_dirty = True + registry_paths = await artifact.materialize(ctx) return registry_paths except Exception as e: if index == len(candidates) - 1: From 88b72943d28eb1de4af827993b3c88c8569db925 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:37:39 -0400 Subject: [PATCH 016/161] fix(executor): persist lru touches for tarball-only cache entries Lease touches refreshed only the SquashFS image mtime, so entries backed solely by a tarball directory kept recency in memory alone. After a restart, measurement fell back to extraction ctime and startup trimming could evict a frequently used older tarball ahead of a newer unused one. Touches now also bump the tarball root mtime and measurement prefers the newest persisted mtime over the ctime fallback. ENG-1568 --- tests/unit/test_registry_artifacts.py | 40 +++++++++++++++++++++++- tracecat/executor/registry_artifacts.py | 41 ++++++++++++++++++------- 2 files changed, 69 insertions(+), 12 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index f28d6d7264..bd28280185 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -788,6 +788,17 @@ async def test_lock_for_different_keys(self, temp_cache_dir): class TestRegistryArtifactCacheLease: """Tests for lease-based pinning of registry artifact cache entries.""" + def test_touch_entry_refreshes_tarball_root_mtime(self, temp_cache_dir): + """Touching a tarball-only entry persists its restart-safe recency.""" + cache = RegistryArtifactCache(temp_cache_dir) + cache_key = "tarball-only" + target_dir = _write_tarball_entry(temp_cache_dir, cache_key) + os.utime(target_dir, (100.0, 100.0)) + + cache._touch_entry(cache_key) + + assert target_dir.stat().st_mtime > 100.0 + @pytest.mark.anyio async def test_lease_refcounts_and_touches_image_mtime(self, temp_cache_dir): """A lease pins its entry and refreshes the restart-safe LRU timestamp.""" @@ -1222,7 +1233,8 @@ async def test_enforce_budget_evicts_least_recently_used_until_under_max_bytes( cache = RegistryArtifactCache(temp_cache_dir) oldest = _write_image_entry(temp_cache_dir, "oldest", size=4096, mtime=100.0) if oldest_has_tarball: - _write_tarball_entry(temp_cache_dir, "oldest") + oldest_tarball = _write_tarball_entry(temp_cache_dir, "oldest") + os.utime(oldest_tarball, (100.0, 100.0)) older = _write_image_entry(temp_cache_dir, "older", size=4096, mtime=200.0) newest = _write_image_entry(temp_cache_dir, "newest", size=4096, mtime=300.0) @@ -1274,6 +1286,7 @@ async def test_auto_non_pool_backend_evicts_tarball_lru(self, temp_cache_dir): temp_cache_dir, "oldest", size=16, mtime=100.0 ) oldest_tarball = _write_tarball_entry(temp_cache_dir, "oldest") + os.utime(oldest_tarball, (100.0, 100.0)) newest = _write_image_entry(temp_cache_dir, "newest", size=16, mtime=200.0) with ( @@ -1718,6 +1731,31 @@ async def test_eviction_skips_busy_key(self, temp_cache_dir): class TestRegistryArtifactCacheStartupSweep: """Tests for the startup sweep that reclaims state from a dead process.""" + @pytest.mark.anyio + async def test_sweep_uses_tarball_root_mtime_for_restart_safe_lru( + self, temp_cache_dir + ): + """Startup trimming preserves a touched older tarball-only entry.""" + previous_cache = RegistryArtifactCache(temp_cache_dir) + old = _write_tarball_entry(temp_cache_dir, "old") + os.utime(old, (100.0, 100.0)) + previous_cache._touch_entry("old") + + await asyncio.sleep(0.01) + new = _write_tarball_entry(temp_cache_dir, "new") + os.utime(new, (200.0, 200.0)) + + with ( + patch(BACKEND_CONFIG, ExecutorBackendType.DIRECT.value), + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + + assert old.is_dir() + assert not new.exists() + @pytest.mark.anyio async def test_sweep_tolerates_missing_cache_dir(self, temp_cache_dir): """A cache directory that does not exist yet is a no-op.""" diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 440b7d30cd..a6a4aa052f 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -1029,7 +1029,7 @@ def _acquire_lease(self, cache_key: str) -> None: lease = self._leases.setdefault(cache_key, RegistryArtifactLease()) lease.refcount += 1 lease.last_used = time.time() - self._touch_image(cache_key) + self._touch_entry(cache_key) def _release_lease(self, cache_key: str) -> None: """Release one pin on a cache entry.""" @@ -1044,18 +1044,31 @@ def _refcount(self, cache_key: str) -> int: lease = self._leases.get(cache_key) return 0 if lease is None else lease.refcount - def _touch_image(self, cache_key: str) -> None: - """Best-effort refresh of an artifact image mtime for restart-safe LRU. + def _touch_entry(self, cache_key: str) -> None: + """Best-effort refresh of artifact entry mtimes for restart-safe LRU. - Only the downloaded image file has a locally meaningful mtime; mount and - extraction directory timestamps come from the artifact build. + Both the image mtime and tarball root mtime are restart-safe recency + signals. """ - image_path = self._paths_for(cache_key).squashfs_image_path + paths = self._paths_for(cache_key) + touched = False try: - os.utime(image_path) + os.utime(paths.squashfs_image_path) except OSError: + pass + else: + touched = True + + try: + os.utime(paths.tarball_target_dir) + except OSError: + pass + else: + touched = True + + if not touched: logger.debug( - "Could not refresh registry artifact image mtime", + "Could not refresh registry artifact entry mtimes", cache_key=cache_key, ) @@ -1461,6 +1474,7 @@ def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: paths = self._paths_for(cache_key) size_bytes = 0 image_mtime = 0.0 + tarball_root_mtime = 0.0 created_at = 0.0 try: @@ -1471,14 +1485,19 @@ def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: size_bytes += image_stat.st_size image_mtime = image_stat.st_mtime + try: + tarball_root_mtime = paths.tarball_target_dir.stat().st_mtime + except OSError: + pass + for directory in (paths.squashfs_extract_dir, paths.tarball_target_dir): directory_bytes, directory_created_at = _directory_footprint(directory) size_bytes += directory_bytes created_at = max(created_at, directory_created_at) - # Image mtimes are refreshed on lease; directory ctimes are only a - # fallback for entries that have no locally downloaded image. - last_used = image_mtime or created_at + # Image and tarball-root mtimes are refreshed on lease; directory ctimes + # remain the fallback for entries that have neither. + last_used = max(image_mtime, tarball_root_mtime) or created_at return RegistryArtifactCacheEntry( cache_key=cache_key, From 73370cb7a25519821380e1f9d5772db119905405 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:53:54 -0400 Subject: [PATCH 017/161] fix(executor): arm budget dirty flag after materialization completes Pre-arming the flag before the awaited materialization left a window where a concurrent lease release consumed the signal and finished its scan before the new entry landed, so the deposit never triggered convergence. The flag is now armed in a finally after the attempt completes, covering success, failure, and cancellation as well as mid-flight consumption. ENG-1568 --- tests/unit/test_registry_artifacts.py | 33 +++++++++++++++++++++++++ tracecat/executor/registry_artifacts.py | 16 +++++++----- 2 files changed, 43 insertions(+), 6 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index bd28280185..ac87c253a0 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1116,6 +1116,39 @@ async def mock_materialize(self, ctx): assert cache._paths_for(cache_key).squashfs_image_path.is_file() assert cache._budget_dirty is True + @pytest.mark.anyio + async def test_materialization_rearms_budget_dirty_consumed_mid_flight( + self, temp_cache_dir + ): + """A convergence pass consuming the signal mid-download cannot unarm it.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + assert cache._budget_dirty is False + + artifact_uri = "s3://bucket/path/site-packages.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) + + async def mock_materialize(self, ctx): + # A concurrent lease release consumes the dirty signal and finishes + # its scan before this attempt lands its canonical bytes. + cache._budget_dirty = False + ctx.paths.squashfs_image_path.write_bytes(b"late image") + return [ctx.paths.squashfs_mount_dir] + + with ( + patch.object( + cache, + "_artifact_candidates", + new_callable=AsyncMock, + return_value=[artifact], + ), + patch.object(SquashfsArtifact, "materialize", mock_materialize), + ): + await cache.materialize(cache_key, artifact_uri) + + assert cache._budget_dirty is True + @pytest.mark.anyio async def test_release_keeps_retrying_while_the_cache_stays_over_budget( self, temp_cache_dir diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index a6a4aa052f..e964d9e712 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -977,12 +977,16 @@ async def materialize(self, cache_key: str, artifact_uri: str) -> list[Path]: candidate=index + 1, candidates=len(candidates), ) - if _is_cache_entry_uri(artifact.uri): - # Any attempt may deposit the canonical image before - # failing or being cancelled, so the budget must be - # re-checked once the entry goes idle. - self._budget_dirty = True - registry_paths = await artifact.materialize(ctx) + try: + registry_paths = await artifact.materialize(ctx) + finally: + if _is_cache_entry_uri(artifact.uri): + # Arm after the attempt completes, whether it + # succeeded, failed, or was cancelled: any attempt + # may deposit canonical bytes, and a concurrent + # convergence pass may have consumed an earlier + # signal before those bytes landed. + self._budget_dirty = True return registry_paths except Exception as e: if index == len(candidates) - 1: From 73bff1f561f790fb5e613527bf82b15fce009022 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:14:37 -0400 Subject: [PATCH 018/161] fix(executor): stop extractors before cancellation unwinds Cancelling a materialization mid-extraction left the extractor running while cleanup removed its scratch directory: unsquashfs kept a live subprocess and the tarball unpack kept a live thread, either of which could recreate scratch that startup-sweep discovery ignores, so repeated cancellations leaked unbounded ephemeral storage. unsquashfs is now killed and reaped, and the tarball thread is rejoined via a shielded task, before cancellation propagates. ENG-1568 --- tests/unit/test_registry_artifacts.py | 134 ++++++++++++++++++++++++ tracecat/executor/registry_artifacts.py | 35 ++++++- 2 files changed, 165 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index ac87c253a0..01130135f1 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -91,6 +91,35 @@ async def wait(self) -> int: return -9 +class _CapturedSubprocess: + """Capture cleanup of a real subprocess used by cancellation tests.""" + + def __init__(self, process: asyncio.subprocess.Process) -> None: + self.process = process + self.killed = False + self.reaped = False + + @property + def returncode(self) -> int | None: + """Return the wrapped subprocess exit status.""" + return self.process.returncode + + async def communicate(self) -> tuple[bytes, bytes]: + """Wait for the wrapped subprocess and collect its output.""" + return await self.process.communicate() + + def kill(self) -> None: + """Kill the wrapped subprocess and record the signal.""" + self.killed = True + self.process.kill() + + async def wait(self) -> int: + """Reap the wrapped subprocess and record completion.""" + returncode = await self.process.wait() + self.reaped = True + return returncode + + @pytest.fixture def temp_cache_dir(): """Create a temporary cache directory.""" @@ -571,6 +600,111 @@ async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): assert target_dir.is_dir() assert not target_dir.is_mount() + @pytest.mark.anyio + async def test_cancelled_squashfs_extract_kills_and_reaps_subprocess( + self, temp_cache_dir + ): + """Cancellation cannot leave unsquashfs writing into scratch.""" + cache_key = "cancelled-extract" + cache = RegistryArtifactCache(temp_cache_dir) + ctx = cache._context_for(cache_key) + artifact = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key=cache_key, + ) + image_path = ctx.paths.squashfs_image_path + target_dir = ctx.paths.squashfs_extract_dir + image_path.write_bytes(b"squashfs") + target_dir.mkdir() + real_create_subprocess_exec = asyncio.create_subprocess_exec + process_started = asyncio.Event() + captured_processes: list[_CapturedSubprocess] = [] + + async def create_sleep_subprocess( + *args: object, **kwargs: object + ) -> _CapturedSubprocess: + del args, kwargs + process = await real_create_subprocess_exec( + "/bin/sleep", + "30", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + captured = _CapturedSubprocess(process) + captured_processes.append(captured) + process_started.set() + return captured + + with ( + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/usr/bin/unsquashfs", + ), + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + side_effect=create_sleep_subprocess, + ), + ): + extracting = asyncio.create_task( + artifact._extract_image(image_path, target_dir) + ) + await process_started.wait() + captured = captured_processes[0] + extracting.cancel() + + try: + with pytest.raises(asyncio.CancelledError): + await extracting + finally: + if captured.returncode is None: + captured.kill() + await captured.wait() + + assert captured.killed is True + assert captured.reaped is True + assert captured.returncode is not None + + @pytest.mark.anyio + async def test_cancelled_tarball_extract_rejoins_thread(self, temp_cache_dir): + """Cancellation waits until the tar extractor stops writing.""" + artifact = TarballArtifact( + uri="s3://bucket/path/site-packages.tar.gz", + cache_key="cancelled-tarball-extract", + ) + tarball_path = temp_cache_dir / "artifact.tar.gz" + target_dir = temp_cache_dir / "target" + target_dir.mkdir() + with tarfile.open(tarball_path, "w:gz"): + pass + + extraction_started = threading.Event() + extraction_release = threading.Event() + extraction_finished = threading.Event() + + def blocking_extractall(*args: object, **kwargs: object) -> None: + del args, kwargs + extraction_started.set() + extraction_release.wait() + extraction_finished.set() + + with patch.object( + tarfile.TarFile, + "extractall", + new=blocking_extractall, + ): + extracting = asyncio.create_task(artifact.extract(tarball_path, target_dir)) + assert await asyncio.to_thread(extraction_started.wait, 1) + extracting.cancel() + done, _ = await asyncio.wait({extracting}, timeout=0.05) + cancellation_propagated_early = bool(done) + extraction_release.set() + + with pytest.raises(asyncio.CancelledError): + await extracting + + assert extraction_finished.is_set() + assert cancellation_propagated_early is False + @pytest.mark.anyio async def test_materialize_extracts_squashfs_when_mount_fails(self, temp_cache_dir): """Test that SquashFS mount failures fall back to unsquashfs extraction.""" diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index e964d9e712..67b658a803 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -498,7 +498,12 @@ async def _mount_image(self, image_path: Path, target_dir: Path) -> None: raise SquashfsMountCommandError(output or "mount command failed") async def _extract_image(self, image_path: Path, target_dir: Path) -> None: - """Extract a SquashFS image to target_dir using unsquashfs.""" + """Extract a SquashFS image to target_dir using unsquashfs. + + Cancellation kills and reaps the extractor before the caller removes + its scratch directory. Otherwise a live extractor could recreate + scratch after cleanup, outside startup-sweep discovery. + """ unsquashfs = shutil.which("unsquashfs") if unsquashfs is None: raise RuntimeError("unsquashfs command is not installed") @@ -512,7 +517,14 @@ async def _extract_image(self, image_path: Path, target_dir: Path) -> None: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - stdout, stderr = await proc.communicate() + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + proc.kill() + await proc.wait() + raise + if proc.returncode == 0: return @@ -604,7 +616,12 @@ async def download( await _download_s3_artifact(self.uri, output_path) async def extract(self, tarball_path: Path, target_dir: Path) -> None: - """Extract a supported registry tarball to target directory.""" + """Extract a supported registry tarball to target directory. + + A cancelled caller rejoins the non-interruptible extraction thread + before propagating cancellation so scratch cleanup cannot race a live + writer and leave undiscoverable ephemeral storage behind. + """ def _do_extract() -> None: if tarball_path.name.endswith(".tar.gz"): @@ -614,7 +631,17 @@ def _do_extract() -> None: raise ValueError(f"Unsupported tarball format: {tarball_path}") - await asyncio.to_thread(_do_extract) + extraction = asyncio.ensure_future(asyncio.to_thread(_do_extract)) + try: + await asyncio.shield(extraction) + except asyncio.CancelledError: + # A thread cannot be killed. Rejoin it before materialize removes + # scratch; a second cancellation may interrupt this best-effort join. + if not extraction.cancelled(): + with contextlib.suppress(Exception): + await asyncio.shield(extraction) + raise + logger.debug( "Tarball extracted", target=str(target_dir), From 71f08be6f3d2f4932bc67fbccdb734ee9488e4e9 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:15:30 -0400 Subject: [PATCH 019/161] fix(executor): single-flight the startup sweep across cancellation Cancelling ensure_swept mid-sweep released the guard lock with the sweep thread still running and the swept flag unset, so a later caller started a second concurrent sweep while the abandoned one kept deleting from a stale snapshot that consults no leases. All callers now join one stored, shielded sweep task: a cancelled waiter leaves the live sweep in place, a completed sweep is reused, and a failed sweep is cleared for retry. ENG-1568 --- tests/unit/test_registry_artifacts.py | 34 +++++++++++++++++++++++ tracecat/executor/registry_artifacts.py | 36 ++++++++++++++++++++----- 2 files changed, 63 insertions(+), 7 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 01130135f1..7ae709c722 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -2024,6 +2024,40 @@ async def test_concurrent_ensure_swept_runs_once(self, temp_cache_dir): assert sweep.call_count == 1 + @pytest.mark.anyio + async def test_cancelled_waiter_joins_same_startup_sweep(self, temp_cache_dir): + """A cancelled waiter cannot start a second concurrent startup sweep.""" + cache = RegistryArtifactCache(temp_cache_dir) + sweep_started = threading.Event() + sweep_release = threading.Event() + invocation_count = 0 + + def blocking_sweep() -> None: + nonlocal invocation_count + invocation_count += 1 + sweep_started.set() + sweep_release.wait() + + with patch.object( + cache, + "_sweep_startup_state", + side_effect=blocking_sweep, + ): + first_waiter = asyncio.create_task(cache.ensure_swept()) + assert await asyncio.to_thread(sweep_started.wait, 1) + first_waiter.cancel() + + with pytest.raises(asyncio.CancelledError): + await first_waiter + + second_waiter = asyncio.create_task(cache.ensure_swept()) + await asyncio.sleep(0) + sweep_release.set() + await second_waiter + + assert invocation_count == 1 + assert cache._swept is True + @pytest.mark.anyio async def test_lease_triggers_startup_sweep(self, temp_cache_dir): """Lease admission reclaims startup scratch before yielding paths.""" diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 67b658a803..6d5285dbee 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -866,6 +866,7 @@ def __init__(self, cache_dir: Path): self._budget_lock = asyncio.Lock() # Guard the off-loop startup sweep independently from cache operations. self._swept: bool = False + self._sweep_task: asyncio.Task[None] | None = None self._sweep_lock = asyncio.Lock() self._squashfs_mount_state = SquashfsMountState() self._leases: dict[str, RegistryArtifactLease] = {} @@ -875,19 +876,40 @@ def __init__(self, cache_dir: Path): self._budget_dirty = True async def ensure_swept(self) -> None: - """Run the startup sweep exactly once, off the event loop. - - Idempotent and safe under concurrency: the first caller performs the - sweep in a thread; concurrent callers wait; later callers return - immediately. Runs before the first lease or materialization so the - sweep never observes in-flight cache entries. + """Run the startup sweep exactly once successfully, off the event loop. + + Idempotent and cancellation-safe under concurrency: every caller joins + one stored sweep task, and cancelling a waiter never abandons or + restarts its live sweep. The lock is deliberately held while awaiting + that shared task so queued callers observe its result before proceeding. + Failures clear the task so the next caller retries. The sweep runs + before the first lease or materialization, so it never observes + in-flight cache entries. """ if self._swept: return async with self._sweep_lock: if self._swept: return - await asyncio.to_thread(self._sweep_startup_state) + sweep_task = self._sweep_task + if sweep_task is None or ( + sweep_task.done() + and (sweep_task.cancelled() or sweep_task.exception() is not None) + ): + sweep_task = asyncio.ensure_future( + asyncio.to_thread(self._sweep_startup_state) + ) + self._sweep_task = sweep_task + try: + await asyncio.shield(sweep_task) + except asyncio.CancelledError: + if sweep_task.cancelled() and self._sweep_task is sweep_task: + self._sweep_task = None + raise + except Exception: + if self._sweep_task is sweep_task: + self._sweep_task = None + raise self._swept = True @asynccontextmanager From 7f1615e2786896ca47ce3d698cc400f9d0f3e8ce Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 17:05:15 -0400 Subject: [PATCH 020/161] fix(executor): retry failed registry cache sweep --- tests/unit/test_registry_artifacts.py | 3 ++- tracecat/executor/registry_artifacts.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 7ae709c722..408a06608d 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1034,6 +1034,7 @@ async def test_lease_is_never_admitted_across_an_in_flight_eviction( ): """A lease must not return a mount an in-flight eviction is deleting.""" cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() artifact_uri = "s3://bucket/path/site-packages.squashfs" cache_key = compute_registry_artifact_cache_key(artifact_uri) paths = cache._paths_for(cache_key) @@ -2079,7 +2080,7 @@ async def test_failed_ensure_swept_retries(self, temp_cache_dir): with ( patch.object( cache, - "_sweep_startup_state", + "_remove_orphaned_temp_paths", side_effect=OSError("simulated sweep failure"), ), pytest.raises(OSError), diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 6d5285dbee..0928a7ad60 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -1587,6 +1587,7 @@ def _sweep_startup_state(self) -> None: cache_dir=str(self.cache_dir), error=str(e), ) + raise def _remove_orphaned_temp_paths(self) -> None: """Delete every materialization scratch path during startup. From 5f2da98965d4333827de232344c120b92b5316b1 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 18:09:39 -0400 Subject: [PATCH 021/161] fix(executor): retain images when reclaiming mounts --- ...registry_artifact_cache_mount_lifecycle.py | 36 +++++++++-- tests/unit/test_registry_artifacts.py | 13 ++-- tracecat/config.py | 5 +- tracecat/executor/registry_artifacts.py | 59 ++++++++++++++++--- 4 files changed, 92 insertions(+), 21 deletions(-) diff --git a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py index 6a40456bcf..be4eda9b66 100644 --- a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py +++ b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py @@ -169,7 +169,14 @@ def test_registry_artifact_cache_mount_lifecycle() -> None: assert payload["remounted"] is True assert payload["squashfs_disabled_after_eviction"] is False - # Releasing a lease converges an over-budget cache, unmounting as it goes. + # Loop-device reclamation unmounts an idle entry but retains its cached image. + assert payload["mount_slot_released"] is True + assert payload["released_target_unmounted"] is True + assert payload["released_loop_device_released"] is True + assert payload["released_image_retained"] is True + assert payload["remounted_from_retained_image"] is True + + # Releasing a lease converges an over-budget cache, destructively evicting it. assert payload["converged_entry_unmounted"] is True assert payload["converged_loop_device_released"] is True assert payload["converged_entries_remaining"] == 3 @@ -307,9 +314,28 @@ async def _run_mount_lifecycle_child() -> None: cache._squashfs_mount_state.disabled ) - # (d) A released lease converges an over-budget cache, unmounting as it - # goes. The fourth entry is materialized under the old budget, so only - # the release-time check can bring the cache back within the new one. + # (d) Loop-device reclamation unmounts the LRU idle artifact while + # retaining its image, and a later lease remounts that image directly. + retained_paths = cache._paths_for(keys[1]) + mounts_before_release = _squashfs_mounts(cache_dir) + retained_device = mounts_before_release[str(retained_paths.squashfs_mount_dir)] + payload["mount_slot_released"] = await cache._release_mounted_slot(keys[0]) + mounts_after_release = _squashfs_mounts(cache_dir) + payload["released_target_unmounted"] = ( + str(retained_paths.squashfs_mount_dir) not in mounts_after_release + ) + payload["released_loop_device_released"] = ( + retained_device not in mounts_after_release.values() + ) + payload["released_image_retained"] = retained_paths.squashfs_image_path.exists() + async with cache.lease([uris[1]]) as registry_paths: + payload["remounted_from_retained_image"] = registry_paths == [ + retained_paths.squashfs_mount_dir + ] + + # (e) A released lease converges an over-budget cache, destructively + # evicting it. The fourth entry is materialized under the old budget, so + # only the release-time check can bring the cache back within the new one. fourth_uri = "s3://bucket/lifecycle/3/site-packages.squashfs" fourth_key = compute_registry_artifact_cache_key(fourth_uri) _build_squashfs_image( @@ -337,7 +363,7 @@ async def _run_mount_lifecycle_child() -> None: ) payload["converged_entries_remaining"] = len(cache._discover_cache_keys()) - # (e) The startup sweep trims to budget and drops stale mount directories. + # (f) The startup sweep trims to budget and drops stale mount directories. sweep_dir = root / "sweep-cache" sweep_dir.mkdir() sweep_keys = ["aaaa1111", "bbbb2222"] diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 408a06608d..956f8354e6 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1531,15 +1531,15 @@ async def test_pool_backend_release_mounted_slot_skips_tarball_entry( patch.object(Path, "is_mount", lambda self: self in mounted), patch.object( cache, - "_evict_entry", + "_unmount_entry", new_callable=AsyncMock, return_value=True, - ) as evict_entry, + ) as unmount_entry, ): released = await cache._release_mounted_slot("protected") assert released is True - evict_entry.assert_awaited_once_with("eligible") + unmount_entry.assert_awaited_once_with("eligible") @pytest.mark.anyio async def test_concurrent_budget_passes_only_evict_once(self, temp_cache_dir): @@ -2267,7 +2267,7 @@ async def mock_extract(self, ctx, image_path): async def test_mount_failure_after_success_reclaims_loop_device_and_retries( self, temp_cache_dir ): - """Loop-device exhaustion evicts an idle mount instead of going sticky.""" + """Loop-device exhaustion unmounts an idle artifact without deleting it.""" cache = RegistryArtifactCache(temp_cache_dir) cache._squashfs_mount_state.mounted_once = True idle = cache._paths_for("idle") @@ -2315,8 +2315,9 @@ async def mock_umount(*args, **kwargs): assert (result[0] / "module.py").read_text() == "VALUE = 1" assert attempts == ["new", "new"] assert cache._squashfs_mount_state.disabled is False - assert not idle.squashfs_image_path.exists() - assert not idle.squashfs_mount_dir.exists() + assert idle.squashfs_image_path.read_bytes() == b"squashfs" + assert idle.squashfs_mount_dir.is_dir() + assert idle.squashfs_mount_dir not in mounted @pytest.mark.anyio async def test_download_failure_does_not_disable_squashfs_process_wide( diff --git a/tracecat/config.py b/tracecat/config.py index 45bb35e548..27d5e548a6 100644 --- a/tracecat/config.py +++ b/tracecat/config.py @@ -191,12 +191,11 @@ class RLSMode(StrEnum): """Prefer SquashFS registry artifacts when sidecars and mount support are available.""" TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES = int( - os.environ.get("TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES") or 8 + os.environ.get("TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES") or 64 ) """Maximum number of registry artifacts kept in the executor-local cache. -Each mounted SquashFS artifact pins one loop device, so this also bounds loop -device usage. Set to 0 to disable entry-count eviction.""" +Set to 0 to disable entry-count eviction.""" TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES = int( os.environ.get("TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES") or 10 * 1024**3 diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 0928a7ad60..21b6f7a3ac 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -60,7 +60,7 @@ class RegistryArtifactFormat(StrEnum): """Matches materialization scratch and doomed eviction paths.""" type MountSlotReleaser = Callable[[str], Awaitable[bool]] -"""Evicts one idle mounted artifact, excluding the given cache key.""" +"""Unmounts one idle artifact, excluding the given cache key.""" class SquashfsMountCommandError(RuntimeError): @@ -138,7 +138,7 @@ def has_mounted_squashfs(self) -> bool: return self.squashfs_mount_state.mounted_once async def release_mounted_slot(self) -> bool: - """Evict one idle mounted artifact to free a loop device.""" + """Unmount one idle artifact to free a loop device.""" if self.mount_slot_releaser is None: return False return await self.mount_slot_releaser(self.cache_key) @@ -250,7 +250,7 @@ async def _try_mount( The first mount-command failure in a process is treated as a capability probe and disables mounting process-wide. Once any mount has succeeded, later failures are attributed to exhausted loop devices instead: one idle - mounted artifact is evicted and the mount is retried once, so a single + artifact is unmounted and the mount is retried once, so a single failure never downgrades the whole process to extraction. The probe lock deliberately covers the first mount's download and mount @@ -259,7 +259,7 @@ async def _try_mount( Only ``SquashfsMountCommandError`` drives this policy. Download and preparation errors propagate to the caller so a transient S3 failure - never disables mounting or evicts an unrelated idle mount. + never disables mounting or unmounts an unrelated idle artifact. Args: ctx: Materialization context for the artifact being mounted. @@ -1335,7 +1335,7 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo return True async def _release_mounted_slot(self, protected_key: str) -> bool: - """Evict one idle mounted artifact so its loop device can be reused. + """Unmount one idle artifact so its loop device can be reused. This path deliberately does not take the budget lock: ``_try_mount`` may call it while holding ``protected_key``'s per-key lock. It excludes that @@ -1346,7 +1346,7 @@ async def _release_mounted_slot(self, protected_key: str) -> bool: protected_key: Cache key that must not be evicted. Returns: - Whether a mounted artifact was unmounted and removed. + Whether a mounted artifact was unmounted. """ entries = await asyncio.to_thread(self._scan_cache_entries) mounted = [ @@ -1359,7 +1359,7 @@ async def _release_mounted_slot(self, protected_key: str) -> bool: while ( candidate := self._least_recently_used(mounted, excluded=skipped) ) is not None: - if await self._evict_entry(candidate.cache_key): + if await self._unmount_entry(candidate.cache_key): return True skipped.add(candidate.cache_key) return False @@ -1407,6 +1407,51 @@ def _recency(self, entry: RegistryArtifactCacheEntry) -> float: return entry.last_used return max(entry.last_used, lease.last_used) + async def _unmount_entry(self, cache_key: str) -> bool: + """Unmount one idle cache entry while retaining its reusable image. + + Loop-device reclamation is independent from disk-budget eviction. The + per-key lock and lease recheck prevent an entry from being unmounted + while an action is importing from it. The image and empty mount + directory remain cached so a later admission can remount without + downloading the artifact again. + + Args: + cache_key: Cache key whose mounted artifact should be released. + + Returns: + Whether a mounted entry was unmounted. + """ + lock = await self._lock_for(cache_key) + if lock.locked(): + logger.debug( + "Skipping unmount of busy registry artifact", + cache_key=cache_key, + ) + return False + + async with lock: + if self._refcount(cache_key) > 0: + return False + + mount_dir = self._paths_for(cache_key).squashfs_mount_dir + if not mount_dir.is_mount(): + return False + if not await self._unmount(mount_dir): + logger.warning( + "Failed to unmount registry artifact for loop-device reclamation", + cache_key=cache_key, + mount_dir=str(mount_dir), + ) + return False + + logger.info( + "Unmounted idle registry artifact", + cache_key=cache_key, + mount_dir=str(mount_dir), + ) + return True + async def _evict_entry(self, cache_key: str) -> bool: """Remove one cache entry from disk, unmounting it first. From a2ad4b47edf9078bbcf34a5dd153fbdffc5c01bc Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 18:44:08 -0400 Subject: [PATCH 022/161] fix(executor): retry failed cache deletions --- tests/unit/test_registry_artifacts.py | 93 ++++++++++++++- tracecat/executor/registry_artifacts.py | 149 +++++++++++++++++++++--- 2 files changed, 221 insertions(+), 21 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 956f8354e6..400b06e56e 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1135,6 +1135,33 @@ async def test_builtin_artifact_is_exempt_from_cache_accounting( class TestRegistryArtifactCacheEviction: """Tests for bounded eviction of registry artifact cache entries.""" + def test_delete_entry_paths_reports_directory_failure(self, temp_cache_dir): + """Physical deletion failures are observable instead of suppressed.""" + paths = RegistryArtifactPaths( + squashfs_image_path=temp_cache_dir / "image.squashfs", + squashfs_mount_dir=temp_cache_dir / "mount", + squashfs_extract_dir=temp_cache_dir / "extract", + tarball_target_dir=temp_cache_dir / "tarball", + ) + paths.squashfs_extract_dir.mkdir() + + with ( + patch( + "tracecat.executor.registry_artifacts.shutil.rmtree", + side_effect=OSError("permission denied"), + ), + patch("tracecat.executor.registry_artifacts.logger.warning") as warning, + ): + deleted = _delete_entry_paths(paths) + + assert deleted is False + assert paths.squashfs_extract_dir.is_dir() + warning.assert_called_once_with( + "Failed to delete registry artifact cache path", + path=str(paths.squashfs_extract_dir), + error="permission denied", + ) + @pytest.mark.anyio async def test_leased_entry_survives_eviction_of_idle_entry(self, temp_cache_dir): """Eviction must never remove an entry a live action is importing from.""" @@ -1199,6 +1226,64 @@ async def mock_extract(self, tarball_path, target_dir): assert (temp_cache_dir / f"tarball-{new_key}").is_dir() assert cache._budget_dirty is False + @pytest.mark.anyio + async def test_deletion_failure_retries_without_blocking_materialization( + self, temp_cache_dir + ): + """Failed cleanup stays retryable while cache admission remains fail-open.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + _write_image_entry(temp_cache_dir, "idle", size=4096, mtime=100.0) + artifact_uri = "s3://bucket/new.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + delete_attempts = 0 + retry_started = threading.Event() + finish_retry = threading.Event() + real_delete_entry_paths = _delete_entry_paths + + def fail_once_then_delete(paths: RegistryArtifactPaths) -> bool: + nonlocal delete_attempts + delete_attempts += 1 + if delete_attempts == 1: + return False + retry_started.set() + finish_retry.wait(timeout=5) + return real_delete_entry_paths(paths) + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifacts._delete_entry_paths", + side_effect=fail_once_then_delete, + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + registry_paths = await cache.materialize(cache_key, artifact_uri) + + assert registry_paths == [temp_cache_dir / f"tarball-{cache_key}"] + assert registry_paths[0].is_dir() + assert cache._pending_deletions + assert cache._budget_dirty is True + + within_budget = await cache._enforce_cache_budget() + assert within_budget is False + assert await asyncio.to_thread(retry_started.wait, 1) + retry_tasks = tuple(cache._deletion_tasks.values()) + finish_retry.set() + await asyncio.gather(*retry_tasks) + await asyncio.sleep(0) + + assert not cache._pending_deletions + assert await cache._enforce_cache_budget() is True + @pytest.mark.anyio async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( self, temp_cache_dir @@ -1756,12 +1841,13 @@ async def test_cancelled_background_deletion_leaves_a_clean_miss( delete_finished = threading.Event() doomed: list[RegistryArtifactPaths] = [] - def blocked_delete(paths: RegistryArtifactPaths) -> None: + def blocked_delete(paths: RegistryArtifactPaths) -> bool: doomed.append(paths) delete_started.set() finish_delete.wait(timeout=5) - _delete_entry_paths(paths) + deleted = _delete_entry_paths(paths) delete_finished.set() + return deleted async def mock_download(self, ctx, path): path.write_bytes(b"fake tarball") @@ -1814,7 +1900,8 @@ async def test_doomed_eviction_names_are_startup_scratch(self, temp_cache_dir): paths.tarball_target_dir.mkdir() with patch( - "tracecat.executor.registry_artifacts._delete_entry_paths" + "tracecat.executor.registry_artifacts._delete_entry_paths", + return_value=True, ) as delete_entry_paths: assert await cache._evict_entry(cache_key) is True diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 21b6f7a3ac..ce323e6796 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -799,16 +799,43 @@ def _directory_footprint(directory: Path) -> tuple[int, float]: return total_bytes, created_at -def _delete_entry_paths(paths: RegistryArtifactPaths) -> None: - """Delete every on-disk path owned by a registry artifact cache key. +def _delete_cache_path(path: Path) -> bool: + """Best-effort delete one cache path while reporting filesystem failures.""" + try: + if path.is_dir(): + shutil.rmtree(path) + else: + path.unlink(missing_ok=True) + except OSError as e: + logger.warning( + "Failed to delete registry artifact cache path", + path=str(path), + error=str(e), + ) + return False + return True + + +def _delete_entry_paths(paths: RegistryArtifactPaths) -> bool: + """Best-effort delete every path owned by a registry artifact cache key. The caller must unmount ``squashfs_mount_dir`` first: unlinking the image file behind a live mount leaves an open-file zombie pinning a loop device. + + Returns: + Whether every path was deleted. Failures are logged and returned instead + of raised so cache cleanup never rejects executor work. """ - shutil.rmtree(paths.squashfs_extract_dir, ignore_errors=True) - shutil.rmtree(paths.tarball_target_dir, ignore_errors=True) - paths.squashfs_image_path.unlink(missing_ok=True) - shutil.rmtree(paths.squashfs_mount_dir, ignore_errors=True) + deleted = True + for path in ( + paths.squashfs_extract_dir, + paths.tarball_target_dir, + paths.squashfs_image_path, + paths.squashfs_mount_dir, + ): + if not _delete_cache_path(path): + deleted = False + return deleted def _rename_entry_paths( @@ -870,6 +897,11 @@ def __init__(self, cache_dir: Path): self._sweep_lock = asyncio.Lock() self._squashfs_mount_state = SquashfsMountState() self._leases: dict[str, RegistryArtifactLease] = {} + # Eviction renames live paths before deleting them so a cancelled or + # failed cleanup can never expose a partial cache hit. Failed physical + # deletions stay here until a later budget pass retries them. + self._pending_deletions: set[RegistryArtifactPaths] = set() + self._deletion_tasks: dict[RegistryArtifactPaths, asyncio.Task[bool]] = {} # Whether the on-disk cache may exceed its budget. Set when a new entry # is materialized and cleared once enforcement measures a cache that # fits, so steady-state cache hits never pay for a disk scan. @@ -1294,10 +1326,11 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo Whether the cache is within budget once eviction has finished. """ async with self._budget_lock: + self._retry_pending_deletions() max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES if max_entries <= 0 and max_bytes <= 0: - return True + return not self._pending_deletions entries = await asyncio.to_thread(self._scan_cache_entries) # The protected key is not on disk yet when it is a fresh entry. @@ -1332,7 +1365,7 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo else: skipped.add(candidate.cache_key) - return True + return not self._pending_deletions async def _release_mounted_slot(self, protected_key: str) -> bool: """Unmount one idle artifact so its loop device can be reused. @@ -1407,6 +1440,62 @@ def _recency(self, entry: RegistryArtifactCacheEntry) -> float: return entry.last_used return max(entry.last_used, lease.last_used) + def _start_pending_deletion( + self, paths: RegistryArtifactPaths + ) -> asyncio.Task[bool]: + """Start or return the background deletion task for renamed cache paths.""" + if task := self._deletion_tasks.get(paths): + return task + + self._pending_deletions.add(paths) + task = asyncio.create_task(asyncio.to_thread(_delete_entry_paths, paths)) + self._deletion_tasks[paths] = task + + def record_result(completed: asyncio.Task[bool]) -> None: + self._record_pending_deletion_result(paths, completed) + + task.add_done_callback(record_result) + return task + + def _record_pending_deletion_result( + self, + paths: RegistryArtifactPaths, + task: asyncio.Task[bool], + ) -> None: + """Update retry bookkeeping after one physical deletion attempt.""" + if self._deletion_tasks.get(paths) is not task: + return + self._deletion_tasks.pop(paths, None) + try: + deleted = task.result() + except asyncio.CancelledError: + self._budget_dirty = True + return + except Exception as e: + logger.error( + "Registry artifact cache deletion task failed", + paths=[ + str(paths.squashfs_extract_dir), + str(paths.tarball_target_dir), + str(paths.squashfs_image_path), + str(paths.squashfs_mount_dir), + ], + error=str(e), + ) + self._budget_dirty = True + return + + if deleted: + self._pending_deletions.discard(paths) + else: + self._budget_dirty = True + + def _retry_pending_deletions(self) -> None: + """Retry failed physical deletions without blocking cache admission.""" + for paths in tuple(self._pending_deletions): + if paths not in self._deletion_tasks: + self._start_pending_deletion(paths) + async def _unmount_entry(self, cache_key: str) -> bool: """Unmount one idle cache entry while retaining its reusable image. @@ -1501,9 +1590,22 @@ async def _evict_entry(self, cache_key: str) -> bool: cache_key=cache_key, ) self._leases.pop(cache_key, None) - logger.info("Evicted registry artifact from cache", cache_key=cache_key) - await asyncio.to_thread(_delete_entry_paths, doomed_paths) + deletion_task = self._start_pending_deletion(doomed_paths) + try: + deleted = await asyncio.shield(deletion_task) + except asyncio.CancelledError: + self._budget_dirty = True + raise + if not deleted: + self._budget_dirty = True + logger.warning( + "Registry artifact eviction remains pending physical deletion", + cache_key=cache_key, + ) + return False + + logger.info("Evicted registry artifact from cache", cache_key=cache_key) return True async def _unmount(self, mount_dir: Path) -> bool: @@ -1623,9 +1725,11 @@ def _sweep_startup_state(self) -> None: return try: - self._remove_orphaned_temp_paths() + orphan_cleanup_succeeded = self._remove_orphaned_temp_paths() self._remove_stale_mount_dirs() self._trim_startup_cache() + if not orphan_cleanup_succeeded: + self._budget_dirty = True except OSError as e: logger.warning( "Failed to sweep registry artifact cache", @@ -1634,23 +1738,29 @@ def _sweep_startup_state(self) -> None: ) raise - def _remove_orphaned_temp_paths(self) -> None: + def _remove_orphaned_temp_paths(self) -> bool: """Delete every materialization scratch path during startup. The sweep runs before the first lease or materialization in this process. Every matching path is therefore interrupted scratch from an earlier process and is safe to remove even when the operating system reused that process's PID. + + Returns: + Whether every orphaned path was removed. """ + deleted = True for name in os.listdir(self.cache_dir): if TEMP_ARTIFACT_PATTERN.match(name) is None: continue path = self.cache_dir / name - if path.is_dir(): - shutil.rmtree(path, ignore_errors=True) + if _delete_cache_path(path): + logger.info( + "Removed orphaned registry artifact scratch path", path=name + ) else: - path.unlink(missing_ok=True) - logger.info("Removed orphaned registry artifact scratch path", path=name) + deleted = False + return deleted def _remove_stale_mount_dirs(self) -> None: """Remove empty mount directories left over from a previous process.""" @@ -1694,6 +1804,7 @@ def _trim_startup_cache(self) -> None: ), key=lambda entry: entry.last_used, ) + deletion_failed = False def within_budget() -> bool: return (max_entries <= 0 or len(entries) <= max_entries) and ( @@ -1703,7 +1814,9 @@ def within_budget() -> bool: for entry in candidates: if within_budget(): break - _delete_entry_paths(self._paths_for(entry.cache_key)) + if not _delete_entry_paths(self._paths_for(entry.cache_key)): + deletion_failed = True + continue del entries[entry.cache_key] total_bytes -= entry.size_bytes logger.info( @@ -1712,4 +1825,4 @@ def within_budget() -> bool: size_bytes=entry.size_bytes, ) - self._budget_dirty = not within_budget() + self._budget_dirty = deletion_failed or not within_budget() From 7abaf238397a7a71cec9fc3e4529e8bde96617fa Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 18:55:48 -0400 Subject: [PATCH 023/161] fix(executor): keep cache cleanup fail-open --- tests/unit/test_registry_artifacts.py | 42 +++------- tracecat/executor/registry_artifacts.py | 100 +++--------------------- 2 files changed, 22 insertions(+), 120 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 400b06e56e..58c44cb52c 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1227,28 +1227,15 @@ async def mock_extract(self, tarball_path, target_dir): assert cache._budget_dirty is False @pytest.mark.anyio - async def test_deletion_failure_retries_without_blocking_materialization( + async def test_deletion_failure_does_not_block_materialization( self, temp_cache_dir ): - """Failed cleanup stays retryable while cache admission remains fail-open.""" + """Failed cleanup is reported while cache admission remains fail-open.""" cache = RegistryArtifactCache(temp_cache_dir) await cache.ensure_swept() _write_image_entry(temp_cache_dir, "idle", size=4096, mtime=100.0) artifact_uri = "s3://bucket/new.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) - delete_attempts = 0 - retry_started = threading.Event() - finish_retry = threading.Event() - real_delete_entry_paths = _delete_entry_paths - - def fail_once_then_delete(paths: RegistryArtifactPaths) -> bool: - nonlocal delete_attempts - delete_attempts += 1 - if delete_attempts == 1: - return False - retry_started.set() - finish_retry.wait(timeout=5) - return real_delete_entry_paths(paths) async def mock_download(self, ctx, path): path.write_bytes(b"fake tarball") @@ -1261,28 +1248,21 @@ async def mock_extract(self, tarball_path, target_dir): patch(MAX_BYTES_CONFIG, 0), patch( "tracecat.executor.registry_artifacts._delete_entry_paths", - side_effect=fail_once_then_delete, + return_value=False, ), patch.object(TarballArtifact, "download", mock_download), patch.object(TarballArtifact, "extract", mock_extract), + patch("tracecat.executor.registry_artifacts.logger.warning") as warning, ): registry_paths = await cache.materialize(cache_key, artifact_uri) - assert registry_paths == [temp_cache_dir / f"tarball-{cache_key}"] - assert registry_paths[0].is_dir() - assert cache._pending_deletions - assert cache._budget_dirty is True - - within_budget = await cache._enforce_cache_budget() - assert within_budget is False - assert await asyncio.to_thread(retry_started.wait, 1) - retry_tasks = tuple(cache._deletion_tasks.values()) - finish_retry.set() - await asyncio.gather(*retry_tasks) - await asyncio.sleep(0) - - assert not cache._pending_deletions - assert await cache._enforce_cache_budget() is True + assert registry_paths == [temp_cache_dir / f"tarball-{cache_key}"] + assert registry_paths[0].is_dir() + assert cache._budget_dirty is True + warning.assert_any_call( + "Registry artifact eviction remains pending physical deletion", + cache_key="idle", + ) @pytest.mark.anyio async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index ce323e6796..52146a4db7 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -897,11 +897,6 @@ def __init__(self, cache_dir: Path): self._sweep_lock = asyncio.Lock() self._squashfs_mount_state = SquashfsMountState() self._leases: dict[str, RegistryArtifactLease] = {} - # Eviction renames live paths before deleting them so a cancelled or - # failed cleanup can never expose a partial cache hit. Failed physical - # deletions stay here until a later budget pass retries them. - self._pending_deletions: set[RegistryArtifactPaths] = set() - self._deletion_tasks: dict[RegistryArtifactPaths, asyncio.Task[bool]] = {} # Whether the on-disk cache may exceed its budget. Set when a new entry # is materialized and cleared once enforcement measures a cache that # fits, so steady-state cache hits never pay for a disk scan. @@ -1326,11 +1321,10 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo Whether the cache is within budget once eviction has finished. """ async with self._budget_lock: - self._retry_pending_deletions() max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES if max_entries <= 0 and max_bytes <= 0: - return not self._pending_deletions + return True entries = await asyncio.to_thread(self._scan_cache_entries) # The protected key is not on disk yet when it is a fresh entry. @@ -1365,7 +1359,7 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo else: skipped.add(candidate.cache_key) - return not self._pending_deletions + return True async def _release_mounted_slot(self, protected_key: str) -> bool: """Unmount one idle artifact so its loop device can be reused. @@ -1440,62 +1434,6 @@ def _recency(self, entry: RegistryArtifactCacheEntry) -> float: return entry.last_used return max(entry.last_used, lease.last_used) - def _start_pending_deletion( - self, paths: RegistryArtifactPaths - ) -> asyncio.Task[bool]: - """Start or return the background deletion task for renamed cache paths.""" - if task := self._deletion_tasks.get(paths): - return task - - self._pending_deletions.add(paths) - task = asyncio.create_task(asyncio.to_thread(_delete_entry_paths, paths)) - self._deletion_tasks[paths] = task - - def record_result(completed: asyncio.Task[bool]) -> None: - self._record_pending_deletion_result(paths, completed) - - task.add_done_callback(record_result) - return task - - def _record_pending_deletion_result( - self, - paths: RegistryArtifactPaths, - task: asyncio.Task[bool], - ) -> None: - """Update retry bookkeeping after one physical deletion attempt.""" - if self._deletion_tasks.get(paths) is not task: - return - self._deletion_tasks.pop(paths, None) - try: - deleted = task.result() - except asyncio.CancelledError: - self._budget_dirty = True - return - except Exception as e: - logger.error( - "Registry artifact cache deletion task failed", - paths=[ - str(paths.squashfs_extract_dir), - str(paths.tarball_target_dir), - str(paths.squashfs_image_path), - str(paths.squashfs_mount_dir), - ], - error=str(e), - ) - self._budget_dirty = True - return - - if deleted: - self._pending_deletions.discard(paths) - else: - self._budget_dirty = True - - def _retry_pending_deletions(self) -> None: - """Retry failed physical deletions without blocking cache admission.""" - for paths in tuple(self._pending_deletions): - if paths not in self._deletion_tasks: - self._start_pending_deletion(paths) - async def _unmount_entry(self, cache_key: str) -> bool: """Unmount one idle cache entry while retaining its reusable image. @@ -1591,12 +1529,7 @@ async def _evict_entry(self, cache_key: str) -> bool: ) self._leases.pop(cache_key, None) - deletion_task = self._start_pending_deletion(doomed_paths) - try: - deleted = await asyncio.shield(deletion_task) - except asyncio.CancelledError: - self._budget_dirty = True - raise + deleted = await asyncio.to_thread(_delete_entry_paths, doomed_paths) if not deleted: self._budget_dirty = True logger.warning( @@ -1725,11 +1658,9 @@ def _sweep_startup_state(self) -> None: return try: - orphan_cleanup_succeeded = self._remove_orphaned_temp_paths() + self._remove_orphaned_temp_paths() self._remove_stale_mount_dirs() self._trim_startup_cache() - if not orphan_cleanup_succeeded: - self._budget_dirty = True except OSError as e: logger.warning( "Failed to sweep registry artifact cache", @@ -1738,29 +1669,23 @@ def _sweep_startup_state(self) -> None: ) raise - def _remove_orphaned_temp_paths(self) -> bool: + def _remove_orphaned_temp_paths(self) -> None: """Delete every materialization scratch path during startup. The sweep runs before the first lease or materialization in this process. Every matching path is therefore interrupted scratch from an earlier process and is safe to remove even when the operating system reused that process's PID. - - Returns: - Whether every orphaned path was removed. """ - deleted = True for name in os.listdir(self.cache_dir): if TEMP_ARTIFACT_PATTERN.match(name) is None: continue path = self.cache_dir / name - if _delete_cache_path(path): - logger.info( - "Removed orphaned registry artifact scratch path", path=name - ) + if path.is_dir(): + shutil.rmtree(path, ignore_errors=True) else: - deleted = False - return deleted + path.unlink(missing_ok=True) + logger.info("Removed orphaned registry artifact scratch path", path=name) def _remove_stale_mount_dirs(self) -> None: """Remove empty mount directories left over from a previous process.""" @@ -1804,7 +1729,6 @@ def _trim_startup_cache(self) -> None: ), key=lambda entry: entry.last_used, ) - deletion_failed = False def within_budget() -> bool: return (max_entries <= 0 or len(entries) <= max_entries) and ( @@ -1814,9 +1738,7 @@ def within_budget() -> bool: for entry in candidates: if within_budget(): break - if not _delete_entry_paths(self._paths_for(entry.cache_key)): - deletion_failed = True - continue + _delete_entry_paths(self._paths_for(entry.cache_key)) del entries[entry.cache_key] total_bytes -= entry.size_bytes logger.info( @@ -1825,4 +1747,4 @@ def within_budget() -> bool: size_bytes=entry.size_bytes, ) - self._budget_dirty = deletion_failed or not within_budget() + self._budget_dirty = not within_budget() From 9542b3d152725cb6bc5f7631fbf24d8a7301310a Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Thu, 30 Jul 2026 12:55:50 -0400 Subject: [PATCH 024/161] refactor(executor): make registry cache entries atomic --- tests/integration/conftest.py | 4 +- tests/integration/test_pool_integration.py | 11 +- ...registry_artifact_cache_mount_lifecycle.py | 21 +- tests/unit/test_action_runner.py | 4 +- tests/unit/test_executor_sandbox_nsjail.py | 9 +- tests/unit/test_multitenant_registry.py | 35 +- tests/unit/test_registry_artifacts.py | 474 +++++++++++---- tracecat/executor/backends/pool/worker.py | 33 +- tracecat/executor/registry_artifacts.py | 575 +++++++++--------- 9 files changed, 726 insertions(+), 440 deletions(-) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index ab9ea8d1d0..4da6c5baec 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -275,8 +275,8 @@ def staged_cache_dirs( Returns tuple of (path_a, path_b) where each path contains the extracted mock modules for that workspace. """ - path_a = temp_registry_cache / "tarball-workspace-a" - path_b = temp_registry_cache / "tarball-workspace-b" + path_a = temp_registry_cache / "entries" / "workspace-a" / "tarball" + path_b = temp_registry_cache / "entries" / "workspace-b" / "tarball" shutil.copytree(mock_modules_dir / "workspace_a", path_a) shutil.copytree(mock_modules_dir / "workspace_b", path_b) diff --git a/tests/integration/test_pool_integration.py b/tests/integration/test_pool_integration.py index e68788f747..782faea6e0 100644 --- a/tests/integration/test_pool_integration.py +++ b/tests/integration/test_pool_integration.py @@ -449,14 +449,17 @@ async def mock_extract(self, tarball_path: Path, target_dir: Path) -> None: patch.object(TarballArtifact, "download", mock_download), patch.object(TarballArtifact, "extract", mock_extract), ): - cache_key = "concurrent-test" tarball_uri = "s3://bucket/concurrent.tar.gz" + async def lease_paths() -> list[Path]: + async with cache.lease([tarball_uri]) as paths: + return paths + # Launch concurrent requests results = await asyncio.gather( - cache.materialize(cache_key, tarball_uri), - cache.materialize(cache_key, tarball_uri), - cache.materialize(cache_key, tarball_uri), + lease_paths(), + lease_paths(), + lease_paths(), ) assert download_count[0] == 1, ( diff --git a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py index be4eda9b66..13c2bef97c 100644 --- a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py +++ b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py @@ -215,6 +215,7 @@ def _build_squashfs_image(source_dir: Path, image_path: Path, module_name: str) """ source_dir.mkdir(parents=True, exist_ok=True) (source_dir / module_name).write_text("VALUE = 1\n") + image_path.parent.mkdir(parents=True, exist_ok=True) image_path.unlink(missing_ok=True) subprocess.run( ["mksquashfs", str(source_dir), str(image_path), "-noappend", "-quiet"], @@ -366,20 +367,28 @@ async def _run_mount_lifecycle_child() -> None: # (f) The startup sweep trims to budget and drops stale mount directories. sweep_dir = root / "sweep-cache" sweep_dir.mkdir() + sweep_cache = RegistryArtifactCache(sweep_dir) sweep_keys = ["aaaa1111", "bbbb2222"] for index, sweep_key in enumerate(sweep_keys): - image_path = sweep_dir / f"squashfs-{sweep_key}.squashfs" + sweep_paths = sweep_cache._paths_for(sweep_key) + sweep_paths.entry_dir.mkdir(parents=True) + image_path = sweep_paths.squashfs_image_path image_path.write_bytes(b"x" * 4096) - os.utime(image_path, (100.0 + index, 100.0 + index)) - stale_mount_dir = sweep_dir / f"squashfs-{sweep_keys[0]}" + os.utime(sweep_paths.entry_dir, (100.0 + index, 100.0 + index)) + stale_mount_dir = sweep_cache._paths_for(sweep_keys[0]).squashfs_mount_dir stale_mount_dir.mkdir() + os.utime( + sweep_cache._paths_for(sweep_keys[0]).entry_dir, + (100.0, 100.0), + ) config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES = 1 - sweep_cache = RegistryArtifactCache(sweep_dir) await sweep_cache.ensure_swept() + evicted_paths = sweep_cache._paths_for(sweep_keys[0]) + retained_paths = sweep_cache._paths_for(sweep_keys[1]) payload["startup_sweep_trimmed"] = ( - not (sweep_dir / f"squashfs-{sweep_keys[0]}.squashfs").exists() - and (sweep_dir / f"squashfs-{sweep_keys[1]}.squashfs").exists() + not evicted_paths.squashfs_image_path.exists() + and retained_paths.squashfs_image_path.exists() ) payload["startup_sweep_removed_stale_mount_dir"] = not stale_mount_dir.exists() diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index d24f79064f..bd17a4b739 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -647,8 +647,8 @@ async def test_execute_action_holds_registry_lease_for_whole_subprocess( runner = ActionRunner(cache_dir=temp_cache_dir) artifact_uri = "s3://bucket/execute.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) - entry_dir = temp_cache_dir / f"tarball-{cache_key}" - entry_dir.mkdir() + entry_dir = runner.registry_artifacts._paths_for(cache_key).tarball_target_dir + entry_dir.mkdir(parents=True) monkeypatch.setattr( action_runner.config, "TRACECAT__EXECUTOR_SANDBOX_ENABLED", False diff --git a/tests/unit/test_executor_sandbox_nsjail.py b/tests/unit/test_executor_sandbox_nsjail.py index 52778e4cdf..f920bf11c2 100644 --- a/tests/unit/test_executor_sandbox_nsjail.py +++ b/tests/unit/test_executor_sandbox_nsjail.py @@ -521,9 +521,10 @@ async def _run_executor_action_smoke_case( runner = ActionRunner(cache_dir=cache_dir) cache_key = compute_registry_artifact_cache_key(_SMOKE_URI) - mount_dir = cache_dir / f"squashfs-{cache_key}" - extract_dir = cache_dir / f"unsquashfs-{cache_key}" - tarball_dir = cache_dir / f"tarball-{cache_key}" + cache_paths = runner.registry_artifacts._paths_for(cache_key) + mount_dir = cache_paths.squashfs_mount_dir + extract_dir = cache_paths.squashfs_extract_dir + tarball_dir = cache_paths.tarball_target_dir async def sidecar_exists( *, @@ -597,7 +598,7 @@ async def download_artifact(self, ctx, output_path: Path) -> float: assert result["marker"] == "registry-artifact" if smoke_case == SmokeCase.DIRECT: assert tarball_dir.exists() - assert f"tarball-{cache_key}" in result["source"] + assert str(tarball_dir) in result["source"] assert result["source"].endswith("/registry_artifact_smoke_action.py") elif smoke_case == SmokeCase.DIRECT_SQUASHFS: assert extract_dir.exists() diff --git a/tests/unit/test_multitenant_registry.py b/tests/unit/test_multitenant_registry.py index 34ff62757f..7857dd6326 100644 --- a/tests/unit/test_multitenant_registry.py +++ b/tests/unit/test_multitenant_registry.py @@ -32,6 +32,14 @@ def temp_cache_dir(): yield Path(tmpdir) +async def _lease_paths( + cache: RegistryArtifactCache, + artifact_uri: str, +) -> list[Path]: + async with cache.lease([artifact_uri]) as paths: + return paths + + # ============================================================================= # Test Class: Tarball Cache Behavior # ============================================================================= @@ -62,7 +70,6 @@ async def test_concurrent_same_uri_single_download(self, temp_cache_dir: Path): - All calls return the same extracted path """ cache = RegistryArtifactCache(temp_cache_dir) - cache_key = "concurrent-test-key" tarball_uri = "s3://bucket/concurrent-test.tar.gz" download_count = [0] # Use list to allow mutation in nested function @@ -81,10 +88,10 @@ async def mock_extract(self, tarball_path: Path, target_dir: Path): ): # Launch multiple concurrent requests results = await asyncio.gather( - cache.materialize(cache_key, tarball_uri), - cache.materialize(cache_key, tarball_uri), - cache.materialize(cache_key, tarball_uri), - cache.materialize(cache_key, tarball_uri), + _lease_paths(cache, tarball_uri), + _lease_paths(cache, tarball_uri), + _lease_paths(cache, tarball_uri), + _lease_paths(cache, tarball_uri), ) # All should return same path @@ -128,8 +135,7 @@ async def mock_extract(self, tarball_path: Path, target_dir: Path): ): results = [] for uri in uris: - cache_key = compute_registry_artifact_cache_key(uri) - result = await cache.materialize(cache_key, uri) + result = await _lease_paths(cache, uri) results.append(result) # All results should be different paths @@ -155,8 +161,8 @@ async def test_failed_extraction_cleans_up_temp_files(self, temp_cache_dir: Path - RuntimeError is raised with appropriate message """ cache = RegistryArtifactCache(temp_cache_dir) - cache_key = "failed-extraction-test" tarball_uri = "s3://bucket/bad.tar.gz" + cache_key = compute_registry_artifact_cache_key(tarball_uri) async def mock_download(self, ctx, path: Path): path.write_bytes(b"corrupt tarball") @@ -169,14 +175,16 @@ async def mock_extract(self, tarball_path: Path, target_dir: Path): patch.object(TarballArtifact, "extract", mock_extract), ): with pytest.raises(RuntimeError, match="Extraction failed"): - await cache.materialize(cache_key, tarball_uri) + await _lease_paths(cache, tarball_uri) # Verify no temp files remain - temp_files = list(temp_cache_dir.glob(f"{cache_key}*")) + temp_files = ( + list(cache.staging_dir.iterdir()) if cache.staging_dir.exists() else [] + ) assert len(temp_files) == 0, f"Temp files not cleaned up: {temp_files}" # Target directory should not exist - target_dir = temp_cache_dir / f"tarball-{cache_key}" + target_dir = cache._paths_for(cache_key).tarball_target_dir assert not target_dir.exists() @pytest.mark.anyio @@ -192,7 +200,6 @@ async def test_cache_reused_on_second_request(self, temp_cache_dir: Path): - Both requests return the same path """ cache = RegistryArtifactCache(temp_cache_dir) - cache_key = "reuse-test" tarball_uri = "s3://bucket/reuse.tar.gz" download_count = [0] @@ -209,11 +216,11 @@ async def mock_extract(self, tarball_path: Path, target_dir: Path): patch.object(TarballArtifact, "extract", mock_extract), ): # First request - result1 = await cache.materialize(cache_key, tarball_uri) + result1 = await _lease_paths(cache, tarball_uri) assert download_count[0] == 1 # Second request (should use cache) - result2 = await cache.materialize(cache_key, tarball_uri) + result2 = await _lease_paths(cache, tarball_uri) assert download_count[0] == 1 # No additional download assert result1 == result2 diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 58c44cb52c..c89b207514 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -16,14 +16,12 @@ from tracecat.executor.registry_artifacts import ( SQUASHFS_MOUNT_OPTIONS, - TEMP_ARTIFACT_PATTERN, RegistryArtifactCache, RegistryArtifactFormat, - RegistryArtifactPaths, SquashfsArtifact, SquashfsMountCommandError, TarballArtifact, - _delete_entry_paths, + _delete_cache_path, bundled_builtin_registry_uri, compute_registry_artifact_cache_key, ) @@ -50,7 +48,7 @@ def _write_tarball_entry(cache_dir: Path, cache_key: str) -> Path: """Create a materialized tarball cache entry on disk.""" - target_dir = cache_dir / f"tarball-{cache_key}" + target_dir = cache_dir / "entries" / cache_key / "tarball" target_dir.mkdir(parents=True) (target_dir / "module.py").write_text("VALUE = 1") return target_dir @@ -60,12 +58,43 @@ def _write_image_entry( cache_dir: Path, cache_key: str, *, size: int, mtime: float ) -> Path: """Create a downloaded SquashFS image cache entry with a fixed mtime.""" - image_path = cache_dir / f"squashfs-{cache_key}.squashfs" + entry_dir = cache_dir / "entries" / cache_key + entry_dir.mkdir(parents=True, exist_ok=True) + image_path = entry_dir / "image.squashfs" image_path.write_bytes(b"x" * size) os.utime(image_path, (mtime, mtime)) + os.utime(entry_dir, (mtime, mtime)) return image_path +async def _materialize( + cache: RegistryArtifactCache, + cache_key: str, + artifact_uri: str, +) -> list[Path]: + """Exercise internal materialization while releasing its test-only lease.""" + await cache.ensure_swept() + ctx = cache._context_for(cache_key) + lock = await cache._lock_for(cache_key) + async with lock: + cache._acquire_lease(cache_key) + try: + candidates = await cache._artifact_candidates(ctx, artifact_uri) + except BaseException: + cache._release_lease(cache_key) + raise + if cached_paths := cache._first_cached_path(candidates, ctx): + cache._release_lease(cache_key) + return cached_paths + + try: + await cache._enforce_cache_budget(protected_key=cache_key) + async with lock: + return await cache._materialize_candidates(ctx, artifact_uri) + finally: + cache._release_lease(cache_key) + + class _BlockingSubprocess: """Fake subprocess that blocks in communicate until it is cancelled.""" @@ -374,7 +403,8 @@ def fake_first_cached_path(candidates, ctx): side_effect=fake_first_cached_path, ), ): - result = await cache.materialize( + result = await _materialize( + cache, cache_key, "s3://bucket/path/site-packages.tar.gz", ) @@ -517,7 +547,8 @@ async def mock_mount(self, ctx, image_path): new_callable=AsyncMock, ) as tarball_materialize, ): - result = await cache.materialize( + result = await _materialize( + cache, "squashfs-key", "s3://bucket/path/site-packages.tar.gz", ) @@ -541,6 +572,7 @@ async def test_mount_squashfs_uses_hardened_read_only_options( ) image_path = ctx.paths.squashfs_image_path target_dir = ctx.paths.squashfs_mount_dir + ctx.paths.entry_dir.mkdir(parents=True) image_path.write_bytes(b"squashfs") target_dir.mkdir() process = AsyncMock() @@ -578,6 +610,7 @@ async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): ) image_path = ctx.paths.squashfs_image_path target_dir = ctx.paths.squashfs_mount_dir + ctx.paths.entry_dir.mkdir(parents=True) image_path.write_bytes(b"squashfs") target_dir.mkdir() process = _BlockingSubprocess() @@ -614,6 +647,7 @@ async def test_cancelled_squashfs_extract_kills_and_reaps_subprocess( ) image_path = ctx.paths.squashfs_image_path target_dir = ctx.paths.squashfs_extract_dir + ctx.paths.entry_dir.mkdir(parents=True) image_path.write_bytes(b"squashfs") target_dir.mkdir() real_create_subprocess_exec = asyncio.create_subprocess_exec @@ -741,14 +775,15 @@ async def mock_extract(self, ctx, image_path): new_callable=AsyncMock, ) as tarball_materialize, ): - result = await cache.materialize( + result = await _materialize( + cache, "fallback-key", "s3://bucket/path/site-packages.tar.gz", ) assert len(result) == 1 assert (result[0] / "module.py").read_text() == "VALUE = 1" - assert result[0].name.startswith("unsquashfs-") + assert result[0].name == "extracted" tarball_materialize.assert_not_awaited() @pytest.mark.anyio @@ -779,14 +814,15 @@ async def mock_extract(self, ctx, image_path): ), patch.object(SquashfsArtifact, "extract", mock_extract), ): - result = await cache.materialize( + result = await _materialize( + cache, "extract-key", "s3://bucket/path/site-packages.tar.gz", ) assert len(result) == 1 assert (result[0] / "module.py").read_text() == "VALUE = 1" - assert result[0].name.startswith("unsquashfs-") + assert result[0].name == "extracted" @pytest.mark.anyio async def test_materialize_falls_back_to_gzip_when_squashfs_extract_fails( @@ -826,14 +862,15 @@ async def mock_extract(self, ctx, image_path): patch.object(SquashfsArtifact, "extract", mock_extract), patch.object(TarballArtifact, "download", mock_tarball_download), ): - result = await cache.materialize( + result = await _materialize( + cache, "gzip-fallback-key", "s3://bucket/path/site-packages.tar.gz", ) assert len(result) == 1 assert (result[0] / "module.py").read_text() == "VALUE = 1" - assert result[0].name.startswith("tarball-") + assert result[0].name == "tarball" @pytest.mark.anyio async def test_materialize_treats_unknown_suffix_as_gzip(self, temp_cache_dir): @@ -849,7 +886,8 @@ async def mock_download(self, ctx, path): tar.add(source / "module.py", arcname="module.py") with patch.object(TarballArtifact, "download", mock_download): - result = await cache.materialize( + result = await _materialize( + cache, "custom-key-test", "s3://bucket/path/custom-key", ) @@ -862,10 +900,14 @@ async def test_materialize_caches_result(self, temp_cache_dir): """Test that tarball extraction is cached.""" cache = RegistryArtifactCache(temp_cache_dir) cache_key = "test-cache-key" - target_dir = temp_cache_dir / f"tarball-{cache_key}" + target_dir = cache._paths_for(cache_key).tarball_target_dir target_dir.mkdir(parents=True) - result = await cache.materialize(cache_key, "s3://bucket/test.tar.gz") + result = await _materialize( + cache, + cache_key, + "s3://bucket/test.tar.gz", + ) assert result == [target_dir] @@ -890,9 +932,9 @@ async def mock_extract(self, tarball_path, target_dir): patch.object(TarballArtifact, "extract", mock_extract), ): results = await asyncio.gather( - cache.materialize(cache_key, "s3://bucket/test.tar.gz"), - cache.materialize(cache_key, "s3://bucket/test.tar.gz"), - cache.materialize(cache_key, "s3://bucket/test.tar.gz"), + _materialize(cache, cache_key, "s3://bucket/test.tar.gz"), + _materialize(cache, cache_key, "s3://bucket/test.tar.gz"), + _materialize(cache, cache_key, "s3://bucket/test.tar.gz"), ) assert all(r == results[0] for r in results) @@ -926,12 +968,13 @@ def test_touch_entry_refreshes_tarball_root_mtime(self, temp_cache_dir): """Touching a tarball-only entry persists its restart-safe recency.""" cache = RegistryArtifactCache(temp_cache_dir) cache_key = "tarball-only" - target_dir = _write_tarball_entry(temp_cache_dir, cache_key) - os.utime(target_dir, (100.0, 100.0)) + _write_tarball_entry(temp_cache_dir, cache_key) + entry_dir = cache._paths_for(cache_key).entry_dir + os.utime(entry_dir, (100.0, 100.0)) cache._touch_entry(cache_key) - assert target_dir.stat().st_mtime > 100.0 + assert entry_dir.stat().st_mtime > 100.0 @pytest.mark.anyio async def test_lease_refcounts_and_touches_image_mtime(self, temp_cache_dir): @@ -941,13 +984,15 @@ async def test_lease_refcounts_and_touches_image_mtime(self, temp_cache_dir): cache_key = compute_registry_artifact_cache_key(artifact_uri) target_dir = _write_tarball_entry(temp_cache_dir, cache_key) image_path = _write_image_entry(temp_cache_dir, cache_key, size=16, mtime=100.0) + entry_dir = cache._paths_for(cache_key).entry_dir async with cache.lease([artifact_uri]) as registry_paths: assert registry_paths == [target_dir] assert cache._refcount(cache_key) == 1 - assert image_path.stat().st_mtime > 100.0 + assert entry_dir.stat().st_mtime > 100.0 assert cache._refcount(cache_key) == 0 + assert image_path.is_file() @pytest.mark.anyio async def test_lease_releases_refcount_when_materialization_fails( @@ -1038,6 +1083,7 @@ async def test_lease_is_never_admitted_across_an_in_flight_eviction( artifact_uri = "s3://bucket/path/site-packages.squashfs" cache_key = compute_registry_artifact_cache_key(artifact_uri) paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) paths.squashfs_image_path.write_bytes(b"squashfs") paths.squashfs_mount_dir.mkdir() mounted = {paths.squashfs_mount_dir} @@ -1135,15 +1181,10 @@ async def test_builtin_artifact_is_exempt_from_cache_accounting( class TestRegistryArtifactCacheEviction: """Tests for bounded eviction of registry artifact cache entries.""" - def test_delete_entry_paths_reports_directory_failure(self, temp_cache_dir): + def test_delete_cache_path_reports_directory_failure(self, temp_cache_dir): """Physical deletion failures are observable instead of suppressed.""" - paths = RegistryArtifactPaths( - squashfs_image_path=temp_cache_dir / "image.squashfs", - squashfs_mount_dir=temp_cache_dir / "mount", - squashfs_extract_dir=temp_cache_dir / "extract", - tarball_target_dir=temp_cache_dir / "tarball", - ) - paths.squashfs_extract_dir.mkdir() + entry_dir = temp_cache_dir / "entry" + entry_dir.mkdir() with ( patch( @@ -1152,13 +1193,13 @@ def test_delete_entry_paths_reports_directory_failure(self, temp_cache_dir): ), patch("tracecat.executor.registry_artifacts.logger.warning") as warning, ): - deleted = _delete_entry_paths(paths) + deleted = _delete_cache_path(entry_dir) assert deleted is False - assert paths.squashfs_extract_dir.is_dir() + assert entry_dir.is_dir() warning.assert_called_once_with( "Failed to delete registry artifact cache path", - path=str(paths.squashfs_extract_dir), + path=str(entry_dir), error="permission denied", ) @@ -1188,8 +1229,8 @@ async def mock_extract(self, tarball_path, target_dir): patch.object(TarballArtifact, "extract", mock_extract), ): async with cache.lease([leased_uri]): - await cache.materialize( - compute_registry_artifact_cache_key(new_uri), new_uri + await _materialize( + cache, compute_registry_artifact_cache_key(new_uri), new_uri ) assert leased_dir.is_dir() @@ -1219,11 +1260,11 @@ async def mock_extract(self, tarball_path, target_dir): ): async with cache.lease([new_uri]) as registry_paths: # Both entries fit only because the new one is still leased. - assert registry_paths == [temp_cache_dir / f"tarball-{new_key}"] + assert registry_paths == [cache._paths_for(new_key).tarball_target_dir] assert idle.exists() assert not idle.exists() - assert (temp_cache_dir / f"tarball-{new_key}").is_dir() + assert cache._paths_for(new_key).tarball_target_dir.is_dir() assert cache._budget_dirty is False @pytest.mark.anyio @@ -1247,16 +1288,16 @@ async def mock_extract(self, tarball_path, target_dir): patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0), patch( - "tracecat.executor.registry_artifacts._delete_entry_paths", + "tracecat.executor.registry_artifacts._delete_cache_path", return_value=False, ), patch.object(TarballArtifact, "download", mock_download), patch.object(TarballArtifact, "extract", mock_extract), patch("tracecat.executor.registry_artifacts.logger.warning") as warning, ): - registry_paths = await cache.materialize(cache_key, artifact_uri) + registry_paths = await _materialize(cache, cache_key, artifact_uri) - assert registry_paths == [temp_cache_dir / f"tarball-{cache_key}"] + assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] assert registry_paths[0].is_dir() assert cache._budget_dirty is True warning.assert_any_call( @@ -1264,6 +1305,85 @@ async def mock_extract(self, tarball_path, target_dir): cache_key="idle", ) + @pytest.mark.anyio + async def test_failed_physical_delete_retries_without_extra_eviction( + self, temp_cache_dir + ): + """A retired entry counts as evicted while its exact trash path retries.""" + cache = RegistryArtifactCache(temp_cache_dir) + oldest = _write_image_entry( + temp_cache_dir, + "oldest", + size=16, + mtime=100.0, + ) + retained = _write_image_entry( + temp_cache_dir, + "retained", + size=16, + mtime=200.0, + ) + real_delete = _delete_cache_path + failed_once = False + + def fail_once(path: Path) -> bool: + nonlocal failed_once + if not failed_once: + failed_once = True + return False + return real_delete(path) + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifacts._delete_cache_path", + side_effect=fail_once, + ), + ): + assert await cache._enforce_cache_budget() is False + assert not oldest.exists() + assert retained.exists() + assert cache._discover_cache_keys() == {"retained"} + assert len(cache._pending_cleanup) == 1 + + assert await cache._enforce_cache_budget() is True + + assert cache._pending_cleanup == set() + assert not any(cache.trash_dir.iterdir()) + + @pytest.mark.anyio + async def test_rename_failure_does_not_block_materialization(self, temp_cache_dir): + """A failed atomic retirement keeps the old entry and admits new work.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + idle = _write_image_entry(temp_cache_dir, "idle", size=16, mtime=100.0) + artifact_uri = "s3://bucket/new-after-rename-failure.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifacts._move_entry_to_trash", + side_effect=OSError("rename failed"), + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + registry_paths = await _materialize(cache, cache_key, artifact_uri) + + assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] + assert registry_paths[0].is_dir() + assert idle.is_file() + assert cache._budget_dirty is True + @pytest.mark.anyio async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( self, temp_cache_dir @@ -1298,6 +1418,7 @@ async def test_failed_materialization_rearms_budget_dirty(self, temp_cache_dir): artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) async def mock_materialize(self, ctx): + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) ctx.paths.squashfs_image_path.write_bytes(b"orphaned image") raise RuntimeError("mount failed") @@ -1311,7 +1432,7 @@ async def mock_materialize(self, ctx): patch.object(SquashfsArtifact, "materialize", mock_materialize), ): with pytest.raises(RuntimeError, match="mount failed"): - await cache.materialize(cache_key, artifact_uri) + await _materialize(cache, cache_key, artifact_uri) assert cache._paths_for(cache_key).squashfs_image_path.is_file() assert cache._budget_dirty is True @@ -1333,6 +1454,7 @@ async def mock_materialize(self, ctx): # A concurrent lease release consumes the dirty signal and finishes # its scan before this attempt lands its canonical bytes. cache._budget_dirty = False + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) ctx.paths.squashfs_image_path.write_bytes(b"late image") return [ctx.paths.squashfs_mount_dir] @@ -1345,7 +1467,7 @@ async def mock_materialize(self, ctx): ), patch.object(SquashfsArtifact, "materialize", mock_materialize), ): - await cache.materialize(cache_key, artifact_uri) + await _materialize(cache, cache_key, artifact_uri) assert cache._budget_dirty is True @@ -1413,11 +1535,11 @@ async def mock_extract(self, tarball_path, target_dir): ): convergence = asyncio.create_task(cache._converge_cache_budget()) await scan_started.wait() - registry_paths = await cache.materialize(cache_key, artifact_uri) + registry_paths = await _materialize(cache, cache_key, artifact_uri) materialized.set() await convergence - assert registry_paths == [temp_cache_dir / f"tarball-{cache_key}"] + assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] assert convergence_scans == 2 assert cache._budget_dirty is False @@ -1466,8 +1588,9 @@ async def test_enforce_budget_evicts_least_recently_used_until_under_max_bytes( cache = RegistryArtifactCache(temp_cache_dir) oldest = _write_image_entry(temp_cache_dir, "oldest", size=4096, mtime=100.0) if oldest_has_tarball: - oldest_tarball = _write_tarball_entry(temp_cache_dir, "oldest") - os.utime(oldest_tarball, (100.0, 100.0)) + _write_tarball_entry(temp_cache_dir, "oldest") + oldest_entry = cache._paths_for("oldest").entry_dir + os.utime(oldest_entry, (100.0, 100.0)) older = _write_image_entry(temp_cache_dir, "older", size=4096, mtime=200.0) newest = _write_image_entry(temp_cache_dir, "newest", size=4096, mtime=300.0) @@ -1519,7 +1642,8 @@ async def test_auto_non_pool_backend_evicts_tarball_lru(self, temp_cache_dir): temp_cache_dir, "oldest", size=16, mtime=100.0 ) oldest_tarball = _write_tarball_entry(temp_cache_dir, "oldest") - os.utime(oldest_tarball, (100.0, 100.0)) + oldest_entry = cache._paths_for("oldest").entry_dir + os.utime(oldest_entry, (100.0, 100.0)) newest = _write_image_entry(temp_cache_dir, "newest", size=16, mtime=200.0) with ( @@ -1578,11 +1702,13 @@ async def test_pool_backend_release_mounted_slot_skips_tarball_entry( """Loop-device recovery must preserve paths visible to warm workers.""" cache = RegistryArtifactCache(temp_cache_dir) pool_visible = cache._paths_for("pool-visible") + pool_visible.entry_dir.mkdir(parents=True) pool_visible.squashfs_image_path.write_bytes(b"squashfs") os.utime(pool_visible.squashfs_image_path, (100.0, 100.0)) pool_visible.squashfs_mount_dir.mkdir() _write_tarball_entry(temp_cache_dir, "pool-visible") eligible = cache._paths_for("eligible") + eligible.entry_dir.mkdir(parents=True) eligible.squashfs_image_path.write_bytes(b"squashfs") os.utime(eligible.squashfs_image_path, (200.0, 200.0)) eligible.squashfs_mount_dir.mkdir() @@ -1645,9 +1771,9 @@ async def controlled_evict(cache_key: str) -> bool: return False eviction_started.set() await finish_eviction.wait() - oldest.unlink(missing_ok=True) + _delete_cache_path(cache._paths_for("oldest").entry_dir) else: - retained.unlink(missing_ok=True) + _delete_cache_path(cache._paths_for("retained").entry_dir) extra_eviction_finished.set() evicted_keys.append(cache_key) return True @@ -1726,6 +1852,7 @@ async def test_eviction_unmounts_before_deleting_the_image(self, temp_cache_dir) """Unlinking a mounted image would strand an open-file zombie.""" cache = RegistryArtifactCache(temp_cache_dir) paths = cache._paths_for("mounted") + paths.entry_dir.mkdir(parents=True) paths.squashfs_image_path.write_bytes(b"squashfs") paths.squashfs_mount_dir.mkdir() mounted = {paths.squashfs_mount_dir} @@ -1772,6 +1899,7 @@ async def test_cancelled_unmount_kills_and_reaps_before_releasing_key_lock( artifact_uri = "s3://bucket/path/cancelled-unmount.squashfs" cache_key = compute_registry_artifact_cache_key(artifact_uri) paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) paths.squashfs_image_path.write_bytes(b"squashfs") paths.squashfs_mount_dir.mkdir() (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") @@ -1819,13 +1947,13 @@ async def test_cancelled_background_deletion_leaves_a_clean_miss( delete_started = threading.Event() finish_delete = threading.Event() delete_finished = threading.Event() - doomed: list[RegistryArtifactPaths] = [] + doomed: list[Path] = [] - def blocked_delete(paths: RegistryArtifactPaths) -> bool: - doomed.append(paths) + def blocked_delete(path: Path) -> bool: + doomed.append(path) delete_started.set() finish_delete.wait(timeout=5) - deleted = _delete_entry_paths(paths) + deleted = _delete_cache_path(path) delete_finished.set() return deleted @@ -1837,7 +1965,7 @@ async def mock_extract(self, tarball_path, target_dir): with ( patch( - "tracecat.executor.registry_artifacts._delete_entry_paths", + "tracecat.executor.registry_artifacts._delete_cache_path", side_effect=blocked_delete, ), patch.object(TarballArtifact, "download", mock_download), @@ -1851,60 +1979,55 @@ async def mock_extract(self, tarball_path, target_dir): try: assert not original_target.exists() - assert doomed[0].tarball_target_dir.is_dir() + assert (doomed[0] / "tarball").is_dir() assert cache._discover_cache_keys() == set() async with cache.lease([artifact_uri]) as registry_paths: assert registry_paths == [original_target] assert original_target.is_dir() - assert doomed[0].tarball_target_dir.is_dir() + assert (doomed[0] / "tarball").is_dir() assert (original_target / "module.py").read_text() == "VALUE = 2" finally: finish_delete.set() assert await asyncio.to_thread(delete_finished.wait, 1) assert original_target.is_dir() - assert not doomed[0].tarball_target_dir.exists() + assert not doomed[0].exists() @pytest.mark.anyio async def test_doomed_eviction_names_are_startup_scratch(self, temp_cache_dir): """Every renamed entry root is invisible and reclaimed on startup.""" cache = RegistryArtifactCache(temp_cache_dir) - # This key makes doomed names look like live image paths unless cache - # discovery rejects the shared scratch pattern first. cache_key = "squashfs-doomed" paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) paths.squashfs_image_path.write_bytes(b"squashfs") paths.squashfs_mount_dir.mkdir() paths.squashfs_extract_dir.mkdir() paths.tarball_target_dir.mkdir() with patch( - "tracecat.executor.registry_artifacts._delete_entry_paths", + "tracecat.executor.registry_artifacts._delete_cache_path", return_value=True, - ) as delete_entry_paths: + ) as delete_cache_path: assert await cache._evict_entry(cache_key) is True - doomed_paths = delete_entry_paths.call_args.args[0] - renamed = ( - doomed_paths.squashfs_image_path, - doomed_paths.squashfs_mount_dir, - doomed_paths.squashfs_extract_dir, - doomed_paths.tarball_target_dir, - ) - assert all(path.exists() for path in renamed) - assert all(TEMP_ARTIFACT_PATTERN.fullmatch(path.name) for path in renamed) + delete_cache_path.assert_called_once() + trash_path = delete_cache_path.call_args.args[0] + assert trash_path.parent == cache.trash_dir + assert trash_path.exists() assert cache._discover_cache_keys() == set() cache._sweep_startup_state() - assert not any(path.exists() for path in renamed) + assert not trash_path.exists() @pytest.mark.anyio async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): """A failed unmount skips the key instead of forcing a lazy detach.""" cache = RegistryArtifactCache(temp_cache_dir) stuck = cache._paths_for("stuck") + stuck.entry_dir.mkdir(parents=True) stuck.squashfs_image_path.write_bytes(b"squashfs") stuck.squashfs_mount_dir.mkdir() os.utime(stuck.squashfs_image_path, (100.0, 100.0)) @@ -1967,18 +2090,18 @@ class TestRegistryArtifactCacheStartupSweep: """Tests for the startup sweep that reclaims state from a dead process.""" @pytest.mark.anyio - async def test_sweep_uses_tarball_root_mtime_for_restart_safe_lru( + async def test_sweep_uses_entry_root_mtime_for_restart_safe_lru( self, temp_cache_dir ): """Startup trimming preserves a touched older tarball-only entry.""" previous_cache = RegistryArtifactCache(temp_cache_dir) old = _write_tarball_entry(temp_cache_dir, "old") - os.utime(old, (100.0, 100.0)) previous_cache._touch_entry("old") - await asyncio.sleep(0.01) + old_mtime = previous_cache._paths_for("old").entry_dir.stat().st_mtime new = _write_tarball_entry(temp_cache_dir, "new") - os.utime(new, (200.0, 200.0)) + new_entry_dir = previous_cache._paths_for("new").entry_dir + os.utime(new_entry_dir, (old_mtime - 1, old_mtime - 1)) with ( patch(BACKEND_CONFIG, ExecutorBackendType.DIRECT.value), @@ -2003,37 +2126,40 @@ async def test_sweep_tolerates_missing_cache_dir(self, temp_cache_dir): assert not cache_dir.exists() @pytest.mark.anyio - async def test_sweep_removes_orphaned_scratch_and_stale_mount_dirs( - self, temp_cache_dir - ): - """Interrupted materializations and dead mount dirs are reclaimed.""" - orphaned = temp_cache_dir / "abc123.999999.4321.squashfs" + async def test_sweep_removes_orphaned_work_and_legacy_paths(self, temp_cache_dir): + """Interrupted work and disposable flat-layout paths are reclaimed.""" + cache = RegistryArtifactCache(temp_cache_dir) + orphaned = cache.staging_dir / "abc123.999999.4321.squashfs" + orphaned.parent.mkdir() orphaned.write_bytes(b"partial") - orphaned_dir = temp_cache_dir / "abc123.999999.4321.tmp" - orphaned_dir.mkdir() - own = temp_cache_dir / f"abc123.{os.getpid()}.4321.tar.gz" - own.write_bytes(b"in flight") - stale_mount_dir = temp_cache_dir / "squashfs-abc123" - stale_mount_dir.mkdir() + orphaned_dir = cache.trash_dir / "abc123.999999.4321" + orphaned_dir.mkdir(parents=True) + legacy = temp_cache_dir / "squashfs-abc123.squashfs" + legacy.write_bytes(b"legacy") + unrelated = temp_cache_dir / "unrelated" + unrelated.mkdir() + unrelated_file = unrelated / "keep.txt" + unrelated_file.write_text("keep") entry_dir = _write_tarball_entry(temp_cache_dir, "abc123") - cache = RegistryArtifactCache(temp_cache_dir) await cache.ensure_swept() assert not orphaned.exists() assert not orphaned_dir.exists() - assert not own.exists() - assert not stale_mount_dir.exists() + assert not legacy.exists() + assert unrelated_file.read_text() == "keep" assert entry_dir.is_dir() @pytest.mark.anyio async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): """A live mountpoint belongs to a running process and must survive.""" - mount_dir = temp_cache_dir / "squashfs-abc123" + cache = RegistryArtifactCache(temp_cache_dir) + paths = cache._paths_for("abc123") + paths.entry_dir.mkdir(parents=True) + mount_dir = paths.squashfs_mount_dir mount_dir.mkdir() with patch.object(Path, "is_mount", lambda self: self == mount_dir): - cache = RegistryArtifactCache(temp_cache_dir) await cache.ensure_swept() assert mount_dir.is_dir() @@ -2129,10 +2255,9 @@ def blocking_sweep() -> None: @pytest.mark.anyio async def test_lease_triggers_startup_sweep(self, temp_cache_dir): """Lease admission reclaims startup scratch before yielding paths.""" - orphaned_dir = temp_cache_dir / "abc123.999999.4321.tmp" - orphaned_dir.mkdir() - assert TEMP_ARTIFACT_PATTERN.match(orphaned_dir.name) is not None cache = RegistryArtifactCache(temp_cache_dir) + orphaned_dir = cache.staging_dir / "abc123.999999.4321" + orphaned_dir.mkdir(parents=True) async with cache.lease(None): assert not orphaned_dir.exists() @@ -2140,14 +2265,14 @@ async def test_lease_triggers_startup_sweep(self, temp_cache_dir): @pytest.mark.anyio async def test_failed_ensure_swept_retries(self, temp_cache_dir): """A failed startup sweep is retried by the next caller.""" - orphaned_dir = temp_cache_dir / "abc123.999999.4321.tmp" - orphaned_dir.mkdir() cache = RegistryArtifactCache(temp_cache_dir) + orphaned_dir = cache.staging_dir / "abc123.999999.4321" + orphaned_dir.mkdir(parents=True) with ( patch.object( cache, - "_remove_orphaned_temp_paths", + "_clear_work_dir", side_effect=OSError("simulated sweep failure"), ), pytest.raises(OSError), @@ -2160,6 +2285,132 @@ async def test_failed_ensure_swept_retries(self, temp_cache_dir): assert not orphaned_dir.exists() + @pytest.mark.anyio + async def test_failed_startup_cleanup_retries_exact_path(self, temp_cache_dir): + """Failed startup scratch cleanup retries without sweeping live staging.""" + cache = RegistryArtifactCache(temp_cache_dir) + orphaned = cache.staging_dir / "orphaned.123.456.tmp" + orphaned.parent.mkdir(parents=True) + orphaned.write_bytes(b"partial") + real_delete = _delete_cache_path + failed_once = False + + def fail_once(path: Path) -> bool: + nonlocal failed_once + if path == orphaned and not failed_once: + failed_once = True + return False + return real_delete(path) + + with patch( + "tracecat.executor.registry_artifacts._delete_cache_path", + side_effect=fail_once, + ): + await cache.ensure_swept() + assert orphaned.is_file() + assert cache._pending_cleanup == {orphaned} + assert cache._budget_dirty is True + + assert await cache._enforce_cache_budget() is True + + assert not orphaned.exists() + assert cache._pending_cleanup == set() + + @pytest.mark.anyio + async def test_failed_startup_retirement_stays_dirty_and_retries( + self, temp_cache_dir + ): + """A startup rename failure preserves entries for later enforcement.""" + oldest = _write_image_entry( + temp_cache_dir, + "oldest", + size=16, + mtime=100.0, + ) + newest = _write_image_entry( + temp_cache_dir, + "newest", + size=16, + mtime=200.0, + ) + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch(BACKEND_CONFIG, ExecutorBackendType.DIRECT.value), + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifacts._move_entry_to_trash", + side_effect=OSError("rename failed"), + ), + ): + await cache.ensure_swept() + + assert oldest.is_file() + assert newest.is_file() + assert cache._budget_dirty is True + + with ( + patch(BACKEND_CONFIG, ExecutorBackendType.DIRECT.value), + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + await cache._converge_cache_budget() + + assert not oldest.exists() + assert newest.is_file() + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_failed_startup_physical_delete_retries_exact_path( + self, temp_cache_dir + ): + """A startup trash delete failure retries without evicting extra entries.""" + oldest = _write_image_entry( + temp_cache_dir, + "oldest", + size=16, + mtime=100.0, + ) + newest = _write_image_entry( + temp_cache_dir, + "newest", + size=16, + mtime=200.0, + ) + cache = RegistryArtifactCache(temp_cache_dir) + real_delete = _delete_cache_path + failed_once = False + + def fail_once(path: Path) -> bool: + nonlocal failed_once + if path.parent == cache.trash_dir and not failed_once: + failed_once = True + return False + return real_delete(path) + + with ( + patch(BACKEND_CONFIG, ExecutorBackendType.DIRECT.value), + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifacts._delete_cache_path", + side_effect=fail_once, + ), + ): + await cache.ensure_swept() + assert not oldest.exists() + assert newest.is_file() + assert len(cache._pending_cleanup) == 1 + assert cache._budget_dirty is True + + await cache._converge_cache_budget() + + assert newest.is_file() + assert cache._pending_cleanup == set() + assert not any(cache.trash_dir.iterdir()) + assert cache._budget_dirty is False + class TestSquashfsMountCapability: """Tests for process-wide SquashFS mount capability tracking.""" @@ -2173,6 +2424,8 @@ async def test_concurrent_first_mount_failure_serializes_capability_probe( first_ctx = cache._context_for("first-probe") second_ctx = cache._context_for("second-probe") assert first_ctx.squashfs_mount_state is second_ctx.squashfs_mount_state + first_ctx.paths.entry_dir.mkdir(parents=True) + second_ctx.paths.entry_dir.mkdir(parents=True) first_ctx.paths.squashfs_image_path.write_bytes(b"squashfs") second_ctx.paths.squashfs_image_path.write_bytes(b"squashfs") first_artifact = SquashfsArtifact( @@ -2237,6 +2490,8 @@ async def test_concurrent_probe_waiter_mounts_after_first_success( first_ctx = cache._context_for("first-success") second_ctx = cache._context_for("second-success") assert first_ctx.squashfs_mount_state is second_ctx.squashfs_mount_state + first_ctx.paths.entry_dir.mkdir(parents=True) + second_ctx.paths.entry_dir.mkdir(parents=True) first_ctx.paths.squashfs_image_path.write_bytes(b"squashfs") second_ctx.paths.squashfs_image_path.write_bytes(b"squashfs") first_artifact = SquashfsArtifact( @@ -2323,8 +2578,8 @@ async def mock_extract(self, ctx, image_path): return_value=False, ) as release_mounted_slot, ): - await cache.materialize( - "probe-key", "s3://bucket/path/site-packages.squashfs" + await _materialize( + cache, "probe-key", "s3://bucket/path/site-packages.squashfs" ) assert cache._squashfs_mount_state.disabled is True @@ -2338,6 +2593,7 @@ async def test_mount_failure_after_success_reclaims_loop_device_and_retries( cache = RegistryArtifactCache(temp_cache_dir) cache._squashfs_mount_state.mounted_once = True idle = cache._paths_for("idle") + idle.entry_dir.mkdir(parents=True) idle.squashfs_image_path.write_bytes(b"squashfs") idle.squashfs_mount_dir.mkdir() mounted = {idle.squashfs_mount_dir} @@ -2374,8 +2630,8 @@ async def mock_umount(*args, **kwargs): ), patch.object(SquashfsArtifact, "mount", mock_mount), ): - result = await cache.materialize( - "new", "s3://bucket/path/site-packages.squashfs" + result = await _materialize( + cache, "new", "s3://bucket/path/site-packages.squashfs" ) assert result == [cache._paths_for("new").squashfs_mount_dir] @@ -2418,8 +2674,8 @@ async def mock_download(self, ctx, image_path): return_value=False, ) as release_mounted_slot, ): - result = await cache.materialize( - "download-failure", "s3://bucket/path/site-packages.squashfs" + result = await _materialize( + cache, "download-failure", "s3://bucket/path/site-packages.squashfs" ) assert result == [tarball_dir] @@ -2460,8 +2716,8 @@ async def mock_download(self, ctx, image_path): return_value=True, ) as release_mounted_slot, ): - result = await cache.materialize( - "download-failure", "s3://bucket/path/site-packages.squashfs" + result = await _materialize( + cache, "download-failure", "s3://bucket/path/site-packages.squashfs" ) assert result == [tarball_dir] @@ -2494,9 +2750,9 @@ async def mock_extract(self, ctx, image_path): patch.object(SquashfsArtifact, "mount", mock_mount), patch.object(SquashfsArtifact, "extract", mock_extract), ): - result = await cache.materialize( - "no-slot", "s3://bucket/path/site-packages.squashfs" + result = await _materialize( + cache, "no-slot", "s3://bucket/path/site-packages.squashfs" ) - assert result[0].name.startswith("unsquashfs-") + assert result[0].name == "extracted" assert cache._squashfs_mount_state.disabled is False diff --git a/tracecat/executor/backends/pool/worker.py b/tracecat/executor/backends/pool/worker.py index 98bbe48cf6..60fce20f30 100644 --- a/tracecat/executor/backends/pool/worker.py +++ b/tracecat/executor/backends/pool/worker.py @@ -56,25 +56,30 @@ def _ensure_tarball_paths_in_sys_path() -> None: """Ensure all tarball extraction directories are in sys.path. - Scans the registry cache directory for tarball-* subdirectories and adds - them to sys.path if not already present. This allows the worker to import - modules from both builtin and custom registries. + Scans atomic cache entries, plus legacy flat-layout entries during rollout, + and adds materialized tarballs to sys.path if not already present. """ cache_dir = Path(config.TRACECAT__EXECUTOR_REGISTRY_CACHE_DIR) if not cache_dir.exists(): return - for path in cache_dir.iterdir(): - if path.is_dir() and path.name.startswith("tarball-"): - path_str = str(path) - if path_str not in _added_tarball_paths and path_str not in sys.path: - sys.path.insert(0, path_str) - _added_tarball_paths.add(path_str) - logger.debug( - "Added tarball path to sys.path", - path=path_str, - worker_id=_worker_id, - ) + tarball_paths = ( + *cache_dir.glob("entries/*/tarball"), + *cache_dir.glob("tarball-*"), + ) + for path in tarball_paths: + if not path.is_dir(): + continue + path_str = str(path) + if path_str in _added_tarball_paths or path_str in sys.path: + continue + sys.path.insert(0, path_str) + _added_tarball_paths.add(path_str) + logger.debug( + "Added tarball path to sys.path", + path=path_str, + worker_id=_worker_id, + ) # Global counters for connection tracking diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 52146a4db7..0ec3cfc449 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -51,13 +51,22 @@ class RegistryArtifactFormat(StrEnum): BASE_PYTHONPATH_DIR_NAME = "base" """Cache subdirectory used as the PYTHONPATH entry when no artifact is requested.""" -CACHE_ENTRY_PREFIXES = ("squashfs-", "unsquashfs-", "tarball-") -"""On-disk name prefixes owned by a registry artifact cache entry.""" +CACHE_ENTRIES_DIR_NAME = "entries" +"""Directory containing one atomic subdirectory per cache key.""" -TEMP_ARTIFACT_PATTERN = re.compile( +CACHE_STAGING_DIR_NAME = "staging" +"""Directory containing in-progress materialization scratch.""" + +CACHE_TRASH_DIR_NAME = "trash" +"""Directory containing atomically retired entries pending physical deletion.""" + +LEGACY_CACHE_ENTRY_PREFIXES = ("squashfs-", "unsquashfs-", "tarball-") +"""Flat-layout prefixes retained only for one-way startup cleanup.""" + +LEGACY_TEMP_ARTIFACT_PATTERN = re.compile( r"^[^.]+\.\d+\.\d+\.(?:squashfs|unsquashfs|tar\.gz|tmp)$" ) -"""Matches materialization scratch and doomed eviction paths.""" +"""Matches scratch paths created by the former flat cache layout.""" type MountSlotReleaser = Callable[[str], Awaitable[bool]] """Unmounts one idle artifact, excluding the given cache key.""" @@ -76,6 +85,7 @@ class SquashfsMountCommandError(RuntimeError): class RegistryArtifactPaths: """Executor-local cache paths for one registry artifact key.""" + entry_dir: Path squashfs_image_path: Path squashfs_mount_dir: Path squashfs_extract_dir: Path @@ -113,7 +123,7 @@ class RegistryArtifactMaterializationContext: """Shared runtime state for artifact materialization.""" cache_key: str - cache_dir: Path + staging_dir: Path paths: RegistryArtifactPaths squashfs_mount_state: SquashfsMountState mount_slot_releaser: MountSlotReleaser | None = None @@ -174,7 +184,8 @@ def _temp_path( suffix: str, ) -> Path: unique_id = id(asyncio.current_task()) - return ctx.cache_dir / f"{self.cache_key}.{os.getpid()}.{unique_id}{suffix}" + ctx.staging_dir.mkdir(parents=True, exist_ok=True) + return ctx.staging_dir / f"{self.cache_key}.{os.getpid()}.{unique_id}{suffix}" @dataclass(frozen=True, slots=True) @@ -334,6 +345,7 @@ async def download( if image_path.exists(): return 0.0 + image_path.parent.mkdir(parents=True, exist_ok=True) temp_image = self._temp_path(ctx, ".squashfs") try: download_start = time.monotonic() @@ -370,7 +382,7 @@ async def mount( ctx.record_squashfs_mount() return target_dir - ctx.cache_dir.mkdir(parents=True, exist_ok=True) + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) target_dir.mkdir(parents=True, exist_ok=True) logger.info( @@ -408,7 +420,7 @@ async def extract( if target_dir.exists(): return target_dir - ctx.cache_dir.mkdir(parents=True, exist_ok=True) + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) logger.info( "Extracting SquashFS registry artifact", @@ -567,7 +579,7 @@ async def materialize( temp_dir = self._temp_path(ctx, ".tmp") try: - ctx.cache_dir.mkdir(parents=True, exist_ok=True) + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) download_start = time.monotonic() await self.download(ctx, temp_tarball) @@ -756,38 +768,18 @@ def _is_cache_entry_uri(artifact_uri: str) -> bool: return _bundled_builtin_registry_version(artifact_uri) is None -def _cache_key_from_entry_name(name: str) -> str | None: - """Return the cache key owning a cache directory entry name, if any.""" - if TEMP_ARTIFACT_PATTERN.fullmatch(name) is not None: - return None - for prefix in CACHE_ENTRY_PREFIXES: - if name.startswith(prefix): - cache_key = name.removeprefix(prefix).removesuffix(".squashfs") - return cache_key or None - return None - - -def _directory_footprint(directory: Path) -> tuple[int, float]: - """Return the total file size and local creation time of a cache directory. - - Directory ``mtime`` values inside extracted artifacts come from the artifact - build, so ``ctime`` (updated when the staging directory is renamed into - place) is used as the local recency signal instead. +def _directory_footprint(directory: Path) -> int: + """Return the total file size of a cache directory. Args: directory: Cache directory to measure. Returns: - Total byte size of contained files and the directory's ctime, or - ``(0, 0.0)`` when the directory is missing. + Total byte size of contained files, or zero when the directory is + missing. """ if not directory.is_dir(): - return 0, 0.0 - - try: - created_at = directory.stat().st_ctime - except OSError: - created_at = 0.0 + return 0 total_bytes = 0 for root, _dirs, files in os.walk(directory): @@ -796,7 +788,7 @@ def _directory_footprint(directory: Path) -> tuple[int, float]: total_bytes += os.lstat(os.path.join(root, file_name)).st_size except OSError: continue - return total_bytes, created_at + return total_bytes def _delete_cache_path(path: Path) -> bool: @@ -816,68 +808,22 @@ def _delete_cache_path(path: Path) -> bool: return True -def _delete_entry_paths(paths: RegistryArtifactPaths) -> bool: - """Best-effort delete every path owned by a registry artifact cache key. - - The caller must unmount ``squashfs_mount_dir`` first: unlinking the image - file behind a live mount leaves an open-file zombie pinning a loop device. - - Returns: - Whether every path was deleted. Failures are logged and returned instead - of raised so cache cleanup never rejects executor work. - """ - deleted = True - for path in ( - paths.squashfs_extract_dir, - paths.tarball_target_dir, - paths.squashfs_image_path, - paths.squashfs_mount_dir, - ): - if not _delete_cache_path(path): - deleted = False - return deleted - - -def _rename_entry_paths( - paths: RegistryArtifactPaths, - *, - cache_dir: Path, - cache_key: str, -) -> RegistryArtifactPaths: - """Synchronously rename live entry paths to unique startup-sweep scratch.""" +def _unique_work_path(root: Path, cache_key: str) -> Path: + """Return a unique path beneath a cache work directory.""" + root.mkdir(parents=True, exist_ok=True) unique_id = time.time_ns() while True: - doomed_paths = RegistryArtifactPaths( - squashfs_image_path=cache_dir - / f"{cache_key}.{os.getpid()}.{unique_id}.squashfs", - squashfs_mount_dir=cache_dir / f"{cache_key}.{os.getpid()}.{unique_id}.tmp", - squashfs_extract_dir=cache_dir - / f"{cache_key}.{os.getpid()}.{unique_id}.unsquashfs", - tarball_target_dir=cache_dir - / f"{cache_key}.{os.getpid()}.{unique_id}.tar.gz", - ) - if not any( - path.exists() - for path in ( - doomed_paths.squashfs_image_path, - doomed_paths.squashfs_mount_dir, - doomed_paths.squashfs_extract_dir, - doomed_paths.tarball_target_dir, - ) - ): - break + path = root / f"{cache_key}.{os.getpid()}.{unique_id}" + if not path.exists(): + return path unique_id += 1 - for source, target in ( - (paths.squashfs_extract_dir, doomed_paths.squashfs_extract_dir), - (paths.tarball_target_dir, doomed_paths.tarball_target_dir), - (paths.squashfs_image_path, doomed_paths.squashfs_image_path), - (paths.squashfs_mount_dir, doomed_paths.squashfs_mount_dir), - ): - if source.exists(): - source.rename(target) - return doomed_paths +def _move_entry_to_trash(entry_dir: Path, trash_dir: Path, cache_key: str) -> Path: + """Atomically retire one cache entry and return its trash path.""" + trash_path = _unique_work_path(trash_dir, cache_key) + entry_dir.rename(trash_path) + return trash_path class RegistryArtifactCache: @@ -885,6 +831,9 @@ class RegistryArtifactCache: def __init__(self, cache_dir: Path): self.cache_dir = cache_dir + self.entries_dir = cache_dir / CACHE_ENTRIES_DIR_NAME + self.staging_dir = cache_dir / CACHE_STAGING_DIR_NAME + self.trash_dir = cache_dir / CACHE_TRASH_DIR_NAME # Per-key locks live for the process lifetime: eviction and lease # admission must serialize on the same object for a given key, so a # lock is never dropped and re-created underneath a waiter. @@ -897,6 +846,10 @@ def __init__(self, cache_dir: Path): self._sweep_lock = asyncio.Lock() self._squashfs_mount_state = SquashfsMountState() self._leases: dict[str, RegistryArtifactLease] = {} + # Only exact paths whose deletion has already failed are retried at + # runtime. Never sweep the whole staging/trash directory while live + # materialization or deletion threads may still own its children. + self._pending_cleanup: set[Path] = set() # Whether the on-disk cache may exceed its budget. Set when a new entry # is materialized and cleared once enforcement measures a cache that # fits, so steady-state cache hits never pay for a disk scan. @@ -964,19 +917,10 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat try: registry_paths: list[Path] = [] for artifact_uri in artifact_uris: - cache_key = compute_registry_artifact_cache_key(artifact_uri) - if not _is_cache_entry_uri(artifact_uri): - registry_paths.extend( - await self.materialize(cache_key, artifact_uri) - ) - continue - - cached_paths = await self._admit_lease(cache_key, artifact_uri) - leased_keys.append(cache_key) - if cached_paths is not None: - registry_paths.extend(cached_paths) - continue - registry_paths.extend(await self.materialize(cache_key, artifact_uri)) + cache_key, artifact_paths = await self._lease_artifact(artifact_uri) + if cache_key is not None: + leased_keys.append(cache_key) + registry_paths.extend(artifact_paths) logger.info( "Using registry artifact environments", count=len(registry_paths), @@ -988,92 +932,77 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat if leased_keys: await self._converge_cache_budget() - async def _admit_lease( - self, cache_key: str, artifact_uri: str - ) -> list[Path] | None: - """Pin a cache entry and return its already-materialized paths, if any. - - The refcount increment and the cached-path check run under the same - per-key lock that eviction holds. An in-flight eviction therefore always - completes before a lease is admitted, and once the refcount is raised no - eviction can delete the entry, so a lease can never be handed a path - that is about to disappear. - - Args: - cache_key: Cache key to pin. - artifact_uri: Registry artifact URI backing the cache key. - - Returns: - Importable paths when the entry is already materialized, else None. + async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Path]]: + """Pin and materialize one artifact, returning its releasable cache key.""" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + ctx = self._context_for(cache_key) + if not _is_cache_entry_uri(artifact_uri): + return None, await self._materialize_candidates(ctx, artifact_uri) - Raises: - BaseException: Any failure or cancellation while resolving artifact - candidates. The lease is released before it propagates. - """ lock = await self._lock_for(cache_key) async with lock: self._acquire_lease(cache_key) try: - ctx = self._context_for(cache_key) candidates = await self._artifact_candidates(ctx, artifact_uri) except BaseException: - # Cancellation is a BaseException and must not leak the pin. self._release_lease(cache_key) raise - return self._first_cached_path(candidates, ctx) + if cached_paths := self._first_cached_path(candidates, ctx): + return cache_key, cached_paths - async def materialize(self, cache_key: str, artifact_uri: str) -> list[Path]: - """Materialize a registry artifact as local importable directories.""" - await self.ensure_swept() + try: + # Make room outside the per-key lock so eviction never nests key locks. + await self._enforce_cache_budget(protected_key=cache_key) + async with lock: + paths = await self._materialize_candidates(ctx, artifact_uri) + self._touch_entry(cache_key) + return cache_key, paths + except BaseException: + self._release_lease(cache_key) + raise - ctx = self._context_for(cache_key) - candidates = await self._artifact_candidates(ctx, artifact_uri) + async def _materialize_candidates( + self, + ctx: RegistryArtifactMaterializationContext, + artifact_uri: str, + ) -> list[Path]: + """Materialize the first viable artifact candidate. + Callers hold the cache key's lock for evictable entries. + """ + candidates = await self._artifact_candidates(ctx, artifact_uri) if cached_paths := self._first_cached_path(candidates, ctx): return cached_paths - if _is_cache_entry_uri(artifact_uri): - # Make room before downloading or expanding a new entry. This runs - # outside the per-key lock so eviction never nests key locks. - await self._enforce_cache_budget(protected_key=cache_key) - - lock = await self._lock_for(cache_key) - async with lock: - candidates = await self._artifact_candidates(ctx, artifact_uri) - if cached_paths := self._first_cached_path(candidates, ctx): - return cached_paths - - for index, artifact in enumerate(candidates): + cache_key = ctx.cache_key + for index, artifact in enumerate(candidates): + try: + logger.info( + "Trying registry artifact candidate", + cache_key=cache_key, + artifact_uri=artifact.uri, + artifact_format=artifact.format.value, + candidate=index + 1, + candidates=len(candidates), + ) try: - logger.info( - "Trying registry artifact candidate", - cache_key=cache_key, - artifact_uri=artifact.uri, - artifact_format=artifact.format.value, - candidate=index + 1, - candidates=len(candidates), - ) - try: - registry_paths = await artifact.materialize(ctx) - finally: - if _is_cache_entry_uri(artifact.uri): - # Arm after the attempt completes, whether it - # succeeded, failed, or was cancelled: any attempt - # may deposit canonical bytes, and a concurrent - # convergence pass may have consumed an earlier - # signal before those bytes landed. - self._budget_dirty = True - return registry_paths - except Exception as e: - if index == len(candidates) - 1: - raise - logger.warning( - "Failed to materialize registry artifact candidate, trying fallback", - cache_key=cache_key, - artifact_uri=artifact.uri, - artifact_format=artifact.format.value, - error=str(e), - ) + registry_paths = await artifact.materialize(ctx) + finally: + if _is_cache_entry_uri(artifact.uri): + # Any attempt may deposit canonical bytes, even when it + # fails or is cancelled. + self._budget_dirty = True + return registry_paths + except Exception as e: + if index == len(candidates) - 1: + raise + logger.warning( + "Failed to materialize registry artifact candidate, trying fallback", + cache_key=cache_key, + artifact_uri=artifact.uri, + artifact_format=artifact.format.value, + error=str(e), + ) raise RuntimeError(f"No registry artifact candidates for {artifact_uri}") @@ -1088,7 +1017,7 @@ def _context_for(self, cache_key: str) -> RegistryArtifactMaterializationContext """Return a materialization context for a registry artifact key.""" return RegistryArtifactMaterializationContext( cache_key=cache_key, - cache_dir=self.cache_dir, + staging_dir=self.staging_dir, paths=self._paths_for(cache_key), squashfs_mount_state=self._squashfs_mount_state, mount_slot_releaser=self._release_mounted_slot, @@ -1125,40 +1054,25 @@ def _refcount(self, cache_key: str) -> int: return 0 if lease is None else lease.refcount def _touch_entry(self, cache_key: str) -> None: - """Best-effort refresh of artifact entry mtimes for restart-safe LRU. - - Both the image mtime and tarball root mtime are restart-safe recency - signals. - """ - paths = self._paths_for(cache_key) - touched = False - try: - os.utime(paths.squashfs_image_path) - except OSError: - pass - else: - touched = True - + """Best-effort refresh of the entry-root mtime for restart-safe LRU.""" + entry_dir = self._paths_for(cache_key).entry_dir try: - os.utime(paths.tarball_target_dir) + os.utime(entry_dir) except OSError: - pass - else: - touched = True - - if not touched: logger.debug( - "Could not refresh registry artifact entry mtimes", + "Could not refresh registry artifact entry mtime", cache_key=cache_key, ) def _paths_for(self, cache_key: str) -> RegistryArtifactPaths: """Return local cache paths for a registry artifact key.""" + entry_dir = self.entries_dir / cache_key return RegistryArtifactPaths( - squashfs_image_path=self.cache_dir / f"squashfs-{cache_key}.squashfs", - squashfs_mount_dir=self.cache_dir / f"squashfs-{cache_key}", - squashfs_extract_dir=self.cache_dir / f"unsquashfs-{cache_key}", - tarball_target_dir=self.cache_dir / f"tarball-{cache_key}", + entry_dir=entry_dir, + squashfs_image_path=entry_dir / "image.squashfs", + squashfs_mount_dir=entry_dir / "mount", + squashfs_extract_dir=entry_dir / "extracted", + tarball_target_dir=entry_dir / "tarball", ) def _first_cached_path( @@ -1321,10 +1235,15 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo Whether the cache is within budget once eviction has finished. """ async with self._budget_lock: + legacy_clean, pending_clean = await asyncio.gather( + asyncio.to_thread(self._remove_legacy_cache_paths), + asyncio.to_thread(self._retry_pending_cleanup), + ) + cleanup_complete = legacy_clean and pending_clean max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES if max_entries <= 0 and max_bytes <= 0: - return True + return cleanup_complete entries = await asyncio.to_thread(self._scan_cache_entries) # The protected key is not on disk yet when it is a fresh entry. @@ -1356,10 +1275,11 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo if await self._evict_entry(candidate.cache_key): del entries[candidate.cache_key] total_bytes -= candidate.size_bytes + cleanup_complete = not self._pending_cleanup and cleanup_complete else: skipped.add(candidate.cache_key) - return True + return cleanup_complete async def _release_mounted_slot(self, protected_key: str) -> bool: """Unmount one idle artifact so its loop device can be reused. @@ -1486,18 +1406,15 @@ async def _evict_entry(self, cache_key: str) -> bool: cannot be unmounted: deleting the image file behind a live mount would leave an open-file zombie holding the loop device. - After unmounting, live paths are synchronously renamed under the per-key - lock to scratch names ignored by cache discovery and budget accounting. - The lease record is then dropped and the lock released before physical - deletion runs in a worker thread. If that await is cancelled, the live - key remains a clean cache miss and the next startup sweep removes any - doomed scratch left behind. + After unmounting, the entry root is atomically renamed into ``trash`` + under the per-key lock. The lock is then released before physical + deletion runs in a worker thread. Args: cache_key: Cache key to evict. Returns: - Whether the entry was removed. + Whether the entry was atomically retired from the active cache. """ lock = await self._lock_for(cache_key) if lock.locked(): @@ -1512,6 +1429,9 @@ async def _evict_entry(self, cache_key: str) -> bool: return False paths = self._paths_for(cache_key) + if not paths.entry_dir.exists(): + self._leases.pop(cache_key, None) + return True if paths.squashfs_mount_dir.is_mount() and not await self._unmount( paths.squashfs_mount_dir ): @@ -1522,21 +1442,37 @@ async def _evict_entry(self, cache_key: str) -> bool: ) return False - doomed_paths = _rename_entry_paths( - paths, - cache_dir=self.cache_dir, - cache_key=cache_key, - ) + try: + trash_path = _move_entry_to_trash( + paths.entry_dir, + self.trash_dir, + cache_key, + ) + except OSError as e: + self._budget_dirty = True + logger.warning( + "Failed to retire registry artifact cache entry", + cache_key=cache_key, + entry_dir=str(paths.entry_dir), + error=str(e), + ) + return False self._leases.pop(cache_key, None) - deleted = await asyncio.to_thread(_delete_entry_paths, doomed_paths) + try: + deleted = await asyncio.to_thread(_delete_cache_path, trash_path) + except BaseException: + self._budget_dirty = True + raise if not deleted: + self._pending_cleanup.add(trash_path) self._budget_dirty = True logger.warning( "Registry artifact eviction remains pending physical deletion", cache_key=cache_key, ) - return False + return True + self._pending_cleanup.discard(trash_path) logger.info("Evicted registry artifact from cache", cache_key=cache_key) return True @@ -1585,17 +1521,17 @@ def _scan_cache_entries(self) -> dict[str, RegistryArtifactCacheEntry]: } def _discover_cache_keys(self) -> set[str]: - """Return the cache keys with at least one path in the cache directory.""" + """Return cache keys represented by atomic entry directories.""" try: - names = os.listdir(self.cache_dir) + entries = list(os.scandir(self.entries_dir)) except OSError: return set() - cache_keys: set[str] = set() - for name in names: - if (cache_key := _cache_key_from_entry_name(name)) is not None: - cache_keys.add(cache_key) - return cache_keys + return { + entry.name + for entry in entries + if entry.name and entry.is_dir(follow_symlinks=False) + } def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: """Measure the on-disk footprint and recency of one cache entry. @@ -1606,9 +1542,6 @@ def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: """ paths = self._paths_for(cache_key) size_bytes = 0 - image_mtime = 0.0 - tarball_root_mtime = 0.0 - created_at = 0.0 try: image_stat = paths.squashfs_image_path.stat() @@ -1616,21 +1549,14 @@ def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: pass else: size_bytes += image_stat.st_size - image_mtime = image_stat.st_mtime - - try: - tarball_root_mtime = paths.tarball_target_dir.stat().st_mtime - except OSError: - pass for directory in (paths.squashfs_extract_dir, paths.tarball_target_dir): - directory_bytes, directory_created_at = _directory_footprint(directory) - size_bytes += directory_bytes - created_at = max(created_at, directory_created_at) + size_bytes += _directory_footprint(directory) - # Image and tarball-root mtimes are refreshed on lease; directory ctimes - # remain the fallback for entries that have neither. - last_used = max(image_mtime, tarball_root_mtime) or created_at + try: + last_used = paths.entry_dir.stat().st_mtime + except OSError: + last_used = 0.0 return RegistryArtifactCacheEntry( cache_key=cache_key, @@ -1641,14 +1567,11 @@ def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: def _sweep_startup_state(self) -> None: """Reclaim orphaned cache state left behind by a previous process. - Mounts never survive a container restart, so any non-mountpoint mount - directory is stale. Scratch paths from interrupted materializations are - removed, and the cache is trimmed to budget using image mtimes as LRU - order. Scratch and stale empty mount directories are never worker import - paths, so they are removed unconditionally. Tarball-bearing entries are - protected for the pool backend because cache construction may happen - after warm workers have already inherited those paths. A missing or - empty cache directory is a no-op. + Scratch and trash paths from interrupted work are removed, legacy flat + cache paths are retired, and active entries are trimmed to budget using + entry-root mtimes as LRU order. Tarball-bearing entries are protected + for the pool backend because cache construction may happen after warm + workers have inherited those paths. The worker warms this sweep before activities can run; lazy first-use sweeping remains a safe fallback. @@ -1658,9 +1581,17 @@ def _sweep_startup_state(self) -> None: return try: - self._remove_orphaned_temp_paths() - self._remove_stale_mount_dirs() - self._trim_startup_cache() + cleanup_complete = self._remove_legacy_cache_paths() + cleanup_complete = ( + self._clear_work_dir(self.staging_dir, remember_failures=True) + and cleanup_complete + ) + cleanup_complete = ( + self._clear_work_dir(self.trash_dir, remember_failures=True) + and cleanup_complete + ) + within_budget = self._trim_startup_cache() + self._budget_dirty = not (cleanup_complete and within_budget) except OSError as e: logger.warning( "Failed to sweep registry artifact cache", @@ -1669,55 +1600,109 @@ def _sweep_startup_state(self) -> None: ) raise - def _remove_orphaned_temp_paths(self) -> None: - """Delete every materialization scratch path during startup. + def _clear_work_dir( + self, + work_dir: Path, + *, + remember_failures: bool = False, + ) -> bool: + """Best-effort remove every child of a staging or trash directory.""" + try: + paths = list(work_dir.iterdir()) + except FileNotFoundError: + return True + except OSError as e: + logger.warning( + "Failed to inspect registry artifact work directory", + path=str(work_dir), + error=str(e), + ) + return False - The sweep runs before the first lease or materialization in this process. - Every matching path is therefore interrupted scratch from an earlier - process and is safe to remove even when the operating system reused that - process's PID. - """ - for name in os.listdir(self.cache_dir): - if TEMP_ARTIFACT_PATTERN.match(name) is None: - continue - path = self.cache_dir / name - if path.is_dir(): - shutil.rmtree(path, ignore_errors=True) + deleted = True + for path in paths: + if _delete_cache_path(path): + if remember_failures: + self._pending_cleanup.discard(path) + logger.info( + "Removed registry artifact work path", + path=str(path), + ) else: - path.unlink(missing_ok=True) - logger.info("Removed orphaned registry artifact scratch path", path=name) + deleted = False + if remember_failures: + self._pending_cleanup.add(path) + return deleted + + def _retry_pending_cleanup(self) -> bool: + """Retry exact failed cleanup paths without touching live work.""" + for path in tuple(self._pending_cleanup): + if _delete_cache_path(path): + self._pending_cleanup.discard(path) + return not self._pending_cleanup + + def _remove_legacy_cache_paths(self) -> bool: + """Best-effort remove flat-layout cache paths from earlier executors.""" + try: + paths = list(self.cache_dir.iterdir()) + except FileNotFoundError: + return True + except OSError as e: + logger.warning( + "Failed to inspect registry artifact cache root", + cache_dir=str(self.cache_dir), + error=str(e), + ) + return False - def _remove_stale_mount_dirs(self) -> None: - """Remove empty mount directories left over from a previous process.""" - for cache_key in self._discover_cache_keys(): - mount_dir = self._paths_for(cache_key).squashfs_mount_dir - if not mount_dir.is_dir() or mount_dir.is_mount(): - continue - try: - mount_dir.rmdir() - except OSError: + current_names = { + BASE_PYTHONPATH_DIR_NAME, + CACHE_ENTRIES_DIR_NAME, + CACHE_STAGING_DIR_NAME, + CACHE_TRASH_DIR_NAME, + } + backend_type = resolve_backend_type() + deleted = True + for path in paths: + if path.name in current_names: continue - logger.debug( - "Removed stale registry artifact mount directory", - cache_key=cache_key, + is_legacy_entry = path.name.startswith(LEGACY_CACHE_ENTRY_PREFIXES) + is_legacy_scratch = ( + LEGACY_TEMP_ARTIFACT_PATTERN.fullmatch(path.name) is not None ) + if not is_legacy_entry and not is_legacy_scratch: + continue + if backend_type == ExecutorBackendType.POOL and path.name.startswith( + "tarball-" + ): + continue + if path.is_mount(): + deleted = False + logger.warning( + "Cannot remove mounted legacy registry artifact path", + path=str(path), + ) + continue + if not _delete_cache_path(path): + deleted = False + return deleted - def _trim_startup_cache(self) -> None: + def _trim_startup_cache(self) -> bool: """Trim the cache to budget before any artifact is leased. Tarball-bearing entries are ineligible when the resolved backend is the pool because existing workers may already import from those paths. - Clears the budget-dirty flag when the cache ends up within budget, so a - healthy cache never rescans until a new entry is materialized. + Returns whether active entries and pending physical deletion fit within + the configured budget. """ max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES if max_entries <= 0 and max_bytes <= 0: - self._budget_dirty = False - return + return True entries = self._scan_cache_entries() total_bytes = sum(entry.size_bytes for entry in entries.values()) + cleanup_complete = True backend_type = resolve_backend_type() # Mounted entries belong to a live process sharing this cache directory. candidates = sorted( @@ -1738,13 +1723,33 @@ def within_budget() -> bool: for entry in candidates: if within_budget(): break - _delete_entry_paths(self._paths_for(entry.cache_key)) + paths = self._paths_for(entry.cache_key) + try: + trash_path = _move_entry_to_trash( + paths.entry_dir, + self.trash_dir, + entry.cache_key, + ) + except OSError as e: + cleanup_complete = False + logger.warning( + "Failed to retire stale registry artifact during startup sweep", + cache_key=entry.cache_key, + entry_dir=str(paths.entry_dir), + error=str(e), + ) + continue + del entries[entry.cache_key] - total_bytes -= entry.size_bytes + if _delete_cache_path(trash_path): + total_bytes -= entry.size_bytes + else: + self._pending_cleanup.add(trash_path) + cleanup_complete = False logger.info( "Evicted stale registry artifact during startup sweep", cache_key=entry.cache_key, size_bytes=entry.size_bytes, ) - self._budget_dirty = not within_budget() + return cleanup_complete and within_budget() From 11350dde654cbeb7aa901177933aa9e0a9c9e468 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:58:43 -0400 Subject: [PATCH 025/161] refactor(executor): simplify registry cache lifecycle --- ...registry_artifact_cache_mount_lifecycle.py | 159 ++-- tests/unit/test_executor_sandbox_nsjail.py | 3 +- tests/unit/test_registry_artifacts.py | 717 ++++++++---------- tracecat/executor/registry_artifacts.py | 460 +++++------ 4 files changed, 556 insertions(+), 783 deletions(-) diff --git a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py index 13c2bef97c..6a3830ef04 100644 --- a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py +++ b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py @@ -1,11 +1,8 @@ """Dockerized lifecycle test for executor registry artifact SquashFS mounts. -The executor caches registry environments as SquashFS images and mounts them, -which consumes one loop device per mounted artifact. Unit tests stub the mount -and umount commands, so this test drives the real thing: it runs inside the -privileged executor image, builds tiny SquashFS images with ``mksquashfs``, and -asserts that materialization mounts them, that eviction unmounts them and -releases their loop devices, and that the startup sweep reclaims stale state. +The executor retains SquashFS images but only keeps them mounted while leased. +Unit tests stub the mount and umount commands, so this test drives the real +mount lifecycle inside the privileged executor image. Run it with the ``integration`` marker; it is skipped when Docker is unavailable. """ @@ -154,32 +151,25 @@ def test_registry_artifact_cache_mount_lifecycle() -> None: if skipped := payload.get("skipped"): pytest.skip(f"SquashFS mounts unsupported in this container: {skipped}") - # Every materialized artifact is mounted and holds its own loop device. - assert payload["mounted_targets"] == 3 - assert payload["mounted_loop_devices"] == 3 + # Each lease mounts its image and its final release frees the loop device. assert payload["module_readable_through_mount"] is True + assert payload["lease_mounts_observed"] == 3 + assert payload["all_targets_unmounted_after_release"] is True + assert payload["all_loop_devices_released"] is True + assert payload["images_retained_after_release"] is True - # Eviction unmounts, frees the loop device, and deletes the entry. + # A later lease remounts from the retained image without downloading again. + assert payload["remounted_from_retained_image"] is True + assert payload["remount_released"] is True + + # Eviction deletes an already-idle entry. assert payload["evicted"] is True - assert payload["evicted_target_unmounted"] is True - assert payload["evicted_loop_device_released"] is True assert payload["evicted_paths_removed"] is True - # Re-materializing the evicted key mounts again: no sticky disable flag. - assert payload["remounted"] is True - assert payload["squashfs_disabled_after_eviction"] is False - - # Loop-device reclamation unmounts an idle entry but retains its cached image. - assert payload["mount_slot_released"] is True - assert payload["released_target_unmounted"] is True - assert payload["released_loop_device_released"] is True - assert payload["released_image_retained"] is True - assert payload["remounted_from_retained_image"] is True - - # Releasing a lease converges an over-budget cache, destructively evicting it. - assert payload["converged_entry_unmounted"] is True + # Releasing a lease unmounts first, then converges the disk cache. + assert payload["converged_lease_unmounted"] is True assert payload["converged_loop_device_released"] is True - assert payload["converged_entries_remaining"] == 3 + assert payload["converged_entries_remaining"] == 2 # The startup sweep trims to budget and removes stale mount directories. assert payload["startup_sweep_trimmed"] is True @@ -255,11 +245,14 @@ async def _run_mount_lifecycle_child() -> None: f"module_{index}.py", ) - # (a) Materialize every artifact through the real cache. + # (a) Every lease mounts its image; final release unmounts it. module_readable = True + lease_mounts_observed = 0 + all_targets_unmounted = True + all_loop_devices_released = True for index, uri in enumerate(uris): + mount_dir = cache._paths_for(keys[index]).squashfs_mount_dir async with cache.lease([uri]) as registry_paths: - mount_dir = cache._paths_for(keys[index]).squashfs_mount_dir if registry_paths != [mount_dir]: if index == 0 and not mount_dir.is_mount(): # Kernels without SquashFS or loop support fall back to @@ -279,92 +272,70 @@ async def _run_mount_lifecycle_child() -> None: module_readable and (mount_dir / f"module_{index}.py").read_text() == "VALUE = 1\n" ) + mounts = _squashfs_mounts(cache_dir) + device = mounts[str(mount_dir)] + lease_mounts_observed += 1 - mounts = _squashfs_mounts(cache_dir) - payload["mounted_targets"] = len(mounts) - payload["mounted_loop_devices"] = len( - {device for device in mounts.values() if device.startswith("/dev/loop")} - ) - payload["module_readable_through_mount"] = module_readable + mounts_after_release = _squashfs_mounts(cache_dir) + all_targets_unmounted = ( + all_targets_unmounted and str(mount_dir) not in mounts_after_release + ) + all_loop_devices_released = ( + all_loop_devices_released + and device not in mounts_after_release.values() + ) - # (b) Evict one entry: it must unmount, free its loop device, and vanish. - evicted_paths = cache._paths_for(keys[0]) - evicted_device = mounts[str(evicted_paths.squashfs_mount_dir)] - payload["evicted"] = await cache._evict_entry(keys[0]) - mounts_after_eviction = _squashfs_mounts(cache_dir) - payload["evicted_target_unmounted"] = ( - str(evicted_paths.squashfs_mount_dir) not in mounts_after_eviction - ) - payload["evicted_loop_device_released"] = ( - evicted_device not in mounts_after_eviction.values() - ) - payload["evicted_paths_removed"] = not ( - evicted_paths.squashfs_image_path.exists() - or evicted_paths.squashfs_mount_dir.exists() + payload["module_readable_through_mount"] = module_readable + payload["lease_mounts_observed"] = lease_mounts_observed + payload["all_targets_unmounted_after_release"] = all_targets_unmounted + payload["all_loop_devices_released"] = all_loop_devices_released + payload["images_retained_after_release"] = all( + cache._paths_for(cache_key).squashfs_image_path.exists() + for cache_key in keys ) - # (c) The evicted key must mount again: eviction is not a capability probe. - _build_squashfs_image( - root / "source-0", - evicted_paths.squashfs_image_path, - "module_0.py", - ) + # (b) A later lease remounts from the retained image. + retained_paths = cache._paths_for(keys[0]) async with cache.lease([uris[0]]) as registry_paths: - payload["remounted"] = registry_paths == [evicted_paths.squashfs_mount_dir] - payload["squashfs_disabled_after_eviction"] = ( - cache._squashfs_mount_state.disabled - ) - - # (d) Loop-device reclamation unmounts the LRU idle artifact while - # retaining its image, and a later lease remounts that image directly. - retained_paths = cache._paths_for(keys[1]) - mounts_before_release = _squashfs_mounts(cache_dir) - retained_device = mounts_before_release[str(retained_paths.squashfs_mount_dir)] - payload["mount_slot_released"] = await cache._release_mounted_slot(keys[0]) - mounts_after_release = _squashfs_mounts(cache_dir) - payload["released_target_unmounted"] = ( - str(retained_paths.squashfs_mount_dir) not in mounts_after_release - ) - payload["released_loop_device_released"] = ( - retained_device not in mounts_after_release.values() - ) - payload["released_image_retained"] = retained_paths.squashfs_image_path.exists() - async with cache.lease([uris[1]]) as registry_paths: payload["remounted_from_retained_image"] = registry_paths == [ retained_paths.squashfs_mount_dir ] + payload["remount_released"] = not retained_paths.squashfs_mount_dir.is_mount() + + # (c) Evict one already-idle entry. + evicted_paths = cache._paths_for(keys[0]) + eviction = await cache._evict_entry(keys[0]) + payload["evicted"] = eviction.retired and eviction.reclaimed + payload["evicted_paths_removed"] = not ( + evicted_paths.squashfs_image_path.exists() + or evicted_paths.squashfs_mount_dir.exists() + ) - # (e) A released lease converges an over-budget cache, destructively - # evicting it. The fourth entry is materialized under the old budget, so - # only the release-time check can bring the cache back within the new one. + # (d) Final release unmounts, then converges an over-budget cache. fourth_uri = "s3://bucket/lifecycle/3/site-packages.squashfs" fourth_key = compute_registry_artifact_cache_key(fourth_uri) + fourth_paths = cache._paths_for(fourth_key) _build_squashfs_image( root / "source-3", - cache._paths_for(fourth_key).squashfs_image_path, + fourth_paths.squashfs_image_path, "module_3.py", ) - # keys[2] is the least recently used idle entry: keys[0] was re-leased - # above and keys[1] is leased again below. - converged_paths = cache._paths_for(keys[2]) async with cache.lease([fourth_uri]): mounts_before_converge = _squashfs_mounts(cache_dir) converged_device = mounts_before_converge[ - str(converged_paths.squashfs_mount_dir) + str(fourth_paths.squashfs_mount_dir) ] - config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES = 3 - async with cache.lease([uris[1]]): - pass - mounts_after_converge = _squashfs_mounts(cache_dir) - payload["converged_entry_unmounted"] = ( - str(converged_paths.squashfs_mount_dir) not in mounts_after_converge - ) - payload["converged_loop_device_released"] = ( - converged_device not in mounts_after_converge.values() - ) - payload["converged_entries_remaining"] = len(cache._discover_cache_keys()) + config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES = 2 + mounts_after_converge = _squashfs_mounts(cache_dir) + payload["converged_lease_unmounted"] = ( + str(fourth_paths.squashfs_mount_dir) not in mounts_after_converge + ) + payload["converged_loop_device_released"] = ( + converged_device not in mounts_after_converge.values() + ) + payload["converged_entries_remaining"] = len(cache._discover_cache_keys()) - # (f) The startup sweep trims to budget and drops stale mount directories. + # (e) The startup sweep trims to budget and drops stale mount directories. sweep_dir = root / "sweep-cache" sweep_dir.mkdir() sweep_cache = RegistryArtifactCache(sweep_dir) diff --git a/tests/unit/test_executor_sandbox_nsjail.py b/tests/unit/test_executor_sandbox_nsjail.py index f920bf11c2..e2025ed634 100644 --- a/tests/unit/test_executor_sandbox_nsjail.py +++ b/tests/unit/test_executor_sandbox_nsjail.py @@ -607,7 +607,8 @@ async def download_artifact(self, ctx, output_path: Path) -> float: else: assert result["source"] == "/packages/0/registry_artifact_smoke_action.py" if smoke_case == SmokeCase.NSJAIL_SQUASHFS: - assert mount_dir.is_mount() + assert not mount_dir.is_mount() + assert cache_paths.squashfs_image_path.is_file() else: assert tarball_dir.exists() finally: diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index c89b207514..4ba58df544 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -8,7 +8,7 @@ import tempfile import threading from pathlib import Path -from unittest.mock import AsyncMock, patch +from unittest.mock import ANY, AsyncMock, patch import httpx import pytest @@ -17,6 +17,7 @@ from tracecat.executor.registry_artifacts import ( SQUASHFS_MOUNT_OPTIONS, RegistryArtifactCache, + RegistryArtifactEviction, RegistryArtifactFormat, SquashfsArtifact, SquashfsMountCommandError, @@ -75,7 +76,7 @@ async def _materialize( """Exercise internal materialization while releasing its test-only lease.""" await cache.ensure_swept() ctx = cache._context_for(cache_key) - lock = await cache._lock_for(cache_key) + lock = cache._runtime_for(cache_key).lock async with lock: cache._acquire_lease(cache_key) try: @@ -940,25 +941,23 @@ async def mock_extract(self, tarball_path, target_dir): assert all(r == results[0] for r in results) assert download_count == 1 - @pytest.mark.anyio - async def test_lock_for_same_key(self, temp_cache_dir): - """Test that same cache key returns same lock.""" + def test_runtime_for_same_key(self, temp_cache_dir): + """The same cache key returns the same runtime state.""" cache = RegistryArtifactCache(temp_cache_dir) - lock1 = await cache._lock_for("key1") - lock2 = await cache._lock_for("key1") + runtime1 = cache._runtime_for("key1") + runtime2 = cache._runtime_for("key1") - assert lock1 is lock2 + assert runtime1 is runtime2 - @pytest.mark.anyio - async def test_lock_for_different_keys(self, temp_cache_dir): - """Test that different cache keys return different locks.""" + def test_runtime_for_different_keys(self, temp_cache_dir): + """Different cache keys return different runtime state.""" cache = RegistryArtifactCache(temp_cache_dir) - lock1 = await cache._lock_for("key1") - lock2 = await cache._lock_for("key2") + runtime1 = cache._runtime_for("key1") + runtime2 = cache._runtime_for("key2") - assert lock1 is not lock2 + assert runtime1 is not runtime2 class TestRegistryArtifactCacheLease: @@ -1044,7 +1043,10 @@ async def take_lease() -> None: await acquisition assert cache._refcount(cache_key) == 0 - assert await cache._evict_entry(cache_key) is True + assert await cache._evict_entry(cache_key) == RegistryArtifactEviction( + retired=True, + reclaimed=True, + ) assert not target_dir.exists() @pytest.mark.anyio @@ -1056,7 +1058,7 @@ async def test_lease_without_uris_returns_base_pythonpath_dir(self, temp_cache_d assert registry_paths == [temp_cache_dir / "base"] assert registry_paths[0].is_dir() - assert cache._leases == {} + assert cache._runtime == {} @pytest.mark.anyio async def test_lease_preserves_uri_order(self, temp_cache_dir): @@ -1138,7 +1140,7 @@ async def take_lease() -> None: finish_umount.set() evicted, _ = await asyncio.gather(eviction, lease) - assert evicted is True + assert evicted == RegistryArtifactEviction(retired=True, reclaimed=True) # The lease waited for the eviction and re-materialized the entry. assert remounts == [cache_key] assert leased_paths == [paths.squashfs_mount_dir] @@ -1173,7 +1175,7 @@ async def test_builtin_artifact_is_exempt_from_cache_accounting( ) as enforce_cache_budget: async with cache.lease([bundled_builtin_registry_uri(version)]) as paths: assert paths == [site_packages.resolve()] - assert cache._leases == {} + assert cache._runtime == {} enforce_cache_budget.assert_not_awaited() @@ -1303,13 +1305,14 @@ async def mock_extract(self, tarball_path, target_dir): warning.assert_any_call( "Registry artifact eviction remains pending physical deletion", cache_key="idle", + trash_path=ANY, ) @pytest.mark.anyio async def test_failed_physical_delete_retries_without_extra_eviction( self, temp_cache_dir ): - """A retired entry counts as evicted while its exact trash path retries.""" + """Failed byte reclamation stops eviction until trash cleanup succeeds.""" cache = RegistryArtifactCache(temp_cache_dir) oldest = _write_image_entry( temp_cache_dir, @@ -1317,12 +1320,18 @@ async def test_failed_physical_delete_retries_without_extra_eviction( size=16, mtime=100.0, ) - retained = _write_image_entry( + older = _write_image_entry( temp_cache_dir, - "retained", + "older", size=16, mtime=200.0, ) + newest = _write_image_entry( + temp_cache_dir, + "newest", + size=16, + mtime=300.0, + ) real_delete = _delete_cache_path failed_once = False @@ -1334,8 +1343,8 @@ def fail_once(path: Path) -> bool: return real_delete(path) with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 16), patch( "tracecat.executor.registry_artifacts._delete_cache_path", side_effect=fail_once, @@ -1343,13 +1352,15 @@ def fail_once(path: Path) -> bool: ): assert await cache._enforce_cache_budget() is False assert not oldest.exists() - assert retained.exists() - assert cache._discover_cache_keys() == {"retained"} - assert len(cache._pending_cleanup) == 1 + assert older.exists() + assert newest.exists() + assert cache._discover_cache_keys() == {"older", "newest"} + assert len(tuple(cache.trash_dir.iterdir())) == 1 assert await cache._enforce_cache_budget() is True - assert cache._pending_cleanup == set() + assert not older.exists() + assert newest.exists() assert not any(cache.trash_dir.iterdir()) @pytest.mark.anyio @@ -1384,6 +1395,36 @@ async def mock_extract(self, tarball_path, target_dir): assert idle.is_file() assert cache._budget_dirty is True + @pytest.mark.anyio + async def test_cache_scan_failure_does_not_block_materialization( + self, temp_cache_dir + ): + """Maintenance errors are observable while artifact admission stays open.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/new-after-scan-failure.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch.object( + cache, + "_enforce_cache_budget", + new_callable=AsyncMock, + side_effect=PermissionError("denied"), + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [ + cache._paths_for(cache_key).tarball_target_dir + ] + @pytest.mark.anyio async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( self, temp_cache_dir @@ -1696,41 +1737,45 @@ async def test_pool_backend_all_tarballs_remain_dirty_when_over_budget( ) @pytest.mark.anyio - async def test_pool_backend_release_mounted_slot_skips_tarball_entry( - self, temp_cache_dir - ): - """Loop-device recovery must preserve paths visible to warm workers.""" - cache = RegistryArtifactCache(temp_cache_dir) - pool_visible = cache._paths_for("pool-visible") - pool_visible.entry_dir.mkdir(parents=True) - pool_visible.squashfs_image_path.write_bytes(b"squashfs") - os.utime(pool_visible.squashfs_image_path, (100.0, 100.0)) - pool_visible.squashfs_mount_dir.mkdir() - _write_tarball_entry(temp_cache_dir, "pool-visible") - eligible = cache._paths_for("eligible") - eligible.entry_dir.mkdir(parents=True) - eligible.squashfs_image_path.write_bytes(b"squashfs") - os.utime(eligible.squashfs_image_path, (200.0, 200.0)) - eligible.squashfs_mount_dir.mkdir() - mounted = { - pool_visible.squashfs_mount_dir, - eligible.squashfs_mount_dir, - } + async def test_final_lease_release_unmounts_and_retains_image(self, temp_cache_dir): + """An idle entry releases its loop device without deleting its image.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/site-packages.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + mounted = {paths.squashfs_mount_dir} + + process = AsyncMock() + process.communicate.return_value = (b"", b"") + process.returncode = 0 + + async def mock_umount(*args, **kwargs): + mounted.discard(paths.squashfs_mount_dir) + return process with ( - patch(BACKEND_CONFIG, ExecutorBackendType.POOL.value), patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/umount", + ), patch.object( - cache, - "_unmount_entry", - new_callable=AsyncMock, - return_value=True, - ) as unmount_entry, + asyncio, + "create_subprocess_exec", + side_effect=mock_umount, + ), ): - released = await cache._release_mounted_slot("protected") + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [paths.squashfs_mount_dir] + assert paths.squashfs_mount_dir in mounted - assert released is True - unmount_entry.assert_awaited_once_with("eligible") + assert paths.squashfs_mount_dir not in mounted + + assert paths.squashfs_image_path.read_bytes() == b"squashfs" + assert paths.squashfs_mount_dir.is_dir() @pytest.mark.anyio async def test_concurrent_budget_passes_only_evict_once(self, temp_cache_dir): @@ -1765,10 +1810,15 @@ def controlled_scan(): extra_eviction_finished = asyncio.Event() evicted_keys: list[str] = [] - async def controlled_evict(cache_key: str) -> bool: + async def controlled_evict( + cache_key: str, + ) -> RegistryArtifactEviction: if cache_key == "oldest": if eviction_started.is_set(): - return False + return RegistryArtifactEviction( + retired=False, + reclaimed=False, + ) eviction_started.set() await finish_eviction.wait() _delete_cache_path(cache._paths_for("oldest").entry_dir) @@ -1776,7 +1826,7 @@ async def controlled_evict(cache_key: str) -> bool: _delete_cache_path(cache._paths_for("retained").entry_dir) extra_eviction_finished.set() evicted_keys.append(cache_key) - return True + return RegistryArtifactEviction(retired=True, reclaimed=True) with ( patch.object(cache, "_scan_cache_entries", side_effect=controlled_scan), @@ -1880,7 +1930,7 @@ async def mock_umount(*args, **kwargs): ): evicted = await cache._evict_entry("mounted") - assert evicted is True + assert evicted == RegistryArtifactEviction(retired=True, reclaimed=True) assert image_present_at_umount == [True] assert not paths.squashfs_image_path.exists() assert not paths.squashfs_mount_dir.exists() @@ -1904,7 +1954,19 @@ async def test_cancelled_unmount_kills_and_reaps_before_releasing_key_lock( paths.squashfs_mount_dir.mkdir() (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") mounted = {paths.squashfs_mount_dir} - process = _BlockingSubprocess() + blocked_process = _BlockingSubprocess() + released_process = AsyncMock() + released_process.communicate.return_value = (b"", b"") + released_process.returncode = 0 + unmount_attempts = 0 + + async def mock_umount(*args, **kwargs): + nonlocal unmount_attempts + unmount_attempts += 1 + if unmount_attempts == 1: + return blocked_process + mounted.discard(paths.squashfs_mount_dir) + return released_process with ( patch.object(Path, "is_mount", lambda self: self in mounted), @@ -1914,18 +1976,17 @@ async def test_cancelled_unmount_kills_and_reaps_before_releasing_key_lock( ), patch( "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", - new_callable=AsyncMock, - return_value=process, + side_effect=mock_umount, ), ): eviction = asyncio.create_task(cache._evict_entry(cache_key)) - await process.communicate_started.wait() + await blocked_process.communicate_started.wait() eviction.cancel() with pytest.raises(asyncio.CancelledError): await eviction - assert process.cleanup_calls == ["kill", "wait"] + assert blocked_process.cleanup_calls == ["kill", "wait"] async with cache.lease([artifact_uri]) as registry_paths: assert registry_paths == [paths.squashfs_mount_dir] assert registry_paths[0].is_dir() @@ -1933,6 +1994,7 @@ async def test_cancelled_unmount_kills_and_reaps_before_releasing_key_lock( assert paths.squashfs_image_path.is_file() assert paths.squashfs_mount_dir.is_dir() + assert paths.squashfs_mount_dir not in mounted @pytest.mark.anyio async def test_cancelled_background_deletion_leaves_a_clean_miss( @@ -1974,25 +2036,25 @@ async def mock_extract(self, tarball_path, target_dir): eviction = asyncio.create_task(cache._evict_entry(cache_key)) assert await asyncio.to_thread(delete_started.wait, 1) eviction.cancel() - with pytest.raises(asyncio.CancelledError): - await eviction - try: + await asyncio.sleep(0) + assert not eviction.done() assert not original_target.exists() assert (doomed[0] / "tarball").is_dir() assert cache._discover_cache_keys() == set() - - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [original_target] - assert original_target.is_dir() - assert (doomed[0] / "tarball").is_dir() - assert (original_target / "module.py").read_text() == "VALUE = 2" finally: finish_delete.set() - assert await asyncio.to_thread(delete_finished.wait, 1) + with pytest.raises(asyncio.CancelledError): + await eviction + assert await asyncio.to_thread(delete_finished.wait, 1) + assert not doomed[0].exists() + + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [original_target] + assert (original_target / "module.py").read_text() == "VALUE = 2" + assert original_target.is_dir() - assert not doomed[0].exists() @pytest.mark.anyio async def test_doomed_eviction_names_are_startup_scratch(self, temp_cache_dir): @@ -2010,7 +2072,9 @@ async def test_doomed_eviction_names_are_startup_scratch(self, temp_cache_dir): "tracecat.executor.registry_artifacts._delete_cache_path", return_value=True, ) as delete_cache_path: - assert await cache._evict_entry(cache_key) is True + assert await cache._evict_entry(cache_key) == RegistryArtifactEviction( + retired=True, reclaimed=True + ) delete_cache_path.assert_called_once() trash_path = delete_cache_path.call_args.args[0] @@ -2059,29 +2123,32 @@ async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): assert not idle.exists() @pytest.mark.anyio - async def test_eviction_drops_lease_records_but_keeps_the_key_lock( - self, temp_cache_dir - ): - """Locks are stable for the process lifetime; lease records are not.""" + async def test_eviction_keeps_stable_runtime_state(self, temp_cache_dir): + """A key keeps one lock and zeroed lease state for the process lifetime.""" cache = RegistryArtifactCache(temp_cache_dir) _write_tarball_entry(temp_cache_dir, "bookkeeping") cache._acquire_lease("bookkeeping") cache._release_lease("bookkeeping") - lock = await cache._lock_for("bookkeeping") + lock = cache._runtime_for("bookkeeping").lock - assert await cache._evict_entry("bookkeeping") is True - assert cache._locks["bookkeeping"] is lock - assert "bookkeeping" not in cache._leases + assert await cache._evict_entry("bookkeeping") == RegistryArtifactEviction( + retired=True, reclaimed=True + ) + runtime = cache._runtime["bookkeeping"] + assert runtime.lock is lock + assert runtime.refcount == 0 @pytest.mark.anyio async def test_eviction_skips_busy_key(self, temp_cache_dir): """A key another task is materializing is never evicted underneath it.""" cache = RegistryArtifactCache(temp_cache_dir) target_dir = _write_tarball_entry(temp_cache_dir, "busy") - lock = await cache._lock_for("busy") + lock = cache._runtime_for("busy").lock async with lock: - assert await cache._evict_entry("busy") is False + assert await cache._evict_entry("busy") == RegistryArtifactEviction( + retired=False, reclaimed=False + ) assert target_dir.is_dir() @@ -2150,6 +2217,33 @@ async def test_sweep_removes_orphaned_work_and_legacy_paths(self, temp_cache_dir assert unrelated_file.read_text() == "keep" assert entry_dir.is_dir() + @pytest.mark.anyio + async def test_sweep_removes_legacy_partial_downloads(self, temp_cache_dir): + """Former ``.part`` downloads are startup scratch, not cache entries.""" + partial = temp_cache_dir / "abc123.999.456.squashfs.part" + partial.write_bytes(b"partial") + + await RegistryArtifactCache(temp_cache_dir).ensure_swept() + + assert not partial.exists() + + @pytest.mark.anyio + async def test_sweep_preserves_backing_image_for_active_legacy_mount( + self, temp_cache_dir + ): + """Legacy cleanup never deletes the image behind a live mount.""" + cache = RegistryArtifactCache(temp_cache_dir) + image = temp_cache_dir / "squashfs-abc123.squashfs" + mount_dir = temp_cache_dir / "squashfs-abc123" + image.write_bytes(b"squashfs") + mount_dir.mkdir() + + with patch.object(Path, "is_mount", lambda self: self == mount_dir): + await cache.ensure_swept() + + assert image.read_bytes() == b"squashfs" + assert mount_dir.is_dir() + @pytest.mark.anyio async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): """A live mountpoint belongs to a running process and must survive.""" @@ -2285,6 +2379,77 @@ async def test_failed_ensure_swept_retries(self, temp_cache_dir): assert not orphaned_dir.exists() + @pytest.mark.parametrize("work_dir_name", ["staging", "trash"]) + def test_work_directory_inspection_errors_propagate( + self, + temp_cache_dir, + work_dir_name: str, + ): + """Unreadable work directories must trigger a later cleanup retry.""" + cache = RegistryArtifactCache(temp_cache_dir) + work_dir = getattr(cache, f"{work_dir_name}_dir") + work_dir.mkdir() + + with ( + patch.object(Path, "iterdir", side_effect=PermissionError("denied")), + pytest.raises(PermissionError, match="denied"), + ): + cache._clear_work_dir(work_dir) + + @pytest.mark.anyio + async def test_active_entry_scan_errors_propagate(self, temp_cache_dir): + """Unreadable active entries cannot be treated as an empty cache.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch( + "tracecat.executor.registry_artifacts.os.scandir", + side_effect=PermissionError("denied"), + ), + pytest.raises(PermissionError, match="denied"), + ): + await cache._enforce_cache_budget() + + @pytest.mark.anyio + async def test_failed_trash_cleanup_skips_startup_trimming(self, temp_cache_dir): + """A failed orphan deletion must not cascade into active retirements.""" + cache = RegistryArtifactCache(temp_cache_dir) + orphan = cache.trash_dir / "orphan" + orphan.mkdir(parents=True) + oldest = _write_image_entry( + temp_cache_dir, + "oldest", + size=16, + mtime=100.0, + ) + newest = _write_image_entry( + temp_cache_dir, + "newest", + size=16, + mtime=200.0, + ) + real_delete = _delete_cache_path + + def fail_orphan(path: Path) -> bool: + if path == orphan: + return False + return real_delete(path) + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifacts._delete_cache_path", + side_effect=fail_orphan, + ), + ): + await cache.ensure_swept() + + assert orphan.is_dir() + assert oldest.is_file() + assert newest.is_file() + assert cache._budget_dirty is True + @pytest.mark.anyio async def test_failed_startup_cleanup_retries_exact_path(self, temp_cache_dir): """Failed startup scratch cleanup retries without sweeping live staging.""" @@ -2308,13 +2473,13 @@ def fail_once(path: Path) -> bool: ): await cache.ensure_swept() assert orphaned.is_file() - assert cache._pending_cleanup == {orphaned} + assert cache._failed_startup_cleanup == {orphaned} assert cache._budget_dirty is True assert await cache._enforce_cache_budget() is True assert not orphaned.exists() - assert cache._pending_cleanup == set() + assert cache._failed_startup_cleanup == set() @pytest.mark.anyio async def test_failed_startup_retirement_stays_dirty_and_retries( @@ -2365,18 +2530,24 @@ async def test_failed_startup_retirement_stays_dirty_and_retries( async def test_failed_startup_physical_delete_retries_exact_path( self, temp_cache_dir ): - """A startup trash delete failure retries without evicting extra entries.""" + """Startup stops retiring entries when bytes were not reclaimed.""" oldest = _write_image_entry( temp_cache_dir, "oldest", size=16, mtime=100.0, ) + older = _write_image_entry( + temp_cache_dir, + "older", + size=16, + mtime=200.0, + ) newest = _write_image_entry( temp_cache_dir, "newest", size=16, - mtime=200.0, + mtime=300.0, ) cache = RegistryArtifactCache(temp_cache_dir) real_delete = _delete_cache_path @@ -2391,8 +2562,8 @@ def fail_once(path: Path) -> bool: with ( patch(BACKEND_CONFIG, ExecutorBackendType.DIRECT.value), - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 16), patch( "tracecat.executor.registry_artifacts._delete_cache_path", side_effect=fail_once, @@ -2400,346 +2571,52 @@ def fail_once(path: Path) -> bool: ): await cache.ensure_swept() assert not oldest.exists() + assert older.is_file() assert newest.is_file() - assert len(cache._pending_cleanup) == 1 + assert len(tuple(cache.trash_dir.iterdir())) == 1 assert cache._budget_dirty is True await cache._converge_cache_budget() + assert not older.exists() assert newest.is_file() - assert cache._pending_cleanup == set() assert not any(cache.trash_dir.iterdir()) assert cache._budget_dirty is False -class TestSquashfsMountCapability: - """Tests for process-wide SquashFS mount capability tracking.""" +class TestSquashfsMountPolicy: + """Tests for per-artifact SquashFS fallback.""" @pytest.mark.anyio - async def test_concurrent_first_mount_failure_serializes_capability_probe( + async def test_mount_failure_does_not_disable_later_artifacts( self, temp_cache_dir ) -> None: - """A failed first probe disables a waiter without racing its mount.""" - cache = RegistryArtifactCache(temp_cache_dir) - first_ctx = cache._context_for("first-probe") - second_ctx = cache._context_for("second-probe") - assert first_ctx.squashfs_mount_state is second_ctx.squashfs_mount_state - first_ctx.paths.entry_dir.mkdir(parents=True) - second_ctx.paths.entry_dir.mkdir(parents=True) - first_ctx.paths.squashfs_image_path.write_bytes(b"squashfs") - second_ctx.paths.squashfs_image_path.write_bytes(b"squashfs") - first_artifact = SquashfsArtifact( - uri="s3://bucket/path/first.squashfs", - cache_key=first_ctx.cache_key, + """One failed mount falls back without changing later mount attempts.""" + cache = RegistryArtifactCache(temp_cache_dir) + first_ctx = cache._context_for("first") + second_ctx = cache._context_for("second") + first = SquashfsArtifact( + uri="s3://bucket/first.squashfs", + cache_key="first", ) - second_artifact = SquashfsArtifact( - uri="s3://bucket/path/second.squashfs", - cache_key=second_ctx.cache_key, + second = SquashfsArtifact( + uri="s3://bucket/second.squashfs", + cache_key="second", ) - first_mount_started = asyncio.Event() - release_first_mount = asyncio.Event() - mount_attempts: list[Path] = [] - - async def mock_mount_image(self, image_path, target_dir): - mount_attempts.append(target_dir) - if target_dir == first_ctx.paths.squashfs_mount_dir: - first_mount_started.set() - await release_first_mount.wait() - raise SquashfsMountCommandError("operation not permitted") - - with patch.object(SquashfsArtifact, "_mount_image", mock_mount_image): - first_mount = asyncio.create_task( - first_artifact._try_mount( - first_ctx, - first_ctx.paths.squashfs_image_path, - ) - ) - await first_mount_started.wait() - second_mount = asyncio.create_task( - second_artifact._try_mount( - second_ctx, - second_ctx.paths.squashfs_image_path, - ) - ) - await asyncio.sleep(0) - - assert first_ctx.squashfs_mount_state.probe_lock.locked() - assert not second_mount.done() - assert mount_attempts == [first_ctx.paths.squashfs_mount_dir] - - release_first_mount.set() - first_result, second_result = await asyncio.gather( - first_mount, - second_mount, - ) - - state = first_ctx.squashfs_mount_state - assert first_result is None - assert second_result is None - assert state.disabled is True - assert state.mounted_once is False - assert not (state.disabled and state.mounted_once) - assert mount_attempts == [first_ctx.paths.squashfs_mount_dir] - - @pytest.mark.anyio - async def test_concurrent_probe_waiter_mounts_after_first_success( - self, temp_cache_dir - ) -> None: - """A waiter mounts after the successful probe releases serialization.""" - cache = RegistryArtifactCache(temp_cache_dir) - first_ctx = cache._context_for("first-success") - second_ctx = cache._context_for("second-success") - assert first_ctx.squashfs_mount_state is second_ctx.squashfs_mount_state - first_ctx.paths.entry_dir.mkdir(parents=True) - second_ctx.paths.entry_dir.mkdir(parents=True) - first_ctx.paths.squashfs_image_path.write_bytes(b"squashfs") - second_ctx.paths.squashfs_image_path.write_bytes(b"squashfs") - first_artifact = SquashfsArtifact( - uri="s3://bucket/path/first.squashfs", - cache_key=first_ctx.cache_key, - ) - second_artifact = SquashfsArtifact( - uri="s3://bucket/path/second.squashfs", - cache_key=second_ctx.cache_key, - ) - first_mount_started = asyncio.Event() - release_first_mount = asyncio.Event() - mount_attempts: list[Path] = [] - - async def mock_mount_image(self, image_path, target_dir): - mount_attempts.append(target_dir) - if target_dir == first_ctx.paths.squashfs_mount_dir: - first_mount_started.set() - await release_first_mount.wait() - - with patch.object(SquashfsArtifact, "_mount_image", mock_mount_image): - first_mount = asyncio.create_task( - first_artifact._try_mount( - first_ctx, - first_ctx.paths.squashfs_image_path, - ) - ) - await first_mount_started.wait() - second_mount = asyncio.create_task( - second_artifact._try_mount( - second_ctx, - second_ctx.paths.squashfs_image_path, - ) - ) - await asyncio.sleep(0) - - assert first_ctx.squashfs_mount_state.probe_lock.locked() - assert not second_mount.done() - assert mount_attempts == [first_ctx.paths.squashfs_mount_dir] - - release_first_mount.set() - first_result, second_result = await asyncio.gather( - first_mount, - second_mount, - ) - - state = first_ctx.squashfs_mount_state - assert first_result == first_ctx.paths.squashfs_mount_dir - assert second_result == second_ctx.paths.squashfs_mount_dir - assert state.disabled is False - assert state.mounted_once is True - assert mount_attempts == [ - first_ctx.paths.squashfs_mount_dir, - second_ctx.paths.squashfs_mount_dir, - ] - - @pytest.mark.anyio - async def test_first_mount_failure_disables_squashfs_process_wide( - self, temp_cache_dir - ): - """With no prior success the failure is a capability probe.""" - cache = RegistryArtifactCache(temp_cache_dir) + mount_attempts: list[str] = [] async def mock_mount(self, ctx, image_path): - raise SquashfsMountCommandError("operation not permitted") - - async def mock_extract(self, ctx, image_path): - target_dir = ctx.paths.squashfs_extract_dir - target_dir.mkdir(parents=True, exist_ok=True) - return target_dir - - with ( - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifacts.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", mock_mount), - patch.object(SquashfsArtifact, "extract", mock_extract), - patch.object( - cache, - "_release_mounted_slot", - new_callable=AsyncMock, - return_value=False, - ) as release_mounted_slot, - ): - await _materialize( - cache, "probe-key", "s3://bucket/path/site-packages.squashfs" - ) - - assert cache._squashfs_mount_state.disabled is True - release_mounted_slot.assert_not_awaited() - - @pytest.mark.anyio - async def test_mount_failure_after_success_reclaims_loop_device_and_retries( - self, temp_cache_dir - ): - """Loop-device exhaustion unmounts an idle artifact without deleting it.""" - cache = RegistryArtifactCache(temp_cache_dir) - cache._squashfs_mount_state.mounted_once = True - idle = cache._paths_for("idle") - idle.entry_dir.mkdir(parents=True) - idle.squashfs_image_path.write_bytes(b"squashfs") - idle.squashfs_mount_dir.mkdir() - mounted = {idle.squashfs_mount_dir} - attempts: list[str] = [] - - async def mock_mount(self, ctx, image_path): - attempts.append(ctx.cache_key) - if mounted: - raise SquashfsMountCommandError("failed to setup loop device") - target_dir = ctx.paths.squashfs_mount_dir - target_dir.mkdir(parents=True, exist_ok=True) - (target_dir / "module.py").write_text("VALUE = 1") - mounted.add(target_dir) - return target_dir - - umount_process = AsyncMock() - umount_process.communicate.return_value = (b"", b"") - umount_process.returncode = 0 - - async def mock_umount(*args, **kwargs): - mounted.discard(idle.squashfs_mount_dir) - return umount_process - - with ( - patch.object(Path, "is_mount", lambda self: self in mounted), - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifacts.shutil.which", - return_value="/sbin/umount", - ), - patch( - "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", - side_effect=mock_umount, - ), - patch.object(SquashfsArtifact, "mount", mock_mount), - ): - result = await _materialize( - cache, "new", "s3://bucket/path/site-packages.squashfs" - ) - - assert result == [cache._paths_for("new").squashfs_mount_dir] - assert (result[0] / "module.py").read_text() == "VALUE = 1" - assert attempts == ["new", "new"] - assert cache._squashfs_mount_state.disabled is False - assert idle.squashfs_image_path.read_bytes() == b"squashfs" - assert idle.squashfs_mount_dir.is_dir() - assert idle.squashfs_mount_dir not in mounted - - @pytest.mark.anyio - async def test_download_failure_does_not_disable_squashfs_process_wide( - self, temp_cache_dir - ): - """A transient download error is not a missing mount capability.""" - cache = RegistryArtifactCache(temp_cache_dir) - tarball_dir = temp_cache_dir / "gzip-fallback" - tarball_dir.mkdir() - - async def mock_download(self, ctx, image_path): - raise RuntimeError("connection reset by peer") - - with ( - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifacts.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "download", mock_download), - patch.object( - TarballArtifact, - "materialize", - new_callable=AsyncMock, - return_value=[tarball_dir], - ) as tarball_materialize, - patch.object( - cache, - "_release_mounted_slot", - new_callable=AsyncMock, - return_value=False, - ) as release_mounted_slot, - ): - result = await _materialize( - cache, "download-failure", "s3://bucket/path/site-packages.squashfs" - ) - - assert result == [tarball_dir] - assert cache._squashfs_mount_state.disabled is False - tarball_materialize.assert_awaited_once() - release_mounted_slot.assert_not_awaited() - - @pytest.mark.anyio - async def test_download_failure_after_a_mount_success_does_not_reclaim_a_slot( - self, temp_cache_dir - ): - """A transient download error must not evict an unrelated idle mount.""" - cache = RegistryArtifactCache(temp_cache_dir) - cache._squashfs_mount_state.mounted_once = True - tarball_dir = temp_cache_dir / "gzip-fallback" - tarball_dir.mkdir() - - async def mock_download(self, ctx, image_path): - raise RuntimeError("connection reset by peer") - - with ( - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifacts.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "download", mock_download), - patch.object( - TarballArtifact, - "materialize", - new_callable=AsyncMock, - return_value=[tarball_dir], - ), - patch.object( - cache, - "_release_mounted_slot", - new_callable=AsyncMock, - return_value=True, - ) as release_mounted_slot, - ): - result = await _materialize( - cache, "download-failure", "s3://bucket/path/site-packages.squashfs" - ) - - assert result == [tarball_dir] - assert cache._squashfs_mount_state.disabled is False - release_mounted_slot.assert_not_awaited() - - @pytest.mark.anyio - async def test_mount_failure_without_reclaimable_slot_falls_back_to_extraction( - self, temp_cache_dir - ): - """Only this artifact degrades when no idle mount can be reclaimed.""" - cache = RegistryArtifactCache(temp_cache_dir) - cache._squashfs_mount_state.mounted_once = True - - async def mock_mount(self, ctx, image_path): - raise SquashfsMountCommandError("failed to setup loop device") + del image_path + mount_attempts.append(ctx.cache_key) + if ctx.cache_key == "first": + raise SquashfsMountCommandError("operation not permitted") + ctx.paths.squashfs_mount_dir.mkdir(parents=True) + return ctx.paths.squashfs_mount_dir async def mock_extract(self, ctx, image_path): - target_dir = ctx.paths.squashfs_extract_dir - target_dir.mkdir(parents=True, exist_ok=True) - (target_dir / "module.py").write_text("VALUE = 1") - return target_dir + del image_path + ctx.paths.squashfs_extract_dir.mkdir(parents=True) + return ctx.paths.squashfs_extract_dir with ( patch(SQUASHFS_ENABLED_CONFIG, True), @@ -2750,9 +2627,11 @@ async def mock_extract(self, ctx, image_path): patch.object(SquashfsArtifact, "mount", mock_mount), patch.object(SquashfsArtifact, "extract", mock_extract), ): - result = await _materialize( - cache, "no-slot", "s3://bucket/path/site-packages.squashfs" - ) - - assert result[0].name == "extracted" - assert cache._squashfs_mount_state.disabled is False + assert await first.materialize(first_ctx) == [ + first_ctx.paths.squashfs_extract_dir + ] + assert await second.materialize(second_ctx) == [ + second_ctx.paths.squashfs_mount_dir + ] + + assert mount_attempts == ["first", "second"] diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 0ec3cfc449..f6d01ad118 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -12,7 +12,7 @@ import tarfile import time from abc import ABC, abstractmethod -from collections.abc import AsyncIterator, Awaitable, Callable, Iterable +from collections.abc import AsyncIterator, Iterable from contextlib import asynccontextmanager from dataclasses import dataclass, field from enum import StrEnum @@ -60,17 +60,11 @@ class RegistryArtifactFormat(StrEnum): CACHE_TRASH_DIR_NAME = "trash" """Directory containing atomically retired entries pending physical deletion.""" -LEGACY_CACHE_ENTRY_PREFIXES = ("squashfs-", "unsquashfs-", "tarball-") -"""Flat-layout prefixes retained only for one-way startup cleanup.""" - LEGACY_TEMP_ARTIFACT_PATTERN = re.compile( - r"^[^.]+\.\d+\.\d+\.(?:squashfs|unsquashfs|tar\.gz|tmp)$" + r"^[^.]+\.\d+\.\d+\.(?:squashfs|unsquashfs|tar\.gz|tmp)(?:\.part)?$" ) """Matches scratch paths created by the former flat cache layout.""" -type MountSlotReleaser = Callable[[str], Awaitable[bool]] -"""Unmounts one idle artifact, excluding the given cache key.""" - class SquashfsMountCommandError(RuntimeError): """The ``mount`` command itself failed for a SquashFS registry artifact. @@ -93,20 +87,20 @@ class RegistryArtifactPaths: @dataclass(slots=True) -class SquashfsMountState: - """Shared process-local SquashFS mount state.""" +class RegistryArtifactRuntimeState: + """Process-local synchronization and lease state for one cache key.""" - disabled: bool = False - mounted_once: bool = False - probe_lock: asyncio.Lock = field(default_factory=asyncio.Lock) + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + refcount: int = 0 + last_used: float = 0.0 -@dataclass(slots=True) -class RegistryArtifactLease: - """In-process lease bookkeeping for one registry artifact cache key.""" +@dataclass(frozen=True, slots=True) +class RegistryArtifactEviction: + """Outcome of atomically retiring and physically deleting one entry.""" - refcount: int = 0 - last_used: float = 0.0 + retired: bool + reclaimed: bool @dataclass(frozen=True, slots=True) @@ -125,34 +119,12 @@ class RegistryArtifactMaterializationContext: cache_key: str staging_dir: Path paths: RegistryArtifactPaths - squashfs_mount_state: SquashfsMountState - mount_slot_releaser: MountSlotReleaser | None = None def can_mount_squashfs(self) -> bool: - return ( - config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED - and not self.squashfs_mount_state.disabled - and (shutil.which("mount") is not None) + return config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED and ( + shutil.which("mount") is not None ) - def disable_squashfs_mount(self) -> None: - """Disable mounting after the serialized first capability probe fails.""" - self.squashfs_mount_state.disabled = True - - def record_squashfs_mount(self) -> None: - """Record that this process has successfully mounted a SquashFS image.""" - self.squashfs_mount_state.mounted_once = True - - def has_mounted_squashfs(self) -> bool: - """Return whether any SquashFS mount has ever succeeded in this process.""" - return self.squashfs_mount_state.mounted_once - - async def release_mounted_slot(self) -> bool: - """Unmount one idle artifact to free a loop device.""" - if self.mount_slot_releaser is None: - return False - return await self.mount_slot_releaser(self.cache_key) - @dataclass(frozen=True, slots=True) class RegistryArtifact(ABC): @@ -246,96 +218,19 @@ async def materialize( ) -> list[Path]: image_path = ctx.paths.squashfs_image_path if ctx.can_mount_squashfs(): - if (mount_dir := await self._try_mount(ctx, image_path)) is not None: - return [mount_dir] + try: + return [await self.mount(ctx, image_path)] + except SquashfsMountCommandError as e: + logger.warning( + "Failed to mount SquashFS registry artifact, trying extraction", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + error=str(e), + ) return [await self.extract(ctx, image_path)] - async def _try_mount( - self, - ctx: RegistryArtifactMaterializationContext, - image_path: Path, - ) -> Path | None: - """Mount the image, retrying once after reclaiming a loop device. - - The first mount-command failure in a process is treated as a capability - probe and disables mounting process-wide. Once any mount has succeeded, - later failures are attributed to exhausted loop devices instead: one idle - artifact is unmounted and the mount is retried once, so a single - failure never downgrades the whole process to extraction. - - The probe lock deliberately covers the first mount's download and mount - command. It nests inside a per-key materialization lock and acquires no - other lock; reclaim and retry remain outside it. - - Only ``SquashfsMountCommandError`` drives this policy. Download and - preparation errors propagate to the caller so a transient S3 failure - never disables mounting or unmounts an unrelated idle artifact. - - Args: - ctx: Materialization context for the artifact being mounted. - image_path: Local path of the SquashFS image. - - Returns: - The mount directory, or None if the caller should extract instead. - - Raises: - Exception: Any non-mount failure raised while preparing the image. - """ - if not ctx.has_mounted_squashfs(): - async with ctx.squashfs_mount_state.probe_lock: - if ctx.squashfs_mount_state.disabled: - return None - if not ctx.has_mounted_squashfs(): - try: - return await self.mount(ctx, image_path) - except SquashfsMountCommandError as e: - # This is the only disable path: the probe lock is held - # and no mount has succeeded, so disabled and mounted_once - # cannot both become true. - ctx.disable_squashfs_mount() - logger.warning( - "Failed to mount SquashFS registry artifact, trying extraction", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - error=str(e), - ) - return None - - try: - return await self.mount(ctx, image_path) - except SquashfsMountCommandError as e: - mount_error = e - - logger.warning( - "Failed to mount SquashFS registry artifact, reclaiming an idle mount", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - error=str(mount_error), - ) - if not await ctx.release_mounted_slot(): - logger.warning( - "No idle SquashFS mount to reclaim, trying extraction", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - ) - return None - - try: - return await self.mount(ctx, image_path) - except SquashfsMountCommandError as e: - logger.warning( - "SquashFS mount retry failed, trying extraction", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - error=str(e), - ) - return None - async def download( self, ctx: RegistryArtifactMaterializationContext, @@ -379,7 +274,6 @@ async def mount( """ target_dir = ctx.paths.squashfs_mount_dir if target_dir.is_mount(): - ctx.record_squashfs_mount() return target_dir ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) @@ -396,7 +290,6 @@ async def mount( mount_start = time.monotonic() await self._mount_image(image_path, target_dir) - ctx.record_squashfs_mount() mount_elapsed = (time.monotonic() - mount_start) * 1000 total_elapsed = (time.monotonic() - start_time) * 1000 @@ -778,19 +671,36 @@ def _directory_footprint(directory: Path) -> int: Total byte size of contained files, or zero when the directory is missing. """ - if not directory.is_dir(): - return 0 + + def raise_walk_error(error: OSError) -> None: + raise error total_bytes = 0 - for root, _dirs, files in os.walk(directory): - for file_name in files: - try: - total_bytes += os.lstat(os.path.join(root, file_name)).st_size - except OSError: - continue + try: + walker = os.walk(directory, onerror=raise_walk_error) + for root, _dirs, files in walker: + for file_name in files: + try: + total_bytes += os.lstat(os.path.join(root, file_name)).st_size + except FileNotFoundError: + continue + except FileNotFoundError: + return 0 return total_bytes +def _legacy_cache_key(path_name: str) -> str | None: + """Return the cache key encoded by a former flat-layout path.""" + if path_name.startswith("squashfs-"): + cache_key = path_name.removeprefix("squashfs-") + return cache_key.removesuffix(".squashfs") or None + if path_name.startswith("unsquashfs-"): + return path_name.removeprefix("unsquashfs-") or None + if path_name.startswith("tarball-"): + return path_name.removeprefix("tarball-") or None + return None + + def _delete_cache_path(path: Path) -> bool: """Best-effort delete one cache path while reporting filesystem failures.""" try: @@ -808,6 +718,20 @@ def _delete_cache_path(path: Path) -> bool: return True +async def _delete_cache_path_off_loop(path: Path) -> bool: + """Delete one path without abandoning its worker thread on cancellation.""" + deletion = asyncio.ensure_future(asyncio.to_thread(_delete_cache_path, path)) + try: + return await asyncio.shield(deletion) + except asyncio.CancelledError: + # A worker thread cannot be killed. Rejoin it so no live deletion can + # race a later trash-directory scan. + if not deletion.cancelled(): + with contextlib.suppress(Exception): + await asyncio.shield(deletion) + raise + + def _unique_work_path(root: Path, cache_key: str) -> Path: """Return a unique path beneath a cache work directory.""" root.mkdir(parents=True, exist_ok=True) @@ -834,22 +758,17 @@ def __init__(self, cache_dir: Path): self.entries_dir = cache_dir / CACHE_ENTRIES_DIR_NAME self.staging_dir = cache_dir / CACHE_STAGING_DIR_NAME self.trash_dir = cache_dir / CACHE_TRASH_DIR_NAME - # Per-key locks live for the process lifetime: eviction and lease - # admission must serialize on the same object for a given key, so a - # lock is never dropped and re-created underneath a waiter. - self._locks: dict[str, asyncio.Lock] = {} - self._locks_lock = asyncio.Lock() + # Runtime states live for the process lifetime so every operation for a + # key always serializes on the same lock. + self._runtime: dict[str, RegistryArtifactRuntimeState] = {} self._budget_lock = asyncio.Lock() # Guard the off-loop startup sweep independently from cache operations. self._swept: bool = False self._sweep_task: asyncio.Task[None] | None = None self._sweep_lock = asyncio.Lock() - self._squashfs_mount_state = SquashfsMountState() - self._leases: dict[str, RegistryArtifactLease] = {} - # Only exact paths whose deletion has already failed are retried at - # runtime. Never sweep the whole staging/trash directory while live - # materialization or deletion threads may still own its children. - self._pending_cleanup: set[Path] = set() + # Startup is the only time the whole staging directory is swept. Exact + # paths that could not be removed are safe to retry later. + self._failed_startup_cleanup: set[Path] = set() # Whether the on-disk cache may exceed its budget. Set when a new entry # is materialized and cleared once enforcement measures a cache that # fits, so steady-state cache hits never pay for a disk scan. @@ -927,8 +846,11 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat ) yield registry_paths finally: - for cache_key in leased_keys: - self._release_lease(cache_key) + idle_keys = [ + cache_key for cache_key in leased_keys if self._release_lease(cache_key) + ] + for cache_key in idle_keys: + await self._unmount_idle_entry(cache_key) if leased_keys: await self._converge_cache_budget() @@ -939,26 +861,31 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat if not _is_cache_entry_uri(artifact_uri): return None, await self._materialize_candidates(ctx, artifact_uri) - lock = await self._lock_for(cache_key) - async with lock: - self._acquire_lease(cache_key) - try: + lock = self._runtime_for(cache_key).lock + try: + async with lock: + self._acquire_lease(cache_key) candidates = await self._artifact_candidates(ctx, artifact_uri) - except BaseException: - self._release_lease(cache_key) - raise - if cached_paths := self._first_cached_path(candidates, ctx): - return cache_key, cached_paths + if cached_paths := self._first_cached_path(candidates, ctx): + return cache_key, cached_paths - try: # Make room outside the per-key lock so eviction never nests key locks. - await self._enforce_cache_budget(protected_key=cache_key) + try: + await self._enforce_cache_budget(protected_key=cache_key) + except OSError as e: + # Cache maintenance must never block artifact admission. + logger.warning( + "Failed to enforce registry artifact cache budget", + cache_dir=str(self.cache_dir), + error=str(e), + ) async with lock: paths = await self._materialize_candidates(ctx, artifact_uri) self._touch_entry(cache_key) return cache_key, paths except BaseException: - self._release_lease(cache_key) + if self._release_lease(cache_key): + await self._unmount_idle_entry(cache_key) raise async def _materialize_candidates( @@ -1006,12 +933,13 @@ async def _materialize_candidates( raise RuntimeError(f"No registry artifact candidates for {artifact_uri}") - async def _lock_for(self, cache_key: str) -> asyncio.Lock: - """Get or create a lock for the given cache key.""" - async with self._locks_lock: - if cache_key not in self._locks: - self._locks[cache_key] = asyncio.Lock() - return self._locks[cache_key] + def _runtime_for(self, cache_key: str) -> RegistryArtifactRuntimeState: + """Return the process-local state for one cache key.""" + if runtime := self._runtime.get(cache_key): + return runtime + runtime = RegistryArtifactRuntimeState() + self._runtime[cache_key] = runtime + return runtime def _context_for(self, cache_key: str) -> RegistryArtifactMaterializationContext: """Return a materialization context for a registry artifact key.""" @@ -1019,8 +947,6 @@ def _context_for(self, cache_key: str) -> RegistryArtifactMaterializationContext cache_key=cache_key, staging_dir=self.staging_dir, paths=self._paths_for(cache_key), - squashfs_mount_state=self._squashfs_mount_state, - mount_slot_releaser=self._release_mounted_slot, ) def _base_pythonpath_dir(self) -> Path: @@ -1035,23 +961,35 @@ def _acquire_lease(self, cache_key: str) -> None: Callers must hold the per-key lock so the increment is ordered against in-flight eviction of the same key. """ - lease = self._leases.setdefault(cache_key, RegistryArtifactLease()) - lease.refcount += 1 - lease.last_used = time.time() + runtime = self._runtime_for(cache_key) + runtime.refcount += 1 + runtime.last_used = time.time() self._touch_entry(cache_key) - def _release_lease(self, cache_key: str) -> None: - """Release one pin on a cache entry.""" - lease = self._leases.get(cache_key) - if lease is None: - return - lease.refcount = max(0, lease.refcount - 1) - lease.last_used = time.time() + def _release_lease(self, cache_key: str) -> bool: + """Release one pin and return whether the entry became idle.""" + runtime = self._runtime.get(cache_key) + if runtime is None or runtime.refcount == 0: + return False + runtime.refcount -= 1 + runtime.last_used = time.time() + return runtime.refcount == 0 + + async def _unmount_idle_entry(self, cache_key: str) -> None: + """Best-effort unmount an entry after its final lease is released.""" + try: + await self._unmount_entry(cache_key) + except OSError as e: + logger.warning( + "Failed to release idle registry artifact mount", + cache_key=cache_key, + error=str(e), + ) def _refcount(self, cache_key: str) -> int: """Return the number of live leases on a cache entry.""" - lease = self._leases.get(cache_key) - return 0 if lease is None else lease.refcount + runtime = self._runtime.get(cache_key) + return 0 if runtime is None else runtime.refcount def _touch_entry(self, cache_key: str) -> None: """Best-effort refresh of the entry-root mtime for restart-safe LRU.""" @@ -1235,15 +1173,18 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo Whether the cache is within budget once eviction has finished. """ async with self._budget_lock: - legacy_clean, pending_clean = await asyncio.gather( - asyncio.to_thread(self._remove_legacy_cache_paths), - asyncio.to_thread(self._retry_pending_cleanup), + trash_clean, startup_clean = await asyncio.gather( + asyncio.to_thread(self._clear_work_dir, self.trash_dir), + asyncio.to_thread(self._retry_failed_startup_cleanup), ) - cleanup_complete = legacy_clean and pending_clean + cleanup_complete = trash_clean and startup_clean + if not cleanup_complete: + return False + max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES if max_entries <= 0 and max_bytes <= 0: - return cleanup_complete + return True entries = await asyncio.to_thread(self._scan_cache_entries) # The protected key is not on disk yet when it is a fresh entry. @@ -1272,44 +1213,16 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo ) return False - if await self._evict_entry(candidate.cache_key): + eviction = await self._evict_entry(candidate.cache_key) + if eviction.retired: del entries[candidate.cache_key] + if not eviction.reclaimed: + return False total_bytes -= candidate.size_bytes - cleanup_complete = not self._pending_cleanup and cleanup_complete else: skipped.add(candidate.cache_key) - return cleanup_complete - - async def _release_mounted_slot(self, protected_key: str) -> bool: - """Unmount one idle artifact so its loop device can be reused. - - This path deliberately does not take the budget lock: ``_try_mount`` may - call it while holding ``protected_key``'s per-key lock. It excludes that - key and only tries candidate per-key locks, so it cannot invert the - budget-lock-to-per-key-lock ordering used by budget enforcement. - - Args: - protected_key: Cache key that must not be evicted. - - Returns: - Whether a mounted artifact was unmounted. - """ - entries = await asyncio.to_thread(self._scan_cache_entries) - mounted = [ - entry - for entry in entries.values() - if entry.cache_key != protected_key - and self._paths_for(entry.cache_key).squashfs_mount_dir.is_mount() - ] - skipped: set[str] = set() - while ( - candidate := self._least_recently_used(mounted, excluded=skipped) - ) is not None: - if await self._unmount_entry(candidate.cache_key): - return True - skipped.add(candidate.cache_key) - return False + return True def _is_pool_worker_visible( self, @@ -1349,10 +1262,10 @@ def _least_recently_used( def _recency(self, entry: RegistryArtifactCacheEntry) -> float: """Return the most recent known use time for a cache entry.""" - lease = self._leases.get(entry.cache_key) - if lease is None: + runtime = self._runtime.get(entry.cache_key) + if runtime is None: return entry.last_used - return max(entry.last_used, lease.last_used) + return max(entry.last_used, runtime.last_used) async def _unmount_entry(self, cache_key: str) -> bool: """Unmount one idle cache entry while retaining its reusable image. @@ -1369,7 +1282,7 @@ async def _unmount_entry(self, cache_key: str) -> bool: Returns: Whether a mounted entry was unmounted. """ - lock = await self._lock_for(cache_key) + lock = self._runtime_for(cache_key).lock if lock.locked(): logger.debug( "Skipping unmount of busy registry artifact", @@ -1399,7 +1312,7 @@ async def _unmount_entry(self, cache_key: str) -> bool: ) return True - async def _evict_entry(self, cache_key: str) -> bool: + async def _evict_entry(self, cache_key: str) -> RegistryArtifactEviction: """Remove one cache entry from disk, unmounting it first. The entry is skipped rather than forced when it is leased, busy, or @@ -1414,24 +1327,23 @@ async def _evict_entry(self, cache_key: str) -> bool: cache_key: Cache key to evict. Returns: - Whether the entry was atomically retired from the active cache. + Whether the entry was retired and its bytes were reclaimed. """ - lock = await self._lock_for(cache_key) + lock = self._runtime_for(cache_key).lock if lock.locked(): logger.debug( "Skipping eviction of busy registry artifact", cache_key=cache_key, ) - return False + return RegistryArtifactEviction(retired=False, reclaimed=False) async with lock: if self._refcount(cache_key) > 0: - return False + return RegistryArtifactEviction(retired=False, reclaimed=False) paths = self._paths_for(cache_key) if not paths.entry_dir.exists(): - self._leases.pop(cache_key, None) - return True + return RegistryArtifactEviction(retired=True, reclaimed=True) if paths.squashfs_mount_dir.is_mount() and not await self._unmount( paths.squashfs_mount_dir ): @@ -1440,7 +1352,7 @@ async def _evict_entry(self, cache_key: str) -> bool: cache_key=cache_key, mount_dir=str(paths.squashfs_mount_dir), ) - return False + return RegistryArtifactEviction(retired=False, reclaimed=False) try: trash_path = _move_entry_to_trash( @@ -1456,26 +1368,24 @@ async def _evict_entry(self, cache_key: str) -> bool: entry_dir=str(paths.entry_dir), error=str(e), ) - return False - self._leases.pop(cache_key, None) + return RegistryArtifactEviction(retired=False, reclaimed=False) try: - deleted = await asyncio.to_thread(_delete_cache_path, trash_path) + deleted = await _delete_cache_path_off_loop(trash_path) except BaseException: self._budget_dirty = True raise if not deleted: - self._pending_cleanup.add(trash_path) self._budget_dirty = True logger.warning( "Registry artifact eviction remains pending physical deletion", cache_key=cache_key, + trash_path=str(trash_path), ) - return True - self._pending_cleanup.discard(trash_path) + return RegistryArtifactEviction(retired=True, reclaimed=False) logger.info("Evicted registry artifact from cache", cache_key=cache_key) - return True + return RegistryArtifactEviction(retired=True, reclaimed=True) async def _unmount(self, mount_dir: Path) -> bool: """Unmount a SquashFS artifact directory, releasing its loop device. @@ -1524,7 +1434,7 @@ def _discover_cache_keys(self) -> set[str]: """Return cache keys represented by atomic entry directories.""" try: entries = list(os.scandir(self.entries_dir)) - except OSError: + except FileNotFoundError: return set() return { @@ -1545,7 +1455,7 @@ def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: try: image_stat = paths.squashfs_image_path.stat() - except OSError: + except FileNotFoundError: pass else: size_bytes += image_stat.st_size @@ -1555,7 +1465,7 @@ def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: try: last_used = paths.entry_dir.stat().st_mtime - except OSError: + except FileNotFoundError: last_used = 0.0 return RegistryArtifactCacheEntry( @@ -1581,16 +1491,14 @@ def _sweep_startup_state(self) -> None: return try: - cleanup_complete = self._remove_legacy_cache_paths() - cleanup_complete = ( - self._clear_work_dir(self.staging_dir, remember_failures=True) - and cleanup_complete + legacy_clean = self._remove_legacy_cache_paths() + staging_clean = self._clear_work_dir( + self.staging_dir, + remember_failures=True, ) - cleanup_complete = ( - self._clear_work_dir(self.trash_dir, remember_failures=True) - and cleanup_complete - ) - within_budget = self._trim_startup_cache() + trash_clean = self._clear_work_dir(self.trash_dir) + cleanup_complete = legacy_clean and staging_clean and trash_clean + within_budget = cleanup_complete and self._trim_startup_cache() self._budget_dirty = not (cleanup_complete and within_budget) except OSError as e: logger.warning( @@ -1617,13 +1525,13 @@ def _clear_work_dir( path=str(work_dir), error=str(e), ) - return False + raise deleted = True for path in paths: if _delete_cache_path(path): if remember_failures: - self._pending_cleanup.discard(path) + self._failed_startup_cleanup.discard(path) logger.info( "Removed registry artifact work path", path=str(path), @@ -1631,15 +1539,15 @@ def _clear_work_dir( else: deleted = False if remember_failures: - self._pending_cleanup.add(path) + self._failed_startup_cleanup.add(path) return deleted - def _retry_pending_cleanup(self) -> bool: - """Retry exact failed cleanup paths without touching live work.""" - for path in tuple(self._pending_cleanup): + def _retry_failed_startup_cleanup(self) -> bool: + """Retry exact startup paths without sweeping live staging work.""" + for path in tuple(self._failed_startup_cleanup): if _delete_cache_path(path): - self._pending_cleanup.discard(path) - return not self._pending_cleanup + self._failed_startup_cleanup.discard(path) + return not self._failed_startup_cleanup def _remove_legacy_cache_paths(self) -> bool: """Best-effort remove flat-layout cache paths from earlier executors.""" @@ -1653,7 +1561,7 @@ def _remove_legacy_cache_paths(self) -> bool: cache_dir=str(self.cache_dir), error=str(e), ) - return False + raise current_names = { BASE_PYTHONPATH_DIR_NAME, @@ -1662,16 +1570,32 @@ def _remove_legacy_cache_paths(self) -> bool: CACHE_TRASH_DIR_NAME, } backend_type = resolve_backend_type() + mounted_keys = { + cache_key + for path in paths + if path.name.startswith("squashfs-") + and not path.name.endswith(".squashfs") + and path.is_mount() + and (cache_key := _legacy_cache_key(path.name)) is not None + } deleted = True for path in paths: if path.name in current_names: continue - is_legacy_entry = path.name.startswith(LEGACY_CACHE_ENTRY_PREFIXES) + cache_key = _legacy_cache_key(path.name) + is_legacy_entry = cache_key is not None is_legacy_scratch = ( LEGACY_TEMP_ARTIFACT_PATTERN.fullmatch(path.name) is not None ) if not is_legacy_entry and not is_legacy_scratch: continue + if cache_key in mounted_keys: + logger.warning( + "Preserving active legacy registry artifact", + cache_key=cache_key, + path=str(path), + ) + continue if backend_type == ExecutorBackendType.POOL and path.name.startswith( "tarball-" ): @@ -1685,6 +1609,7 @@ def _remove_legacy_cache_paths(self) -> bool: continue if not _delete_cache_path(path): deleted = False + self._failed_startup_cleanup.add(path) return deleted def _trim_startup_cache(self) -> bool: @@ -1702,7 +1627,6 @@ def _trim_startup_cache(self) -> bool: entries = self._scan_cache_entries() total_bytes = sum(entry.size_bytes for entry in entries.values()) - cleanup_complete = True backend_type = resolve_backend_type() # Mounted entries belong to a live process sharing this cache directory. candidates = sorted( @@ -1731,25 +1655,23 @@ def within_budget() -> bool: entry.cache_key, ) except OSError as e: - cleanup_complete = False logger.warning( "Failed to retire stale registry artifact during startup sweep", cache_key=entry.cache_key, entry_dir=str(paths.entry_dir), error=str(e), ) - continue + return False del entries[entry.cache_key] if _delete_cache_path(trash_path): total_bytes -= entry.size_bytes else: - self._pending_cleanup.add(trash_path) - cleanup_complete = False + return False logger.info( "Evicted stale registry artifact during startup sweep", cache_key=entry.cache_key, size_bytes=entry.size_bytes, ) - return cleanup_complete and within_budget() + return within_budget() From b6fde73667a37c1d732a74abfc94bf0e661a611b Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:56:14 -0400 Subject: [PATCH 026/161] fix(executor): clean failed cache materializations --- tests/unit/test_registry_artifacts.py | 4 ++- tests/unit/test_storage_blob.py | 40 +++++++++++++++++++++++++ tracecat/executor/registry_artifacts.py | 24 +++++++++++++++ tracecat/storage/blob.py | 2 +- 4 files changed, 68 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 4ba58df544..d96d379447 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -997,7 +997,7 @@ async def test_lease_refcounts_and_touches_image_mtime(self, temp_cache_dir): async def test_lease_releases_refcount_when_materialization_fails( self, temp_cache_dir ): - """A failed materialization must not leak a permanent pin.""" + """A failed materialization must not leak a pin or empty cache entry.""" cache = RegistryArtifactCache(temp_cache_dir) artifact_uri = "s3://bucket/broken.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) @@ -1011,6 +1011,8 @@ async def mock_download(self, ctx, path): pass assert cache._refcount(cache_key) == 0 + assert not cache._paths_for(cache_key).entry_dir.exists() + assert cache_key not in cache._discover_cache_keys() @pytest.mark.anyio async def test_cancelled_lease_admission_releases_refcount(self, temp_cache_dir): diff --git a/tests/unit/test_storage_blob.py b/tests/unit/test_storage_blob.py index c8df650a66..d9b9f34f87 100644 --- a/tests/unit/test_storage_blob.py +++ b/tests/unit/test_storage_blob.py @@ -803,6 +803,46 @@ async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 assert not out.exists() assert not (tmp_path / "out.bin.part").exists() + @pytest.mark.anyio + async def test_download_file_to_path_cancellation_cleans_partial( + self, tmp_path: Path, monkeypatch + ): + """Cancellation removes the downloader-owned partial file.""" + download_blocked = asyncio.Event() + + class DummyStream: + async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 + yield b"partial" + download_blocked.set() + await asyncio.Event().wait() + + @asynccontextmanager + async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 + yield DummyStream(), None + + monkeypatch.setattr( + "tracecat.storage.blob.open_download_stream", + _fake_open_download_stream, + ) + + out = tmp_path / "out.bin" + download = asyncio.create_task( + download_file_to_path( + key="k", + bucket="b", + output_path=out, + ) + ) + await download_blocked.wait() + assert (tmp_path / "out.bin.part").exists() + + download.cancel() + with pytest.raises(asyncio.CancelledError): + await download + + assert not out.exists() + assert not (tmp_path / "out.bin.part").exists() + @pytest.mark.anyio @patch("tracecat.storage.blob.get_storage_client") async def test_ensure_bucket_exists_create_error_propagates(self, mock_get_client): diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index f6d01ad118..23f14d3ad6 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -912,13 +912,17 @@ async def _materialize_candidates( candidate=index + 1, candidates=len(candidates), ) + materialized = False try: registry_paths = await artifact.materialize(ctx) + materialized = True finally: if _is_cache_entry_uri(artifact.uri): # Any attempt may deposit canonical bytes, even when it # fails or is cancelled. self._budget_dirty = True + if not materialized: + self._remove_unpublished_entry(ctx) return registry_paths except Exception as e: if index == len(candidates) - 1: @@ -1024,6 +1028,26 @@ def _first_cached_path( return cached_paths return None + def _remove_unpublished_entry( + self, + ctx: RegistryArtifactMaterializationContext, + ) -> None: + """Remove an entry shell when an attempt published no reusable artifact. + + Callers hold the cache key lock. ``rmdir`` only removes empty + directories, so canonical artifacts and unknown contents are preserved. + """ + paths = ctx.paths + try: + if paths.squashfs_mount_dir.is_mount(): + return + except OSError: + return + + for directory in (paths.squashfs_mount_dir, paths.entry_dir): + with contextlib.suppress(OSError): + directory.rmdir() + async def _artifact_candidates( self, ctx: RegistryArtifactMaterializationContext, diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index a66c801cfc..e9d160c3c1 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -819,7 +819,7 @@ async def download_file_to_path( ) os.replace(temp_path, output_path) - except Exception: + except BaseException: try: temp_path.unlink(missing_ok=True) except Exception: From 634480f3feb0c5ca85328fdb9ce54aad8c3e1f90 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Fri, 31 Jul 2026 14:51:32 -0400 Subject: [PATCH 027/161] fix(executor): harden registry cache ownership Bind the process-wide cache to one event loop, expand cancellation and concurrency coverage across unit, Temporal, and real mount tests, and drop obsolete flat-layout handling for pod-local caches. --- ...registry_artifact_cache_mount_lifecycle.py | 74 ++- ...registry_artifact_cache_temporal_worker.py | 175 ++++++ tests/unit/test_action_runner.py | 104 ++++ tests/unit/test_registry_artifacts.py | 567 +++++++++++++++++- tracecat/executor/backends/pool/worker.py | 10 +- tracecat/executor/registry_artifacts.py | 129 ++-- 6 files changed, 928 insertions(+), 131 deletions(-) create mode 100644 tests/integration/test_registry_artifact_cache_temporal_worker.py diff --git a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py index 6a3830ef04..6982282480 100644 --- a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py +++ b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py @@ -145,7 +145,12 @@ def _run_mount_lifecycle_in_docker_or_skip() -> dict[str, Any]: @pytest.mark.integration def test_registry_artifact_cache_mount_lifecycle() -> None: - """Real mounts, evictions, and loop devices behave as the cache assumes.""" + """Protect lease invariants against the real kernel mount lifecycle. + + Unit mocks cannot prove that overlapping holders share one loop device or + that final release actually returns it to the kernel. This privileged test + guards those assumptions alongside retained-image remount and eviction. + """ payload = _run_mount_lifecycle_in_docker_or_skip() if skipped := payload.get("skipped"): @@ -162,6 +167,14 @@ def test_registry_artifact_cache_mount_lifecycle() -> None: assert payload["remounted_from_retained_image"] is True assert payload["remount_released"] is True + # Overlapping holders share one real mount and only the last one releases it. + assert payload["concurrent_paths_shared"] is True + assert payload["concurrent_peak_refcount"] == 3 + assert payload["concurrent_single_mount"] is True + assert payload["concurrent_intermediate_release_preserved_mount"] is True + assert payload["concurrent_final_release_unmounted"] is True + assert payload["concurrent_loop_device_released"] is True + # Eviction deletes an already-idle entry. assert payload["evicted"] is True assert payload["evicted_paths_removed"] is True @@ -302,7 +315,60 @@ async def _run_mount_lifecycle_child() -> None: ] payload["remount_released"] = not retained_paths.squashfs_mount_dir.is_mount() - # (c) Evict one already-idle entry. + # (c) Overlapping leases share one mount until final release. + concurrent_holders = 3 + concurrent_entered = 0 + all_concurrent_entered = asyncio.Event() + concurrent_releases = [asyncio.Event() for _ in range(concurrent_holders)] + concurrent_paths_shared: list[bool] = [] + + async def hold_concurrent_lease(index: int) -> None: + nonlocal concurrent_entered + async with cache.lease([uris[0]]) as registry_paths: + concurrent_paths_shared.append( + registry_paths == [retained_paths.squashfs_mount_dir] + and (registry_paths[0] / "module_0.py").read_text() == "VALUE = 1\n" + ) + concurrent_entered += 1 + if concurrent_entered == concurrent_holders: + all_concurrent_entered.set() + await concurrent_releases[index].wait() + + holders = [ + asyncio.create_task(hold_concurrent_lease(index)) + for index in range(concurrent_holders) + ] + await asyncio.wait_for(all_concurrent_entered.wait(), timeout=10) + concurrent_mounts = _squashfs_mounts(cache_dir) + concurrent_device = concurrent_mounts[str(retained_paths.squashfs_mount_dir)] + payload["concurrent_paths_shared"] = all(concurrent_paths_shared) + payload["concurrent_peak_refcount"] = cache._refcount(keys[0]) + payload["concurrent_single_mount"] = ( + list(concurrent_mounts).count(str(retained_paths.squashfs_mount_dir)) == 1 + ) + + for release in concurrent_releases[:-1]: + release.set() + await asyncio.gather(*holders[:-1]) + mounts_after_intermediate_releases = _squashfs_mounts(cache_dir) + payload["concurrent_intermediate_release_preserved_mount"] = ( + cache._refcount(keys[0]) == 1 + and str(retained_paths.squashfs_mount_dir) + in mounts_after_intermediate_releases + ) + + concurrent_releases[-1].set() + await holders[-1] + mounts_after_concurrent_release = _squashfs_mounts(cache_dir) + payload["concurrent_final_release_unmounted"] = ( + str(retained_paths.squashfs_mount_dir) + not in mounts_after_concurrent_release + ) + payload["concurrent_loop_device_released"] = ( + concurrent_device not in mounts_after_concurrent_release.values() + ) + + # (d) Evict one already-idle entry. evicted_paths = cache._paths_for(keys[0]) eviction = await cache._evict_entry(keys[0]) payload["evicted"] = eviction.retired and eviction.reclaimed @@ -311,7 +377,7 @@ async def _run_mount_lifecycle_child() -> None: or evicted_paths.squashfs_mount_dir.exists() ) - # (d) Final release unmounts, then converges an over-budget cache. + # (e) Final release unmounts, then converges an over-budget cache. fourth_uri = "s3://bucket/lifecycle/3/site-packages.squashfs" fourth_key = compute_registry_artifact_cache_key(fourth_uri) fourth_paths = cache._paths_for(fourth_key) @@ -335,7 +401,7 @@ async def _run_mount_lifecycle_child() -> None: ) payload["converged_entries_remaining"] = len(cache._discover_cache_keys()) - # (e) The startup sweep trims to budget and drops stale mount directories. + # (f) The startup sweep trims to budget and drops stale mount directories. sweep_dir = root / "sweep-cache" sweep_dir.mkdir() sweep_cache = RegistryArtifactCache(sweep_dir) diff --git a/tests/integration/test_registry_artifact_cache_temporal_worker.py b/tests/integration/test_registry_artifact_cache_temporal_worker.py new file mode 100644 index 0000000000..47fae9f691 --- /dev/null +++ b/tests/integration/test_registry_artifact_cache_temporal_worker.py @@ -0,0 +1,175 @@ +"""Temporal-worker concurrency coverage for the registry artifact cache.""" + +from __future__ import annotations + +import asyncio +import inspect +import threading +import uuid +from collections.abc import AsyncGenerator +from dataclasses import dataclass +from datetime import timedelta +from pathlib import Path + +import pytest +from temporalio import activity, workflow +from temporalio.common import RetryPolicy +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import UnsandboxedWorkflowRunner, Worker + +from tracecat.executor.activities import ExecutorActivities +from tracecat.executor.registry_artifacts import ( + RegistryArtifactCache, + compute_registry_artifact_cache_key, +) + +_PROBE_ACTIVITY_NAME = "registry_artifact_cache_concurrency_probe" + + +@pytest.fixture +async def temporal_env() -> AsyncGenerator[WorkflowEnvironment, None]: + """Run this worker test against Temporal's self-contained test server.""" + async with await WorkflowEnvironment.start_time_skipping() as environment: + yield environment + + +@dataclass(frozen=True, slots=True) +class _ProbeResult: + """Runtime identity observed by one Temporal activity.""" + + cache_instance_id: int + event_loop_id: int + thread_id: int + registry_path: str + + +class _RegistryCacheProbe: + """Hold overlapping activity leases against one cache instance.""" + + def __init__( + self, + *, + cache: RegistryArtifactCache, + artifact_uri: str, + expected_holders: int, + ) -> None: + self.cache = cache + self.artifact_uri = artifact_uri + self.expected_holders = expected_holders + self.all_entered = asyncio.Event() + self.release = asyncio.Event() + self.active_holders = 0 + self.peak_holders = 0 + + @activity.defn(name=_PROBE_ACTIVITY_NAME) + async def run(self, index: int) -> _ProbeResult: + """Lease the shared artifact until every scheduled activity overlaps.""" + del index + async with self.cache.lease([self.artifact_uri]) as registry_paths: + self.active_holders += 1 + self.peak_holders = max(self.peak_holders, self.active_holders) + if self.active_holders == self.expected_holders: + self.all_entered.set() + try: + await self.release.wait() + return _ProbeResult( + cache_instance_id=id(self.cache), + event_loop_id=id(asyncio.get_running_loop()), + thread_id=threading.get_ident(), + registry_path=str(registry_paths[0]), + ) + finally: + self.active_holders -= 1 + + +@workflow.defn +class _RegistryCacheConcurrencyWorkflow: + """Fan out enough activities to force overlapping cache leases.""" + + @workflow.run + async def run(self, activity_count: int) -> list[_ProbeResult]: + """Run the cache probe activities concurrently without retries.""" + handles = [ + workflow.start_activity( + _PROBE_ACTIVITY_NAME, + index, + result_type=_ProbeResult, + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(maximum_attempts=1), + ) + for index in range(activity_count) + ] + return await asyncio.gather(*handles) + + +@pytest.mark.anyio +@pytest.mark.integration +@pytest.mark.temporal +async def test_one_temporal_worker_uses_one_cache_loop_and_thread( + temporal_env: WorkflowEnvironment, + tmp_path: Path, +) -> None: + """Protect the production async-activity ownership contract. + + A real Temporal worker must schedule overlapping cache users on one event + loop and thread while sharing one process-wide cache instance. The explicit + production-activity assertion also prevents action execution from quietly + moving into Temporal's synchronous thread pool, the historical failure mode + for process-wide async storage state. + """ + assert inspect.iscoroutinefunction(ExecutorActivities.execute_action_activity) + + activity_count = 32 + artifact_uri = "s3://bucket/temporal-shared.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + cache_dir = tmp_path / "registry-cache" + cache_dir.mkdir() + cache = RegistryArtifactCache(cache_dir) + registry_path = cache._paths_for(cache_key).tarball_target_dir + registry_path.mkdir(parents=True) + (registry_path / "module.py").write_text("VALUE = 1") + + probe = _RegistryCacheProbe( + cache=cache, + artifact_uri=artifact_uri, + expected_holders=activity_count, + ) + task_queue = f"registry-cache-concurrency-{uuid.uuid4()}" + + async with Worker( + client=temporal_env.client, + task_queue=task_queue, + activities=[probe.run], + workflows=[_RegistryCacheConcurrencyWorkflow], + workflow_runner=UnsandboxedWorkflowRunner(), + max_concurrent_activities=activity_count, + ): + handle = await temporal_env.client.start_workflow( + _RegistryCacheConcurrencyWorkflow.run, + activity_count, + id=f"registry-cache-concurrency-{uuid.uuid4()}", + task_queue=task_queue, + execution_timeout=timedelta(seconds=45), + ) + results: list[_ProbeResult] = [] + try: + await asyncio.wait_for(probe.all_entered.wait(), timeout=20) + assert cache._refcount(cache_key) == activity_count + probe.release.set() + results = await handle.result() + except BaseException: + probe.release.set() + await handle.terminate(reason="Registry cache concurrency test failed") + raise + finally: + probe.release.set() + + assert probe.peak_holders == activity_count + assert probe.active_holders == 0 + assert {result.cache_instance_id for result in results} == {id(cache)} + assert len({result.event_loop_id for result in results}) == 1 + assert len({result.thread_id for result in results}) == 1 + assert {result.registry_path for result in results} == {str(registry_path)} + assert cache._refcount(cache_key) == 0 + assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) + assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index bd17a4b739..ac749d35e8 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -700,3 +700,107 @@ async def create_subprocess_exec_side_effect(*args, **kwargs): # noqa: ARG001 assert refcounts == [1] assert registry_paths[0].startswith(str(entry_dir)) assert runner.registry_artifacts._refcount(cache_key) == 0 + + @pytest.mark.anyio + async def test_cancelled_action_reaps_child_before_releasing_mounted_artifact( + self, + temp_cache_dir: Path, + mock_run_action_input: RunActionInput, + mock_role: Role, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Protect the subprocess-to-cache ownership boundary on cancellation. + + Cancellation must kill and reap the action before dropping its registry + pin; only then may final release unmount the artifact. This prevents a + child from importing through a reclaimed mount while still retaining + the reusable SquashFS image for the next action. + """ + runner = ActionRunner(cache_dir=temp_cache_dir) + cache = runner.registry_artifacts + artifact_uri = "s3://bucket/cancelled-action.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") + mounted = {paths.squashfs_mount_dir} + + monkeypatch.setattr( + action_runner.config, "TRACECAT__EXECUTOR_SANDBOX_ENABLED", False + ) + + resolved_context = ResolvedContext( + action_impl=ActionImplementation( + type="udf", + action_name="core.table.search_rows", + module="tracecat_registry.core.table", + name="search_rows", + ), + evaluated_args={"table": "customers"}, + workspace_id=str(mock_role.workspace_id), + workflow_id=str(mock_run_action_input.run_context.wf_id), + run_id=str(mock_run_action_input.run_context.wf_run_id), + executor_token="test-executor-token", + secret_projection=_empty_secret_projection(), + ) + + real_create_subprocess_exec = asyncio.create_subprocess_exec + process_started = asyncio.Event() + process: asyncio.subprocess.Process | None = None + reaped_before_unmount: list[bool] = [] + + async def capture_subprocess(*args, **kwargs): + nonlocal process + process = await real_create_subprocess_exec(*args, **kwargs) + process_started.set() + return process + + async def release_mount(mount_dir: Path) -> bool: + reaped_before_unmount.append( + process is not None and process.returncode is not None + ) + mounted.discard(mount_dir) + return True + + with ( + patch.object(Path, "is_mount", lambda path: path in mounted), + patch.object( + action_runner, + "_direct_subprocess_command", + return_value=["/bin/sleep", "30"], + ), + patch( + "tracecat.executor.action_runner.asyncio.create_subprocess_exec", + side_effect=capture_subprocess, + ), + patch.object(cache, "_unmount", side_effect=release_mount), + ): + execution = asyncio.create_task( + runner.execute_action( + input=mock_run_action_input, + role=mock_role, + resolved_context=resolved_context, + artifact_uris=[artifact_uri], + timeout=60.0, + ) + ) + try: + await process_started.wait() + assert cache._refcount(cache_key) == 1 + execution.cancel() + + with pytest.raises(asyncio.CancelledError): + await execution + finally: + if process is not None and process.returncode is None: + process.kill() + await process.wait() + + assert process is not None + assert process.returncode is not None + assert reaped_before_unmount == [True] + assert cache._refcount(cache_key) == 0 + assert paths.squashfs_mount_dir not in mounted + assert paths.squashfs_image_path.is_file() diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index d96d379447..aa8e184ec9 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -7,6 +7,7 @@ import tarfile import tempfile import threading +from dataclasses import dataclass, field from pathlib import Path from unittest.mock import ANY, AsyncMock, patch @@ -17,8 +18,10 @@ from tracecat.executor.registry_artifacts import ( SQUASHFS_MOUNT_OPTIONS, RegistryArtifactCache, + RegistryArtifactCacheLoopError, RegistryArtifactEviction, RegistryArtifactFormat, + RegistryArtifactMaterializationContext, SquashfsArtifact, SquashfsMountCommandError, TarballArtifact, @@ -68,6 +71,62 @@ def _write_image_entry( return image_path +@dataclass(slots=True) +class _SquashfsMountHarness: + """Model mount ownership without consuming host loop devices.""" + + cache: RegistryArtifactCache + failed_mount_keys: set[str] = field(default_factory=set) + mounted: set[Path] = field(default_factory=set) + mount_attempts: list[str] = field(default_factory=list) + extraction_attempts: list[str] = field(default_factory=list) + unmounts: list[Path] = field(default_factory=list) + + def seed_mount(self, cache_key: str) -> Path: + """Create one reusable image and mark its mount directory active.""" + paths = self.cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True, exist_ok=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir(exist_ok=True) + (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") + self.mounted.add(paths.squashfs_mount_dir) + return paths.squashfs_mount_dir + + async def mount( + self, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> Path: + """Record a mount attempt, failing selected keys like loop exhaustion.""" + self.mount_attempts.append(ctx.cache_key) + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + image_path.write_bytes(b"squashfs") + if ctx.cache_key in self.failed_mount_keys: + raise SquashfsMountCommandError("no free loop device") + ctx.paths.squashfs_mount_dir.mkdir(exist_ok=True) + (ctx.paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") + self.mounted.add(ctx.paths.squashfs_mount_dir) + return ctx.paths.squashfs_mount_dir + + async def extract( + self, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> Path: + """Publish an extracted fallback from the image retained by mount.""" + assert image_path.is_file() + self.extraction_attempts.append(ctx.cache_key) + ctx.paths.squashfs_extract_dir.mkdir(parents=True, exist_ok=True) + (ctx.paths.squashfs_extract_dir / "module.py").write_text("VALUE = 1") + return ctx.paths.squashfs_extract_dir + + async def unmount(self, mount_dir: Path) -> bool: + """Record final-release unmounts and release the modeled loop device.""" + self.unmounts.append(mount_dir) + self.mounted.discard(mount_dir) + return True + + async def _materialize( cache: RegistryArtifactCache, cache_key: str, @@ -963,6 +1022,30 @@ def test_runtime_for_different_keys(self, temp_cache_dir): class TestRegistryArtifactCacheLease: """Tests for lease-based pinning of registry artifact cache entries.""" + @pytest.mark.anyio + async def test_cache_rejects_use_from_a_second_event_loop( + self, temp_cache_dir: Path + ) -> None: + """Protect the process-wide cache from thread-local Temporal loops. + + The cache owns asyncio locks, tasks, and refcounts as one unit. A typed + ownership failure prevents a future synchronous activity from silently + sharing only part of that state across thread-local event loops, which + previously caused intermittent cross-loop failures in storage caches. + """ + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + + def use_cache_from_another_thread() -> None: + asyncio.run(cache.ensure_swept()) + + with pytest.raises(RegistryArtifactCacheLoopError): + await asyncio.to_thread(use_cache_from_another_thread) + + # A rejected caller must not poison the owning loop's cache. + async with cache.lease(None) as registry_paths: + assert registry_paths == [temp_cache_dir / "base"] + def test_touch_entry_refreshes_tarball_root_mtime(self, temp_cache_dir): """Touching a tarball-only entry persists its restart-safe recency.""" cache = RegistryArtifactCache(temp_cache_dir) @@ -1077,6 +1160,390 @@ async def test_lease_preserves_uri_order(self, temp_cache_dir): async with cache.lease(uris) as registry_paths: assert registry_paths == expected + @pytest.mark.anyio + async def test_overlapping_same_key_leases_share_one_mount_until_final_release( + self, temp_cache_dir: Path + ) -> None: + """Protect refcount and loop-device lifetime under heavy fan-in. + + Every holder must observe one published mount, intermediate releases + must leave that mount usable, and exactly the final release may reclaim + its loop device while retaining the downloaded image for a remount. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/shared.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + harness = _SquashfsMountHarness(cache) + holder_count = 32 + entered = 0 + all_entered = asyncio.Event() + releases = [asyncio.Event() for _ in range(holder_count)] + + async def hold_lease(index: int) -> None: + nonlocal entered + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [ + cache._paths_for(cache_key).squashfs_mount_dir + ] + entered += 1 + if entered == holder_count: + all_entered.set() + await releases[index].wait() + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", harness.mount), + patch.object(SquashfsArtifact, "extract", harness.extract), + patch.object(cache, "_unmount", harness.unmount), + ): + holders = [ + asyncio.create_task(hold_lease(index)) for index in range(holder_count) + ] + await asyncio.wait_for(all_entered.wait(), timeout=5) + + mount_dir = cache._paths_for(cache_key).squashfs_mount_dir + assert cache._refcount(cache_key) == holder_count + assert harness.mount_attempts == [cache_key] + assert harness.extraction_attempts == [] + assert mount_dir in harness.mounted + + for release in releases[:-1]: + release.set() + await asyncio.gather(*holders[:-1]) + + assert cache._refcount(cache_key) == 1 + assert mount_dir in harness.mounted + assert harness.unmounts == [] + + releases[-1].set() + await holders[-1] + + assert cache._refcount(cache_key) == 0 + assert harness.unmounts == [mount_dir] + assert mount_dir not in harness.mounted + assert cache._paths_for(cache_key).squashfs_image_path.is_file() + assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) + assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) + + @pytest.mark.anyio + async def test_cancelling_one_holder_preserves_sibling_leases( + self, temp_cache_dir: Path + ) -> None: + """Protect sibling actions from another holder's cancellation. + + Cancellation must release exactly one pin without unmounting the shared + artifact beneath surviving actions; the last surviving holder remains + solely responsible for loop-device reclamation. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/cancel-one.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + harness = _SquashfsMountHarness(cache) + entered = [asyncio.Event() for _ in range(3)] + releases = [asyncio.Event() for _ in range(3)] + + async def hold_lease(index: int) -> None: + async with cache.lease([artifact_uri]): + entered[index].set() + await releases[index].wait() + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", harness.mount), + patch.object(SquashfsArtifact, "extract", harness.extract), + patch.object(cache, "_unmount", harness.unmount), + ): + holders = [asyncio.create_task(hold_lease(index)) for index in range(3)] + await asyncio.wait_for( + asyncio.gather(*(event.wait() for event in entered)), + timeout=5, + ) + + holders[0].cancel() + with pytest.raises(asyncio.CancelledError): + await holders[0] + + mount_dir = cache._paths_for(cache_key).squashfs_mount_dir + assert cache._refcount(cache_key) == 2 + assert mount_dir in harness.mounted + assert harness.unmounts == [] + + releases[1].set() + await holders[1] + assert cache._refcount(cache_key) == 1 + assert mount_dir in harness.mounted + assert harness.unmounts == [] + + releases[2].set() + await holders[2] + + assert cache._refcount(cache_key) == 0 + assert harness.mount_attempts == [cache_key] + assert harness.unmounts == [mount_dir] + + @pytest.mark.anyio + async def test_new_lease_racing_final_release_prevents_stale_unmount( + self, temp_cache_dir: Path + ) -> None: + """Protect the zero-refcount-to-unmount handoff from a new acquisition. + + A lease can arrive after the old holder decrements to zero but before + unmount takes the key lock. The lock-time refcount recheck must preserve + the mount for that newcomer instead of returning a path being torn down. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/release-race.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + harness = _SquashfsMountHarness(cache) + mount_dir = harness.seed_mount(cache_key) + first_entered = asyncio.Event() + release_first = asyncio.Event() + release_reached_unmount = asyncio.Event() + allow_unmount_recheck = asyncio.Event() + newcomer_entered = asyncio.Event() + release_newcomer = asyncio.Event() + original_unmount_idle_entry = cache._unmount_idle_entry + unmount_requests = 0 + + async def pause_first_unmount_request(requested_key: str) -> None: + nonlocal unmount_requests + unmount_requests += 1 + if unmount_requests == 1: + release_reached_unmount.set() + await allow_unmount_recheck.wait() + await original_unmount_idle_entry(requested_key) + + async def old_holder() -> None: + async with cache.lease([artifact_uri]): + first_entered.set() + await release_first.wait() + + async def new_holder() -> None: + async with cache.lease([artifact_uri]): + newcomer_entered.set() + await release_newcomer.wait() + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch.object(cache, "_unmount", harness.unmount), + patch.object( + cache, + "_unmount_idle_entry", + pause_first_unmount_request, + ), + ): + old_task = asyncio.create_task(old_holder()) + await first_entered.wait() + release_first.set() + await release_reached_unmount.wait() + assert cache._refcount(cache_key) == 0 + + new_task = asyncio.create_task(new_holder()) + await newcomer_entered.wait() + assert cache._refcount(cache_key) == 1 + + allow_unmount_recheck.set() + await old_task + + assert mount_dir in harness.mounted + assert harness.unmounts == [] + + release_newcomer.set() + await new_task + + assert cache._refcount(cache_key) == 0 + assert harness.unmounts == [mount_dir] + + @pytest.mark.anyio + async def test_partial_multi_artifact_failure_rolls_back_prior_leases( + self, temp_cache_dir: Path + ) -> None: + """Protect sequential multi-artifact admission from partial pin leaks. + + If a middle artifact fails, every earlier acquisition must be released + and unmounted, the failed key must leave no shell, later artifacts must + never be acquired, and the release path must still request convergence. + """ + cache = RegistryArtifactCache(temp_cache_dir) + first_uri = "s3://bucket/first.squashfs" + failed_uri = "s3://bucket/failed.tar.gz" + untouched_uri = "s3://bucket/untouched.tar.gz" + first_key = compute_registry_artifact_cache_key(first_uri) + failed_key = compute_registry_artifact_cache_key(failed_uri) + untouched_key = compute_registry_artifact_cache_key(untouched_uri) + harness = _SquashfsMountHarness(cache) + first_mount = harness.seed_mount(first_key) + untouched_path = _write_tarball_entry(temp_cache_dir, untouched_key) + + async def fail_download( + artifact: TarballArtifact, + ctx: RegistryArtifactMaterializationContext, + output_path: Path, + ) -> None: + del artifact, ctx, output_path + raise RuntimeError("download failed") + + tracked_lease_artifact = AsyncMock(wraps=cache._lease_artifact) + converge_cache_budget = AsyncMock() + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "download", fail_download), + patch.object(cache, "_unmount", harness.unmount), + patch.object(cache, "_lease_artifact", tracked_lease_artifact), + patch.object( + cache, + "_converge_cache_budget", + converge_cache_budget, + ), + ): + with pytest.raises(RuntimeError, match="download failed"): + async with cache.lease([first_uri, failed_uri, untouched_uri]): + pass + + requested_uris = [ + await_call.args[0] for await_call in tracked_lease_artifact.await_args_list + ] + assert requested_uris == [first_uri, failed_uri] + assert cache._refcount(first_key) == 0 + assert cache._refcount(failed_key) == 0 + assert cache._refcount(untouched_key) == 0 + assert harness.unmounts == [first_mount] + assert cache._paths_for(first_key).squashfs_image_path.is_file() + assert not cache._paths_for(failed_key).entry_dir.exists() + assert untouched_path.is_dir() + converge_cache_budget.assert_awaited_once_with() + assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) + assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) + + @pytest.mark.anyio + async def test_waiter_retries_after_first_same_key_materializer_fails( + self, temp_cache_dir: Path + ) -> None: + """Protect same-key waiters from a failed cold-cache publisher. + + The first materializer may fail after writing scratch. Its waiter must + retry against a clean staging area, publish the sole valid entry, and + leave neither a leaked pin nor an empty LRU-visible cache shell. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/retry-after-failure.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + first_download_started = asyncio.Event() + fail_first_download = asyncio.Event() + retry_download_started = asyncio.Event() + allow_retry_download = asyncio.Event() + download_attempts = 0 + retry_saw_clean_staging = False + + async def controlled_download( + artifact: TarballArtifact, + ctx: RegistryArtifactMaterializationContext, + output_path: Path, + ) -> None: + del artifact, ctx + nonlocal download_attempts, retry_saw_clean_staging + download_attempts += 1 + if download_attempts == 1: + output_path.write_bytes(b"partial") + first_download_started.set() + await fail_first_download.wait() + raise RuntimeError("first publisher failed") + + retry_saw_clean_staging = not any(cache.staging_dir.iterdir()) + retry_download_started.set() + await allow_retry_download.wait() + output_path.write_bytes(b"complete") + + async def mock_extract( + artifact: TarballArtifact, + tarball_path: Path, + target_dir: Path, + ) -> None: + del artifact + assert tarball_path.read_bytes() == b"complete" + (target_dir / "module.py").write_text("VALUE = 1") + + async def take_lease() -> list[Path]: + async with cache.lease([artifact_uri]) as registry_paths: + return registry_paths + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "download", controlled_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + first = asyncio.create_task(take_lease()) + await first_download_started.wait() + waiter = asyncio.create_task(take_lease()) + await asyncio.sleep(0) + assert download_attempts == 1 + + fail_first_download.set() + with pytest.raises(RuntimeError, match="first publisher failed"): + await first + + await retry_download_started.wait() + allow_retry_download.set() + registry_paths = await waiter + + target_dir = cache._paths_for(cache_key).tarball_target_dir + assert registry_paths == [target_dir] + assert (target_dir / "module.py").read_text() == "VALUE = 1" + assert download_attempts == 2 + assert retry_saw_clean_staging is True + assert cache._discover_cache_keys() == {cache_key} + assert cache._refcount(cache_key) == 0 + assert not any(cache.staging_dir.iterdir()) + assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) + + @pytest.mark.anyio + async def test_duplicate_uri_balances_each_acquisition_and_release( + self, temp_cache_dir: Path + ) -> None: + """Protect list bookkeeping when one lease requests the same URI twice. + + Duplicate PYTHONPATH entries intentionally acquire two pins. Both must + be released, while materialization and final unmount still occur once; + accidental deduplication on only one side would leak or underflow pins. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/duplicate.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + harness = _SquashfsMountHarness(cache) + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", harness.mount), + patch.object(SquashfsArtifact, "extract", harness.extract), + patch.object(cache, "_unmount", harness.unmount), + ): + async with cache.lease([artifact_uri, artifact_uri]) as registry_paths: + mount_dir = cache._paths_for(cache_key).squashfs_mount_dir + assert registry_paths == [mount_dir, mount_dir] + assert cache._refcount(cache_key) == 2 + assert harness.mount_attempts == [cache_key] + + assert cache._refcount(cache_key) == 0 + assert harness.unmounts == [mount_dir] + assert cache._paths_for(cache_key).squashfs_image_path.is_file() + @pytest.mark.anyio async def test_lease_is_never_admitted_across_an_in_flight_eviction( self, temp_cache_dir @@ -2195,16 +2662,14 @@ async def test_sweep_tolerates_missing_cache_dir(self, temp_cache_dir): assert not cache_dir.exists() @pytest.mark.anyio - async def test_sweep_removes_orphaned_work_and_legacy_paths(self, temp_cache_dir): - """Interrupted work and disposable flat-layout paths are reclaimed.""" + async def test_sweep_removes_orphaned_work(self, temp_cache_dir): + """Interrupted work is reclaimed without touching unrelated paths.""" cache = RegistryArtifactCache(temp_cache_dir) orphaned = cache.staging_dir / "abc123.999999.4321.squashfs" orphaned.parent.mkdir() orphaned.write_bytes(b"partial") orphaned_dir = cache.trash_dir / "abc123.999999.4321" orphaned_dir.mkdir(parents=True) - legacy = temp_cache_dir / "squashfs-abc123.squashfs" - legacy.write_bytes(b"legacy") unrelated = temp_cache_dir / "unrelated" unrelated.mkdir() unrelated_file = unrelated / "keep.txt" @@ -2215,37 +2680,9 @@ async def test_sweep_removes_orphaned_work_and_legacy_paths(self, temp_cache_dir assert not orphaned.exists() assert not orphaned_dir.exists() - assert not legacy.exists() assert unrelated_file.read_text() == "keep" assert entry_dir.is_dir() - @pytest.mark.anyio - async def test_sweep_removes_legacy_partial_downloads(self, temp_cache_dir): - """Former ``.part`` downloads are startup scratch, not cache entries.""" - partial = temp_cache_dir / "abc123.999.456.squashfs.part" - partial.write_bytes(b"partial") - - await RegistryArtifactCache(temp_cache_dir).ensure_swept() - - assert not partial.exists() - - @pytest.mark.anyio - async def test_sweep_preserves_backing_image_for_active_legacy_mount( - self, temp_cache_dir - ): - """Legacy cleanup never deletes the image behind a live mount.""" - cache = RegistryArtifactCache(temp_cache_dir) - image = temp_cache_dir / "squashfs-abc123.squashfs" - mount_dir = temp_cache_dir / "squashfs-abc123" - image.write_bytes(b"squashfs") - mount_dir.mkdir() - - with patch.object(Path, "is_mount", lambda self: self == mount_dir): - await cache.ensure_swept() - - assert image.read_bytes() == b"squashfs" - assert mount_dir.is_dir() - @pytest.mark.anyio async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): """A live mountpoint belongs to a running process and must survive.""" @@ -2589,6 +3026,72 @@ def fail_once(path: Path) -> bool: class TestSquashfsMountPolicy: """Tests for per-artifact SquashFS fallback.""" + @pytest.mark.anyio + async def test_loop_device_exhaustion_isolated_sticky_extraction_fallback( + self, temp_cache_dir: Path + ) -> None: + """Protect the cache's fail-open policy when loop devices are saturated. + + A mount-command failure must extract only the affected cold artifact, + preserve already-leased mounts, reuse that extraction on later leases, + and still let unrelated artifacts attempt mounting. This models loop + exhaustion deterministically without consuming host-global devices. + """ + cache = RegistryArtifactCache(temp_cache_dir) + held_uri = "s3://bucket/already-mounted.squashfs" + saturated_uri = "s3://bucket/no-loop-available.squashfs" + later_uri = "s3://bucket/later-artifact.squashfs" + held_key = compute_registry_artifact_cache_key(held_uri) + saturated_key = compute_registry_artifact_cache_key(saturated_uri) + later_key = compute_registry_artifact_cache_key(later_uri) + harness = _SquashfsMountHarness( + cache, + failed_mount_keys={saturated_key}, + ) + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", harness.mount), + patch.object(SquashfsArtifact, "extract", harness.extract), + patch.object(cache, "_unmount", harness.unmount), + ): + async with cache.lease([held_uri]) as held_paths: + held_mount = cache._paths_for(held_key).squashfs_mount_dir + assert held_paths == [held_mount] + + async with cache.lease([saturated_uri]) as saturated_paths: + extracted = cache._paths_for(saturated_key).squashfs_extract_dir + assert saturated_paths == [extracted] + assert held_mount in harness.mounted + assert cache._refcount(held_key) == 1 + assert harness.unmounts == [] + + # The extracted directory is the cache entry's sticky path; + # freeing a loop elsewhere does not trigger a remount attempt. + async with cache.lease([saturated_uri]) as cached_paths: + assert cached_paths == [extracted] + + async with cache.lease([later_uri]) as later_paths: + later_mount = cache._paths_for(later_key).squashfs_mount_dir + assert later_paths == [later_mount] + assert held_mount in harness.mounted + + assert held_mount in harness.mounted + + assert harness.mount_attempts == [held_key, saturated_key, later_key] + assert harness.extraction_attempts == [saturated_key] + assert harness.unmounts == [later_mount, held_mount] + assert cache._paths_for(saturated_key).squashfs_image_path.is_file() + assert cache._paths_for(saturated_key).squashfs_extract_dir.is_dir() + assert cache._refcount(held_key) == 0 + assert cache._refcount(saturated_key) == 0 + assert cache._refcount(later_key) == 0 + @pytest.mark.anyio async def test_mount_failure_does_not_disable_later_artifacts( self, temp_cache_dir diff --git a/tracecat/executor/backends/pool/worker.py b/tracecat/executor/backends/pool/worker.py index 60fce20f30..9c2ed05773 100644 --- a/tracecat/executor/backends/pool/worker.py +++ b/tracecat/executor/backends/pool/worker.py @@ -56,18 +56,14 @@ def _ensure_tarball_paths_in_sys_path() -> None: """Ensure all tarball extraction directories are in sys.path. - Scans atomic cache entries, plus legacy flat-layout entries during rollout, - and adds materialized tarballs to sys.path if not already present. + Scans atomic cache entries and adds materialized tarballs to sys.path if not + already present. """ cache_dir = Path(config.TRACECAT__EXECUTOR_REGISTRY_CACHE_DIR) if not cache_dir.exists(): return - tarball_paths = ( - *cache_dir.glob("entries/*/tarball"), - *cache_dir.glob("tarball-*"), - ) - for path in tarball_paths: + for path in cache_dir.glob("entries/*/tarball"): if not path.is_dir(): continue path_str = str(path) diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 23f14d3ad6..5dea7f1f18 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -6,10 +6,10 @@ import contextlib import hashlib import os -import re import shutil import sysconfig import tarfile +import threading import time from abc import ABC, abstractmethod from collections.abc import AsyncIterator, Iterable @@ -60,11 +60,6 @@ class RegistryArtifactFormat(StrEnum): CACHE_TRASH_DIR_NAME = "trash" """Directory containing atomically retired entries pending physical deletion.""" -LEGACY_TEMP_ARTIFACT_PATTERN = re.compile( - r"^[^.]+\.\d+\.\d+\.(?:squashfs|unsquashfs|tar\.gz|tmp)(?:\.part)?$" -) -"""Matches scratch paths created by the former flat cache layout.""" - class SquashfsMountCommandError(RuntimeError): """The ``mount`` command itself failed for a SquashFS registry artifact. @@ -75,6 +70,10 @@ class SquashfsMountCommandError(RuntimeError): """ +class RegistryArtifactCacheLoopError(RuntimeError): + """A registry artifact cache was used outside its owning event loop.""" + + @dataclass(frozen=True, slots=True) class RegistryArtifactPaths: """Executor-local cache paths for one registry artifact key.""" @@ -689,18 +688,6 @@ def raise_walk_error(error: OSError) -> None: return total_bytes -def _legacy_cache_key(path_name: str) -> str | None: - """Return the cache key encoded by a former flat-layout path.""" - if path_name.startswith("squashfs-"): - cache_key = path_name.removeprefix("squashfs-") - return cache_key.removesuffix(".squashfs") or None - if path_name.startswith("unsquashfs-"): - return path_name.removeprefix("unsquashfs-") or None - if path_name.startswith("tarball-"): - return path_name.removeprefix("tarball-") or None - return None - - def _delete_cache_path(path: Path) -> bool: """Best-effort delete one cache path while reporting filesystem failures.""" try: @@ -761,6 +748,13 @@ def __init__(self, cache_dir: Path): # Runtime states live for the process lifetime so every operation for a # key always serializes on the same lock. self._runtime: dict[str, RegistryArtifactRuntimeState] = {} + # The cache contains asyncio locks, tasks, and multi-step lease state. + # Bind the public API to one loop/thread so a future synchronous + # Temporal activity fails immediately instead of corrupting that state + # through a thread-local event loop. + self._owner_binding_lock = threading.Lock() + self._owner_loop: asyncio.AbstractEventLoop | None = None + self._owner_thread_id: int | None = None self._budget_lock = asyncio.Lock() # Guard the off-loop startup sweep independently from cache operations. self._swept: bool = False @@ -785,6 +779,7 @@ async def ensure_swept(self) -> None: before the first lease or materialization, so it never observes in-flight cache entries. """ + self._assert_owner_loop() if self._swept: return async with self._sweep_lock: @@ -811,6 +806,29 @@ async def ensure_swept(self) -> None: raise self._swept = True + def _assert_owner_loop(self) -> None: + """Bind to the current loop or reject use from another loop/thread.""" + current_loop = asyncio.get_running_loop() + current_thread_id = threading.get_ident() + with self._owner_binding_lock: + if self._owner_loop is None: + self._owner_loop = current_loop + self._owner_thread_id = current_thread_id + return + if ( + self._owner_loop is current_loop + and self._owner_thread_id == current_thread_id + ): + return + + raise RegistryArtifactCacheLoopError( + "RegistryArtifactCache is bound to one event loop and thread " + f"(owner_loop={id(self._owner_loop)}, " + f"owner_thread={self._owner_thread_id}, " + f"current_loop={id(current_loop)}, " + f"current_thread={current_thread_id})" + ) + @asynccontextmanager async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Path]]: """Materialize registry artifacts and pin them for the life of the context. @@ -1501,11 +1519,10 @@ def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: def _sweep_startup_state(self) -> None: """Reclaim orphaned cache state left behind by a previous process. - Scratch and trash paths from interrupted work are removed, legacy flat - cache paths are retired, and active entries are trimmed to budget using - entry-root mtimes as LRU order. Tarball-bearing entries are protected - for the pool backend because cache construction may happen after warm - workers have inherited those paths. + Scratch and trash paths from interrupted work are removed, and active + entries are trimmed to budget using entry-root mtimes as LRU order. + Tarball-bearing entries are protected for the pool backend because cache + construction may happen after warm workers have inherited those paths. The worker warms this sweep before activities can run; lazy first-use sweeping remains a safe fallback. @@ -1515,13 +1532,12 @@ def _sweep_startup_state(self) -> None: return try: - legacy_clean = self._remove_legacy_cache_paths() staging_clean = self._clear_work_dir( self.staging_dir, remember_failures=True, ) trash_clean = self._clear_work_dir(self.trash_dir) - cleanup_complete = legacy_clean and staging_clean and trash_clean + cleanup_complete = staging_clean and trash_clean within_budget = cleanup_complete and self._trim_startup_cache() self._budget_dirty = not (cleanup_complete and within_budget) except OSError as e: @@ -1573,69 +1589,6 @@ def _retry_failed_startup_cleanup(self) -> bool: self._failed_startup_cleanup.discard(path) return not self._failed_startup_cleanup - def _remove_legacy_cache_paths(self) -> bool: - """Best-effort remove flat-layout cache paths from earlier executors.""" - try: - paths = list(self.cache_dir.iterdir()) - except FileNotFoundError: - return True - except OSError as e: - logger.warning( - "Failed to inspect registry artifact cache root", - cache_dir=str(self.cache_dir), - error=str(e), - ) - raise - - current_names = { - BASE_PYTHONPATH_DIR_NAME, - CACHE_ENTRIES_DIR_NAME, - CACHE_STAGING_DIR_NAME, - CACHE_TRASH_DIR_NAME, - } - backend_type = resolve_backend_type() - mounted_keys = { - cache_key - for path in paths - if path.name.startswith("squashfs-") - and not path.name.endswith(".squashfs") - and path.is_mount() - and (cache_key := _legacy_cache_key(path.name)) is not None - } - deleted = True - for path in paths: - if path.name in current_names: - continue - cache_key = _legacy_cache_key(path.name) - is_legacy_entry = cache_key is not None - is_legacy_scratch = ( - LEGACY_TEMP_ARTIFACT_PATTERN.fullmatch(path.name) is not None - ) - if not is_legacy_entry and not is_legacy_scratch: - continue - if cache_key in mounted_keys: - logger.warning( - "Preserving active legacy registry artifact", - cache_key=cache_key, - path=str(path), - ) - continue - if backend_type == ExecutorBackendType.POOL and path.name.startswith( - "tarball-" - ): - continue - if path.is_mount(): - deleted = False - logger.warning( - "Cannot remove mounted legacy registry artifact path", - path=str(path), - ) - continue - if not _delete_cache_path(path): - deleted = False - self._failed_startup_cleanup.add(path) - return deleted - def _trim_startup_cache(self) -> bool: """Trim the cache to budget before any artifact is leased. From f2967ae61020d7621d56eda2eedf71c323799748 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Fri, 31 Jul 2026 14:59:35 -0400 Subject: [PATCH 028/161] fix(executor): converge cache after failed admission --- tests/unit/test_registry_artifacts.py | 41 +++++++++++++++++++++++++ tracecat/executor/registry_artifacts.py | 3 +- 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index aa8e184ec9..61dbaad236 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1097,6 +1097,47 @@ async def mock_download(self, ctx, path): assert not cache._paths_for(cache_key).entry_dir.exists() assert cache_key not in cache._discover_cache_keys() + @pytest.mark.anyio + async def test_failed_first_admission_converges_deposited_image( + self, temp_cache_dir: Path + ) -> None: + """A failed first artifact must not strand an over-budget image.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/failed-first.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) + + async def fail_after_deposit( + artifact: SquashfsArtifact, + ctx: RegistryArtifactMaterializationContext, + ) -> list[Path]: + del artifact + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + ctx.paths.squashfs_image_path.write_bytes(b"reusable image") + raise RuntimeError("mount failed") + + with ( + patch(BACKEND_CONFIG, ExecutorBackendType.DIRECT.value), + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 1), + patch.object( + cache, + "_artifact_candidates", + new_callable=AsyncMock, + return_value=[artifact], + ), + patch.object(SquashfsArtifact, "materialize", fail_after_deposit), + ): + with pytest.raises(RuntimeError, match="mount failed"): + async with cache.lease([artifact_uri]): + pass + + assert cache._refcount(cache_key) == 0 + assert not cache._paths_for(cache_key).entry_dir.exists() + assert cache._discover_cache_keys() == set() + assert cache._budget_dirty is False + @pytest.mark.anyio async def test_cancelled_lease_admission_releases_refcount(self, temp_cache_dir): """Cancellation during candidate lookup must not leak a permanent pin.""" diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 5dea7f1f18..1fa68f69ec 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -869,8 +869,7 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat ] for cache_key in idle_keys: await self._unmount_idle_entry(cache_key) - if leased_keys: - await self._converge_cache_budget() + await self._converge_cache_budget() async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Path]]: """Pin and materialize one artifact, returning its releasable cache key.""" From 5dcaf6bda5dc34f321a4a378a1132fe88f86abe3 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:37:25 -0400 Subject: [PATCH 029/161] test(executor): isolate registry smoke cache runner --- tests/unit/test_executor_sandbox_nsjail.py | 28 +++++++++++++++------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/tests/unit/test_executor_sandbox_nsjail.py b/tests/unit/test_executor_sandbox_nsjail.py index e2025ed634..3845b40661 100644 --- a/tests/unit/test_executor_sandbox_nsjail.py +++ b/tests/unit/test_executor_sandbox_nsjail.py @@ -620,6 +620,7 @@ async def _run_current_builtin_smoke_case( *, smoke_case: SmokeCase, monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, ) -> None: if reason := _missing_prerequisite(smoke_case): _skip_smoke(reason) @@ -652,16 +653,23 @@ async def _run_current_builtin_smoke_case( role=role, ) + runner = ActionRunner(cache_dir=tmp_path / "registry-cache") backend = EphemeralBackend() if smoke_case.force_sandbox else DirectBackend() + get_action_runner_path = ( + "tracecat.executor.backends.ephemeral.get_action_runner" + if smoke_case.force_sandbox + else "tracecat.executor.backends.direct.get_action_runner" + ) action_gateway = ActionGateway() try: await action_gateway.start() - result = await backend.execute( - input=action_input, - role=role, - resolved_context=resolved_context, - timeout=30, - ) + with patch(get_action_runner_path, return_value=runner): + result = await backend.execute( + input=action_input, + role=role, + resolved_context=resolved_context, + timeout=30, + ) finally: await action_gateway.stop() @@ -809,7 +817,9 @@ async def test_action_runner_executes_registry_action_smoke( _run_executor_action_smoke_in_docker_or_skip(smoke_case) return await _run_current_builtin_smoke_case( - smoke_case=smoke_case, monkeypatch=monkeypatch + smoke_case=smoke_case, + monkeypatch=monkeypatch, + tmp_path=tmp_path, ) return @@ -831,7 +841,9 @@ async def run() -> None: try: if smoke_case in _CURRENT_BUILTIN_CASES: await _run_current_builtin_smoke_case( - smoke_case=smoke_case, monkeypatch=monkeypatch + smoke_case=smoke_case, + monkeypatch=monkeypatch, + tmp_path=tmp_path, ) else: await _run_executor_action_smoke_case( From 0658e4b2cc4c9733212e254dcc49881462696a5b Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Fri, 31 Jul 2026 16:22:32 -0400 Subject: [PATCH 030/161] fix(executor): preserve warm cache on failed admission --- tests/unit/test_registry_artifacts.py | 55 +++++++++++++++++++------ tracecat/executor/registry_artifacts.py | 41 +++++++++--------- 2 files changed, 62 insertions(+), 34 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 57151c9f5b..b36fc6c9a7 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -143,9 +143,10 @@ async def _materialize( return cached_paths try: - await cache._enforce_cache_budget(protected_key=cache_key) async with lock: - return await cache._materialize_candidates(ctx, artifact_uri) + paths = await cache._materialize_candidates(ctx, artifact_uri) + await cache._enforce_cache_budget(protected_key=cache_key) + return paths finally: cache._release_lease(cache_key) @@ -1743,10 +1744,10 @@ async def mock_extract(self, tarball_path, target_dir): assert not idle_dir.exists() @pytest.mark.anyio - async def test_releasing_a_lease_converges_the_cache_to_budget( + async def test_successful_admission_enforces_actual_size_before_yield( self, temp_cache_dir ): - """Enforcement before materialization cannot see the new entry's size.""" + """Post-publication enforcement sees the new entry's actual size.""" cache = RegistryArtifactCache(temp_cache_dir) idle = _write_image_entry(temp_cache_dir, "idle", size=4096, mtime=100.0) new_uri = "s3://bucket/new.tar.gz" @@ -1765,9 +1766,8 @@ async def mock_extract(self, tarball_path, target_dir): patch.object(TarballArtifact, "extract", mock_extract), ): async with cache.lease([new_uri]) as registry_paths: - # Both entries fit only because the new one is still leased. assert registry_paths == [cache._paths_for(new_key).tarball_target_dir] - assert idle.exists() + assert not idle.exists() assert not idle.exists() assert cache._paths_for(new_key).tarball_target_dir.is_dir() @@ -1951,6 +1951,33 @@ async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( assert cache._budget_dirty is False + @pytest.mark.anyio + async def test_failed_cold_admission_keeps_warm_lru(self, temp_cache_dir): + """A missing cold artifact must not evict an existing warm entry.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + warm_uri = "s3://bucket/warm.tar.gz" + warm_key = compute_registry_artifact_cache_key(warm_uri) + warm_dir = _write_tarball_entry(temp_cache_dir, warm_key) + missing_uri = "s3://bucket/missing.tar.gz" + missing_key = compute_registry_artifact_cache_key(missing_uri) + + async def mock_download(self, ctx, path): + raise FileNotFoundError("missing artifact") + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "download", mock_download), + ): + with pytest.raises(FileNotFoundError, match="missing artifact"): + async with cache.lease([missing_uri]): + pytest.fail("failed admission must not yield a lease") + + assert warm_dir.is_dir() + assert not cache._paths_for(missing_key).entry_dir.exists() + @pytest.mark.anyio async def test_failed_materialization_rearms_budget_dirty(self, temp_cache_dir): """A failed materialization may leave a canonical image to evict.""" @@ -2273,19 +2300,20 @@ async def controlled_evict( assert retained.exists() @pytest.mark.anyio - async def test_enforce_budget_counts_the_pending_entry(self, temp_cache_dir): - """The entry about to be materialized counts against the entry budget.""" + async def test_enforce_budget_ignores_missing_protected_key(self, temp_cache_dir): + """Only successfully published entries count against the entry budget.""" cache = RegistryArtifactCache(temp_cache_dir) existing = _write_image_entry(temp_cache_dir, "existing", size=16, mtime=100.0) with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): - await cache._enforce_cache_budget(protected_key="pending") + within_budget = await cache._enforce_cache_budget(protected_key="missing") - assert not existing.exists() + assert within_budget is True + assert existing.exists() @pytest.mark.anyio async def test_enforce_budget_never_evicts_the_protected_key(self, temp_cache_dir): - """The key being materialized is exempt even when it is the LRU entry.""" + """The newly materialized key is exempt even when it is the LRU entry.""" cache = RegistryArtifactCache(temp_cache_dir) protected = _write_image_entry( temp_cache_dir, "protected", size=4096, mtime=100.0 @@ -2307,9 +2335,10 @@ async def test_enforce_budget_proceeds_over_budget_when_everything_is_leased( leased = _write_image_entry(temp_cache_dir, "leased", size=4096, mtime=100.0) cache._acquire_lease("leased") - with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): - await cache._enforce_cache_budget(protected_key="pending") + with patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 1): + within_budget = await cache._enforce_cache_budget(protected_key="missing") + assert within_budget is False assert leased.exists() @pytest.mark.anyio diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 3af3be723f..e729a12351 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -885,7 +885,13 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat if cached_paths := self._first_cached_path(candidates, ctx): return cache_key, cached_paths - # Make room outside the per-key lock so eviction never nests key locks. + async with lock: + paths = await self._materialize_candidates(ctx, artifact_uri) + self._touch_entry(cache_key) + + # Enforce only after publication, and outside the per-key lock so + # eviction never nests key locks. A failed cold admission must not + # discard a usable warm entry before a replacement exists. try: await self._enforce_cache_budget(protected_key=cache_key) except OSError as e: @@ -895,10 +901,7 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat cache_dir=str(self.cache_dir), error=str(e), ) - async with lock: - paths = await self._materialize_candidates(ctx, artifact_uri) - self._touch_entry(cache_key) - return cache_key, paths + return cache_key, paths except BaseException: if self._release_lease(cache_key): await self._unmount_idle_entry(cache_key) @@ -1163,11 +1166,11 @@ def _can_try_squashfs(self) -> bool: async def _converge_cache_budget(self) -> None: """Bring an idle cache back under budget after a lease is released. - Materialization enforces the budget before a new entry exists, so the - cache can legitimately sit over budget while that entry is leased. This - runs on release, when the real on-disk size is known and the entry is - evictable. The scan is skipped entirely unless a new entry has landed - since the last successful enforcement. + Successful materialization enforces the budget after publication while + protecting the new entry. The cache can still sit over budget while + entries are leased. This runs on release, when every newly idle entry is + evictable. The scan is skipped entirely unless a materialization attempt + has occurred since the last successful enforcement. Each successful pass consumes the dirty signal before its awaited scan. A follow-up pass therefore occurs only when a concurrent materialization @@ -1205,9 +1208,9 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo invoke enforcement without holding a per-key lock. Args: - protected_key: Cache key about to be materialized. It is counted - against the budget but never evicted. None when enforcing - against the entries already on disk. + protected_key: Newly materialized cache key. It is counted against + the budget when present but never evicted. None when enforcing + against idle entries after leases are released. Returns: Whether the cache is within budget once eviction has finished. @@ -1227,17 +1230,13 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo return True entries = await asyncio.to_thread(self._scan_cache_entries) - # The protected key is not on disk yet when it is a fresh entry. - pending_entries = ( - 1 if protected_key is not None and protected_key not in entries else 0 - ) total_bytes = sum(entry.size_bytes for entry in entries.values()) protected = set() if protected_key is None else {protected_key} skipped: set[str] = set() - while ( - max_entries > 0 and len(entries) + pending_entries > max_entries - ) or (max_bytes > 0 and total_bytes > max_bytes): + while (max_entries > 0 and len(entries) > max_entries) or ( + max_bytes > 0 and total_bytes > max_bytes + ): candidate = self._least_recently_used( entries.values(), excluded=skipped | protected, @@ -1246,7 +1245,7 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo logger.warning( "Registry artifact cache is over budget but every entry is in use", cache_dir=str(self.cache_dir), - entries=len(entries) + pending_entries, + entries=len(entries), max_entries=max_entries, total_bytes=total_bytes, max_bytes=max_bytes, From 58246ecba7f402c38cc23e91341bb30aeb4411a0 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Fri, 31 Jul 2026 17:50:15 -0400 Subject: [PATCH 031/161] fix(executor): preserve leases on cancelled admission --- tests/unit/test_registry_artifacts.py | 33 +++++++++++++++++++++++++ tracecat/executor/registry_artifacts.py | 4 ++- 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index b36fc6c9a7..df798df1b2 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1170,6 +1170,39 @@ async def take_lease() -> None: ) assert not target_dir.exists() + @pytest.mark.anyio + async def test_cancelled_waiter_preserves_existing_same_key_lease( + self, temp_cache_dir: Path + ) -> None: + """A waiter cancelled before admission cannot release another holder.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/path/shared.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + lock = cache._runtime_for(cache_key).lock + + async def take_lease() -> None: + async with cache.lease([artifact_uri]): + pytest.fail("cancelled waiter must not enter the lease context") + + unmount_idle_entry = AsyncMock() + with patch.object(cache, "_unmount_idle_entry", unmount_idle_entry): + async with lock: + cache._acquire_lease(cache_key) + try: + waiter = asyncio.create_task(take_lease()) + await asyncio.sleep(0) + assert not waiter.done() + + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + + assert cache._refcount(cache_key) == 1 + unmount_idle_entry.assert_not_awaited() + finally: + cache._release_lease(cache_key) + @pytest.mark.anyio async def test_lease_without_uris_returns_base_pythonpath_dir(self, temp_cache_dir): """No artifact URIs still yields the base PYTHONPATH directory.""" diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index e729a12351..2a517afcc3 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -878,9 +878,11 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat return None, await self._materialize_candidates(ctx, artifact_uri) lock = self._runtime_for(cache_key).lock + lease_acquired = False try: async with lock: self._acquire_lease(cache_key) + lease_acquired = True candidates = await self._artifact_candidates(ctx, artifact_uri) if cached_paths := self._first_cached_path(candidates, ctx): return cache_key, cached_paths @@ -903,7 +905,7 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat ) return cache_key, paths except BaseException: - if self._release_lease(cache_key): + if lease_acquired and self._release_lease(cache_key): await self._unmount_idle_entry(cache_key) raise From 096b1268f0977a662f03e34bbfaa7ca6af4df33e Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:58:49 -0400 Subject: [PATCH 032/161] fix(executor): rejoin extraction after repeated cancellation --- tests/unit/test_registry_artifacts.py | 15 +++++++++++---- tracecat/executor/registry_artifacts.py | 13 +++++++++++-- 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index df798df1b2..5e802af088 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -755,8 +755,10 @@ async def create_sleep_subprocess( assert captured.returncode is not None @pytest.mark.anyio - async def test_cancelled_tarball_extract_rejoins_thread(self, temp_cache_dir): - """Cancellation waits until the tar extractor stops writing.""" + async def test_repeatedly_cancelled_tarball_extract_rejoins_thread( + self, temp_cache_dir + ): + """Repeated cancellation waits until the tar extractor stops writing.""" artifact = TarballArtifact( uri="s3://bucket/path/site-packages.tar.gz", cache_key="cancelled-tarball-extract", @@ -786,14 +788,19 @@ def blocking_extractall(*args: object, **kwargs: object) -> None: assert await asyncio.to_thread(extraction_started.wait, 1) extracting.cancel() done, _ = await asyncio.wait({extracting}, timeout=0.05) - cancellation_propagated_early = bool(done) + first_cancellation_propagated_early = bool(done) + + extracting.cancel() + done, _ = await asyncio.wait({extracting}, timeout=0.05) + second_cancellation_propagated_early = bool(done) extraction_release.set() with pytest.raises(asyncio.CancelledError): await extracting assert extraction_finished.is_set() - assert cancellation_propagated_early is False + assert first_cancellation_propagated_early is False + assert second_cancellation_propagated_early is False @pytest.mark.anyio async def test_materialize_extracts_squashfs_when_mount_fails(self, temp_cache_dir): diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 2a517afcc3..dabf7225da 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -539,10 +539,19 @@ def _do_extract() -> None: await asyncio.shield(extraction) except asyncio.CancelledError: # A thread cannot be killed. Rejoin it before materialize removes - # scratch; a second cancellation may interrupt this best-effort join. + # scratch. Each cancellation can interrupt shield without stopping + # the thread, so keep waiting until extraction reaches a terminal + # state before propagating the original cancellation. + while not extraction.done(): + try: + await asyncio.shield(extraction) + except asyncio.CancelledError: + continue + except Exception: + break if not extraction.cancelled(): with contextlib.suppress(Exception): - await asyncio.shield(extraction) + extraction.result() raise logger.debug( From 8369ae682a6cf9aa0b12d066713d7dac5d9626c4 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Fri, 31 Jul 2026 23:18:34 -0400 Subject: [PATCH 033/161] fix(executor): keep cache warmup fail-open --- .../unit/test_worker_activity_registration.py | 63 +++++++++++++++++++ tracecat/executor/worker.py | 10 ++- 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_worker_activity_registration.py b/tests/unit/test_worker_activity_registration.py index 1c1476da82..c68ff8a4b7 100644 --- a/tests/unit/test_worker_activity_registration.py +++ b/tests/unit/test_worker_activity_registration.py @@ -62,6 +62,69 @@ def test_agent_executor_worker_registers_runtime_execution_activities() -> None: } +@pytest.mark.anyio +async def test_executor_worker_continues_after_registry_cache_warmup_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from tracecat.executor import worker + + shutdown_event = asyncio.Event() + warmup = AsyncMock(side_effect=OSError("transient cache failure")) + action_runner = Mock() + action_runner.registry_artifacts.ensure_swept = warmup + action_gateway = Mock() + action_gateway.start = AsyncMock() + action_gateway.stop = AsyncMock() + initialize_backend = AsyncMock() + shutdown_backend = AsyncMock() + close_storage_cache = AsyncMock() + get_temporal_client = AsyncMock(return_value=object()) + warning = Mock() + worker_constructed = False + + class _FakeWorker: + def __init__(self, *args: object, **kwargs: object) -> None: + nonlocal worker_constructed + del args, kwargs + worker_constructed = True + + async def __aenter__(self) -> _FakeWorker: + shutdown_event.set() + return self + + async def __aexit__( + self, + exc_type: object, + exc: object, + tb: object, + ) -> None: + del exc_type, exc, tb + + monkeypatch.setattr(worker, "ActionGateway", lambda: action_gateway) + monkeypatch.setattr(worker, "get_action_runner", lambda: action_runner) + monkeypatch.setattr(worker, "initialize_executor_backend", initialize_backend) + monkeypatch.setattr(worker, "shutdown_executor_backend", shutdown_backend) + monkeypatch.setattr(worker, "close_storage_client_cache", close_storage_cache) + monkeypatch.setattr(worker, "get_temporal_client", get_temporal_client) + monkeypatch.setattr(worker, "Worker", _FakeWorker) + monkeypatch.setattr(worker, "new_sandbox_runner", lambda: object()) + monkeypatch.setattr(worker.logger, "warning", warning) + + await worker.main(shutdown_event=shutdown_event) + + warmup.assert_awaited_once() + initialize_backend.assert_awaited_once() + get_temporal_client.assert_awaited_once() + assert worker_constructed is True + warning.assert_called_once_with( + "Registry artifact cache warmup failed; continuing worker startup", + error="transient cache failure", + ) + shutdown_backend.assert_awaited_once() + close_storage_cache.assert_awaited_once() + action_gateway.stop.assert_awaited_once() + + @pytest.mark.anyio async def test_dsl_worker_treats_empty_concurrency_env_vars_as_defaults( monkeypatch: pytest.MonkeyPatch, diff --git a/tracecat/executor/worker.py b/tracecat/executor/worker.py index 08e5114786..d5df125935 100644 --- a/tracecat/executor/worker.py +++ b/tracecat/executor/worker.py @@ -138,7 +138,15 @@ async def main(shutdown_event: asyncio.Event | None = None) -> None: # Warm the registry artifact cache sweep before the backend spawns # workers or activities run; cache construction itself is cheap. - await get_action_runner().registry_artifacts.ensure_swept() + try: + await get_action_runner().registry_artifacts.ensure_swept() + except OSError as e: + # Cache cleanup is best-effort. The failed sweep remains retryable + # at the first lease or materialization boundary. + logger.warning( + "Registry artifact cache warmup failed; continuing worker startup", + error=str(e), + ) # Initialize the executor backend before accepting tasks await initialize_executor_backend() From 0f2180dc0ccf65168e6fce39eea4bff00362b1ff Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 14:43:26 -0400 Subject: [PATCH 034/161] fix(executor): finish lease cleanup before cancellation --- tests/unit/test_registry_artifacts.py | 53 +++++++++++++++++++++++++ tracecat/executor/registry_artifacts.py | 22 ++++++++-- 2 files changed, 72 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 5e802af088..354283499d 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1367,6 +1367,59 @@ async def hold_lease(index: int) -> None: assert harness.mount_attempts == [cache_key] assert harness.unmounts == [mount_dir] + @pytest.mark.anyio + async def test_repeated_cancellation_finishes_all_lease_cleanup( + self, temp_cache_dir: Path + ) -> None: + """Every unmount and budget pass finishes before cancellation propagates.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uris = [ + "s3://bucket/first.tar.gz", + "s3://bucket/second.tar.gz", + ] + cache_keys = [compute_registry_artifact_cache_key(uri) for uri in artifact_uris] + for cache_key in cache_keys: + _write_tarball_entry(temp_cache_dir, cache_key) + + lease_entered = asyncio.Event() + first_unmount_started = asyncio.Event() + finish_first_unmount = asyncio.Event() + cleanup_calls: list[str] = [] + + async def mock_unmount_idle_entry(cache_key: str) -> None: + cleanup_calls.append(cache_key) + if cache_key == cache_keys[0]: + first_unmount_started.set() + await finish_first_unmount.wait() + + async def mock_converge_cache_budget() -> None: + cleanup_calls.append("converge") + + async def hold_lease() -> None: + async with cache.lease(artifact_uris): + lease_entered.set() + await asyncio.Event().wait() + + with ( + patch.object(cache, "_unmount_idle_entry", mock_unmount_idle_entry), + patch.object(cache, "_converge_cache_budget", mock_converge_cache_budget), + ): + holder = asyncio.create_task(hold_lease()) + await lease_entered.wait() + holder.cancel() + await first_unmount_started.wait() + + holder.cancel() + await asyncio.sleep(0) + assert not holder.done() + + finish_first_unmount.set() + with pytest.raises(asyncio.CancelledError): + await holder + + assert cleanup_calls == [*cache_keys, "converge"] + assert all(cache._refcount(cache_key) == 0 for cache_key in cache_keys) + @pytest.mark.anyio async def test_new_lease_racing_final_release_prevents_stale_unmount( self, temp_cache_dir: Path diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index dabf7225da..1d264e03e4 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -875,9 +875,25 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat idle_keys = [ cache_key for cache_key in leased_keys if self._release_lease(cache_key) ] - for cache_key in idle_keys: - await self._unmount_idle_entry(cache_key) - await self._converge_cache_budget() + cleanup_task = asyncio.ensure_future(self._finish_lease_cleanup(idle_keys)) + pending_cancellation: asyncio.CancelledError | None = None + while True: + try: + await asyncio.shield(cleanup_task) + break + except asyncio.CancelledError as e: + if cleanup_task.cancelled(): + raise + pending_cancellation = e + + if pending_cancellation is not None: + raise pending_cancellation + + async def _finish_lease_cleanup(self, idle_keys: list[str]) -> None: + """Unmount every newly idle entry and converge the cache budget.""" + for cache_key in idle_keys: + await self._unmount_idle_entry(cache_key) + await self._converge_cache_budget() async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Path]]: """Pin and materialize one artifact, returning its releasable cache key.""" From 1e87df3dd382041d7322baaa53fc5b5eb578eb40 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 14:49:27 -0400 Subject: [PATCH 035/161] fix(sandbox): terminate nsjail process groups --- tests/unit/test_executor_sandbox_nsjail.py | 114 +++++++++++---------- tracecat/sandbox/executor.py | 49 ++------- 2 files changed, 68 insertions(+), 95 deletions(-) diff --git a/tests/unit/test_executor_sandbox_nsjail.py b/tests/unit/test_executor_sandbox_nsjail.py index 3845b40661..c451035e9e 100644 --- a/tests/unit/test_executor_sandbox_nsjail.py +++ b/tests/unit/test_executor_sandbox_nsjail.py @@ -7,6 +7,7 @@ import json import os import shutil +import signal import subprocess import sys import tarfile @@ -78,6 +79,14 @@ def preferred_format(self) -> RegistryArtifactFormat: return RegistryArtifactFormat.TAR_GZ +class CancelledNsjailOperation(StrEnum): + """NsJail entry points whose cancellation must terminate descendants.""" + + EXECUTE = "execute" + INSTALL = "install" + ACTION = "action" + + def _executor_nsjail_available() -> bool: return ( Path(config.TRACECAT__SANDBOX_NSJAIL_PATH).is_file() @@ -683,61 +692,13 @@ async def _run_current_builtin_smoke_case( } +@pytest.mark.parametrize("operation", list(CancelledNsjailOperation)) @pytest.mark.anyio -async def test_cancelled_nsjail_execute_kills_and_reaps_subprocess( +async def test_cancelled_nsjail_operation_kills_process_group( tmp_path: Path, + operation: CancelledNsjailOperation, ) -> None: - """Cancellation propagates only after the nsjail child is reaped.""" - job_dir = tmp_path / "job" - job_dir.mkdir() - rootfs_dir = tmp_path / "rootfs" - rootfs_dir.mkdir() - runner = NsjailExecutor( - nsjail_path=str(tmp_path / "nsjail"), - rootfs_path=str(rootfs_dir), - cache_dir=str(tmp_path / "cache"), - ) - real_create_subprocess_exec = asyncio.create_subprocess_exec - process_started = asyncio.Event() - process: asyncio.subprocess.Process | None = None - - async def capture_subprocess(*args, **kwargs): - nonlocal process - process = await real_create_subprocess_exec( - "/bin/sleep", - "30", - **kwargs, - ) - process_started.set() - return process - - with patch( - "tracecat.sandbox.executor.asyncio.create_subprocess_exec", - side_effect=capture_subprocess, - ): - execution = asyncio.create_task(runner.execute(job_dir, SandboxConfig())) - try: - await process_started.wait() - await asyncio.sleep(0) - execution.cancel() - - with pytest.raises(asyncio.CancelledError): - await execution - - assert process is not None - assert process.returncode is not None - assert not (job_dir / "nsjail.cfg").exists() - finally: - if process is not None and process.returncode is None: - process.kill() - await process.wait() - - -@pytest.mark.anyio -async def test_cancelled_nsjail_action_kills_and_reaps_subprocess( - tmp_path: Path, -) -> None: - """Cancellation propagates only after the nsjail child is reaped.""" + """Cancellation propagates only after nsjail and its child are gone.""" job_dir = tmp_path / "job" job_dir.mkdir() rootfs_dir = tmp_path / "rootfs" @@ -755,12 +716,23 @@ async def test_cancelled_nsjail_action_kills_and_reaps_subprocess( real_create_subprocess_exec = asyncio.create_subprocess_exec process_started = asyncio.Event() process: asyncio.subprocess.Process | None = None + descendant_pid_path = tmp_path / "descendant.pid" + descendant_pid: int | None = None async def capture_subprocess(*args, **kwargs): nonlocal process + assert kwargs["start_new_session"] is True process = await real_create_subprocess_exec( - "/bin/sleep", - "30", + sys.executable, + "-c", + ( + "import subprocess, sys, time; " + "child = subprocess.Popen([sys.executable, '-c', " + "'import time; time.sleep(30)']); " + "open(sys.argv[1], 'w').write(str(child.pid)); " + "time.sleep(30)" + ), + str(descendant_pid_path), **kwargs, ) process_started.set() @@ -770,10 +742,29 @@ async def capture_subprocess(*args, **kwargs): "tracecat.sandbox.executor.asyncio.create_subprocess_exec", side_effect=capture_subprocess, ): - execution = asyncio.create_task(runner.execute_action(job_dir, sandbox_config)) + match operation: + case CancelledNsjailOperation.EXECUTE: + execution = asyncio.create_task( + runner.execute(job_dir, SandboxConfig()) + ) + case CancelledNsjailOperation.INSTALL: + execution = asyncio.create_task( + runner.execute_install(job_dir, "deadbeef") + ) + case CancelledNsjailOperation.ACTION: + execution = asyncio.create_task( + runner.execute_action(job_dir, sandbox_config) + ) + try: await process_started.wait() - await asyncio.sleep(0) + for _ in range(100): + if descendant_pid_path.exists(): + break + await asyncio.sleep(0.01) + assert descendant_pid_path.is_file() + descendant_pid = int(descendant_pid_path.read_text()) + execution.cancel() with pytest.raises(asyncio.CancelledError): @@ -782,10 +773,21 @@ async def capture_subprocess(*args, **kwargs): assert process is not None assert process.returncode is not None assert not (job_dir / "nsjail.cfg").exists() + for _ in range(100): + try: + os.kill(descendant_pid, 0) + except ProcessLookupError: + break + await asyncio.sleep(0.01) + else: + pytest.fail("nsjail descendant survived process-group cleanup") finally: if process is not None and process.returncode is None: process.kill() await process.wait() + if descendant_pid is not None: + with contextlib.suppress(ProcessLookupError): + os.kill(descendant_pid, signal.SIGKILL) @pytest.mark.parametrize( diff --git a/tracecat/sandbox/executor.py b/tracecat/sandbox/executor.py index af861a9988..faf72df5ec 100644 --- a/tracecat/sandbox/executor.py +++ b/tracecat/sandbox/executor.py @@ -1,7 +1,6 @@ """nsjail executor for sandboxed Python execution.""" import asyncio -import contextlib import json import os import re @@ -31,6 +30,7 @@ SandboxErrorCode, SandboxResult, ) +from tracecat.sandbox.utils import communicate_process_group RUN_PYTHON_ACTION_GATEWAY_SOCKET = Path("/var/run/tracecat/action-gateway.sock") """Path visible inside run_python nsjail for executor-owned SDK calls.""" @@ -467,28 +467,17 @@ async def execute( stderr=asyncio.subprocess.PIPE, cwd=str(job_dir), env=env_map, + start_new_session=True, ) try: # Wait with timeout (add buffer for nsjail overhead) timeout = config.resources.timeout_seconds + 10 - stdout_bytes, stderr_bytes = await asyncio.wait_for( - process.communicate(), + stdout_bytes, stderr_bytes = await communicate_process_group( + process, timeout=timeout, ) - - except asyncio.CancelledError: - # The registry-path lease is released as cancellation unwinds, so - # the importing child must be dead and reaped before propagation. - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - raise - except TimeoutError as e: - # Kill the process if it times out - process.kill() - await process.wait() raise SandboxTimeoutError( f"Execution timed out after {config.resources.timeout_seconds}s" ) from e @@ -621,24 +610,16 @@ async def execute_install( stderr=asyncio.subprocess.PIPE, cwd=str(job_dir), env=env_map, + start_new_session=True, ) try: timeout = timeout_seconds + 30 # Extra buffer for package downloads - stdout_bytes, stderr_bytes = await asyncio.wait_for( - process.communicate(), + stdout_bytes, stderr_bytes = await communicate_process_group( + process, timeout=timeout, ) - - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - raise - except TimeoutError as e: - process.kill() - await process.wait() raise SandboxTimeoutError( f"Package installation timed out after {timeout_seconds}s" ) from e @@ -894,27 +875,17 @@ async def execute_action( stderr=asyncio.subprocess.PIPE, cwd=str(job_dir), env=env_map, + start_new_session=True, ) try: # Wait with timeout (add buffer for nsjail overhead) timeout = config.timeout_seconds + 10 - stdout_bytes, stderr_bytes = await asyncio.wait_for( - process.communicate(), + stdout_bytes, stderr_bytes = await communicate_process_group( + process, timeout=timeout, ) - - except asyncio.CancelledError: - # The registry-path lease is released as cancellation unwinds, so - # the importing child must be dead and reaped before propagation. - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - raise - except TimeoutError as e: - process.kill() - await process.wait() raise SandboxTimeoutError( f"Action execution timed out after {config.timeout_seconds}s" ) from e From c1678072ee4fe0403dc634a2b260d21e68affede Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 14:52:32 -0400 Subject: [PATCH 036/161] fix(executor): retain test leases through sync timeout --- .../test_test_backend_no_registry_action.py | 100 +++++++++++++++++- tracecat/executor/backends/test.py | 33 +++++- 2 files changed, 130 insertions(+), 3 deletions(-) diff --git a/tests/unit/executor/test_test_backend_no_registry_action.py b/tests/unit/executor/test_test_backend_no_registry_action.py index faf96571f4..50363e03fc 100644 --- a/tests/unit/executor/test_test_backend_no_registry_action.py +++ b/tests/unit/executor/test_test_backend_no_registry_action.py @@ -8,7 +8,9 @@ from __future__ import annotations +import asyncio import sys +import threading import uuid from collections.abc import AsyncIterator from contextlib import asynccontextmanager @@ -27,7 +29,11 @@ RunContext, ) from tracecat.executor.backends.test import TestBackend -from tracecat.executor.schemas import ActionImplementation, ResolvedContext +from tracecat.executor.schemas import ( + ActionImplementation, + ExecutorResult, + ResolvedContext, +) from tracecat.executor.secret_preprocessors import SecretEnvProjection from tracecat.identifiers.workflow import ExecutionUUID, WorkflowUUID from tracecat.registry.lock.types import RegistryLock @@ -328,6 +334,98 @@ async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: finally: await backend.shutdown() + @pytest.mark.anyio + async def test_timed_out_sync_udf_keeps_artifact_lease_until_thread_finishes( + self, + test_role: Role, + test_resolved_context: ResolvedContext, + test_run_action_input: RunActionInput, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + """A soft timeout cannot release paths beneath a live UDF thread.""" + artifact_path = tmp_path / "artifact" + artifact_path.mkdir() + artifact_uri = "s3://bucket/sync-timeout.tar.gz" + worker_started = threading.Event() + finish_worker = threading.Event() + + class FakeRegistryArtifacts: + def __init__(self) -> None: + self.active = 0 + + @asynccontextmanager + async def lease( + self, artifact_uris: list[str] | None = None + ) -> AsyncIterator[list[Path]]: + assert artifact_uris == [artifact_uri] + self.active += 1 + try: + yield [artifact_path] + finally: + self.active -= 1 + + class FakeActionRunner: + def __init__(self) -> None: + self.registry_artifacts = FakeRegistryArtifacts() + + fake_runner = FakeActionRunner() + + async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: + return [artifact_uri] + + def blocking_udf(**_kwargs: object) -> str: + worker_started.set() + assert finish_worker.wait(timeout=5) + return "finished" + + backend = TestBackend() + await backend.start() + execution: asyncio.Task[ExecutorResult] | None = None + try: + monkeypatch.setattr( + "tracecat.executor.backends.test.config" + ".TRACECAT__LOCAL_REPOSITORY_ENABLED", + False, + ) + monkeypatch.setattr( + "tracecat.executor.backends.test.get_action_runner", + lambda: fake_runner, + ) + monkeypatch.setattr(backend, "_get_artifact_uris", _get_artifact_uris) + monkeypatch.setattr( + backend, + "_load_udf_callable", + lambda _action_impl: blocking_udf, + ) + + execution = asyncio.create_task( + backend.execute( + input=test_run_action_input, + role=test_role, + resolved_context=test_resolved_context, + timeout=0.01, + ) + ) + assert await asyncio.to_thread(worker_started.wait, 1) + await asyncio.sleep(0.05) + + assert not execution.done() + assert fake_runner.registry_artifacts.active == 1 + assert str(artifact_path) in sys.path + + finish_worker.set() + result = await execution + assert result.type == "failure" + assert result.error.type == "TimeoutError" + assert fake_runner.registry_artifacts.active == 0 + assert str(artifact_path) not in sys.path + finally: + finish_worker.set() + if execution is not None: + await execution + await backend.shutdown() + @pytest.mark.anyio async def test_execute_udf_reuses_cached_secret_projection( self, diff --git a/tracecat/executor/backends/test.py b/tracecat/executor/backends/test.py index 1f7cae6952..696141180c 100644 --- a/tracecat/executor/backends/test.py +++ b/tracecat/executor/backends/test.py @@ -18,6 +18,7 @@ from __future__ import annotations import asyncio +import contextlib import sys import threading from contextlib import AsyncExitStack, contextmanager @@ -54,7 +55,7 @@ from tracecat.secrets import secrets_manager if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Callable, Iterator from tracecat.auth.types import Role from tracecat.dsl.schemas import RunActionInput @@ -228,13 +229,41 @@ async def _execute_with_context( if asyncio.iscoroutinefunction(fn): result = await fn(**args) else: - result = await asyncio.to_thread(fn, **args) + result = await self._run_sync_udf(fn, args) log.trace("Result", result=result) return result finally: registry_secrets.reset_context(secrets_token) + async def _run_sync_udf( + self, + fn: Callable[..., Any], + args: dict[str, Any], + ) -> Any: + """Keep a non-interruptible UDF thread joined through cancellation. + + TestBackend timeouts are necessarily soft for synchronous functions: a + Python thread cannot be killed safely. Rejoining it keeps registry + leases, temporary ``sys.path`` entries, and secret contexts alive until + the function actually stops. + """ + worker = asyncio.ensure_future(asyncio.to_thread(fn, **args)) + try: + return await asyncio.shield(worker) + except asyncio.CancelledError: + while not worker.done(): + try: + await asyncio.shield(worker) + except asyncio.CancelledError: + continue + except Exception: + break + if not worker.cancelled(): + with contextlib.suppress(Exception): + worker.result() + raise + def _load_udf_callable(self, action_impl: ActionImplementation): """Load the UDF callable from action_impl metadata.""" if not action_impl.module or not action_impl.name: From 6df96a11440dc1f304948762cf6e498807876f60 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 14:54:02 -0400 Subject: [PATCH 037/161] fix(executor): rejoin deletion after repeated cancellation --- tests/unit/test_registry_artifacts.py | 2 ++ tracecat/executor/registry_artifacts.py | 12 ++++++++++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 354283499d..f04a4d89d9 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -2574,6 +2574,8 @@ async def mock_extract(self, tarball_path, target_dir): assert await asyncio.to_thread(delete_started.wait, 1) eviction.cancel() try: + await asyncio.sleep(0) + eviction.cancel() await asyncio.sleep(0) assert not eviction.done() assert not original_target.exists() diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 1d264e03e4..0aa8a9e212 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -720,10 +720,18 @@ async def _delete_cache_path_off_loop(path: Path) -> bool: return await asyncio.shield(deletion) except asyncio.CancelledError: # A worker thread cannot be killed. Rejoin it so no live deletion can - # race a later trash-directory scan. + # race a later trash-directory scan. Repeated cancellation can interrupt + # shield without stopping the thread, so keep waiting for termination. + while not deletion.done(): + try: + await asyncio.shield(deletion) + except asyncio.CancelledError: + continue + except Exception: + break if not deletion.cancelled(): with contextlib.suppress(Exception): - await asyncio.shield(deletion) + deletion.result() raise From 988fc6f93ac0e271f09b08f8caec2243f51c7378 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 15:03:34 -0400 Subject: [PATCH 038/161] fix(executor): single-flight cold cache admission --- tests/unit/test_registry_artifacts.py | 113 +++++++++++------------- tracecat/executor/registry_artifacts.py | 42 +++++++-- 2 files changed, 87 insertions(+), 68 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index f04a4d89d9..2bb9e51fed 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -131,24 +131,21 @@ async def _materialize( await cache.ensure_swept() ctx = cache._context_for(cache_key) lock = cache._runtime_for(cache_key).lock - async with lock: - cache._acquire_lease(cache_key) - try: - candidates = await cache._artifact_candidates(ctx, artifact_uri) - except BaseException: - cache._release_lease(cache_key) - raise - if cached_paths := cache._first_cached_path(candidates, ctx): - cache._release_lease(cache_key) - return cached_paths - + lease_acquired = False try: async with lock: - paths = await cache._materialize_candidates(ctx, artifact_uri) + cache._acquire_lease(cache_key) + lease_acquired = True + candidates = await cache._artifact_candidates(ctx, artifact_uri) + if cached_paths := cache._first_cached_path(candidates, ctx): + return cached_paths + paths = await cache._materialize_candidates(ctx, candidates) + cache._touch_entry(cache_key) await cache._enforce_cache_budget(protected_key=cache_key) return paths finally: - cache._release_lease(cache_key) + if lease_acquired: + cache._release_lease(cache_key) class _BlockingSubprocess: @@ -421,56 +418,56 @@ async def test_artifact_candidates_prefer_squashfs_sidecar(self, temp_cache_dir) ) @pytest.mark.anyio - async def test_materialize_recomputes_candidates_after_lock(self, temp_cache_dir): - """Test that lock waiters re-check preferred artifact candidates.""" + async def test_same_key_cold_fan_in_materializes_and_enforces_once( + self, temp_cache_dir + ): + """Same-key waiters share candidate lookup, materialization, and enforcement.""" cache = RegistryArtifactCache(temp_cache_dir) - cache_key = "recompute-key" - cached_path = temp_cache_dir / "cached-squashfs" - cached_path.mkdir() - tarball = TarballArtifact( - uri="s3://bucket/path/site-packages.tar.gz", - cache_key=cache_key, - ) - squashfs = SquashfsArtifact( - uri="s3://bucket/path/site-packages.squashfs", - cache_key=cache_key, - ) - pre_lock_candidates = [tarball] - post_lock_candidates = [squashfs, tarball] - seen_candidates: list[list[RegistryArtifactFormat]] = [] + await cache.ensure_swept() + artifact_uri = "s3://bucket/path/site-packages.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + materialization_started = asyncio.Event() + finish_materialization = asyncio.Event() + + async def mock_materialize( + self: TarballArtifact, + ctx: RegistryArtifactMaterializationContext, + ) -> list[Path]: + materialization_started.set() + await finish_materialization.wait() + ctx.paths.tarball_target_dir.mkdir(parents=True) + return [ctx.paths.tarball_target_dir] - def fake_first_cached_path(candidates, ctx): - del ctx - seen_candidates.append([artifact.format for artifact in candidates]) - if candidates is post_lock_candidates: - return [cached_path] - return None + async def take_lease() -> list[Path]: + async with cache.lease([artifact_uri]) as paths: + return paths with ( patch.object( cache, - "_artifact_candidates", + "_sidecar_exists", new_callable=AsyncMock, - side_effect=[pre_lock_candidates, post_lock_candidates], - ) as artifact_candidates, + return_value=False, + ) as sidecar_exists, patch.object( cache, - "_first_cached_path", - side_effect=fake_first_cached_path, - ), + "_enforce_cache_budget", + new_callable=AsyncMock, + return_value=True, + ) as enforce_cache_budget, + patch.object(cache, "_converge_cache_budget", new_callable=AsyncMock), + patch.object(TarballArtifact, "materialize", mock_materialize), ): - result = await _materialize( - cache, - cache_key, - "s3://bucket/path/site-packages.tar.gz", - ) + leases = [asyncio.create_task(take_lease()) for _ in range(5)] + await materialization_started.wait() + await asyncio.sleep(0) + finish_materialization.set() + results = await asyncio.gather(*leases) - assert result == [cached_path] - assert artifact_candidates.await_count == 2 - assert seen_candidates == [ - [RegistryArtifactFormat.TAR_GZ], - [RegistryArtifactFormat.SQUASHFS, RegistryArtifactFormat.TAR_GZ], - ] + expected_paths = [cache._paths_for(cache_key).tarball_target_dir] + assert results == [expected_paths] * 5 + sidecar_exists.assert_awaited_once() + enforce_cache_budget.assert_awaited_once_with(protected_key=cache_key) @pytest.mark.anyio async def test_artifact_candidates_direct_squashfs_include_gzip_fallback( @@ -1146,7 +1143,6 @@ async def test_cancelled_lease_admission_releases_refcount(self, temp_cache_dir) cache = RegistryArtifactCache(temp_cache_dir) artifact_uri = "s3://bucket/path/site-packages.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) - target_dir = _write_tarball_entry(temp_cache_dir, cache_key) lookup_started = asyncio.Event() finish_lookup = asyncio.Event() @@ -1171,11 +1167,7 @@ async def take_lease() -> None: await acquisition assert cache._refcount(cache_key) == 0 - assert await cache._evict_entry(cache_key) == RegistryArtifactEviction( - retired=True, - reclaimed=True, - ) - assert not target_dir.exists() + assert not cache._paths_for(cache_key).entry_dir.exists() @pytest.mark.anyio async def test_cancelled_waiter_preserves_existing_same_key_lease( @@ -1620,12 +1612,11 @@ async def take_lease() -> list[Path]: assert download_attempts == 1 fail_first_download.set() - with pytest.raises(RuntimeError, match="first publisher failed"): - await first - await retry_download_started.wait() allow_retry_download.set() registry_paths = await waiter + with pytest.raises(RuntimeError, match="first publisher failed"): + await first target_dir = cache._paths_for(cache_key).tarball_target_dir assert registry_paths == [target_dir] diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 0aa8a9e212..cbc31c6062 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -908,7 +908,8 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat cache_key = compute_registry_artifact_cache_key(artifact_uri) ctx = self._context_for(cache_key) if not _is_cache_entry_uri(artifact_uri): - return None, await self._materialize_candidates(ctx, artifact_uri) + candidates = await self._artifact_candidates(ctx, artifact_uri) + return None, await self._materialize_candidates(ctx, candidates) lock = self._runtime_for(cache_key).lock lease_acquired = False @@ -916,12 +917,12 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat async with lock: self._acquire_lease(cache_key) lease_acquired = True + if cached_paths := self._locally_cached_path(ctx, artifact_uri): + return cache_key, cached_paths candidates = await self._artifact_candidates(ctx, artifact_uri) if cached_paths := self._first_cached_path(candidates, ctx): return cache_key, cached_paths - - async with lock: - paths = await self._materialize_candidates(ctx, artifact_uri) + paths = await self._materialize_candidates(ctx, candidates) self._touch_entry(cache_key) # Enforce only after publication, and outside the per-key lock so @@ -945,13 +946,12 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat async def _materialize_candidates( self, ctx: RegistryArtifactMaterializationContext, - artifact_uri: str, + candidates: list[RegistryArtifact], ) -> list[Path]: """Materialize the first viable artifact candidate. Callers hold the cache key's lock for evictable entries. """ - candidates = await self._artifact_candidates(ctx, artifact_uri) if cached_paths := self._first_cached_path(candidates, ctx): return cached_paths @@ -989,7 +989,7 @@ async def _materialize_candidates( error=str(e), ) - raise RuntimeError(f"No registry artifact candidates for {artifact_uri}") + raise RuntimeError(f"No registry artifact candidates for {ctx.cache_key}") def _runtime_for(self, cache_key: str) -> RegistryArtifactRuntimeState: """Return the process-local state for one cache key.""" @@ -1082,6 +1082,34 @@ def _first_cached_path( return cached_paths return None + def _locally_cached_path( + self, + ctx: RegistryArtifactMaterializationContext, + artifact_uri: str, + ) -> list[Path] | None: + """Return a reusable local candidate without probing remote sidecars.""" + artifact_format = _artifact_format(artifact_uri) + candidates: list[RegistryArtifact] = [] + if artifact_format == RegistryArtifactFormat.SQUASHFS: + candidates.append( + SquashfsArtifact(uri=artifact_uri, cache_key=ctx.cache_key) + ) + if tarball_uri := _tarball_uri_for_squashfs(artifact_uri): + candidates.append( + TarballArtifact(uri=tarball_uri, cache_key=ctx.cache_key) + ) + else: + if self._can_try_squashfs() and ( + squashfs_uri := _squashfs_sidecar_uri(artifact_uri) + ): + candidates.append( + SquashfsArtifact(uri=squashfs_uri, cache_key=ctx.cache_key) + ) + candidates.append( + TarballArtifact(uri=artifact_uri, cache_key=ctx.cache_key) + ) + return self._first_cached_path(candidates, ctx) + def _remove_unpublished_entry( self, ctx: RegistryArtifactMaterializationContext, From cb64624806801ed83cf21f3824356f95ba5ef941 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 15:22:49 -0400 Subject: [PATCH 039/161] fix(executor): bound registry cache peak disk use --- tests/unit/test_registry_artifacts.py | 244 ++++++++++++++++++ tests/unit/test_storage_blob.py | 75 ++++++ tracecat/config.py | 5 +- tracecat/executor/registry_artifacts.py | 322 ++++++++++++++++++++---- tracecat/storage/blob.py | 28 ++- 5 files changed, 621 insertions(+), 53 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 2bb9e51fed..6f2f2e5d73 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -3,10 +3,12 @@ from __future__ import annotations import asyncio +import io import os import tarfile import tempfile import threading +from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from pathlib import Path from unittest.mock import ANY, AsyncMock, patch @@ -18,6 +20,7 @@ from tracecat.executor.registry_artifacts import ( SQUASHFS_MOUNT_OPTIONS, RegistryArtifactCache, + RegistryArtifactCacheCapacityError, RegistryArtifactCacheLoopError, RegistryArtifactEviction, RegistryArtifactFormat, @@ -26,6 +29,7 @@ SquashfsMountCommandError, TarballArtifact, _delete_cache_path, + _squashfs_listing_size, bundled_builtin_registry_uri, compute_registry_artifact_cache_key, ) @@ -66,6 +70,17 @@ def _write_image_entry( return image_path +def _tarball_payload(*, size: int) -> bytes: + """Return a gzip tarball containing one synthetic regular file.""" + payload = b"x" * size + output = io.BytesIO() + with tarfile.open(fileobj=output, mode="w:gz") as tar: + member = tarfile.TarInfo("module.py") + member.size = len(payload) + tar.addfile(member, io.BytesIO(payload)) + return output.getvalue() + + @dataclass(slots=True) class _SquashfsMountHarness: """Model mount ownership without consuming host loop devices.""" @@ -469,6 +484,72 @@ async def take_lease() -> list[Path]: sidecar_exists.assert_awaited_once() enforce_cache_budget.assert_awaited_once_with(protected_key=cache_key) + @pytest.mark.anyio + async def test_different_cold_keys_serialize_materialization( + self, temp_cache_dir: Path + ) -> None: + """Distinct cold keys cannot multiply staging and extraction peaks.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + uris = ["s3://bucket/first.tar.gz", "s3://bucket/second.tar.gz"] + first_started = asyncio.Event() + release_first = asyncio.Event() + second_started = asyncio.Event() + started_keys: list[str] = [] + + async def controlled_materialize( + self: TarballArtifact, + ctx: RegistryArtifactMaterializationContext, + ) -> list[Path]: + del self + started_keys.append(ctx.cache_key) + if len(started_keys) == 1: + first_started.set() + await release_first.wait() + else: + second_started.set() + ctx.paths.tarball_target_dir.mkdir(parents=True) + return [ctx.paths.tarball_target_dir] + + async def take_lease(uri: str) -> None: + async with cache.lease([uri]): + pass + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "materialize", controlled_materialize), + patch.object(cache, "_enforce_cache_budget", new_callable=AsyncMock), + patch.object(cache, "_converge_cache_budget", new_callable=AsyncMock), + ): + first = asyncio.create_task(take_lease(uris[0])) + await first_started.wait() + second = asyncio.create_task(take_lease(uris[1])) + await asyncio.sleep(0) + assert not second_started.is_set() + + release_first.set() + await second_started.wait() + await asyncio.gather(first, second) + + assert started_keys == [ + compute_registry_artifact_cache_key(uri) for uri in uris + ] + + def test_squashfs_listing_size_sums_files_and_symlinks(self) -> None: + listing = b"\n".join( + [ + b"drwxr-xr-x 0/0 64 2026-01-01 00:00 squashfs-root", + b"-rw-r--r-- 0/0 123 2026-01-01 00:00 squashfs-root/module.py", + b"lrwxrwxrwx 0/0 9 2026-01-01 00:00 squashfs-root/current -> module.py", + ] + ) + + assert _squashfs_listing_size(listing) == 132 + + def test_squashfs_listing_size_rejects_unparseable_files(self) -> None: + with pytest.raises(ValueError, match="Could not parse SquashFS listing"): + _squashfs_listing_size(b"-rw-r--r-- malformed") + @pytest.mark.anyio async def test_artifact_candidates_direct_squashfs_include_gzip_fallback( self, temp_cache_dir @@ -1602,6 +1683,10 @@ async def take_lease() -> list[Path]: with ( patch(SQUASHFS_ENABLED_CONFIG, False), + patch( + "tracecat.executor.registry_artifacts._tarball_extracted_size", + return_value=1, + ), patch.object(TarballArtifact, "download", controlled_download), patch.object(TarballArtifact, "extract", mock_extract), ): @@ -1827,6 +1912,153 @@ async def mock_extract(self, tarball_path, target_dir): assert leased_dir.is_dir() assert not idle_dir.exists() + @pytest.mark.anyio + async def test_cold_download_reserves_space_before_writing( + self, temp_cache_dir: Path + ) -> None: + """Admission evicts idle bytes before a new download enters staging.""" + cache = RegistryArtifactCache(temp_cache_dir) + idle = _write_image_entry(temp_cache_dir, "idle", size=80, mtime=100.0) + artifact_uri = "s3://bucket/new.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + payload = _tarball_payload(size=32) + max_bytes = len(payload) + 32 + capacity_checked = False + + async def download_file_to_path( + *, + key: str, + bucket: str, + output_path: Path, + max_bytes: int, + ensure_capacity: Callable[[int], Awaitable[None]], + ) -> int: + del key, bucket + nonlocal capacity_checked + assert max_bytes == len(payload) + 32 + await ensure_capacity(len(payload)) + capacity_checked = True + assert not idle.exists() + output_path.write_bytes(payload) + return len(payload) + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, max_bytes), + patch( + "tracecat.executor.registry_artifacts.blob.download_file_to_path", + side_effect=download_file_to_path, + ), + ): + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [ + cache._paths_for(cache_key).tarball_target_dir + ] + + assert capacity_checked is True + assert not idle.exists() + assert (registry_paths[0] / "module.py").read_bytes() == b"x" * 32 + + @pytest.mark.anyio + async def test_compression_heavy_tarball_is_rejected_before_extraction( + self, temp_cache_dir: Path + ) -> None: + """Compressed bytes plus declared extraction cannot exceed the cache cap.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/compression-heavy.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + payload = _tarball_payload(size=4096) + max_bytes = len(payload) + 256 + + async def download_file_to_path( + *, + key: str, + bucket: str, + output_path: Path, + max_bytes: int, + ensure_capacity: Callable[[int], Awaitable[None]], + ) -> int: + del key, bucket + assert max_bytes == len(payload) + 256 + await ensure_capacity(len(payload)) + output_path.write_bytes(payload) + return len(payload) + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, max_bytes), + patch( + "tracecat.executor.registry_artifacts.blob.download_file_to_path", + side_effect=download_file_to_path, + ), + patch.object( + TarballArtifact, + "extract", + new_callable=AsyncMock, + ) as extract, + ): + with pytest.raises(RegistryArtifactCacheCapacityError) as raised: + async with cache.lease([artifact_uri]): + pass + + assert raised.value.additional_bytes == 4096 + assert raised.value.max_bytes == max_bytes + extract.assert_not_awaited() + assert not cache._paths_for(cache_key).entry_dir.exists() + assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) + + @pytest.mark.anyio + async def test_squashfs_expansion_is_rejected_before_extraction( + self, temp_cache_dir: Path + ) -> None: + """SquashFS metadata is accounted before unsquashfs writes scratch.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/oversized.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def download( + self: SquashfsArtifact, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> float: + del self, ctx + image_path.parent.mkdir(parents=True, exist_ok=True) + image_path.write_bytes(b"image") + return 0.0 + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 100), + patch.object( + RegistryArtifactMaterializationContext, + "can_mount_squashfs", + return_value=False, + ), + patch.object(SquashfsArtifact, "download", download), + patch.object( + SquashfsArtifact, + "_squashfs_extracted_size", + new_callable=AsyncMock, + return_value=101, + ), + patch.object( + SquashfsArtifact, + "_extract_image", + new_callable=AsyncMock, + ) as extract_image, + ): + with pytest.raises(RegistryArtifactCacheCapacityError) as raised: + async with cache.lease([artifact_uri]): + pass + + assert raised.value.additional_bytes == 101 + extract_image.assert_not_awaited() + paths = cache._paths_for(cache_key) + assert paths.squashfs_image_path.read_bytes() == b"image" + assert not paths.squashfs_extract_dir.exists() + @pytest.mark.anyio async def test_successful_admission_enforces_actual_size_before_yield( self, temp_cache_dir @@ -1846,6 +2078,10 @@ async def mock_extract(self, tarball_path, target_dir): with ( patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 6000), + patch( + "tracecat.executor.registry_artifacts._tarball_extracted_size", + return_value=4096, + ), patch.object(TarballArtifact, "download", mock_download), patch.object(TarballArtifact, "extract", mock_extract), ): @@ -2005,6 +2241,10 @@ async def mock_extract(self, tarball_path, target_dir): new_callable=AsyncMock, side_effect=PermissionError("denied"), ), + patch( + "tracecat.executor.registry_artifacts._tarball_extracted_size", + return_value=9, + ), patch.object(TarballArtifact, "download", mock_download), patch.object(TarballArtifact, "extract", mock_extract), ): @@ -2558,6 +2798,10 @@ async def mock_extract(self, tarball_path, target_dir): "tracecat.executor.registry_artifacts._delete_cache_path", side_effect=blocked_delete, ), + patch( + "tracecat.executor.registry_artifacts._tarball_extracted_size", + return_value=9, + ), patch.object(TarballArtifact, "download", mock_download), patch.object(TarballArtifact, "extract", mock_extract), ): diff --git a/tests/unit/test_storage_blob.py b/tests/unit/test_storage_blob.py index d9b9f34f87..c960a2cf5b 100644 --- a/tests/unit/test_storage_blob.py +++ b/tests/unit/test_storage_blob.py @@ -771,6 +771,81 @@ async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 assert not out.exists() assert not (tmp_path / "out.bin.part").exists() + @pytest.mark.anyio + async def test_download_file_to_path_reserves_content_length_before_writing( + self, tmp_path: Path, monkeypatch + ): + """Capacity is reserved before a downloader creates its partial file.""" + + class DummyStream: + async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 + yield b"payload" + + @asynccontextmanager + async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 + yield DummyStream(), 7 + + monkeypatch.setattr( + "tracecat.storage.blob.open_download_stream", + _fake_open_download_stream, + ) + + out = tmp_path / "out.bin" + reservations: list[int] = [] + + async def ensure_capacity(size_bytes: int) -> None: + assert not (tmp_path / "out.bin.part").exists() + reservations.append(size_bytes) + + await download_file_to_path( + key="k", + bucket="b", + output_path=out, + max_bytes=10, + ensure_capacity=ensure_capacity, + ) + + assert reservations == [7] + assert out.read_bytes() == b"payload" + + @pytest.mark.anyio + async def test_download_file_to_path_never_exceeds_reserved_length( + self, tmp_path: Path, monkeypatch + ): + """A body larger than ContentLength is rejected before its excess write.""" + + class DummyStream: + async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 + yield b"too-large" + + @asynccontextmanager + async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 + yield DummyStream(), 5 + + monkeypatch.setattr( + "tracecat.storage.blob.open_download_stream", + _fake_open_download_stream, + ) + + out = tmp_path / "out.bin" + reservations: list[int] = [] + + async def ensure_capacity(size_bytes: int) -> None: + reservations.append(size_bytes) + + with pytest.raises(ValueError, match="exceeds max_bytes=5"): + await download_file_to_path( + key="k", + bucket="b", + output_path=out, + max_bytes=10, + ensure_capacity=ensure_capacity, + ) + + assert reservations == [5] + assert not out.exists() + assert not (tmp_path / "out.bin.part").exists() + @pytest.mark.anyio async def test_download_file_to_path_sha256_mismatch_cleans_partial( self, tmp_path: Path, monkeypatch diff --git a/tracecat/config.py b/tracecat/config.py index cbbceedc8f..b92988332d 100644 --- a/tracecat/config.py +++ b/tracecat/config.py @@ -202,8 +202,9 @@ class RLSMode(StrEnum): ) """Maximum on-disk size of the executor-local registry artifact cache, in bytes. -Mounted artifacts only account for their backing image file. Set to 0 to -disable size-based eviction.""" +Cold downloads and extraction scratch are admitted within this bound. Mounted +artifacts only account for their backing image file. Set to 0 to disable +size-based eviction and materialization limits.""" TRACECAT__AGENT_SKILL_CACHE_DIR = os.environ.get( "TRACECAT__AGENT_SKILL_CACHE_DIR", "/tmp/tracecat/agent-skill-cache" diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index cbc31c6062..58e8ad2129 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -12,7 +12,7 @@ import threading import time from abc import ABC, abstractmethod -from collections.abc import AsyncIterator, Iterable +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable from contextlib import asynccontextmanager from dataclasses import dataclass, field from enum import StrEnum @@ -73,6 +73,26 @@ class RegistryArtifactCacheLoopError(RuntimeError): """A registry artifact cache was used outside its owning event loop.""" +class RegistryArtifactCacheCapacityError(RuntimeError): + """A cold artifact cannot fit within the configured cache byte budget.""" + + def __init__( + self, + *, + current_bytes: int, + additional_bytes: int, + max_bytes: int, + ) -> None: + super().__init__( + "Registry artifact admission exceeds the cache byte budget: " + f"current_bytes={current_bytes}, additional_bytes={additional_bytes}, " + f"max_bytes={max_bytes}" + ) + self.current_bytes = current_bytes + self.additional_bytes = additional_bytes + self.max_bytes = max_bytes + + @dataclass(frozen=True, slots=True) class RegistryArtifactPaths: """Executor-local cache paths for one registry artifact key.""" @@ -110,6 +130,14 @@ class RegistryArtifactCacheEntry: last_used: float +@dataclass(frozen=True, slots=True) +class RegistryArtifactAdmission: + """Byte-bound admission hook shared by one cold materialization.""" + + max_bytes: int + ensure_capacity: Callable[[int], Awaitable[None]] + + @dataclass(slots=True) class RegistryArtifactMaterializationContext: """Shared runtime state for artifact materialization.""" @@ -117,6 +145,7 @@ class RegistryArtifactMaterializationContext: cache_key: str staging_dir: Path paths: RegistryArtifactPaths + admission: RegistryArtifactAdmission | None = None def can_mount_squashfs(self) -> bool: return config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED and ( @@ -242,7 +271,11 @@ async def download( temp_image = self._temp_path(ctx, ".squashfs") try: download_start = time.monotonic() - await _download_s3_artifact(self.uri, temp_image) + await _download_s3_artifact( + self.uri, + temp_image, + admission=ctx.admission, + ) try: temp_image.rename(image_path) except OSError: @@ -324,6 +357,9 @@ async def extract( temp_dir = self._temp_path(ctx, ".unsquashfs") try: + if ctx.admission is not None: + extracted_size = await self._squashfs_extracted_size(image_path) + await ctx.admission.ensure_capacity(extracted_size) extract_start = time.monotonic() temp_dir.mkdir(parents=True, exist_ok=True) await self._extract_image(image_path, temp_dir) @@ -434,6 +470,32 @@ async def _extract_image(self, image_path: Path, target_dir: Path) -> None: output = (stderr or stdout).decode(errors="replace").strip() raise RuntimeError(output or "unsquashfs command failed") + async def _squashfs_extracted_size(self, image_path: Path) -> int: + """Return a conservative logical size for an extracted SquashFS image.""" + unsquashfs = shutil.which("unsquashfs") + if unsquashfs is None: + raise RuntimeError("unsquashfs command is not installed") + + proc = await asyncio.create_subprocess_exec( + unsquashfs, + "-lln", + str(image_path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + proc.kill() + await proc.wait() + raise + + if proc.returncode != 0: + output = (stderr or stdout).decode(errors="replace").strip() + raise RuntimeError(output or "unsquashfs listing failed") + return _squashfs_listing_size(stdout) + @dataclass(frozen=True, slots=True) class TarballArtifact(RegistryArtifact): @@ -476,6 +538,13 @@ async def materialize( await self.download(ctx, temp_tarball) download_elapsed = (time.monotonic() - download_start) * 1000 + if ctx.admission is not None: + extracted_size = await asyncio.to_thread( + _tarball_extracted_size, + temp_tarball, + ) + await ctx.admission.ensure_capacity(extracted_size) + extract_start = time.monotonic() temp_dir.mkdir(parents=True, exist_ok=True) await self.extract(temp_tarball, temp_dir) @@ -516,7 +585,11 @@ async def download( ctx: RegistryArtifactMaterializationContext, output_path: Path, ) -> None: - await _download_s3_artifact(self.uri, output_path) + await _download_s3_artifact( + self.uri, + output_path, + admission=ctx.admission, + ) async def extract(self, tarball_path: Path, target_dir: Path) -> None: """Extract a supported registry tarball to target directory. @@ -561,15 +634,29 @@ def _do_extract() -> None: ) -async def _download_s3_artifact(artifact_uri: str, output_path: Path) -> None: +async def _download_s3_artifact( + artifact_uri: str, + output_path: Path, + *, + admission: RegistryArtifactAdmission | None = None, +) -> None: """Download an S3 registry artifact to a local path.""" bucket, key = parse_s3_uri(artifact_uri) try: - await blob.download_file_to_path( - key=key, - bucket=bucket, - output_path=output_path, - ) + if admission is None: + await blob.download_file_to_path( + key=key, + bucket=bucket, + output_path=output_path, + ) + else: + await blob.download_file_to_path( + key=key, + bucket=bucket, + output_path=output_path, + max_bytes=admission.max_bytes, + ensure_capacity=admission.ensure_capacity, + ) except FileNotFoundError as e: request = httpx.Request("GET", artifact_uri) response = httpx.Response(status_code=404, request=request) @@ -668,6 +755,38 @@ def _is_cache_entry_uri(artifact_uri: str) -> bool: return _bundled_builtin_registry_version(artifact_uri) is None +def _tarball_extracted_size(tarball_path: Path) -> int: + """Return a conservative logical size for all tarball members.""" + total_bytes = 0 + with tarfile.open(tarball_path, "r:gz") as tar: + for member in tar: + if member.size < 0: + raise ValueError( + f"Registry tarball member has a negative size: {member.name}" + ) + total_bytes += member.size + return total_bytes + + +def _squashfs_listing_size(output: bytes) -> int: + """Sum file sizes from ``unsquashfs -lln`` output, failing closed.""" + total_bytes = 0 + for raw_line in output.decode(errors="strict").splitlines(): + line = raw_line.strip() + if not line: + continue + fields = line.split(maxsplit=4) + mode = fields[0] + if len(mode) != 10 or mode[0] not in "bcdlps-": + continue + if mode[0] not in "-l": + continue + if len(fields) < 5 or "/" not in fields[1] or not fields[2].isdigit(): + raise ValueError(f"Could not parse SquashFS listing line: {line}") + total_bytes += int(fields[2]) + return total_bytes + + def _directory_footprint(directory: Path) -> int: """Return the total file size of a cache directory. @@ -771,6 +890,9 @@ def __init__(self, cache_dir: Path): self._owner_binding_lock = threading.Lock() self._owner_loop: asyncio.AbstractEventLoop | None = None self._owner_thread_id: int | None = None + # Cold materializations and budget passes share this outer lock. It + # keeps byte reservations stable while downloads and extraction write. + self._admission_lock = asyncio.Lock() self._budget_lock = asyncio.Lock() # Guard the off-loop startup sweep independently from cache operations. self._swept: bool = False @@ -919,15 +1041,25 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat lease_acquired = True if cached_paths := self._locally_cached_path(ctx, artifact_uri): return cache_key, cached_paths - candidates = await self._artifact_candidates(ctx, artifact_uri) - if cached_paths := self._first_cached_path(candidates, ctx): - return cache_key, cached_paths - paths = await self._materialize_candidates(ctx, candidates) - self._touch_entry(cache_key) - # Enforce only after publication, and outside the per-key lock so - # eviction never nests key locks. A failed cold admission must not - # discard a usable warm entry before a replacement exists. + async with self._admission_lock: + async with lock: + if cached_paths := self._locally_cached_path(ctx, artifact_uri): + return cache_key, cached_paths + ctx = self._context_for( + cache_key, + admission=self._admission_for(cache_key), + ) + candidates = await self._artifact_candidates(ctx, artifact_uri) + if cached_paths := self._first_cached_path(candidates, ctx): + return cache_key, cached_paths + paths = await self._materialize_candidates(ctx, candidates) + self._touch_entry(cache_key) + + # Enforce entry count and final measured size after publication, + # outside the per-key lock. Peak-byte reservations may already + # have reclaimed idle LRU entries when that was required to keep + # staging and extraction within the hard byte cap. try: await self._enforce_cache_budget(protected_key=cache_key) except OSError as e: @@ -999,12 +1131,36 @@ def _runtime_for(self, cache_key: str) -> RegistryArtifactRuntimeState: self._runtime[cache_key] = runtime return runtime - def _context_for(self, cache_key: str) -> RegistryArtifactMaterializationContext: + def _context_for( + self, + cache_key: str, + *, + admission: RegistryArtifactAdmission | None = None, + ) -> RegistryArtifactMaterializationContext: """Return a materialization context for a registry artifact key.""" return RegistryArtifactMaterializationContext( cache_key=cache_key, staging_dir=self.staging_dir, paths=self._paths_for(cache_key), + admission=admission, + ) + + def _admission_for(self, cache_key: str) -> RegistryArtifactAdmission | None: + """Return byte-bound admission controls for one cold cache key.""" + max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES + if max_bytes <= 0: + return None + + async def ensure_capacity(additional_bytes: int) -> None: + await self._ensure_cache_capacity( + additional_bytes=additional_bytes, + protected_key=cache_key, + max_bytes=max_bytes, + ) + + return RegistryArtifactAdmission( + max_bytes=max_bytes, + ensure_capacity=ensure_capacity, ) def _base_pythonpath_dir(self) -> Path: @@ -1266,9 +1422,10 @@ async def _converge_cache_budget(self) -> None: async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bool: """Evict least-recently-used idle entries until the cache fits its budget. - The budget lock serializes the complete scan/select/evict pass. It is - always acquired before any candidate's per-key lock, and callers must - invoke enforcement without holding a per-key lock. + The admission lock excludes cold writers before the budget lock begins + a scan/select/evict pass. Callers invoke enforcement without holding a + per-key lock. Cold writers already hold the admission lock and use + ``_ensure_cache_capacity`` for their staged reservations instead. Args: protected_key: Newly materialized cache key. It is counted against @@ -1278,54 +1435,123 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo Returns: Whether the cache is within budget once eviction has finished. """ - async with self._budget_lock: - trash_clean, startup_clean = await asyncio.gather( - asyncio.to_thread(self._clear_work_dir, self.trash_dir), - asyncio.to_thread(self._retry_failed_startup_cleanup), + async with self._admission_lock: + async with self._budget_lock: + return await self._enforce_cache_budget_locked( + protected_key=protected_key + ) + + async def _enforce_cache_budget_locked( + self, + *, + protected_key: str | None, + ) -> bool: + """Enforce entry and byte limits while both cache-wide locks are held.""" + trash_clean, startup_clean = await asyncio.gather( + asyncio.to_thread(self._clear_work_dir, self.trash_dir), + asyncio.to_thread(self._retry_failed_startup_cleanup), + ) + cleanup_complete = trash_clean and startup_clean + if not cleanup_complete: + return False + + max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES + max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES + if max_entries <= 0 and max_bytes <= 0: + return True + + entries = await asyncio.to_thread(self._scan_cache_entries) + total_bytes = sum(entry.size_bytes for entry in entries.values()) + protected = set() if protected_key is None else {protected_key} + skipped: set[str] = set() + + while (max_entries > 0 and len(entries) > max_entries) or ( + max_bytes > 0 and total_bytes > max_bytes + ): + candidate = self._least_recently_used( + entries.values(), + excluded=skipped | protected, ) - cleanup_complete = trash_clean and startup_clean - if not cleanup_complete: + if candidate is None: + logger.warning( + "Registry artifact cache is over budget but every entry is in use", + cache_dir=str(self.cache_dir), + entries=len(entries), + max_entries=max_entries, + total_bytes=total_bytes, + max_bytes=max_bytes, + ) return False - max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES - max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES - if max_entries <= 0 and max_bytes <= 0: - return True + eviction = await self._evict_entry(candidate.cache_key) + if eviction.retired: + del entries[candidate.cache_key] + if not eviction.reclaimed: + return False + total_bytes -= candidate.size_bytes + else: + skipped.add(candidate.cache_key) + + return True + + async def _ensure_cache_capacity( + self, + *, + additional_bytes: int, + protected_key: str, + max_bytes: int, + ) -> None: + """Reserve peak bytes for a cold writer without exceeding the cap. + + The caller holds the admission lock and its key lock. Every normal + budget pass takes the admission lock first, so acquiring the budget + lock here cannot deadlock with eviction of the protected key. + """ + if additional_bytes < 0: + raise ValueError("additional_bytes must be non-negative") + async with self._budget_lock: + await asyncio.gather( + asyncio.to_thread(self._clear_work_dir, self.trash_dir), + asyncio.to_thread(self._retry_failed_startup_cleanup), + ) entries = await asyncio.to_thread(self._scan_cache_entries) - total_bytes = sum(entry.size_bytes for entry in entries.values()) - protected = set() if protected_key is None else {protected_key} - skipped: set[str] = set() + staging_bytes, trash_bytes = await asyncio.gather( + asyncio.to_thread(_directory_footprint, self.staging_dir), + asyncio.to_thread(_directory_footprint, self.trash_dir), + ) + total_bytes = ( + sum(entry.size_bytes for entry in entries.values()) + + staging_bytes + + trash_bytes + ) + skipped = {protected_key} - while (max_entries > 0 and len(entries) > max_entries) or ( - max_bytes > 0 and total_bytes > max_bytes - ): + while total_bytes + additional_bytes > max_bytes: candidate = self._least_recently_used( entries.values(), - excluded=skipped | protected, + excluded=skipped, ) if candidate is None: - logger.warning( - "Registry artifact cache is over budget but every entry is in use", - cache_dir=str(self.cache_dir), - entries=len(entries), - max_entries=max_entries, - total_bytes=total_bytes, + raise RegistryArtifactCacheCapacityError( + current_bytes=total_bytes, + additional_bytes=additional_bytes, max_bytes=max_bytes, ) - return False eviction = await self._evict_entry(candidate.cache_key) if eviction.retired: del entries[candidate.cache_key] if not eviction.reclaimed: - return False + raise RegistryArtifactCacheCapacityError( + current_bytes=total_bytes, + additional_bytes=additional_bytes, + max_bytes=max_bytes, + ) total_bytes -= candidate.size_bytes else: skipped.add(candidate.cache_key) - return True - def _least_recently_used( self, entries: Iterable[RegistryArtifactCacheEntry], diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index e9d160c3c1..79fb6e5e22 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -7,7 +7,7 @@ import os import threading import weakref -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass, field from pathlib import Path @@ -758,6 +758,7 @@ async def download_file_to_path( chunk_size: int = DEFAULT_DOWNLOAD_CHUNK_SIZE_BYTES, max_bytes: int | None = None, expected_sha256: str | None = None, + ensure_capacity: Callable[[int], Awaitable[None]] | None = None, ) -> int: """Stream an S3/MinIO object to a local file. @@ -771,6 +772,9 @@ async def download_file_to_path( chunk_size: Chunk size for streaming reads (default: 8MB). max_bytes: Optional guardrail; raise if the object exceeds this size. expected_sha256: Optional integrity check; raise if computed SHA-256 differs. + ensure_capacity: Optional callback invoked before the first disk write with + the maximum number of bytes the download may occupy. When the server + omits ContentLength, max_bytes is required to provide that bound. Returns: Total bytes written. @@ -796,15 +800,33 @@ async def download_file_to_path( f"ContentLength={content_length} exceeds max_bytes={max_bytes}" ) + download_limit = max_bytes + if ensure_capacity is not None: + reserved_bytes = content_length + if reserved_bytes is None: + if max_bytes is None: + raise ValueError( + "Cannot reserve disk capacity for a download without " + f"ContentLength or max_bytes: {bucket}/{key}" + ) + reserved_bytes = max_bytes + await ensure_capacity(reserved_bytes) + download_limit = ( + reserved_bytes + if download_limit is None + else min(download_limit, reserved_bytes) + ) + async with aiofiles.open(temp_path, "wb") as f: async for chunk in stream.iter_chunks(chunk_size=chunk_size): if not chunk: continue bytes_written += len(chunk) - if max_bytes is not None and bytes_written > max_bytes: + if download_limit is not None and bytes_written > download_limit: raise ValueError( f"Refusing to download {bucket}/{key} to disk: " - f"bytes_written={bytes_written} exceeds max_bytes={max_bytes}" + f"bytes_written={bytes_written} exceeds " + f"max_bytes={download_limit}" ) if hasher is not None: hasher.update(chunk) From b05e2d2bbb8f90ac0ef8cb971e3f359ab0950207 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 15:32:43 -0400 Subject: [PATCH 040/161] refactor(executor): split registry cache responsibilities --- tests/unit/test_executor_sandbox_nsjail.py | 2 +- tests/unit/test_registry_artifacts.py | 91 +- .../executor/registry_artifact_cache_state.py | 175 ++ .../registry_artifact_materialization.py | 729 +++++++ .../executor/registry_artifact_storage.py | 745 ++++++++ tracecat/executor/registry_artifacts.py | 1676 +---------------- 6 files changed, 1785 insertions(+), 1633 deletions(-) create mode 100644 tracecat/executor/registry_artifact_cache_state.py create mode 100644 tracecat/executor/registry_artifact_materialization.py create mode 100644 tracecat/executor/registry_artifact_storage.py diff --git a/tests/unit/test_executor_sandbox_nsjail.py b/tests/unit/test_executor_sandbox_nsjail.py index c451035e9e..36389a6e0f 100644 --- a/tests/unit/test_executor_sandbox_nsjail.py +++ b/tests/unit/test_executor_sandbox_nsjail.py @@ -573,7 +573,7 @@ async def download_artifact(self, ctx, output_path: Path) -> float: # path instead of attempting a loopback mount. patches.append( patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value=None, ) ) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 6f2f2e5d73..f0b8ea5e68 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -329,7 +329,7 @@ async def test_lease_uses_bundled_current_builtin( monkeypatch.setattr(tracecat_registry, "__version__", version) monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) monkeypatch.setattr( - "tracecat.executor.registry_artifacts.sysconfig.get_path", + "tracecat.executor.registry_artifact_materialization.sysconfig.get_path", lambda name: str(site_packages) if name == "purelib" else None, ) @@ -356,7 +356,7 @@ async def test_lease_exposes_editable_builtin_parent( monkeypatch.setattr(tracecat_registry, "__version__", version) monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) monkeypatch.setattr( - "tracecat.executor.registry_artifacts.sysconfig.get_path", + "tracecat.executor.registry_artifact_materialization.sysconfig.get_path", lambda name: str(site_packages) if name == "purelib" else None, ) @@ -612,7 +612,7 @@ def test_can_try_squashfs_does_not_require_mount_binary(self, temp_cache_dir): with ( patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value=None, ), patch( @@ -671,7 +671,7 @@ async def mock_mount(self, ctx, image_path): True, ), patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/mount", ), patch.object(SquashfsArtifact, "mount", mock_mount), @@ -714,7 +714,7 @@ async def test_mount_squashfs_uses_hardened_read_only_options( process.returncode = 0 with patch( - "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process, ) as create_subprocess_exec: @@ -750,7 +750,7 @@ async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): process = _BlockingSubprocess() with patch( - "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process, ): @@ -805,11 +805,11 @@ async def create_sleep_subprocess( with ( patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/usr/bin/unsquashfs", ), patch( - "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", side_effect=create_sleep_subprocess, ), ): @@ -905,7 +905,7 @@ async def mock_extract(self, ctx, image_path): True, ), patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/mount", ), patch.object(SquashfsArtifact, "mount", mock_mount), @@ -951,7 +951,8 @@ async def mock_extract(self, ctx, image_path): True, ), patch( - "tracecat.executor.registry_artifacts.shutil.which", return_value=None + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value=None, ), patch.object(SquashfsArtifact, "extract", mock_extract), ): @@ -996,7 +997,7 @@ async def mock_extract(self, ctx, image_path): True, ), patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/mount", ), patch.object(SquashfsArtifact, "mount", mock_mount), @@ -1343,7 +1344,7 @@ async def hold_lease(index: int) -> None: patch.object(Path, "is_mount", lambda path: path in harness.mounted), patch(SQUASHFS_ENABLED_CONFIG, True), patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/mount", ), patch.object(SquashfsArtifact, "mount", harness.mount), @@ -1405,7 +1406,7 @@ async def hold_lease(index: int) -> None: patch.object(Path, "is_mount", lambda path: path in harness.mounted), patch(SQUASHFS_ENABLED_CONFIG, True), patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/mount", ), patch.object(SquashfsArtifact, "mount", harness.mount), @@ -1684,7 +1685,7 @@ async def take_lease() -> list[Path]: with ( patch(SQUASHFS_ENABLED_CONFIG, False), patch( - "tracecat.executor.registry_artifacts._tarball_extracted_size", + "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", return_value=1, ), patch.object(TarballArtifact, "download", controlled_download), @@ -1732,7 +1733,7 @@ async def test_duplicate_uri_balances_each_acquisition_and_release( patch.object(Path, "is_mount", lambda path: path in harness.mounted), patch(SQUASHFS_ENABLED_CONFIG, True), patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/mount", ), patch.object(SquashfsArtifact, "mount", harness.mount), @@ -1797,11 +1798,11 @@ async def take_lease() -> None: patch.object(Path, "is_mount", lambda self: self in mounted), patch(SQUASHFS_ENABLED_CONFIG, True), patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/umount", ), patch( - "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", side_effect=mock_umount, ), patch.object(SquashfsArtifact, "mount", mock_mount), @@ -1836,7 +1837,7 @@ async def test_builtin_artifact_is_exempt_from_cache_accounting( monkeypatch.setattr(tracecat_registry, "__version__", version) monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) monkeypatch.setattr( - "tracecat.executor.registry_artifacts.sysconfig.get_path", + "tracecat.executor.registry_artifact_materialization.sysconfig.get_path", lambda name: str(site_packages) if name == "purelib" else None, ) @@ -1864,10 +1865,12 @@ def test_delete_cache_path_reports_directory_failure(self, temp_cache_dir): with ( patch( - "tracecat.executor.registry_artifacts.shutil.rmtree", + "tracecat.executor.registry_artifact_storage.shutil.rmtree", side_effect=OSError("permission denied"), ), - patch("tracecat.executor.registry_artifacts.logger.warning") as warning, + patch( + "tracecat.executor.registry_artifact_storage.logger.warning" + ) as warning, ): deleted = _delete_cache_path(entry_dir) @@ -2079,7 +2082,7 @@ async def mock_extract(self, tarball_path, target_dir): patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 6000), patch( - "tracecat.executor.registry_artifacts._tarball_extracted_size", + "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", return_value=4096, ), patch.object(TarballArtifact, "download", mock_download), @@ -2114,12 +2117,14 @@ async def mock_extract(self, tarball_path, target_dir): patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0), patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", return_value=False, ), patch.object(TarballArtifact, "download", mock_download), patch.object(TarballArtifact, "extract", mock_extract), - patch("tracecat.executor.registry_artifacts.logger.warning") as warning, + patch( + "tracecat.executor.registry_artifact_storage.logger.warning" + ) as warning, ): registry_paths = await _materialize(cache, cache_key, artifact_uri) @@ -2170,7 +2175,7 @@ def fail_once(path: Path) -> bool: patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 16), patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", side_effect=fail_once, ), ): @@ -2206,7 +2211,7 @@ async def mock_extract(self, tarball_path, target_dir): patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0), patch( - "tracecat.executor.registry_artifacts._move_entry_to_trash", + "tracecat.executor.registry_artifact_storage._move_entry_to_trash", side_effect=OSError("rename failed"), ), patch.object(TarballArtifact, "download", mock_download), @@ -2242,7 +2247,7 @@ async def mock_extract(self, tarball_path, target_dir): side_effect=PermissionError("denied"), ), patch( - "tracecat.executor.registry_artifacts._tarball_extracted_size", + "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", return_value=9, ), patch.object(TarballArtifact, "download", mock_download), @@ -2525,7 +2530,7 @@ async def mock_umount(*args, **kwargs): with ( patch.object(Path, "is_mount", lambda self: self in mounted), patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/umount", ), patch.object( @@ -2688,11 +2693,11 @@ async def mock_umount(*args, **kwargs): with ( patch.object(Path, "is_mount", lambda self: self in mounted), patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/umount", ), patch( - "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", side_effect=mock_umount, ) as create_subprocess_exec, ): @@ -2739,11 +2744,11 @@ async def mock_umount(*args, **kwargs): with ( patch.object(Path, "is_mount", lambda self: self in mounted), patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/umount", ), patch( - "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", side_effect=mock_umount, ), ): @@ -2795,11 +2800,11 @@ async def mock_extract(self, tarball_path, target_dir): with ( patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", side_effect=blocked_delete, ), patch( - "tracecat.executor.registry_artifacts._tarball_extracted_size", + "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", return_value=9, ), patch.object(TarballArtifact, "download", mock_download), @@ -2843,7 +2848,7 @@ async def test_doomed_eviction_names_are_startup_scratch(self, temp_cache_dir): paths.tarball_target_dir.mkdir() with patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", return_value=True, ) as delete_cache_path: assert await cache._evict_entry(cache_key) == RegistryArtifactEviction( @@ -2879,11 +2884,11 @@ async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): with ( patch.object(Path, "is_mount", lambda self: self in mounted), patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/umount", ), patch( - "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process, ), @@ -3121,7 +3126,7 @@ async def test_active_entry_scan_errors_propagate(self, temp_cache_dir): with ( patch( - "tracecat.executor.registry_artifacts.os.scandir", + "tracecat.executor.registry_artifact_storage.os.scandir", side_effect=PermissionError("denied"), ), pytest.raises(PermissionError, match="denied"), @@ -3157,7 +3162,7 @@ def fail_orphan(path: Path) -> bool: patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0), patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", side_effect=fail_orphan, ), ): @@ -3186,7 +3191,7 @@ def fail_once(path: Path) -> bool: return real_delete(path) with patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", side_effect=fail_once, ): await cache.ensure_swept() @@ -3222,7 +3227,7 @@ async def test_failed_startup_retirement_stays_dirty_and_retries( patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0), patch( - "tracecat.executor.registry_artifacts._move_entry_to_trash", + "tracecat.executor.registry_artifact_storage._move_entry_to_trash", side_effect=OSError("rename failed"), ), ): @@ -3280,7 +3285,7 @@ def fail_once(path: Path) -> bool: patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 16), patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", side_effect=fail_once, ), ): @@ -3329,7 +3334,7 @@ async def test_loop_device_exhaustion_isolated_sticky_extraction_fallback( patch.object(Path, "is_mount", lambda path: path in harness.mounted), patch(SQUASHFS_ENABLED_CONFIG, True), patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/mount", ), patch.object(SquashfsArtifact, "mount", harness.mount), @@ -3402,7 +3407,7 @@ async def mock_extract(self, ctx, image_path): with ( patch(SQUASHFS_ENABLED_CONFIG, True), patch( - "tracecat.executor.registry_artifacts.shutil.which", + "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/mount", ), patch.object(SquashfsArtifact, "mount", mock_mount), diff --git a/tracecat/executor/registry_artifact_cache_state.py b/tracecat/executor/registry_artifact_cache_state.py new file mode 100644 index 0000000000..5f921bc1c5 --- /dev/null +++ b/tracecat/executor/registry_artifact_cache_state.py @@ -0,0 +1,175 @@ +"""Process-local state and lease bookkeeping for registry artifact caches.""" + +from __future__ import annotations + +import asyncio +import os +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path + +from tracecat.executor.registry_artifact_materialization import ( + RegistryArtifactAdmission, + RegistryArtifactMaterializationContext, + RegistryArtifactPaths, +) +from tracecat.logger import logger + +BASE_PYTHONPATH_DIR_NAME = "base" +"""Cache subdirectory used as the PYTHONPATH entry when no artifact is requested.""" + +CACHE_ENTRIES_DIR_NAME = "entries" +"""Directory containing one atomic subdirectory per cache key.""" + +CACHE_STAGING_DIR_NAME = "staging" +"""Directory containing in-progress materialization scratch.""" + +CACHE_TRASH_DIR_NAME = "trash" +"""Directory containing atomically retired entries pending physical deletion.""" + + +class RegistryArtifactCacheLoopError(RuntimeError): + """A registry artifact cache was used outside its owning event loop.""" + + +@dataclass(slots=True) +class RegistryArtifactRuntimeState: + """Process-local synchronization and lease state for one cache key.""" + + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + refcount: int = 0 + last_used: float = 0.0 + + +class _RegistryArtifactCacheState: + """Owns cache paths, event-loop affinity, and per-key lease state.""" + + def __init__(self, cache_dir: Path): + self.cache_dir = cache_dir + self.entries_dir = cache_dir / CACHE_ENTRIES_DIR_NAME + self.staging_dir = cache_dir / CACHE_STAGING_DIR_NAME + self.trash_dir = cache_dir / CACHE_TRASH_DIR_NAME + # Runtime states live for the process lifetime so every operation for a + # key always serializes on the same lock. + self._runtime: dict[str, RegistryArtifactRuntimeState] = {} + # The cache contains asyncio locks, tasks, and multi-step lease state. + # Bind the public API to one loop/thread so a future synchronous + # Temporal activity fails immediately instead of corrupting that state + # through a thread-local event loop. + self._owner_binding_lock = threading.Lock() + self._owner_loop: asyncio.AbstractEventLoop | None = None + self._owner_thread_id: int | None = None + # Cold materializations and budget passes share this outer lock. It + # keeps byte reservations stable while downloads and extraction write. + self._admission_lock = asyncio.Lock() + self._budget_lock = asyncio.Lock() + # Guard the off-loop startup sweep independently from cache operations. + self._swept: bool = False + self._sweep_task: asyncio.Task[None] | None = None + self._sweep_lock = asyncio.Lock() + # Startup is the only time the whole staging directory is swept. Exact + # paths that could not be removed are safe to retry later. + self._failed_startup_cleanup: set[Path] = set() + # Whether the on-disk cache may exceed its budget. Set when a new entry + # is materialized and cleared once enforcement measures a cache that + # fits, so steady-state cache hits never pay for a disk scan. + self._budget_dirty = True + + def _assert_owner_loop(self) -> None: + """Bind to the current loop or reject use from another loop/thread.""" + current_loop = asyncio.get_running_loop() + current_thread_id = threading.get_ident() + with self._owner_binding_lock: + if self._owner_loop is None: + self._owner_loop = current_loop + self._owner_thread_id = current_thread_id + return + if ( + self._owner_loop is current_loop + and self._owner_thread_id == current_thread_id + ): + return + + raise RegistryArtifactCacheLoopError( + "RegistryArtifactCache is bound to one event loop and thread " + f"(owner_loop={id(self._owner_loop)}, " + f"owner_thread={self._owner_thread_id}, " + f"current_loop={id(current_loop)}, " + f"current_thread={current_thread_id})" + ) + + def _runtime_for(self, cache_key: str) -> RegistryArtifactRuntimeState: + """Return the process-local state for one cache key.""" + if runtime := self._runtime.get(cache_key): + return runtime + runtime = RegistryArtifactRuntimeState() + self._runtime[cache_key] = runtime + return runtime + + def _context_for( + self, + cache_key: str, + *, + admission: RegistryArtifactAdmission | None = None, + ) -> RegistryArtifactMaterializationContext: + """Return a materialization context for a registry artifact key.""" + return RegistryArtifactMaterializationContext( + cache_key=cache_key, + staging_dir=self.staging_dir, + paths=self._paths_for(cache_key), + admission=admission, + ) + + def _base_pythonpath_dir(self) -> Path: + """Return the base PYTHONPATH directory used when no artifact is requested.""" + base_dir = self.cache_dir / BASE_PYTHONPATH_DIR_NAME + base_dir.mkdir(parents=True, exist_ok=True) + return base_dir + + def _acquire_lease(self, cache_key: str) -> None: + """Pin a cache entry against eviction and mark it as recently used. + + Callers must hold the per-key lock so the increment is ordered against + in-flight eviction of the same key. + """ + runtime = self._runtime_for(cache_key) + runtime.refcount += 1 + runtime.last_used = time.time() + self._touch_entry(cache_key) + + def _release_lease(self, cache_key: str) -> bool: + """Release one pin and return whether the entry became idle.""" + runtime = self._runtime.get(cache_key) + if runtime is None or runtime.refcount == 0: + return False + runtime.refcount -= 1 + runtime.last_used = time.time() + return runtime.refcount == 0 + + def _refcount(self, cache_key: str) -> int: + """Return the number of live leases on a cache entry.""" + runtime = self._runtime.get(cache_key) + return 0 if runtime is None else runtime.refcount + + def _touch_entry(self, cache_key: str) -> None: + """Best-effort refresh of the entry-root mtime for restart-safe LRU.""" + entry_dir = self._paths_for(cache_key).entry_dir + try: + os.utime(entry_dir) + except OSError: + logger.debug( + "Could not refresh registry artifact entry mtime", + cache_key=cache_key, + ) + + def _paths_for(self, cache_key: str) -> RegistryArtifactPaths: + """Return local cache paths for a registry artifact key.""" + entry_dir = self.entries_dir / cache_key + return RegistryArtifactPaths( + entry_dir=entry_dir, + squashfs_image_path=entry_dir / "image.squashfs", + squashfs_mount_dir=entry_dir / "mount", + squashfs_extract_dir=entry_dir / "extracted", + tarball_target_dir=entry_dir / "tarball", + ) diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py new file mode 100644 index 0000000000..f2f1f68bf1 --- /dev/null +++ b/tracecat/executor/registry_artifact_materialization.py @@ -0,0 +1,729 @@ +"""Registry artifact formats and local materialization primitives.""" + +from __future__ import annotations + +import asyncio +import contextlib +import hashlib +import os +import shutil +import sysconfig +import tarfile +import time +from abc import ABC, abstractmethod +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from enum import StrEnum +from pathlib import Path + +import httpx +import tracecat_registry + +from tracecat import config +from tracecat.logger import logger +from tracecat.registry.artifact_keys import parse_s3_uri +from tracecat.registry.constants import DEFAULT_REGISTRY_ORIGIN +from tracecat.storage import blob + +__all__ = [ + "_is_cache_entry_uri", + "_squashfs_sidecar_uri", + "_tarball_uri_for_squashfs", +] + + +class RegistryArtifactFormat(StrEnum): + """Executor-supported registry artifact encodings.""" + + BUILTIN = "builtin" + SQUASHFS = "squashfs" + TAR_GZ = "tar.gz" + + +SQUASHFS_MOUNT_OPTIONS = "loop,ro,nodev,nosuid" +"""Mount options for executor-managed SquashFS registry artifacts. + +The image must stay read-only and should not expose device nodes or setuid bits +from registry package contents. Avoid noexec because Python packages may include +native extension modules that need to be loaded from the mounted artifact. +""" + +BUNDLED_BUILTIN_REGISTRY_URI_PREFIX = f"tracecat-builtin://{DEFAULT_REGISTRY_ORIGIN}/" +"""Pseudo-URI for the builtin registry already installed in the executor image.""" + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactPaths: + """Executor-local cache paths for one registry artifact key.""" + + entry_dir: Path + squashfs_image_path: Path + squashfs_mount_dir: Path + squashfs_extract_dir: Path + tarball_target_dir: Path + + +class SquashfsMountCommandError(RuntimeError): + """The ``mount`` command itself failed for a SquashFS registry artifact. + + Only this error drives SquashFS mount policy. Download, mkdir, and other + preparation failures must not be mistaken for a missing mount capability or + for loop-device exhaustion. + """ + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactAdmission: + """Byte-bound admission hook shared by one cold materialization.""" + + max_bytes: int + ensure_capacity: Callable[[int], Awaitable[None]] + + +@dataclass(slots=True) +class RegistryArtifactMaterializationContext: + """Shared runtime state for artifact materialization.""" + + cache_key: str + staging_dir: Path + paths: RegistryArtifactPaths + admission: RegistryArtifactAdmission | None = None + + def can_mount_squashfs(self) -> bool: + return config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED and ( + shutil.which("mount") is not None + ) + + +@dataclass(frozen=True, slots=True) +class RegistryArtifact(ABC): + """An executor-local materializable registry artifact.""" + + uri: str + cache_key: str + + @property + @abstractmethod + def format(self) -> RegistryArtifactFormat: + """Artifact format used for logging and dispatch.""" + + @abstractmethod + def cached_path( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path] | None: + """Return already-materialized import paths for this artifact, if present.""" + + @abstractmethod + async def materialize( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path]: + """Return importable Python paths, materializing the artifact if needed.""" + + def _temp_path( + self, + ctx: RegistryArtifactMaterializationContext, + suffix: str, + ) -> Path: + unique_id = id(asyncio.current_task()) + ctx.staging_dir.mkdir(parents=True, exist_ok=True) + return ctx.staging_dir / f"{self.cache_key}.{os.getpid()}.{unique_id}{suffix}" + + +@dataclass(frozen=True, slots=True) +class BuiltinArtifact(RegistryArtifact): + """Current builtin registry package already installed in the executor image.""" + + version: str + + @property + def format(self) -> RegistryArtifactFormat: + return RegistryArtifactFormat.BUILTIN + + def cached_path( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path] | None: + return None + + async def materialize( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path]: + del ctx + import_paths = _bundled_builtin_registry_import_paths(self.version) + logger.info( + "Using bundled builtin registry environment", + registry_version=self.version, + paths=[str(p) for p in import_paths], + ) + return import_paths + + +@dataclass(frozen=True, slots=True) +class SquashfsArtifact(RegistryArtifact): + """SquashFS registry environment image.""" + + @property + def format(self) -> RegistryArtifactFormat: + return RegistryArtifactFormat.SQUASHFS + + def cached_path( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path] | None: + if ctx.paths.squashfs_mount_dir.is_mount(): + logger.debug( + "Using cached SquashFS registry mount", + cache_key=ctx.cache_key, + ) + return [ctx.paths.squashfs_mount_dir] + if ctx.paths.squashfs_extract_dir.exists(): + logger.debug( + "Using cached SquashFS registry extraction", + cache_key=ctx.cache_key, + ) + return [ctx.paths.squashfs_extract_dir] + return None + + async def materialize( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path]: + image_path = ctx.paths.squashfs_image_path + if ctx.can_mount_squashfs(): + try: + return [await self.mount(ctx, image_path)] + except SquashfsMountCommandError as e: + logger.warning( + "Failed to mount SquashFS registry artifact, trying extraction", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + error=str(e), + ) + + return [await self.extract(ctx, image_path)] + + async def download( + self, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> float: + """Ensure the SquashFS image exists locally and return download time.""" + if image_path.exists(): + return 0.0 + + image_path.parent.mkdir(parents=True, exist_ok=True) + temp_image = self._temp_path(ctx, ".squashfs") + try: + download_start = time.monotonic() + await _download_s3_artifact( + self.uri, + temp_image, + admission=ctx.admission, + ) + try: + temp_image.rename(image_path) + except OSError: + if not image_path.exists(): + raise + return (time.monotonic() - download_start) * 1000 + finally: + temp_image.unlink(missing_ok=True) + + async def mount( + self, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> Path: + """Download the image if needed and mount it read-only. + + Args: + ctx: Materialization context for the artifact being mounted. + image_path: Local path of the SquashFS image. + + Returns: + The mount directory. + + Raises: + SquashfsMountCommandError: The ``mount`` command failed. + Exception: The image could not be downloaded or prepared. + """ + target_dir = ctx.paths.squashfs_mount_dir + if target_dir.is_mount(): + return target_dir + + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + target_dir.mkdir(parents=True, exist_ok=True) + + logger.info( + "Materializing SquashFS registry artifact", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + ) + start_time = time.monotonic() + download_elapsed = await self.download(ctx, image_path) + + mount_start = time.monotonic() + await self._mount_image(image_path, target_dir) + mount_elapsed = (time.monotonic() - mount_start) * 1000 + total_elapsed = (time.monotonic() - start_time) * 1000 + + logger.info( + "SquashFS registry artifact mounted", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + download_ms=f"{download_elapsed:.1f}", + mount_ms=f"{mount_elapsed:.1f}", + total_ms=f"{total_elapsed:.1f}", + ) + return target_dir + + async def extract( + self, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> Path: + target_dir = ctx.paths.squashfs_extract_dir + if target_dir.exists(): + return target_dir + + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + + logger.info( + "Extracting SquashFS registry artifact", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + ) + start_time = time.monotonic() + download_elapsed = await self.download(ctx, image_path) + + temp_dir = self._temp_path(ctx, ".unsquashfs") + try: + if ctx.admission is not None: + extracted_size = await self._squashfs_extracted_size(image_path) + await ctx.admission.ensure_capacity(extracted_size) + extract_start = time.monotonic() + temp_dir.mkdir(parents=True, exist_ok=True) + await self._extract_image(image_path, temp_dir) + extract_elapsed = (time.monotonic() - extract_start) * 1000 + + try: + temp_dir.rename(target_dir) + total_elapsed = (time.monotonic() - start_time) * 1000 + logger.info( + "SquashFS registry artifact extracted", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + download_ms=f"{download_elapsed:.1f}", + extract_ms=f"{extract_elapsed:.1f}", + total_ms=f"{total_elapsed:.1f}", + ) + except OSError: + if target_dir.exists(): + logger.debug( + "SquashFS already extracted by another process", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + ) + else: + raise + finally: + if temp_dir.exists(): + shutil.rmtree(temp_dir, ignore_errors=True) + + return target_dir + + async def _mount_image(self, image_path: Path, target_dir: Path) -> None: + """Mount a SquashFS image read-only at target_dir. + + Cancellation kills and reaps the mount subprocess before propagating, + so the caller's per-key lock covers the complete mount lifecycle. The + target therefore remains an unmounted cache miss if mount never took + effect, or is already mounted and reusable by the next admission. + + Args: + image_path: Local path of the SquashFS image. + target_dir: Existing directory to mount the image at. + + Raises: + SquashfsMountCommandError: The ``mount`` command failed. + """ + if target_dir.is_mount(): + return + + proc = await asyncio.create_subprocess_exec( + "mount", + "-t", + "squashfs", + "-o", + SQUASHFS_MOUNT_OPTIONS, + str(image_path), + str(target_dir), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + proc.kill() + await proc.wait() + raise + + if proc.returncode == 0 or target_dir.is_mount(): + return + + output = (stderr or stdout).decode(errors="replace").strip() + raise SquashfsMountCommandError(output or "mount command failed") + + async def _extract_image(self, image_path: Path, target_dir: Path) -> None: + """Extract a SquashFS image to target_dir using unsquashfs. + + Cancellation kills and reaps the extractor before the caller removes + its scratch directory. Otherwise a live extractor could recreate + scratch after cleanup, outside startup-sweep discovery. + """ + unsquashfs = shutil.which("unsquashfs") + if unsquashfs is None: + raise RuntimeError("unsquashfs command is not installed") + + proc = await asyncio.create_subprocess_exec( + unsquashfs, + "-f", + "-d", + str(target_dir), + str(image_path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + proc.kill() + await proc.wait() + raise + + if proc.returncode == 0: + return + + output = (stderr or stdout).decode(errors="replace").strip() + raise RuntimeError(output or "unsquashfs command failed") + + async def _squashfs_extracted_size(self, image_path: Path) -> int: + """Return a conservative logical size for an extracted SquashFS image.""" + unsquashfs = shutil.which("unsquashfs") + if unsquashfs is None: + raise RuntimeError("unsquashfs command is not installed") + + proc = await asyncio.create_subprocess_exec( + unsquashfs, + "-lln", + str(image_path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + proc.kill() + await proc.wait() + raise + + if proc.returncode != 0: + output = (stderr or stdout).decode(errors="replace").strip() + raise RuntimeError(output or "unsquashfs listing failed") + return _squashfs_listing_size(stdout) + + +@dataclass(frozen=True, slots=True) +class TarballArtifact(RegistryArtifact): + """Legacy gzip tarball registry environment.""" + + @property + def format(self) -> RegistryArtifactFormat: + return RegistryArtifactFormat.TAR_GZ + + def cached_path( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path] | None: + if ctx.paths.tarball_target_dir.exists(): + logger.debug( + "Using cached tarball extraction", + cache_key=ctx.cache_key, + ) + return [ctx.paths.tarball_target_dir] + return None + + async def materialize( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path]: + target_dir = ctx.paths.tarball_target_dir + logger.info( + "Materializing tarball registry artifact", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + ) + start_time = time.monotonic() + + temp_tarball = self._temp_path(ctx, ".tar.gz") + temp_dir = self._temp_path(ctx, ".tmp") + + try: + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + + download_start = time.monotonic() + await self.download(ctx, temp_tarball) + download_elapsed = (time.monotonic() - download_start) * 1000 + + if ctx.admission is not None: + extracted_size = await asyncio.to_thread( + _tarball_extracted_size, + temp_tarball, + ) + await ctx.admission.ensure_capacity(extracted_size) + + extract_start = time.monotonic() + temp_dir.mkdir(parents=True, exist_ok=True) + await self.extract(temp_tarball, temp_dir) + extract_elapsed = (time.monotonic() - extract_start) * 1000 + + try: + temp_dir.rename(target_dir) + total_elapsed = (time.monotonic() - start_time) * 1000 + logger.info( + "Tarball extracted and cached", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + download_ms=f"{download_elapsed:.1f}", + extract_ms=f"{extract_elapsed:.1f}", + total_ms=f"{total_elapsed:.1f}", + ) + except OSError: + if target_dir.exists(): + logger.debug( + "Tarball already extracted by another process", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + ) + else: + raise + finally: + if temp_dir.exists(): + shutil.rmtree(temp_dir, ignore_errors=True) + if temp_tarball.exists(): + temp_tarball.unlink(missing_ok=True) + + return [target_dir] + + async def download( + self, + ctx: RegistryArtifactMaterializationContext, + output_path: Path, + ) -> None: + await _download_s3_artifact( + self.uri, + output_path, + admission=ctx.admission, + ) + + async def extract(self, tarball_path: Path, target_dir: Path) -> None: + """Extract a supported registry tarball to target directory. + + A cancelled caller rejoins the non-interruptible extraction thread + before propagating cancellation so scratch cleanup cannot race a live + writer and leave undiscoverable ephemeral storage behind. + """ + + def _do_extract() -> None: + if tarball_path.name.endswith(".tar.gz"): + with tarfile.open(tarball_path, "r:gz") as tar: + tar.extractall(path=target_dir, filter="data") + return + + raise ValueError(f"Unsupported tarball format: {tarball_path}") + + extraction = asyncio.ensure_future(asyncio.to_thread(_do_extract)) + try: + await asyncio.shield(extraction) + except asyncio.CancelledError: + # A thread cannot be killed. Rejoin it before materialize removes + # scratch. Each cancellation can interrupt shield without stopping + # the thread, so keep waiting until extraction reaches a terminal + # state before propagating the original cancellation. + while not extraction.done(): + try: + await asyncio.shield(extraction) + except asyncio.CancelledError: + continue + except Exception: + break + if not extraction.cancelled(): + with contextlib.suppress(Exception): + extraction.result() + raise + + logger.debug( + "Tarball extracted", + target=str(target_dir), + artifact_format=_artifact_format(str(tarball_path)).value, + ) + + +async def _download_s3_artifact( + artifact_uri: str, + output_path: Path, + *, + admission: RegistryArtifactAdmission | None = None, +) -> None: + """Download an S3 registry artifact to a local path.""" + bucket, key = parse_s3_uri(artifact_uri) + try: + if admission is None: + await blob.download_file_to_path( + key=key, + bucket=bucket, + output_path=output_path, + ) + else: + await blob.download_file_to_path( + key=key, + bucket=bucket, + output_path=output_path, + max_bytes=admission.max_bytes, + ensure_capacity=admission.ensure_capacity, + ) + except FileNotFoundError as e: + request = httpx.Request("GET", artifact_uri) + response = httpx.Response(status_code=404, request=request) + raise httpx.HTTPStatusError( + f"Registry artifact not found: {artifact_uri}", + request=request, + response=response, + ) from e + + +def compute_registry_artifact_cache_key(artifact_uri: str) -> str: + """Compute the local cache key for a registry artifact URI.""" + if not artifact_uri: + return "base" + # S3 keys are case-sensitive, so preserve URI case when hashing. + content = artifact_uri.strip() + return hashlib.sha256(content.encode()).hexdigest()[:16] + + +def bundled_builtin_registry_uri(version: str) -> str: + """Return the pseudo-URI for the installed builtin registry package.""" + return f"{BUNDLED_BUILTIN_REGISTRY_URI_PREFIX}{version}" + + +def _bundled_builtin_registry_version(artifact_uri: str) -> str | None: + """Return the builtin registry version encoded in a bundled pseudo-URI.""" + if not artifact_uri.startswith(BUNDLED_BUILTIN_REGISTRY_URI_PREFIX): + return None + version = artifact_uri.removeprefix(BUNDLED_BUILTIN_REGISTRY_URI_PREFIX) + return version or None + + +def _bundled_builtin_registry_import_paths(version: str) -> list[Path]: + """Return import paths for the current builtin registry and its dependencies. + + Dependencies always live in the executor's site-packages. For editable + installs the parent of ``package_dir`` (the package wrapper, e.g. + ``packages/tracecat-registry/``) is exposed first so its ``tracecat_registry/`` + shadows any stale copy in site-packages. + """ + installed_version = tracecat_registry.__version__ + if version != installed_version: + raise RuntimeError( + "Bundled builtin registry version does not match installed version: " + f"requested={version!r}, installed={installed_version!r}" + ) + + package_file = tracecat_registry.__file__ + if package_file is None: + raise RuntimeError("Installed tracecat_registry package has no __file__") + + site_packages_path = sysconfig.get_path("purelib") + if site_packages_path is None: + raise RuntimeError("Could not resolve installed Python site-packages path") + + site_packages = Path(site_packages_path).resolve() + if not site_packages.exists(): + raise RuntimeError( + f"Installed Python site-packages path does not exist: {site_packages}" + ) + + package_dir = Path(package_file).resolve().parent + if package_dir.is_relative_to(site_packages): + return [site_packages] + + return [package_dir.parent, site_packages] + + +def _squashfs_sidecar_uri(tarball_uri: str) -> str | None: + """Return the sibling SquashFS URI for registry site-packages tarballs.""" + if not tarball_uri.endswith("site-packages.tar.gz"): + return None + return tarball_uri.removesuffix(".tar.gz") + ".squashfs" + + +def _tarball_uri_for_squashfs(squashfs_uri: str) -> str | None: + """Return the sibling gzip tarball URI for registry SquashFS artifacts.""" + if not squashfs_uri.endswith("site-packages.squashfs"): + return None + return squashfs_uri.removesuffix(".squashfs") + ".tar.gz" + + +def _artifact_format(artifact_uri: str) -> RegistryArtifactFormat: + """Return the materialization format for an artifact URI.""" + if artifact_uri.endswith(".squashfs"): + return RegistryArtifactFormat.SQUASHFS + return RegistryArtifactFormat.TAR_GZ + + +def _is_cache_entry_uri(artifact_uri: str) -> bool: + """Return whether an artifact URI materializes into an evictable cache entry. + + The bundled builtin registry is served from the executor image and never + writes into the cache directory, so it is exempt from eviction accounting. + """ + return _bundled_builtin_registry_version(artifact_uri) is None + + +def _tarball_extracted_size(tarball_path: Path) -> int: + """Return a conservative logical size for all tarball members.""" + total_bytes = 0 + with tarfile.open(tarball_path, "r:gz") as tar: + for member in tar: + if member.size < 0: + raise ValueError( + f"Registry tarball member has a negative size: {member.name}" + ) + total_bytes += member.size + return total_bytes + + +def _squashfs_listing_size(output: bytes) -> int: + """Sum file sizes from ``unsquashfs -lln`` output, failing closed.""" + total_bytes = 0 + for raw_line in output.decode(errors="strict").splitlines(): + line = raw_line.strip() + if not line: + continue + fields = line.split(maxsplit=4) + mode = fields[0] + if len(mode) != 10 or mode[0] not in "bcdlps-": + continue + if mode[0] not in "-l": + continue + if len(fields) < 5 or "/" not in fields[1] or not fields[2].isdigit(): + raise ValueError(f"Could not parse SquashFS listing line: {line}") + total_bytes += int(fields[2]) + return total_bytes diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py new file mode 100644 index 0000000000..35c930fab4 --- /dev/null +++ b/tracecat/executor/registry_artifact_storage.py @@ -0,0 +1,745 @@ +"""Disk-budget enforcement and cleanup for registry artifact caches.""" + +from __future__ import annotations + +import asyncio +import contextlib +import os +import shutil +import time +from collections.abc import Iterable +from dataclasses import dataclass +from pathlib import Path + +from tracecat import config +from tracecat.executor.registry_artifact_cache_state import ( + _RegistryArtifactCacheState, +) +from tracecat.executor.registry_artifact_materialization import ( + RegistryArtifactAdmission, +) +from tracecat.logger import logger + + +class RegistryArtifactCacheCapacityError(RuntimeError): + """A cold artifact cannot fit within the configured cache byte budget.""" + + def __init__( + self, + *, + current_bytes: int, + additional_bytes: int, + max_bytes: int, + ) -> None: + super().__init__( + "Registry artifact admission exceeds the cache byte budget: " + f"current_bytes={current_bytes}, additional_bytes={additional_bytes}, " + f"max_bytes={max_bytes}" + ) + self.current_bytes = current_bytes + self.additional_bytes = additional_bytes + self.max_bytes = max_bytes + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactEviction: + """Outcome of atomically retiring and physically deleting one entry.""" + + retired: bool + reclaimed: bool + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactCacheEntry: + """Measured on-disk footprint and recency for one registry artifact key.""" + + cache_key: str + size_bytes: int + last_used: float + + +def _directory_footprint(directory: Path) -> int: + """Return the total file size of a cache directory. + + Args: + directory: Cache directory to measure. + + Returns: + Total byte size of contained files, or zero when the directory is + missing. + """ + + def raise_walk_error(error: OSError) -> None: + raise error + + total_bytes = 0 + try: + walker = os.walk(directory, onerror=raise_walk_error) + for root, _dirs, files in walker: + for file_name in files: + try: + total_bytes += os.lstat(os.path.join(root, file_name)).st_size + except FileNotFoundError: + continue + except FileNotFoundError: + return 0 + return total_bytes + + +def _delete_cache_path(path: Path) -> bool: + """Best-effort delete one cache path while reporting filesystem failures.""" + try: + if path.is_dir(): + shutil.rmtree(path) + else: + path.unlink(missing_ok=True) + except OSError as e: + logger.warning( + "Failed to delete registry artifact cache path", + path=str(path), + error=str(e), + ) + return False + return True + + +async def _delete_cache_path_off_loop(path: Path) -> bool: + """Delete one path without abandoning its worker thread on cancellation.""" + deletion = asyncio.ensure_future(asyncio.to_thread(_delete_cache_path, path)) + try: + return await asyncio.shield(deletion) + except asyncio.CancelledError: + # A worker thread cannot be killed. Rejoin it so no live deletion can + # race a later trash-directory scan. Repeated cancellation can interrupt + # shield without stopping the thread, so keep waiting for termination. + while not deletion.done(): + try: + await asyncio.shield(deletion) + except asyncio.CancelledError: + continue + except Exception: + break + if not deletion.cancelled(): + with contextlib.suppress(Exception): + deletion.result() + raise + + +def _unique_work_path(root: Path, cache_key: str) -> Path: + """Return a unique path beneath a cache work directory.""" + root.mkdir(parents=True, exist_ok=True) + unique_id = time.time_ns() + while True: + path = root / f"{cache_key}.{os.getpid()}.{unique_id}" + if not path.exists(): + return path + unique_id += 1 + + +def _move_entry_to_trash(entry_dir: Path, trash_dir: Path, cache_key: str) -> Path: + """Atomically retire one cache entry and return its trash path.""" + trash_path = _unique_work_path(trash_dir, cache_key) + entry_dir.rename(trash_path) + return trash_path + + +class _RegistryArtifactCacheStorage(_RegistryArtifactCacheState): + """Adds startup recovery, eviction, and byte-budget enforcement.""" + + async def ensure_swept(self) -> None: + """Run the startup sweep exactly once successfully, off the event loop. + + Idempotent and cancellation-safe under concurrency: every caller joins + one stored sweep task, and cancelling a waiter never abandons or + restarts its live sweep. The lock is deliberately held while awaiting + that shared task so queued callers observe its result before proceeding. + Failures clear the task so the next caller retries. The sweep runs + before the first lease or materialization, so it never observes + in-flight cache entries. + """ + self._assert_owner_loop() + if self._swept: + return + async with self._sweep_lock: + if self._swept: + return + sweep_task = self._sweep_task + if sweep_task is None or ( + sweep_task.done() + and (sweep_task.cancelled() or sweep_task.exception() is not None) + ): + sweep_task = asyncio.ensure_future( + asyncio.to_thread(self._sweep_startup_state) + ) + self._sweep_task = sweep_task + try: + await asyncio.shield(sweep_task) + except asyncio.CancelledError: + if sweep_task.cancelled() and self._sweep_task is sweep_task: + self._sweep_task = None + raise + except Exception: + if self._sweep_task is sweep_task: + self._sweep_task = None + raise + self._swept = True + + def _admission_for(self, cache_key: str) -> RegistryArtifactAdmission | None: + """Return byte-bound admission controls for one cold cache key.""" + max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES + if max_bytes <= 0: + return None + + async def ensure_capacity(additional_bytes: int) -> None: + await self._ensure_cache_capacity( + additional_bytes=additional_bytes, + protected_key=cache_key, + max_bytes=max_bytes, + ) + + return RegistryArtifactAdmission( + max_bytes=max_bytes, + ensure_capacity=ensure_capacity, + ) + + async def _unmount_idle_entry(self, cache_key: str) -> None: + """Best-effort unmount an entry after its final lease is released.""" + try: + await self._unmount_entry(cache_key) + except OSError as e: + logger.warning( + "Failed to release idle registry artifact mount", + cache_key=cache_key, + error=str(e), + ) + + async def _converge_cache_budget(self) -> None: + """Bring an idle cache back under budget after a lease is released. + + Successful materialization enforces the budget after publication while + protecting the new entry. The cache can still sit over budget while + entries are leased. This runs on release, when every newly idle entry is + evictable. The scan is skipped entirely unless a materialization attempt + has occurred since the last successful enforcement. + + Each successful pass consumes the dirty signal before its awaited scan. + A follow-up pass therefore occurs only when a concurrent materialization + sets the flag again. Without new materializations the loop terminates, + while an over-budget or failed scan restores the flag and breaks so it + cannot spin while entries remain leased. Cancellation also restores the + consumed flag before propagating. + """ + while self._budget_dirty: + self._budget_dirty = False + try: + within_budget = await self._enforce_cache_budget() + except OSError as e: + logger.warning( + "Failed to converge registry artifact cache to budget", + cache_dir=str(self.cache_dir), + error=str(e), + ) + self._budget_dirty = True + break + except BaseException: + # Cancellation must re-arm the consumed dirty signal before propagating. + self._budget_dirty = True + raise + + if not within_budget: + self._budget_dirty = True + break + + async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bool: + """Evict least-recently-used idle entries until the cache fits its budget. + + The admission lock excludes cold writers before the budget lock begins + a scan/select/evict pass. Callers invoke enforcement without holding a + per-key lock. Cold writers already hold the admission lock and use + ``_ensure_cache_capacity`` for their staged reservations instead. + + Args: + protected_key: Newly materialized cache key. It is counted against + the budget when present but never evicted. None when enforcing + against idle entries after leases are released. + + Returns: + Whether the cache is within budget once eviction has finished. + """ + async with self._admission_lock: + async with self._budget_lock: + return await self._enforce_cache_budget_locked( + protected_key=protected_key + ) + + async def _enforce_cache_budget_locked( + self, + *, + protected_key: str | None, + ) -> bool: + """Enforce entry and byte limits while both cache-wide locks are held.""" + trash_clean, startup_clean = await asyncio.gather( + asyncio.to_thread(self._clear_work_dir, self.trash_dir), + asyncio.to_thread(self._retry_failed_startup_cleanup), + ) + cleanup_complete = trash_clean and startup_clean + if not cleanup_complete: + return False + + max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES + max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES + if max_entries <= 0 and max_bytes <= 0: + return True + + entries = await asyncio.to_thread(self._scan_cache_entries) + total_bytes = sum(entry.size_bytes for entry in entries.values()) + protected = set() if protected_key is None else {protected_key} + skipped: set[str] = set() + + while (max_entries > 0 and len(entries) > max_entries) or ( + max_bytes > 0 and total_bytes > max_bytes + ): + candidate = self._least_recently_used( + entries.values(), + excluded=skipped | protected, + ) + if candidate is None: + logger.warning( + "Registry artifact cache is over budget but every entry is in use", + cache_dir=str(self.cache_dir), + entries=len(entries), + max_entries=max_entries, + total_bytes=total_bytes, + max_bytes=max_bytes, + ) + return False + + eviction = await self._evict_entry(candidate.cache_key) + if eviction.retired: + del entries[candidate.cache_key] + if not eviction.reclaimed: + return False + total_bytes -= candidate.size_bytes + else: + skipped.add(candidate.cache_key) + + return True + + async def _ensure_cache_capacity( + self, + *, + additional_bytes: int, + protected_key: str, + max_bytes: int, + ) -> None: + """Reserve peak bytes for a cold writer without exceeding the cap. + + The caller holds the admission lock and its key lock. Every normal + budget pass takes the admission lock first, so acquiring the budget + lock here cannot deadlock with eviction of the protected key. + """ + if additional_bytes < 0: + raise ValueError("additional_bytes must be non-negative") + + async with self._budget_lock: + await asyncio.gather( + asyncio.to_thread(self._clear_work_dir, self.trash_dir), + asyncio.to_thread(self._retry_failed_startup_cleanup), + ) + entries = await asyncio.to_thread(self._scan_cache_entries) + staging_bytes, trash_bytes = await asyncio.gather( + asyncio.to_thread(_directory_footprint, self.staging_dir), + asyncio.to_thread(_directory_footprint, self.trash_dir), + ) + total_bytes = ( + sum(entry.size_bytes for entry in entries.values()) + + staging_bytes + + trash_bytes + ) + skipped = {protected_key} + + while total_bytes + additional_bytes > max_bytes: + candidate = self._least_recently_used( + entries.values(), + excluded=skipped, + ) + if candidate is None: + raise RegistryArtifactCacheCapacityError( + current_bytes=total_bytes, + additional_bytes=additional_bytes, + max_bytes=max_bytes, + ) + + eviction = await self._evict_entry(candidate.cache_key) + if eviction.retired: + del entries[candidate.cache_key] + if not eviction.reclaimed: + raise RegistryArtifactCacheCapacityError( + current_bytes=total_bytes, + additional_bytes=additional_bytes, + max_bytes=max_bytes, + ) + total_bytes -= candidate.size_bytes + else: + skipped.add(candidate.cache_key) + + def _least_recently_used( + self, + entries: Iterable[RegistryArtifactCacheEntry], + *, + excluded: set[str], + ) -> RegistryArtifactCacheEntry | None: + """Return the least recently used idle entry eligible for eviction.""" + eligible = [ + entry + for entry in entries + if entry.cache_key not in excluded and self._refcount(entry.cache_key) == 0 + ] + if not eligible: + return None + return min(eligible, key=self._recency) + + def _recency(self, entry: RegistryArtifactCacheEntry) -> float: + """Return the most recent known use time for a cache entry.""" + runtime = self._runtime.get(entry.cache_key) + if runtime is None: + return entry.last_used + return max(entry.last_used, runtime.last_used) + + async def _unmount_entry(self, cache_key: str) -> bool: + """Unmount one idle cache entry while retaining its reusable image. + + Loop-device reclamation is independent from disk-budget eviction. The + per-key lock and lease recheck prevent an entry from being unmounted + while an action is importing from it. The image and empty mount + directory remain cached so a later admission can remount without + downloading the artifact again. + + Args: + cache_key: Cache key whose mounted artifact should be released. + + Returns: + Whether a mounted entry was unmounted. + """ + lock = self._runtime_for(cache_key).lock + if lock.locked(): + logger.debug( + "Skipping unmount of busy registry artifact", + cache_key=cache_key, + ) + return False + + async with lock: + if self._refcount(cache_key) > 0: + return False + + mount_dir = self._paths_for(cache_key).squashfs_mount_dir + if not mount_dir.is_mount(): + return False + if not await self._unmount(mount_dir): + logger.warning( + "Failed to unmount registry artifact for loop-device reclamation", + cache_key=cache_key, + mount_dir=str(mount_dir), + ) + return False + + logger.info( + "Unmounted idle registry artifact", + cache_key=cache_key, + mount_dir=str(mount_dir), + ) + return True + + async def _evict_entry(self, cache_key: str) -> RegistryArtifactEviction: + """Remove one cache entry from disk, unmounting it first. + + The entry is skipped rather than forced when it is leased, busy, or + cannot be unmounted: deleting the image file behind a live mount would + leave an open-file zombie holding the loop device. + + After unmounting, the entry root is atomically renamed into ``trash`` + under the per-key lock. The lock is then released before physical + deletion runs in a worker thread. + + Args: + cache_key: Cache key to evict. + + Returns: + Whether the entry was retired and its bytes were reclaimed. + """ + lock = self._runtime_for(cache_key).lock + if lock.locked(): + logger.debug( + "Skipping eviction of busy registry artifact", + cache_key=cache_key, + ) + return RegistryArtifactEviction(retired=False, reclaimed=False) + + async with lock: + if self._refcount(cache_key) > 0: + return RegistryArtifactEviction(retired=False, reclaimed=False) + + paths = self._paths_for(cache_key) + if not paths.entry_dir.exists(): + return RegistryArtifactEviction(retired=True, reclaimed=True) + if paths.squashfs_mount_dir.is_mount() and not await self._unmount( + paths.squashfs_mount_dir + ): + logger.warning( + "Failed to unmount registry artifact, skipping eviction", + cache_key=cache_key, + mount_dir=str(paths.squashfs_mount_dir), + ) + return RegistryArtifactEviction(retired=False, reclaimed=False) + + try: + trash_path = _move_entry_to_trash( + paths.entry_dir, + self.trash_dir, + cache_key, + ) + except OSError as e: + self._budget_dirty = True + logger.warning( + "Failed to retire registry artifact cache entry", + cache_key=cache_key, + entry_dir=str(paths.entry_dir), + error=str(e), + ) + return RegistryArtifactEviction(retired=False, reclaimed=False) + + try: + deleted = await _delete_cache_path_off_loop(trash_path) + except BaseException: + self._budget_dirty = True + raise + if not deleted: + self._budget_dirty = True + logger.warning( + "Registry artifact eviction remains pending physical deletion", + cache_key=cache_key, + trash_path=str(trash_path), + ) + return RegistryArtifactEviction(retired=True, reclaimed=False) + + logger.info("Evicted registry artifact from cache", cache_key=cache_key) + return RegistryArtifactEviction(retired=True, reclaimed=True) + + async def _unmount(self, mount_dir: Path) -> bool: + """Unmount a SquashFS artifact directory, releasing its loop device. + + Cancellation kills and reaps the umount subprocess before propagating, + so the caller's per-key lock covers the complete unmount lifecycle. If + umount never took effect, the mounted entry stays consistent and can be + reused; if it already took effect, the missing extraction directory + makes the entry a plain cache miss on the next admission. + """ + umount = shutil.which("umount") + if umount is None: + return False + + proc = await asyncio.create_subprocess_exec( + umount, + str(mount_dir), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + proc.kill() + await proc.wait() + raise + if proc.returncode == 0 or not mount_dir.is_mount(): + return True + + logger.warning( + "umount command failed", + mount_dir=str(mount_dir), + output=(stderr or stdout).decode(errors="replace").strip(), + ) + return False + + def _scan_cache_entries(self) -> dict[str, RegistryArtifactCacheEntry]: + """Measure every registry artifact entry currently on disk.""" + return { + cache_key: self._measure_entry(cache_key) + for cache_key in self._discover_cache_keys() + } + + def _discover_cache_keys(self) -> set[str]: + """Return cache keys represented by atomic entry directories.""" + try: + entries = list(os.scandir(self.entries_dir)) + except FileNotFoundError: + return set() + + return { + entry.name + for entry in entries + if entry.name and entry.is_dir(follow_symlinks=False) + } + + def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: + """Measure the on-disk footprint and recency of one cache entry. + + The mount directory is excluded because a mounted view only costs the + image file that backs it. The image is measured with a single ``stat`` + so a concurrent eviction deleting it cannot fail the scan. + """ + paths = self._paths_for(cache_key) + size_bytes = 0 + + try: + image_stat = paths.squashfs_image_path.stat() + except FileNotFoundError: + pass + else: + size_bytes += image_stat.st_size + + for directory in (paths.squashfs_extract_dir, paths.tarball_target_dir): + size_bytes += _directory_footprint(directory) + + try: + last_used = paths.entry_dir.stat().st_mtime + except FileNotFoundError: + last_used = 0.0 + + return RegistryArtifactCacheEntry( + cache_key=cache_key, + size_bytes=size_bytes, + last_used=last_used, + ) + + def _sweep_startup_state(self) -> None: + """Reclaim orphaned cache state left behind by a previous process. + + Scratch and trash paths from interrupted work are removed, and active + entries are trimmed to budget using entry-root mtimes as LRU order. + + The worker warms this sweep before activities can run; lazy first-use + sweeping remains a safe fallback. + """ + if not self.cache_dir.is_dir(): + self._budget_dirty = False + return + + try: + staging_clean = self._clear_work_dir( + self.staging_dir, + remember_failures=True, + ) + trash_clean = self._clear_work_dir(self.trash_dir) + cleanup_complete = staging_clean and trash_clean + within_budget = cleanup_complete and self._trim_startup_cache() + self._budget_dirty = not (cleanup_complete and within_budget) + except OSError as e: + logger.warning( + "Failed to sweep registry artifact cache", + cache_dir=str(self.cache_dir), + error=str(e), + ) + raise + + def _clear_work_dir( + self, + work_dir: Path, + *, + remember_failures: bool = False, + ) -> bool: + """Best-effort remove every child of a staging or trash directory.""" + try: + paths = list(work_dir.iterdir()) + except FileNotFoundError: + return True + except OSError as e: + logger.warning( + "Failed to inspect registry artifact work directory", + path=str(work_dir), + error=str(e), + ) + raise + + deleted = True + for path in paths: + if _delete_cache_path(path): + if remember_failures: + self._failed_startup_cleanup.discard(path) + logger.info( + "Removed registry artifact work path", + path=str(path), + ) + else: + deleted = False + if remember_failures: + self._failed_startup_cleanup.add(path) + return deleted + + def _retry_failed_startup_cleanup(self) -> bool: + """Retry exact startup paths without sweeping live staging work.""" + for path in tuple(self._failed_startup_cleanup): + if _delete_cache_path(path): + self._failed_startup_cleanup.discard(path) + return not self._failed_startup_cleanup + + def _trim_startup_cache(self) -> bool: + """Trim the cache to budget before any artifact is leased. + + Returns whether active entries and pending physical deletion fit within + the configured budget. + """ + max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES + max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES + if max_entries <= 0 and max_bytes <= 0: + return True + + entries = self._scan_cache_entries() + total_bytes = sum(entry.size_bytes for entry in entries.values()) + # Mounted entries belong to a live process sharing this cache directory. + candidates = sorted( + ( + entry + for entry in entries.values() + if not self._paths_for(entry.cache_key).squashfs_mount_dir.is_mount() + ), + key=lambda entry: entry.last_used, + ) + + def within_budget() -> bool: + return (max_entries <= 0 or len(entries) <= max_entries) and ( + max_bytes <= 0 or total_bytes <= max_bytes + ) + + for entry in candidates: + if within_budget(): + break + paths = self._paths_for(entry.cache_key) + try: + trash_path = _move_entry_to_trash( + paths.entry_dir, + self.trash_dir, + entry.cache_key, + ) + except OSError as e: + logger.warning( + "Failed to retire stale registry artifact during startup sweep", + cache_key=entry.cache_key, + entry_dir=str(paths.entry_dir), + error=str(e), + ) + return False + + del entries[entry.cache_key] + if _delete_cache_path(trash_path): + total_bytes -= entry.size_bytes + else: + return False + logger.info( + "Evicted stale registry artifact during startup sweep", + cache_key=entry.cache_key, + size_bytes=entry.size_bytes, + ) + + return within_budget() diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 58e8ad2129..f038418bd8 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -4,968 +4,101 @@ import asyncio import contextlib -import hashlib -import os -import shutil -import sysconfig -import tarfile -import threading -import time -from abc import ABC, abstractmethod -from collections.abc import AsyncIterator, Awaitable, Callable, Iterable +from collections.abc import AsyncIterator from contextlib import asynccontextmanager -from dataclasses import dataclass, field -from enum import StrEnum from pathlib import Path -import httpx -import tracecat_registry - from tracecat import config +from tracecat.executor.registry_artifact_cache_state import ( + BASE_PYTHONPATH_DIR_NAME, + CACHE_ENTRIES_DIR_NAME, + CACHE_STAGING_DIR_NAME, + CACHE_TRASH_DIR_NAME, + RegistryArtifactCacheLoopError, + RegistryArtifactRuntimeState, +) +from tracecat.executor.registry_artifact_materialization import ( + BUNDLED_BUILTIN_REGISTRY_URI_PREFIX, + SQUASHFS_MOUNT_OPTIONS, + BuiltinArtifact, + RegistryArtifact, + RegistryArtifactAdmission, + RegistryArtifactFormat, + RegistryArtifactMaterializationContext, + RegistryArtifactPaths, + SquashfsArtifact, + SquashfsMountCommandError, + TarballArtifact, + _artifact_format, + _bundled_builtin_registry_import_paths, + _bundled_builtin_registry_version, + _download_s3_artifact, + _is_cache_entry_uri, + _squashfs_listing_size, + _squashfs_sidecar_uri, + _tarball_extracted_size, + _tarball_uri_for_squashfs, + bundled_builtin_registry_uri, + compute_registry_artifact_cache_key, +) +from tracecat.executor.registry_artifact_storage import ( + RegistryArtifactCacheCapacityError, + RegistryArtifactCacheEntry, + RegistryArtifactEviction, + _delete_cache_path, + _delete_cache_path_off_loop, + _directory_footprint, + _move_entry_to_trash, + _RegistryArtifactCacheStorage, + _unique_work_path, +) from tracecat.logger import logger from tracecat.registry.artifact_keys import parse_s3_uri -from tracecat.registry.constants import DEFAULT_REGISTRY_ORIGIN from tracecat.storage import blob - -class RegistryArtifactFormat(StrEnum): - """Executor-supported registry artifact encodings.""" - - BUILTIN = "builtin" - SQUASHFS = "squashfs" - TAR_GZ = "tar.gz" - - -SQUASHFS_MOUNT_OPTIONS = "loop,ro,nodev,nosuid" -"""Mount options for executor-managed SquashFS registry artifacts. - -The image must stay read-only and should not expose device nodes or setuid bits -from registry package contents. Avoid noexec because Python packages may include -native extension modules that need to be loaded from the mounted artifact. -""" - -BUNDLED_BUILTIN_REGISTRY_URI_PREFIX = f"tracecat-builtin://{DEFAULT_REGISTRY_ORIGIN}/" -"""Pseudo-URI for the builtin registry already installed in the executor image.""" - -BASE_PYTHONPATH_DIR_NAME = "base" -"""Cache subdirectory used as the PYTHONPATH entry when no artifact is requested.""" - -CACHE_ENTRIES_DIR_NAME = "entries" -"""Directory containing one atomic subdirectory per cache key.""" - -CACHE_STAGING_DIR_NAME = "staging" -"""Directory containing in-progress materialization scratch.""" - -CACHE_TRASH_DIR_NAME = "trash" -"""Directory containing atomically retired entries pending physical deletion.""" - - -class SquashfsMountCommandError(RuntimeError): - """The ``mount`` command itself failed for a SquashFS registry artifact. - - Only this error drives SquashFS mount policy. Download, mkdir, and other - preparation failures must not be mistaken for a missing mount capability or - for loop-device exhaustion. - """ - - -class RegistryArtifactCacheLoopError(RuntimeError): - """A registry artifact cache was used outside its owning event loop.""" - - -class RegistryArtifactCacheCapacityError(RuntimeError): - """A cold artifact cannot fit within the configured cache byte budget.""" - - def __init__( - self, - *, - current_bytes: int, - additional_bytes: int, - max_bytes: int, - ) -> None: - super().__init__( - "Registry artifact admission exceeds the cache byte budget: " - f"current_bytes={current_bytes}, additional_bytes={additional_bytes}, " - f"max_bytes={max_bytes}" - ) - self.current_bytes = current_bytes - self.additional_bytes = additional_bytes - self.max_bytes = max_bytes - - -@dataclass(frozen=True, slots=True) -class RegistryArtifactPaths: - """Executor-local cache paths for one registry artifact key.""" - - entry_dir: Path - squashfs_image_path: Path - squashfs_mount_dir: Path - squashfs_extract_dir: Path - tarball_target_dir: Path - - -@dataclass(slots=True) -class RegistryArtifactRuntimeState: - """Process-local synchronization and lease state for one cache key.""" - - lock: asyncio.Lock = field(default_factory=asyncio.Lock) - refcount: int = 0 - last_used: float = 0.0 - - -@dataclass(frozen=True, slots=True) -class RegistryArtifactEviction: - """Outcome of atomically retiring and physically deleting one entry.""" - - retired: bool - reclaimed: bool - - -@dataclass(frozen=True, slots=True) -class RegistryArtifactCacheEntry: - """Measured on-disk footprint and recency for one registry artifact key.""" - - cache_key: str - size_bytes: int - last_used: float - - -@dataclass(frozen=True, slots=True) -class RegistryArtifactAdmission: - """Byte-bound admission hook shared by one cold materialization.""" - - max_bytes: int - ensure_capacity: Callable[[int], Awaitable[None]] - - -@dataclass(slots=True) -class RegistryArtifactMaterializationContext: - """Shared runtime state for artifact materialization.""" - - cache_key: str - staging_dir: Path - paths: RegistryArtifactPaths - admission: RegistryArtifactAdmission | None = None - - def can_mount_squashfs(self) -> bool: - return config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED and ( - shutil.which("mount") is not None - ) - - -@dataclass(frozen=True, slots=True) -class RegistryArtifact(ABC): - """An executor-local materializable registry artifact.""" - - uri: str - cache_key: str - - @property - @abstractmethod - def format(self) -> RegistryArtifactFormat: - """Artifact format used for logging and dispatch.""" - - @abstractmethod - def cached_path( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path] | None: - """Return already-materialized import paths for this artifact, if present.""" - - @abstractmethod - async def materialize( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path]: - """Return importable Python paths, materializing the artifact if needed.""" - - def _temp_path( - self, - ctx: RegistryArtifactMaterializationContext, - suffix: str, - ) -> Path: - unique_id = id(asyncio.current_task()) - ctx.staging_dir.mkdir(parents=True, exist_ok=True) - return ctx.staging_dir / f"{self.cache_key}.{os.getpid()}.{unique_id}{suffix}" - - -@dataclass(frozen=True, slots=True) -class BuiltinArtifact(RegistryArtifact): - """Current builtin registry package already installed in the executor image.""" - - version: str - - @property - def format(self) -> RegistryArtifactFormat: - return RegistryArtifactFormat.BUILTIN - - def cached_path( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path] | None: - return None - - async def materialize( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path]: - del ctx - import_paths = _bundled_builtin_registry_import_paths(self.version) - logger.info( - "Using bundled builtin registry environment", - registry_version=self.version, - paths=[str(p) for p in import_paths], - ) - return import_paths - - -@dataclass(frozen=True, slots=True) -class SquashfsArtifact(RegistryArtifact): - """SquashFS registry environment image.""" - - @property - def format(self) -> RegistryArtifactFormat: - return RegistryArtifactFormat.SQUASHFS - - def cached_path( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path] | None: - if ctx.paths.squashfs_mount_dir.is_mount(): - logger.debug( - "Using cached SquashFS registry mount", - cache_key=ctx.cache_key, - ) - return [ctx.paths.squashfs_mount_dir] - if ctx.paths.squashfs_extract_dir.exists(): - logger.debug( - "Using cached SquashFS registry extraction", - cache_key=ctx.cache_key, - ) - return [ctx.paths.squashfs_extract_dir] - return None - - async def materialize( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path]: - image_path = ctx.paths.squashfs_image_path - if ctx.can_mount_squashfs(): - try: - return [await self.mount(ctx, image_path)] - except SquashfsMountCommandError as e: - logger.warning( - "Failed to mount SquashFS registry artifact, trying extraction", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - error=str(e), - ) - - return [await self.extract(ctx, image_path)] - - async def download( - self, - ctx: RegistryArtifactMaterializationContext, - image_path: Path, - ) -> float: - """Ensure the SquashFS image exists locally and return download time.""" - if image_path.exists(): - return 0.0 - - image_path.parent.mkdir(parents=True, exist_ok=True) - temp_image = self._temp_path(ctx, ".squashfs") - try: - download_start = time.monotonic() - await _download_s3_artifact( - self.uri, - temp_image, - admission=ctx.admission, - ) - try: - temp_image.rename(image_path) - except OSError: - if not image_path.exists(): - raise - return (time.monotonic() - download_start) * 1000 - finally: - temp_image.unlink(missing_ok=True) - - async def mount( - self, - ctx: RegistryArtifactMaterializationContext, - image_path: Path, - ) -> Path: - """Download the image if needed and mount it read-only. - - Args: - ctx: Materialization context for the artifact being mounted. - image_path: Local path of the SquashFS image. - - Returns: - The mount directory. - - Raises: - SquashfsMountCommandError: The ``mount`` command failed. - Exception: The image could not be downloaded or prepared. - """ - target_dir = ctx.paths.squashfs_mount_dir - if target_dir.is_mount(): - return target_dir - - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - target_dir.mkdir(parents=True, exist_ok=True) - - logger.info( - "Materializing SquashFS registry artifact", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - ) - start_time = time.monotonic() - download_elapsed = await self.download(ctx, image_path) - - mount_start = time.monotonic() - await self._mount_image(image_path, target_dir) - mount_elapsed = (time.monotonic() - mount_start) * 1000 - total_elapsed = (time.monotonic() - start_time) * 1000 - - logger.info( - "SquashFS registry artifact mounted", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - download_ms=f"{download_elapsed:.1f}", - mount_ms=f"{mount_elapsed:.1f}", - total_ms=f"{total_elapsed:.1f}", - ) - return target_dir - - async def extract( - self, - ctx: RegistryArtifactMaterializationContext, - image_path: Path, - ) -> Path: - target_dir = ctx.paths.squashfs_extract_dir - if target_dir.exists(): - return target_dir - - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - - logger.info( - "Extracting SquashFS registry artifact", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - ) - start_time = time.monotonic() - download_elapsed = await self.download(ctx, image_path) - - temp_dir = self._temp_path(ctx, ".unsquashfs") - try: - if ctx.admission is not None: - extracted_size = await self._squashfs_extracted_size(image_path) - await ctx.admission.ensure_capacity(extracted_size) - extract_start = time.monotonic() - temp_dir.mkdir(parents=True, exist_ok=True) - await self._extract_image(image_path, temp_dir) - extract_elapsed = (time.monotonic() - extract_start) * 1000 - - try: - temp_dir.rename(target_dir) - total_elapsed = (time.monotonic() - start_time) * 1000 - logger.info( - "SquashFS registry artifact extracted", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - download_ms=f"{download_elapsed:.1f}", - extract_ms=f"{extract_elapsed:.1f}", - total_ms=f"{total_elapsed:.1f}", - ) - except OSError: - if target_dir.exists(): - logger.debug( - "SquashFS already extracted by another process", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - ) - else: - raise - finally: - if temp_dir.exists(): - shutil.rmtree(temp_dir, ignore_errors=True) - - return target_dir - - async def _mount_image(self, image_path: Path, target_dir: Path) -> None: - """Mount a SquashFS image read-only at target_dir. - - Cancellation kills and reaps the mount subprocess before propagating, - so the caller's per-key lock covers the complete mount lifecycle. The - target therefore remains an unmounted cache miss if mount never took - effect, or is already mounted and reusable by the next admission. - - Args: - image_path: Local path of the SquashFS image. - target_dir: Existing directory to mount the image at. - - Raises: - SquashfsMountCommandError: The ``mount`` command failed. - """ - if target_dir.is_mount(): - return - - proc = await asyncio.create_subprocess_exec( - "mount", - "-t", - "squashfs", - "-o", - SQUASHFS_MOUNT_OPTIONS, - str(image_path), - str(target_dir), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - try: - stdout, stderr = await proc.communicate() - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - proc.kill() - await proc.wait() - raise - - if proc.returncode == 0 or target_dir.is_mount(): - return - - output = (stderr or stdout).decode(errors="replace").strip() - raise SquashfsMountCommandError(output or "mount command failed") - - async def _extract_image(self, image_path: Path, target_dir: Path) -> None: - """Extract a SquashFS image to target_dir using unsquashfs. - - Cancellation kills and reaps the extractor before the caller removes - its scratch directory. Otherwise a live extractor could recreate - scratch after cleanup, outside startup-sweep discovery. - """ - unsquashfs = shutil.which("unsquashfs") - if unsquashfs is None: - raise RuntimeError("unsquashfs command is not installed") - - proc = await asyncio.create_subprocess_exec( - unsquashfs, - "-f", - "-d", - str(target_dir), - str(image_path), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - try: - stdout, stderr = await proc.communicate() - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - proc.kill() - await proc.wait() - raise - - if proc.returncode == 0: - return - - output = (stderr or stdout).decode(errors="replace").strip() - raise RuntimeError(output or "unsquashfs command failed") - - async def _squashfs_extracted_size(self, image_path: Path) -> int: - """Return a conservative logical size for an extracted SquashFS image.""" - unsquashfs = shutil.which("unsquashfs") - if unsquashfs is None: - raise RuntimeError("unsquashfs command is not installed") - - proc = await asyncio.create_subprocess_exec( - unsquashfs, - "-lln", - str(image_path), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - try: - stdout, stderr = await proc.communicate() - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - proc.kill() - await proc.wait() - raise - - if proc.returncode != 0: - output = (stderr or stdout).decode(errors="replace").strip() - raise RuntimeError(output or "unsquashfs listing failed") - return _squashfs_listing_size(stdout) - - -@dataclass(frozen=True, slots=True) -class TarballArtifact(RegistryArtifact): - """Legacy gzip tarball registry environment.""" - - @property - def format(self) -> RegistryArtifactFormat: - return RegistryArtifactFormat.TAR_GZ - - def cached_path( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path] | None: - if ctx.paths.tarball_target_dir.exists(): - logger.debug( - "Using cached tarball extraction", - cache_key=ctx.cache_key, - ) - return [ctx.paths.tarball_target_dir] - return None - - async def materialize( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path]: - target_dir = ctx.paths.tarball_target_dir - logger.info( - "Materializing tarball registry artifact", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - ) - start_time = time.monotonic() - - temp_tarball = self._temp_path(ctx, ".tar.gz") - temp_dir = self._temp_path(ctx, ".tmp") - - try: - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - - download_start = time.monotonic() - await self.download(ctx, temp_tarball) - download_elapsed = (time.monotonic() - download_start) * 1000 - - if ctx.admission is not None: - extracted_size = await asyncio.to_thread( - _tarball_extracted_size, - temp_tarball, - ) - await ctx.admission.ensure_capacity(extracted_size) - - extract_start = time.monotonic() - temp_dir.mkdir(parents=True, exist_ok=True) - await self.extract(temp_tarball, temp_dir) - extract_elapsed = (time.monotonic() - extract_start) * 1000 - - try: - temp_dir.rename(target_dir) - total_elapsed = (time.monotonic() - start_time) * 1000 - logger.info( - "Tarball extracted and cached", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - download_ms=f"{download_elapsed:.1f}", - extract_ms=f"{extract_elapsed:.1f}", - total_ms=f"{total_elapsed:.1f}", - ) - except OSError: - if target_dir.exists(): - logger.debug( - "Tarball already extracted by another process", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - artifact_format=self.format.value, - ) - else: - raise - finally: - if temp_dir.exists(): - shutil.rmtree(temp_dir, ignore_errors=True) - if temp_tarball.exists(): - temp_tarball.unlink(missing_ok=True) - - return [target_dir] - - async def download( - self, - ctx: RegistryArtifactMaterializationContext, - output_path: Path, - ) -> None: - await _download_s3_artifact( - self.uri, - output_path, - admission=ctx.admission, - ) - - async def extract(self, tarball_path: Path, target_dir: Path) -> None: - """Extract a supported registry tarball to target directory. - - A cancelled caller rejoins the non-interruptible extraction thread - before propagating cancellation so scratch cleanup cannot race a live - writer and leave undiscoverable ephemeral storage behind. - """ - - def _do_extract() -> None: - if tarball_path.name.endswith(".tar.gz"): - with tarfile.open(tarball_path, "r:gz") as tar: - tar.extractall(path=target_dir, filter="data") - return - - raise ValueError(f"Unsupported tarball format: {tarball_path}") - - extraction = asyncio.ensure_future(asyncio.to_thread(_do_extract)) - try: - await asyncio.shield(extraction) - except asyncio.CancelledError: - # A thread cannot be killed. Rejoin it before materialize removes - # scratch. Each cancellation can interrupt shield without stopping - # the thread, so keep waiting until extraction reaches a terminal - # state before propagating the original cancellation. - while not extraction.done(): - try: - await asyncio.shield(extraction) - except asyncio.CancelledError: - continue - except Exception: - break - if not extraction.cancelled(): - with contextlib.suppress(Exception): - extraction.result() - raise - - logger.debug( - "Tarball extracted", - target=str(target_dir), - artifact_format=_artifact_format(str(tarball_path)).value, - ) - - -async def _download_s3_artifact( - artifact_uri: str, - output_path: Path, - *, - admission: RegistryArtifactAdmission | None = None, -) -> None: - """Download an S3 registry artifact to a local path.""" - bucket, key = parse_s3_uri(artifact_uri) - try: - if admission is None: - await blob.download_file_to_path( - key=key, - bucket=bucket, - output_path=output_path, - ) - else: - await blob.download_file_to_path( - key=key, - bucket=bucket, - output_path=output_path, - max_bytes=admission.max_bytes, - ensure_capacity=admission.ensure_capacity, - ) - except FileNotFoundError as e: - request = httpx.Request("GET", artifact_uri) - response = httpx.Response(status_code=404, request=request) - raise httpx.HTTPStatusError( - f"Registry artifact not found: {artifact_uri}", - request=request, - response=response, - ) from e - - -def compute_registry_artifact_cache_key(artifact_uri: str) -> str: - """Compute the local cache key for a registry artifact URI.""" - if not artifact_uri: - return "base" - # S3 keys are case-sensitive, so preserve URI case when hashing. - content = artifact_uri.strip() - return hashlib.sha256(content.encode()).hexdigest()[:16] - - -def bundled_builtin_registry_uri(version: str) -> str: - """Return the pseudo-URI for the installed builtin registry package.""" - return f"{BUNDLED_BUILTIN_REGISTRY_URI_PREFIX}{version}" - - -def _bundled_builtin_registry_version(artifact_uri: str) -> str | None: - """Return the builtin registry version encoded in a bundled pseudo-URI.""" - if not artifact_uri.startswith(BUNDLED_BUILTIN_REGISTRY_URI_PREFIX): - return None - version = artifact_uri.removeprefix(BUNDLED_BUILTIN_REGISTRY_URI_PREFIX) - return version or None - - -def _bundled_builtin_registry_import_paths(version: str) -> list[Path]: - """Return import paths for the current builtin registry and its dependencies. - - Dependencies always live in the executor's site-packages. For editable - installs the parent of ``package_dir`` (the package wrapper, e.g. - ``packages/tracecat-registry/``) is exposed first so its ``tracecat_registry/`` - shadows any stale copy in site-packages. - """ - installed_version = tracecat_registry.__version__ - if version != installed_version: - raise RuntimeError( - "Bundled builtin registry version does not match installed version: " - f"requested={version!r}, installed={installed_version!r}" - ) - - package_file = tracecat_registry.__file__ - if package_file is None: - raise RuntimeError("Installed tracecat_registry package has no __file__") - - site_packages_path = sysconfig.get_path("purelib") - if site_packages_path is None: - raise RuntimeError("Could not resolve installed Python site-packages path") - - site_packages = Path(site_packages_path).resolve() - if not site_packages.exists(): - raise RuntimeError( - f"Installed Python site-packages path does not exist: {site_packages}" - ) - - package_dir = Path(package_file).resolve().parent - if package_dir.is_relative_to(site_packages): - return [site_packages] - - return [package_dir.parent, site_packages] - - -def _squashfs_sidecar_uri(tarball_uri: str) -> str | None: - """Return the sibling SquashFS URI for registry site-packages tarballs.""" - if not tarball_uri.endswith("site-packages.tar.gz"): - return None - return tarball_uri.removesuffix(".tar.gz") + ".squashfs" - - -def _tarball_uri_for_squashfs(squashfs_uri: str) -> str | None: - """Return the sibling gzip tarball URI for registry SquashFS artifacts.""" - if not squashfs_uri.endswith("site-packages.squashfs"): - return None - return squashfs_uri.removesuffix(".squashfs") + ".tar.gz" - - -def _artifact_format(artifact_uri: str) -> RegistryArtifactFormat: - """Return the materialization format for an artifact URI.""" - if artifact_uri.endswith(".squashfs"): - return RegistryArtifactFormat.SQUASHFS - return RegistryArtifactFormat.TAR_GZ - - -def _is_cache_entry_uri(artifact_uri: str) -> bool: - """Return whether an artifact URI materializes into an evictable cache entry. - - The bundled builtin registry is served from the executor image and never - writes into the cache directory, so it is exempt from eviction accounting. - """ - return _bundled_builtin_registry_version(artifact_uri) is None - - -def _tarball_extracted_size(tarball_path: Path) -> int: - """Return a conservative logical size for all tarball members.""" - total_bytes = 0 - with tarfile.open(tarball_path, "r:gz") as tar: - for member in tar: - if member.size < 0: - raise ValueError( - f"Registry tarball member has a negative size: {member.name}" - ) - total_bytes += member.size - return total_bytes - - -def _squashfs_listing_size(output: bytes) -> int: - """Sum file sizes from ``unsquashfs -lln`` output, failing closed.""" - total_bytes = 0 - for raw_line in output.decode(errors="strict").splitlines(): - line = raw_line.strip() - if not line: - continue - fields = line.split(maxsplit=4) - mode = fields[0] - if len(mode) != 10 or mode[0] not in "bcdlps-": - continue - if mode[0] not in "-l": - continue - if len(fields) < 5 or "/" not in fields[1] or not fields[2].isdigit(): - raise ValueError(f"Could not parse SquashFS listing line: {line}") - total_bytes += int(fields[2]) - return total_bytes - - -def _directory_footprint(directory: Path) -> int: - """Return the total file size of a cache directory. - - Args: - directory: Cache directory to measure. - - Returns: - Total byte size of contained files, or zero when the directory is - missing. - """ - - def raise_walk_error(error: OSError) -> None: - raise error - - total_bytes = 0 - try: - walker = os.walk(directory, onerror=raise_walk_error) - for root, _dirs, files in walker: - for file_name in files: - try: - total_bytes += os.lstat(os.path.join(root, file_name)).st_size - except FileNotFoundError: - continue - except FileNotFoundError: - return 0 - return total_bytes - - -def _delete_cache_path(path: Path) -> bool: - """Best-effort delete one cache path while reporting filesystem failures.""" - try: - if path.is_dir(): - shutil.rmtree(path) - else: - path.unlink(missing_ok=True) - except OSError as e: - logger.warning( - "Failed to delete registry artifact cache path", - path=str(path), - error=str(e), - ) - return False - return True - - -async def _delete_cache_path_off_loop(path: Path) -> bool: - """Delete one path without abandoning its worker thread on cancellation.""" - deletion = asyncio.ensure_future(asyncio.to_thread(_delete_cache_path, path)) - try: - return await asyncio.shield(deletion) - except asyncio.CancelledError: - # A worker thread cannot be killed. Rejoin it so no live deletion can - # race a later trash-directory scan. Repeated cancellation can interrupt - # shield without stopping the thread, so keep waiting for termination. - while not deletion.done(): - try: - await asyncio.shield(deletion) - except asyncio.CancelledError: - continue - except Exception: - break - if not deletion.cancelled(): - with contextlib.suppress(Exception): - deletion.result() - raise - - -def _unique_work_path(root: Path, cache_key: str) -> Path: - """Return a unique path beneath a cache work directory.""" - root.mkdir(parents=True, exist_ok=True) - unique_id = time.time_ns() - while True: - path = root / f"{cache_key}.{os.getpid()}.{unique_id}" - if not path.exists(): - return path - unique_id += 1 - - -def _move_entry_to_trash(entry_dir: Path, trash_dir: Path, cache_key: str) -> Path: - """Atomically retire one cache entry and return its trash path.""" - trash_path = _unique_work_path(trash_dir, cache_key) - entry_dir.rename(trash_path) - return trash_path - - -class RegistryArtifactCache: - """Materializes registry artifacts into executor-local Python paths.""" - - def __init__(self, cache_dir: Path): - self.cache_dir = cache_dir - self.entries_dir = cache_dir / CACHE_ENTRIES_DIR_NAME - self.staging_dir = cache_dir / CACHE_STAGING_DIR_NAME - self.trash_dir = cache_dir / CACHE_TRASH_DIR_NAME - # Runtime states live for the process lifetime so every operation for a - # key always serializes on the same lock. - self._runtime: dict[str, RegistryArtifactRuntimeState] = {} - # The cache contains asyncio locks, tasks, and multi-step lease state. - # Bind the public API to one loop/thread so a future synchronous - # Temporal activity fails immediately instead of corrupting that state - # through a thread-local event loop. - self._owner_binding_lock = threading.Lock() - self._owner_loop: asyncio.AbstractEventLoop | None = None - self._owner_thread_id: int | None = None - # Cold materializations and budget passes share this outer lock. It - # keeps byte reservations stable while downloads and extraction write. - self._admission_lock = asyncio.Lock() - self._budget_lock = asyncio.Lock() - # Guard the off-loop startup sweep independently from cache operations. - self._swept: bool = False - self._sweep_task: asyncio.Task[None] | None = None - self._sweep_lock = asyncio.Lock() - # Startup is the only time the whole staging directory is swept. Exact - # paths that could not be removed are safe to retry later. - self._failed_startup_cleanup: set[Path] = set() - # Whether the on-disk cache may exceed its budget. Set when a new entry - # is materialized and cleared once enforcement measures a cache that - # fits, so steady-state cache hits never pay for a disk scan. - self._budget_dirty = True - - async def ensure_swept(self) -> None: - """Run the startup sweep exactly once successfully, off the event loop. - - Idempotent and cancellation-safe under concurrency: every caller joins - one stored sweep task, and cancelling a waiter never abandons or - restarts its live sweep. The lock is deliberately held while awaiting - that shared task so queued callers observe its result before proceeding. - Failures clear the task so the next caller retries. The sweep runs - before the first lease or materialization, so it never observes - in-flight cache entries. - """ - self._assert_owner_loop() - if self._swept: - return - async with self._sweep_lock: - if self._swept: - return - sweep_task = self._sweep_task - if sweep_task is None or ( - sweep_task.done() - and (sweep_task.cancelled() or sweep_task.exception() is not None) - ): - sweep_task = asyncio.ensure_future( - asyncio.to_thread(self._sweep_startup_state) - ) - self._sweep_task = sweep_task - try: - await asyncio.shield(sweep_task) - except asyncio.CancelledError: - if sweep_task.cancelled() and self._sweep_task is sweep_task: - self._sweep_task = None - raise - except Exception: - if self._sweep_task is sweep_task: - self._sweep_task = None - raise - self._swept = True - - def _assert_owner_loop(self) -> None: - """Bind to the current loop or reject use from another loop/thread.""" - current_loop = asyncio.get_running_loop() - current_thread_id = threading.get_ident() - with self._owner_binding_lock: - if self._owner_loop is None: - self._owner_loop = current_loop - self._owner_thread_id = current_thread_id - return - if ( - self._owner_loop is current_loop - and self._owner_thread_id == current_thread_id - ): - return - - raise RegistryArtifactCacheLoopError( - "RegistryArtifactCache is bound to one event loop and thread " - f"(owner_loop={id(self._owner_loop)}, " - f"owner_thread={self._owner_thread_id}, " - f"current_loop={id(current_loop)}, " - f"current_thread={current_thread_id})" - ) +__all__ = [ + "BASE_PYTHONPATH_DIR_NAME", + "BUNDLED_BUILTIN_REGISTRY_URI_PREFIX", + "CACHE_ENTRIES_DIR_NAME", + "CACHE_STAGING_DIR_NAME", + "CACHE_TRASH_DIR_NAME", + "SQUASHFS_MOUNT_OPTIONS", + "BuiltinArtifact", + "RegistryArtifact", + "RegistryArtifactAdmission", + "RegistryArtifactCache", + "RegistryArtifactCacheCapacityError", + "RegistryArtifactCacheEntry", + "RegistryArtifactCacheLoopError", + "RegistryArtifactEviction", + "RegistryArtifactFormat", + "RegistryArtifactMaterializationContext", + "RegistryArtifactPaths", + "RegistryArtifactRuntimeState", + "SquashfsArtifact", + "SquashfsMountCommandError", + "TarballArtifact", + "_artifact_format", + "_bundled_builtin_registry_import_paths", + "_bundled_builtin_registry_version", + "_delete_cache_path", + "_delete_cache_path_off_loop", + "_directory_footprint", + "_download_s3_artifact", + "_is_cache_entry_uri", + "_move_entry_to_trash", + "_squashfs_listing_size", + "_squashfs_sidecar_uri", + "_tarball_extracted_size", + "_tarball_uri_for_squashfs", + "_unique_work_path", + "bundled_builtin_registry_uri", + "compute_registry_artifact_cache_key", +] + + +class RegistryArtifactCache(_RegistryArtifactCacheStorage): + """Materializes and leases executor-local registry artifact paths.""" @asynccontextmanager async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Path]]: @@ -1123,110 +256,6 @@ async def _materialize_candidates( raise RuntimeError(f"No registry artifact candidates for {ctx.cache_key}") - def _runtime_for(self, cache_key: str) -> RegistryArtifactRuntimeState: - """Return the process-local state for one cache key.""" - if runtime := self._runtime.get(cache_key): - return runtime - runtime = RegistryArtifactRuntimeState() - self._runtime[cache_key] = runtime - return runtime - - def _context_for( - self, - cache_key: str, - *, - admission: RegistryArtifactAdmission | None = None, - ) -> RegistryArtifactMaterializationContext: - """Return a materialization context for a registry artifact key.""" - return RegistryArtifactMaterializationContext( - cache_key=cache_key, - staging_dir=self.staging_dir, - paths=self._paths_for(cache_key), - admission=admission, - ) - - def _admission_for(self, cache_key: str) -> RegistryArtifactAdmission | None: - """Return byte-bound admission controls for one cold cache key.""" - max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES - if max_bytes <= 0: - return None - - async def ensure_capacity(additional_bytes: int) -> None: - await self._ensure_cache_capacity( - additional_bytes=additional_bytes, - protected_key=cache_key, - max_bytes=max_bytes, - ) - - return RegistryArtifactAdmission( - max_bytes=max_bytes, - ensure_capacity=ensure_capacity, - ) - - def _base_pythonpath_dir(self) -> Path: - """Return the base PYTHONPATH directory used when no artifact is requested.""" - base_dir = self.cache_dir / BASE_PYTHONPATH_DIR_NAME - base_dir.mkdir(parents=True, exist_ok=True) - return base_dir - - def _acquire_lease(self, cache_key: str) -> None: - """Pin a cache entry against eviction and mark it as recently used. - - Callers must hold the per-key lock so the increment is ordered against - in-flight eviction of the same key. - """ - runtime = self._runtime_for(cache_key) - runtime.refcount += 1 - runtime.last_used = time.time() - self._touch_entry(cache_key) - - def _release_lease(self, cache_key: str) -> bool: - """Release one pin and return whether the entry became idle.""" - runtime = self._runtime.get(cache_key) - if runtime is None or runtime.refcount == 0: - return False - runtime.refcount -= 1 - runtime.last_used = time.time() - return runtime.refcount == 0 - - async def _unmount_idle_entry(self, cache_key: str) -> None: - """Best-effort unmount an entry after its final lease is released.""" - try: - await self._unmount_entry(cache_key) - except OSError as e: - logger.warning( - "Failed to release idle registry artifact mount", - cache_key=cache_key, - error=str(e), - ) - - def _refcount(self, cache_key: str) -> int: - """Return the number of live leases on a cache entry.""" - runtime = self._runtime.get(cache_key) - return 0 if runtime is None else runtime.refcount - - def _touch_entry(self, cache_key: str) -> None: - """Best-effort refresh of the entry-root mtime for restart-safe LRU.""" - entry_dir = self._paths_for(cache_key).entry_dir - try: - os.utime(entry_dir) - except OSError: - logger.debug( - "Could not refresh registry artifact entry mtime", - cache_key=cache_key, - ) - - def _paths_for(self, cache_key: str) -> RegistryArtifactPaths: - """Return local cache paths for a registry artifact key.""" - entry_dir = self.entries_dir / cache_key - return RegistryArtifactPaths( - entry_dir=entry_dir, - squashfs_image_path=entry_dir / "image.squashfs", - squashfs_mount_dir=entry_dir / "mount", - squashfs_extract_dir=entry_dir / "extracted", - tarball_target_dir=entry_dir / "tarball", - ) - def _first_cached_path( self, candidates: list[RegistryArtifact], @@ -1381,534 +410,3 @@ async def _sidecar_exists( def _can_try_squashfs(self) -> bool: """Return whether this process should prefer SquashFS artifacts.""" return config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED - - async def _converge_cache_budget(self) -> None: - """Bring an idle cache back under budget after a lease is released. - - Successful materialization enforces the budget after publication while - protecting the new entry. The cache can still sit over budget while - entries are leased. This runs on release, when every newly idle entry is - evictable. The scan is skipped entirely unless a materialization attempt - has occurred since the last successful enforcement. - - Each successful pass consumes the dirty signal before its awaited scan. - A follow-up pass therefore occurs only when a concurrent materialization - sets the flag again. Without new materializations the loop terminates, - while an over-budget or failed scan restores the flag and breaks so it - cannot spin while entries remain leased. Cancellation also restores the - consumed flag before propagating. - """ - while self._budget_dirty: - self._budget_dirty = False - try: - within_budget = await self._enforce_cache_budget() - except OSError as e: - logger.warning( - "Failed to converge registry artifact cache to budget", - cache_dir=str(self.cache_dir), - error=str(e), - ) - self._budget_dirty = True - break - except BaseException: - # Cancellation must re-arm the consumed dirty signal before propagating. - self._budget_dirty = True - raise - - if not within_budget: - self._budget_dirty = True - break - - async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bool: - """Evict least-recently-used idle entries until the cache fits its budget. - - The admission lock excludes cold writers before the budget lock begins - a scan/select/evict pass. Callers invoke enforcement without holding a - per-key lock. Cold writers already hold the admission lock and use - ``_ensure_cache_capacity`` for their staged reservations instead. - - Args: - protected_key: Newly materialized cache key. It is counted against - the budget when present but never evicted. None when enforcing - against idle entries after leases are released. - - Returns: - Whether the cache is within budget once eviction has finished. - """ - async with self._admission_lock: - async with self._budget_lock: - return await self._enforce_cache_budget_locked( - protected_key=protected_key - ) - - async def _enforce_cache_budget_locked( - self, - *, - protected_key: str | None, - ) -> bool: - """Enforce entry and byte limits while both cache-wide locks are held.""" - trash_clean, startup_clean = await asyncio.gather( - asyncio.to_thread(self._clear_work_dir, self.trash_dir), - asyncio.to_thread(self._retry_failed_startup_cleanup), - ) - cleanup_complete = trash_clean and startup_clean - if not cleanup_complete: - return False - - max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES - max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES - if max_entries <= 0 and max_bytes <= 0: - return True - - entries = await asyncio.to_thread(self._scan_cache_entries) - total_bytes = sum(entry.size_bytes for entry in entries.values()) - protected = set() if protected_key is None else {protected_key} - skipped: set[str] = set() - - while (max_entries > 0 and len(entries) > max_entries) or ( - max_bytes > 0 and total_bytes > max_bytes - ): - candidate = self._least_recently_used( - entries.values(), - excluded=skipped | protected, - ) - if candidate is None: - logger.warning( - "Registry artifact cache is over budget but every entry is in use", - cache_dir=str(self.cache_dir), - entries=len(entries), - max_entries=max_entries, - total_bytes=total_bytes, - max_bytes=max_bytes, - ) - return False - - eviction = await self._evict_entry(candidate.cache_key) - if eviction.retired: - del entries[candidate.cache_key] - if not eviction.reclaimed: - return False - total_bytes -= candidate.size_bytes - else: - skipped.add(candidate.cache_key) - - return True - - async def _ensure_cache_capacity( - self, - *, - additional_bytes: int, - protected_key: str, - max_bytes: int, - ) -> None: - """Reserve peak bytes for a cold writer without exceeding the cap. - - The caller holds the admission lock and its key lock. Every normal - budget pass takes the admission lock first, so acquiring the budget - lock here cannot deadlock with eviction of the protected key. - """ - if additional_bytes < 0: - raise ValueError("additional_bytes must be non-negative") - - async with self._budget_lock: - await asyncio.gather( - asyncio.to_thread(self._clear_work_dir, self.trash_dir), - asyncio.to_thread(self._retry_failed_startup_cleanup), - ) - entries = await asyncio.to_thread(self._scan_cache_entries) - staging_bytes, trash_bytes = await asyncio.gather( - asyncio.to_thread(_directory_footprint, self.staging_dir), - asyncio.to_thread(_directory_footprint, self.trash_dir), - ) - total_bytes = ( - sum(entry.size_bytes for entry in entries.values()) - + staging_bytes - + trash_bytes - ) - skipped = {protected_key} - - while total_bytes + additional_bytes > max_bytes: - candidate = self._least_recently_used( - entries.values(), - excluded=skipped, - ) - if candidate is None: - raise RegistryArtifactCacheCapacityError( - current_bytes=total_bytes, - additional_bytes=additional_bytes, - max_bytes=max_bytes, - ) - - eviction = await self._evict_entry(candidate.cache_key) - if eviction.retired: - del entries[candidate.cache_key] - if not eviction.reclaimed: - raise RegistryArtifactCacheCapacityError( - current_bytes=total_bytes, - additional_bytes=additional_bytes, - max_bytes=max_bytes, - ) - total_bytes -= candidate.size_bytes - else: - skipped.add(candidate.cache_key) - - def _least_recently_used( - self, - entries: Iterable[RegistryArtifactCacheEntry], - *, - excluded: set[str], - ) -> RegistryArtifactCacheEntry | None: - """Return the least recently used idle entry eligible for eviction.""" - eligible = [ - entry - for entry in entries - if entry.cache_key not in excluded and self._refcount(entry.cache_key) == 0 - ] - if not eligible: - return None - return min(eligible, key=self._recency) - - def _recency(self, entry: RegistryArtifactCacheEntry) -> float: - """Return the most recent known use time for a cache entry.""" - runtime = self._runtime.get(entry.cache_key) - if runtime is None: - return entry.last_used - return max(entry.last_used, runtime.last_used) - - async def _unmount_entry(self, cache_key: str) -> bool: - """Unmount one idle cache entry while retaining its reusable image. - - Loop-device reclamation is independent from disk-budget eviction. The - per-key lock and lease recheck prevent an entry from being unmounted - while an action is importing from it. The image and empty mount - directory remain cached so a later admission can remount without - downloading the artifact again. - - Args: - cache_key: Cache key whose mounted artifact should be released. - - Returns: - Whether a mounted entry was unmounted. - """ - lock = self._runtime_for(cache_key).lock - if lock.locked(): - logger.debug( - "Skipping unmount of busy registry artifact", - cache_key=cache_key, - ) - return False - - async with lock: - if self._refcount(cache_key) > 0: - return False - - mount_dir = self._paths_for(cache_key).squashfs_mount_dir - if not mount_dir.is_mount(): - return False - if not await self._unmount(mount_dir): - logger.warning( - "Failed to unmount registry artifact for loop-device reclamation", - cache_key=cache_key, - mount_dir=str(mount_dir), - ) - return False - - logger.info( - "Unmounted idle registry artifact", - cache_key=cache_key, - mount_dir=str(mount_dir), - ) - return True - - async def _evict_entry(self, cache_key: str) -> RegistryArtifactEviction: - """Remove one cache entry from disk, unmounting it first. - - The entry is skipped rather than forced when it is leased, busy, or - cannot be unmounted: deleting the image file behind a live mount would - leave an open-file zombie holding the loop device. - - After unmounting, the entry root is atomically renamed into ``trash`` - under the per-key lock. The lock is then released before physical - deletion runs in a worker thread. - - Args: - cache_key: Cache key to evict. - - Returns: - Whether the entry was retired and its bytes were reclaimed. - """ - lock = self._runtime_for(cache_key).lock - if lock.locked(): - logger.debug( - "Skipping eviction of busy registry artifact", - cache_key=cache_key, - ) - return RegistryArtifactEviction(retired=False, reclaimed=False) - - async with lock: - if self._refcount(cache_key) > 0: - return RegistryArtifactEviction(retired=False, reclaimed=False) - - paths = self._paths_for(cache_key) - if not paths.entry_dir.exists(): - return RegistryArtifactEviction(retired=True, reclaimed=True) - if paths.squashfs_mount_dir.is_mount() and not await self._unmount( - paths.squashfs_mount_dir - ): - logger.warning( - "Failed to unmount registry artifact, skipping eviction", - cache_key=cache_key, - mount_dir=str(paths.squashfs_mount_dir), - ) - return RegistryArtifactEviction(retired=False, reclaimed=False) - - try: - trash_path = _move_entry_to_trash( - paths.entry_dir, - self.trash_dir, - cache_key, - ) - except OSError as e: - self._budget_dirty = True - logger.warning( - "Failed to retire registry artifact cache entry", - cache_key=cache_key, - entry_dir=str(paths.entry_dir), - error=str(e), - ) - return RegistryArtifactEviction(retired=False, reclaimed=False) - - try: - deleted = await _delete_cache_path_off_loop(trash_path) - except BaseException: - self._budget_dirty = True - raise - if not deleted: - self._budget_dirty = True - logger.warning( - "Registry artifact eviction remains pending physical deletion", - cache_key=cache_key, - trash_path=str(trash_path), - ) - return RegistryArtifactEviction(retired=True, reclaimed=False) - - logger.info("Evicted registry artifact from cache", cache_key=cache_key) - return RegistryArtifactEviction(retired=True, reclaimed=True) - - async def _unmount(self, mount_dir: Path) -> bool: - """Unmount a SquashFS artifact directory, releasing its loop device. - - Cancellation kills and reaps the umount subprocess before propagating, - so the caller's per-key lock covers the complete unmount lifecycle. If - umount never took effect, the mounted entry stays consistent and can be - reused; if it already took effect, the missing extraction directory - makes the entry a plain cache miss on the next admission. - """ - umount = shutil.which("umount") - if umount is None: - return False - - proc = await asyncio.create_subprocess_exec( - umount, - str(mount_dir), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - try: - stdout, stderr = await proc.communicate() - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - proc.kill() - await proc.wait() - raise - if proc.returncode == 0 or not mount_dir.is_mount(): - return True - - logger.warning( - "umount command failed", - mount_dir=str(mount_dir), - output=(stderr or stdout).decode(errors="replace").strip(), - ) - return False - - def _scan_cache_entries(self) -> dict[str, RegistryArtifactCacheEntry]: - """Measure every registry artifact entry currently on disk.""" - return { - cache_key: self._measure_entry(cache_key) - for cache_key in self._discover_cache_keys() - } - - def _discover_cache_keys(self) -> set[str]: - """Return cache keys represented by atomic entry directories.""" - try: - entries = list(os.scandir(self.entries_dir)) - except FileNotFoundError: - return set() - - return { - entry.name - for entry in entries - if entry.name and entry.is_dir(follow_symlinks=False) - } - - def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: - """Measure the on-disk footprint and recency of one cache entry. - - The mount directory is excluded because a mounted view only costs the - image file that backs it. The image is measured with a single ``stat`` - so a concurrent eviction deleting it cannot fail the scan. - """ - paths = self._paths_for(cache_key) - size_bytes = 0 - - try: - image_stat = paths.squashfs_image_path.stat() - except FileNotFoundError: - pass - else: - size_bytes += image_stat.st_size - - for directory in (paths.squashfs_extract_dir, paths.tarball_target_dir): - size_bytes += _directory_footprint(directory) - - try: - last_used = paths.entry_dir.stat().st_mtime - except FileNotFoundError: - last_used = 0.0 - - return RegistryArtifactCacheEntry( - cache_key=cache_key, - size_bytes=size_bytes, - last_used=last_used, - ) - - def _sweep_startup_state(self) -> None: - """Reclaim orphaned cache state left behind by a previous process. - - Scratch and trash paths from interrupted work are removed, and active - entries are trimmed to budget using entry-root mtimes as LRU order. - - The worker warms this sweep before activities can run; lazy first-use - sweeping remains a safe fallback. - """ - if not self.cache_dir.is_dir(): - self._budget_dirty = False - return - - try: - staging_clean = self._clear_work_dir( - self.staging_dir, - remember_failures=True, - ) - trash_clean = self._clear_work_dir(self.trash_dir) - cleanup_complete = staging_clean and trash_clean - within_budget = cleanup_complete and self._trim_startup_cache() - self._budget_dirty = not (cleanup_complete and within_budget) - except OSError as e: - logger.warning( - "Failed to sweep registry artifact cache", - cache_dir=str(self.cache_dir), - error=str(e), - ) - raise - - def _clear_work_dir( - self, - work_dir: Path, - *, - remember_failures: bool = False, - ) -> bool: - """Best-effort remove every child of a staging or trash directory.""" - try: - paths = list(work_dir.iterdir()) - except FileNotFoundError: - return True - except OSError as e: - logger.warning( - "Failed to inspect registry artifact work directory", - path=str(work_dir), - error=str(e), - ) - raise - - deleted = True - for path in paths: - if _delete_cache_path(path): - if remember_failures: - self._failed_startup_cleanup.discard(path) - logger.info( - "Removed registry artifact work path", - path=str(path), - ) - else: - deleted = False - if remember_failures: - self._failed_startup_cleanup.add(path) - return deleted - - def _retry_failed_startup_cleanup(self) -> bool: - """Retry exact startup paths without sweeping live staging work.""" - for path in tuple(self._failed_startup_cleanup): - if _delete_cache_path(path): - self._failed_startup_cleanup.discard(path) - return not self._failed_startup_cleanup - - def _trim_startup_cache(self) -> bool: - """Trim the cache to budget before any artifact is leased. - - Returns whether active entries and pending physical deletion fit within - the configured budget. - """ - max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES - max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES - if max_entries <= 0 and max_bytes <= 0: - return True - - entries = self._scan_cache_entries() - total_bytes = sum(entry.size_bytes for entry in entries.values()) - # Mounted entries belong to a live process sharing this cache directory. - candidates = sorted( - ( - entry - for entry in entries.values() - if not self._paths_for(entry.cache_key).squashfs_mount_dir.is_mount() - ), - key=lambda entry: entry.last_used, - ) - - def within_budget() -> bool: - return (max_entries <= 0 or len(entries) <= max_entries) and ( - max_bytes <= 0 or total_bytes <= max_bytes - ) - - for entry in candidates: - if within_budget(): - break - paths = self._paths_for(entry.cache_key) - try: - trash_path = _move_entry_to_trash( - paths.entry_dir, - self.trash_dir, - entry.cache_key, - ) - except OSError as e: - logger.warning( - "Failed to retire stale registry artifact during startup sweep", - cache_key=entry.cache_key, - entry_dir=str(paths.entry_dir), - error=str(e), - ) - return False - - del entries[entry.cache_key] - if _delete_cache_path(trash_path): - total_bytes -= entry.size_bytes - else: - return False - logger.info( - "Evicted stale registry artifact during startup sweep", - cache_key=entry.cache_key, - size_bytes=entry.size_bytes, - ) - - return within_budget() From 10033b3baa10a057659fd27d7d20584bc6e29b5d Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 15:41:53 -0400 Subject: [PATCH 041/161] test(executor): split registry cache coverage --- tests/unit/executor/conftest.py | 16 + .../registry_artifact_test_helpers.py | 180 + .../executor/test_registry_artifact_budget.py | 817 ++++ .../test_registry_artifact_eviction.py | 293 ++ .../executor/test_registry_artifact_leases.py | 785 ++++ .../test_registry_artifact_materialization.py | 698 ++++ .../test_registry_artifact_resolution.py | 345 ++ .../test_registry_artifact_startup.py | 395 ++ .../test_registry_artifact_tarball.py} | 59 +- tests/unit/test_registry_artifacts.py | 3423 ----------------- 10 files changed, 3546 insertions(+), 3465 deletions(-) create mode 100644 tests/unit/executor/conftest.py create mode 100644 tests/unit/executor/registry_artifact_test_helpers.py create mode 100644 tests/unit/executor/test_registry_artifact_budget.py create mode 100644 tests/unit/executor/test_registry_artifact_eviction.py create mode 100644 tests/unit/executor/test_registry_artifact_leases.py create mode 100644 tests/unit/executor/test_registry_artifact_materialization.py create mode 100644 tests/unit/executor/test_registry_artifact_resolution.py create mode 100644 tests/unit/executor/test_registry_artifact_startup.py rename tests/unit/{test_multitenant_registry.py => executor/test_registry_artifact_tarball.py} (83%) delete mode 100644 tests/unit/test_registry_artifacts.py diff --git a/tests/unit/executor/conftest.py b/tests/unit/executor/conftest.py new file mode 100644 index 0000000000..5a218e3f97 --- /dev/null +++ b/tests/unit/executor/conftest.py @@ -0,0 +1,16 @@ +"""Fixtures for executor registry artifact tests.""" + +from __future__ import annotations + +import tempfile +from collections.abc import Iterator +from pathlib import Path + +import pytest + + +@pytest.fixture +def temp_cache_dir() -> Iterator[Path]: + """Create an isolated registry artifact cache directory.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield Path(tmpdir) diff --git a/tests/unit/executor/registry_artifact_test_helpers.py b/tests/unit/executor/registry_artifact_test_helpers.py new file mode 100644 index 0000000000..062a0d0c71 --- /dev/null +++ b/tests/unit/executor/registry_artifact_test_helpers.py @@ -0,0 +1,180 @@ +"""Shared fixtures and fakes for registry artifact cache tests.""" + +from __future__ import annotations + +import asyncio +import io +import os +import tarfile +from dataclasses import dataclass, field +from pathlib import Path + +from tracecat.executor.registry_artifacts import ( + RegistryArtifactCache, + RegistryArtifactMaterializationContext, + SquashfsMountCommandError, +) + +MAX_ENTRIES_CONFIG = ( + "tracecat.executor.registry_artifacts.config" + ".TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES" +) +MAX_BYTES_CONFIG = ( + "tracecat.executor.registry_artifacts.config" + ".TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES" +) +SQUASHFS_ENABLED_CONFIG = ( + "tracecat.executor.registry_artifacts.config" + ".TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED" +) + + +def write_tarball_entry(cache_dir: Path, cache_key: str) -> Path: + """Create a materialized tarball cache entry on disk.""" + target_dir = cache_dir / "entries" / cache_key / "tarball" + target_dir.mkdir(parents=True) + (target_dir / "module.py").write_text("VALUE = 1") + return target_dir + + +def write_image_entry( + cache_dir: Path, cache_key: str, *, size: int, mtime: float +) -> Path: + """Create a downloaded SquashFS image cache entry with a fixed mtime.""" + entry_dir = cache_dir / "entries" / cache_key + entry_dir.mkdir(parents=True, exist_ok=True) + image_path = entry_dir / "image.squashfs" + image_path.write_bytes(b"x" * size) + os.utime(image_path, (mtime, mtime)) + os.utime(entry_dir, (mtime, mtime)) + return image_path + + +def tarball_payload(*, size: int) -> bytes: + """Return a gzip tarball containing one synthetic regular file.""" + payload = b"x" * size + output = io.BytesIO() + with tarfile.open(fileobj=output, mode="w:gz") as tar: + member = tarfile.TarInfo("module.py") + member.size = len(payload) + tar.addfile(member, io.BytesIO(payload)) + return output.getvalue() + + +@dataclass(slots=True) +class SquashfsMountHarness: + """Model mount ownership without consuming host loop devices.""" + + cache: RegistryArtifactCache + failed_mount_keys: set[str] = field(default_factory=set) + mounted: set[Path] = field(default_factory=set) + mount_attempts: list[str] = field(default_factory=list) + extraction_attempts: list[str] = field(default_factory=list) + unmounts: list[Path] = field(default_factory=list) + + def seed_mount(self, cache_key: str) -> Path: + """Create one reusable image and mark its mount directory active.""" + paths = self.cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True, exist_ok=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir(exist_ok=True) + (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") + self.mounted.add(paths.squashfs_mount_dir) + return paths.squashfs_mount_dir + + async def mount( + self, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> Path: + """Record a mount attempt, failing selected keys like loop exhaustion.""" + self.mount_attempts.append(ctx.cache_key) + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + image_path.write_bytes(b"squashfs") + if ctx.cache_key in self.failed_mount_keys: + raise SquashfsMountCommandError("no free loop device") + ctx.paths.squashfs_mount_dir.mkdir(exist_ok=True) + (ctx.paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") + self.mounted.add(ctx.paths.squashfs_mount_dir) + return ctx.paths.squashfs_mount_dir + + async def extract( + self, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> Path: + """Publish an extracted fallback from the image retained by mount.""" + assert image_path.is_file() + self.extraction_attempts.append(ctx.cache_key) + ctx.paths.squashfs_extract_dir.mkdir(parents=True, exist_ok=True) + (ctx.paths.squashfs_extract_dir / "module.py").write_text("VALUE = 1") + return ctx.paths.squashfs_extract_dir + + async def unmount(self, mount_dir: Path) -> bool: + """Record final-release unmounts and release the modeled loop device.""" + self.unmounts.append(mount_dir) + self.mounted.discard(mount_dir) + return True + + +async def lease_paths( + cache: RegistryArtifactCache, + artifact_uri: str, +) -> list[Path]: + """Return paths from the cache's public lease API.""" + async with cache.lease([artifact_uri]) as paths: + return paths + + +class BlockingSubprocess: + """Fake subprocess that blocks in communicate until it is cancelled.""" + + def __init__(self) -> None: + self.communicate_started = asyncio.Event() + self.cleanup_calls: list[str] = [] + self.returncode: int | None = None + + async def communicate(self) -> tuple[bytes, bytes]: + """Block until the task awaiting subprocess completion is cancelled.""" + self.communicate_started.set() + await asyncio.Event().wait() + return b"", b"" + + def kill(self) -> None: + """Record that the subprocess was killed.""" + self.cleanup_calls.append("kill") + self.returncode = -9 + + async def wait(self) -> int: + """Record that the killed subprocess was reaped.""" + self.cleanup_calls.append("wait") + return -9 + + +class CapturedSubprocess: + """Capture cleanup of a real subprocess used by cancellation tests.""" + + def __init__(self, process: asyncio.subprocess.Process) -> None: + self.process = process + self.killed = False + self.reaped = False + + @property + def returncode(self) -> int | None: + """Return the wrapped subprocess exit status.""" + return self.process.returncode + + async def communicate(self) -> tuple[bytes, bytes]: + """Wait for the wrapped subprocess and collect its output.""" + return await self.process.communicate() + + def kill(self) -> None: + """Kill the wrapped subprocess and record the signal.""" + self.killed = True + self.process.kill() + + async def wait(self) -> int: + """Reap the wrapped subprocess and record completion.""" + returncode = await self.process.wait() + self.reaped = True + return returncode diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py new file mode 100644 index 0000000000..e3438a4a90 --- /dev/null +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -0,0 +1,817 @@ +"""Registry artifact byte admission and budget convergence tests.""" + +from __future__ import annotations + +import asyncio +import os +import threading +from collections.abc import Awaitable, Callable +from pathlib import Path +from unittest.mock import ANY, AsyncMock, patch + +import pytest + +from tracecat.executor.registry_artifacts import ( + RegistryArtifactCache, + RegistryArtifactCacheCapacityError, + RegistryArtifactEviction, + RegistryArtifactMaterializationContext, + SquashfsArtifact, + TarballArtifact, + _delete_cache_path, + compute_registry_artifact_cache_key, +) + +from .registry_artifact_test_helpers import ( + MAX_BYTES_CONFIG, + MAX_ENTRIES_CONFIG, + SQUASHFS_ENABLED_CONFIG, + lease_paths, + tarball_payload, + write_image_entry, + write_tarball_entry, +) + + +class TestRegistryArtifactCacheBudget: + """Enforce peak and steady-state cache capacity.""" + + def test_delete_cache_path_reports_directory_failure(self, temp_cache_dir): + """Physical deletion failures are observable instead of suppressed.""" + entry_dir = temp_cache_dir / "entry" + entry_dir.mkdir() + + with ( + patch( + "tracecat.executor.registry_artifact_storage.shutil.rmtree", + side_effect=OSError("permission denied"), + ), + patch( + "tracecat.executor.registry_artifact_storage.logger.warning" + ) as warning, + ): + deleted = _delete_cache_path(entry_dir) + + assert deleted is False + assert entry_dir.is_dir() + warning.assert_called_once_with( + "Failed to delete registry artifact cache path", + path=str(entry_dir), + error="permission denied", + ) + + @pytest.mark.anyio + async def test_leased_entry_survives_eviction_of_idle_entry(self, temp_cache_dir): + """Eviction must never remove an entry a live action is importing from.""" + cache = RegistryArtifactCache(temp_cache_dir) + leased_uri = "s3://bucket/leased.tar.gz" + idle_uri = "s3://bucket/idle.tar.gz" + new_uri = "s3://bucket/new.tar.gz" + leased_dir = write_tarball_entry( + temp_cache_dir, compute_registry_artifact_cache_key(leased_uri) + ) + idle_dir = write_tarball_entry( + temp_cache_dir, compute_registry_artifact_cache_key(idle_uri) + ) + + async def mock_download(self, ctx, path): + path.write_bytes(tarball_payload(size=1)) + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch(MAX_ENTRIES_CONFIG, 2), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + async with cache.lease([leased_uri]): + await lease_paths(cache, new_uri) + + assert leased_dir.is_dir() + assert not idle_dir.exists() + + @pytest.mark.anyio + async def test_cold_download_reserves_space_before_writing( + self, temp_cache_dir: Path + ) -> None: + """Admission evicts idle bytes before a new download enters staging.""" + cache = RegistryArtifactCache(temp_cache_dir) + idle = write_image_entry(temp_cache_dir, "idle", size=80, mtime=100.0) + artifact_uri = "s3://bucket/new.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + payload = tarball_payload(size=32) + max_bytes = len(payload) + 32 + capacity_checked = False + + async def download_file_to_path( + *, + key: str, + bucket: str, + output_path: Path, + max_bytes: int, + ensure_capacity: Callable[[int], Awaitable[None]], + ) -> int: + del key, bucket + nonlocal capacity_checked + assert max_bytes == len(payload) + 32 + await ensure_capacity(len(payload)) + capacity_checked = True + assert not idle.exists() + output_path.write_bytes(payload) + return len(payload) + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, max_bytes), + patch( + "tracecat.executor.registry_artifacts.blob.download_file_to_path", + side_effect=download_file_to_path, + ), + ): + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [ + cache._paths_for(cache_key).tarball_target_dir + ] + + assert capacity_checked is True + assert not idle.exists() + assert (registry_paths[0] / "module.py").read_bytes() == b"x" * 32 + + @pytest.mark.anyio + async def test_compression_heavy_tarball_is_rejected_before_extraction( + self, temp_cache_dir: Path + ) -> None: + """Compressed bytes plus declared extraction cannot exceed the cache cap.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/compression-heavy.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + payload = tarball_payload(size=4096) + max_bytes = len(payload) + 256 + + async def download_file_to_path( + *, + key: str, + bucket: str, + output_path: Path, + max_bytes: int, + ensure_capacity: Callable[[int], Awaitable[None]], + ) -> int: + del key, bucket + assert max_bytes == len(payload) + 256 + await ensure_capacity(len(payload)) + output_path.write_bytes(payload) + return len(payload) + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, max_bytes), + patch( + "tracecat.executor.registry_artifacts.blob.download_file_to_path", + side_effect=download_file_to_path, + ), + patch.object( + TarballArtifact, + "extract", + new_callable=AsyncMock, + ) as extract, + ): + with pytest.raises(RegistryArtifactCacheCapacityError) as raised: + async with cache.lease([artifact_uri]): + pass + + assert raised.value.additional_bytes == 4096 + assert raised.value.max_bytes == max_bytes + extract.assert_not_awaited() + assert not cache._paths_for(cache_key).entry_dir.exists() + assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) + + @pytest.mark.anyio + async def test_squashfs_expansion_is_rejected_before_extraction( + self, temp_cache_dir: Path + ) -> None: + """SquashFS metadata is accounted before unsquashfs writes scratch.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/oversized.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def download( + self: SquashfsArtifact, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> float: + del self, ctx + image_path.parent.mkdir(parents=True, exist_ok=True) + image_path.write_bytes(b"image") + return 0.0 + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 100), + patch.object( + RegistryArtifactMaterializationContext, + "can_mount_squashfs", + return_value=False, + ), + patch.object(SquashfsArtifact, "download", download), + patch.object( + SquashfsArtifact, + "_squashfs_extracted_size", + new_callable=AsyncMock, + return_value=101, + ), + patch.object( + SquashfsArtifact, + "_extract_image", + new_callable=AsyncMock, + ) as extract_image, + ): + with pytest.raises(RegistryArtifactCacheCapacityError) as raised: + async with cache.lease([artifact_uri]): + pass + + assert raised.value.additional_bytes == 101 + extract_image.assert_not_awaited() + paths = cache._paths_for(cache_key) + assert paths.squashfs_image_path.read_bytes() == b"image" + assert not paths.squashfs_extract_dir.exists() + + @pytest.mark.anyio + async def test_successful_admission_enforces_actual_size_before_yield( + self, temp_cache_dir + ): + """Post-publication enforcement sees the new entry's actual size.""" + cache = RegistryArtifactCache(temp_cache_dir) + idle = write_image_entry(temp_cache_dir, "idle", size=4096, mtime=100.0) + new_uri = "s3://bucket/new.tar.gz" + new_key = compute_registry_artifact_cache_key(new_uri) + + async def mock_download(self, ctx, path): + path.write_bytes(tarball_payload(size=1)) + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_bytes(b"x" * 4096) + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 6000), + patch( + "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", + return_value=4096, + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + async with cache.lease([new_uri]) as registry_paths: + assert registry_paths == [cache._paths_for(new_key).tarball_target_dir] + assert not idle.exists() + + assert not idle.exists() + assert cache._paths_for(new_key).tarball_target_dir.is_dir() + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_deletion_failure_does_not_block_materialization( + self, temp_cache_dir + ): + """Failed cleanup is reported while cache admission remains fail-open.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + write_image_entry(temp_cache_dir, "idle", size=4096, mtime=100.0) + artifact_uri = "s3://bucket/new.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def mock_download(self, ctx, path): + path.write_bytes(tarball_payload(size=1)) + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifact_storage._delete_cache_path", + return_value=False, + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + patch( + "tracecat.executor.registry_artifact_storage.logger.warning" + ) as warning, + ): + registry_paths = await lease_paths(cache, artifact_uri) + + assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] + assert registry_paths[0].is_dir() + assert cache._budget_dirty is True + warning.assert_any_call( + "Registry artifact eviction remains pending physical deletion", + cache_key="idle", + trash_path=ANY, + ) + + @pytest.mark.anyio + async def test_failed_physical_delete_retries_without_extra_eviction( + self, temp_cache_dir + ): + """Failed byte reclamation stops eviction until trash cleanup succeeds.""" + cache = RegistryArtifactCache(temp_cache_dir) + oldest = write_image_entry( + temp_cache_dir, + "oldest", + size=16, + mtime=100.0, + ) + older = write_image_entry( + temp_cache_dir, + "older", + size=16, + mtime=200.0, + ) + newest = write_image_entry( + temp_cache_dir, + "newest", + size=16, + mtime=300.0, + ) + real_delete = _delete_cache_path + failed_once = False + + def fail_once(path: Path) -> bool: + nonlocal failed_once + if not failed_once: + failed_once = True + return False + return real_delete(path) + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 16), + patch( + "tracecat.executor.registry_artifact_storage._delete_cache_path", + side_effect=fail_once, + ), + ): + assert await cache._enforce_cache_budget() is False + assert not oldest.exists() + assert older.exists() + assert newest.exists() + assert cache._discover_cache_keys() == {"older", "newest"} + assert len(tuple(cache.trash_dir.iterdir())) == 1 + + assert await cache._enforce_cache_budget() is True + + assert not older.exists() + assert newest.exists() + assert not any(cache.trash_dir.iterdir()) + + @pytest.mark.anyio + async def test_rename_failure_does_not_block_materialization(self, temp_cache_dir): + """A failed atomic retirement keeps the old entry and admits new work.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + idle = write_image_entry(temp_cache_dir, "idle", size=16, mtime=100.0) + artifact_uri = "s3://bucket/new-after-rename-failure.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def mock_download(self, ctx, path): + path.write_bytes(tarball_payload(size=1)) + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifact_storage._move_entry_to_trash", + side_effect=OSError("rename failed"), + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + registry_paths = await lease_paths(cache, artifact_uri) + + assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] + assert registry_paths[0].is_dir() + assert idle.is_file() + assert cache._budget_dirty is True + + @pytest.mark.anyio + async def test_cache_scan_failure_does_not_block_materialization( + self, temp_cache_dir + ): + """Maintenance errors are observable while artifact admission stays open.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/new-after-scan-failure.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def mock_download(self, ctx, path): + path.write_bytes(tarball_payload(size=1)) + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch.object( + cache, + "_enforce_cache_budget", + new_callable=AsyncMock, + side_effect=PermissionError("denied"), + ), + patch( + "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", + return_value=9, + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [ + cache._paths_for(cache_key).tarball_target_dir + ] + + @pytest.mark.anyio + async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( + self, temp_cache_dir + ): + """Steady-state cache hits must not pay for a full cache scan.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/cached.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + write_tarball_entry(temp_cache_dir, cache_key) + cache._budget_dirty = False + + with patch.object( + cache, + "_scan_cache_entries", + side_effect=AssertionError("cache hits must not scan the cache dir"), + ): + async with cache.lease([artifact_uri]): + pass + + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_failed_cold_admission_keeps_warm_lru(self, temp_cache_dir): + """A missing cold artifact must not evict an existing warm entry.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + warm_uri = "s3://bucket/warm.tar.gz" + warm_key = compute_registry_artifact_cache_key(warm_uri) + warm_dir = write_tarball_entry(temp_cache_dir, warm_key) + missing_uri = "s3://bucket/missing.tar.gz" + missing_key = compute_registry_artifact_cache_key(missing_uri) + + async def mock_download(self, ctx, path): + raise FileNotFoundError("missing artifact") + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "download", mock_download), + ): + with pytest.raises(FileNotFoundError, match="missing artifact"): + async with cache.lease([missing_uri]): + pytest.fail("failed admission must not yield a lease") + + assert warm_dir.is_dir() + assert not cache._paths_for(missing_key).entry_dir.exists() + + @pytest.mark.anyio + async def test_materialization_rearms_budget_dirty_consumed_mid_flight( + self, temp_cache_dir + ): + """A convergence pass consuming the signal mid-download cannot unarm it.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + assert cache._budget_dirty is False + + artifact_uri = "s3://bucket/path/site-packages.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) + + async def mock_materialize(self, ctx): + # A concurrent lease release consumes the dirty signal and finishes + # its scan before this attempt lands its canonical bytes. + cache._budget_dirty = False + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + ctx.paths.squashfs_image_path.write_bytes(b"late image") + return [ctx.paths.squashfs_mount_dir] + + with ( + patch.object( + cache, + "_artifact_candidates", + new_callable=AsyncMock, + return_value=[artifact], + ), + patch.object(SquashfsArtifact, "materialize", mock_materialize), + ): + async with cache.lease([artifact_uri]): + assert cache._budget_dirty is True + + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_release_keeps_retrying_while_the_cache_stays_over_budget( + self, temp_cache_dir + ): + """A cache that cannot shrink yet must stay marked for re-enforcement.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/pinned.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + write_image_entry(temp_cache_dir, cache_key, size=4096, mtime=100.0) + write_tarball_entry(temp_cache_dir, cache_key) + cache._budget_dirty = True + # A second holder keeps the entry pinned past the inner lease. + cache._acquire_lease(cache_key) + + with patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 1): + async with cache.lease([artifact_uri]): + pass + + assert cache._budget_dirty is True + + @pytest.mark.anyio + async def test_convergence_rescans_after_concurrent_materialization( + self, temp_cache_dir + ): + """A materialization during a budget scan must schedule another scan.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/concurrent.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + scan_started = asyncio.Event() + materialized = asyncio.Event() + convergence_scans = 0 + + async def mock_enforce_cache_budget( + *, protected_key: str | None = None + ) -> bool: + nonlocal convergence_scans + if protected_key is not None: + return True + + convergence_scans += 1 + if convergence_scans == 1: + scan_started.set() + await materialized.wait() + return True + + async def mock_download(self, ctx, path): + path.write_bytes(tarball_payload(size=1)) + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 1") + + cache._budget_dirty = True + with ( + patch.object( + cache, + "_enforce_cache_budget", + side_effect=mock_enforce_cache_budget, + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + convergence = asyncio.create_task(cache._converge_cache_budget()) + await scan_started.wait() + registry_paths = await lease_paths(cache, artifact_uri) + materialized.set() + await convergence + + assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] + assert convergence_scans == 2 + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_cancelled_convergence_rearms_budget_dirty( + self, temp_cache_dir + ) -> None: + """Cancellation must preserve the need for a later budget pass.""" + cache = RegistryArtifactCache(temp_cache_dir) + enforcement_started = asyncio.Event() + finish_enforcement = asyncio.Event() + + async def mock_enforce_cache_budget( + *, protected_key: str | None = None + ) -> bool: + del protected_key + enforcement_started.set() + await finish_enforcement.wait() + return True + + cache._budget_dirty = True + with patch.object( + cache, + "_enforce_cache_budget", + side_effect=mock_enforce_cache_budget, + ): + convergence = asyncio.create_task(cache._converge_cache_budget()) + await enforcement_started.wait() + convergence.cancel() + + with pytest.raises(asyncio.CancelledError): + await convergence + + assert cache._budget_dirty is True + + @pytest.mark.parametrize( + "oldest_has_tarball", + [False, True], + ids=["squashfs-only", "tarball-bearing"], + ) + @pytest.mark.anyio + async def test_enforce_budget_evicts_least_recently_used_until_under_max_bytes( + self, temp_cache_dir, oldest_has_tarball: bool + ): + """Size eviction stops once the cache is within budget.""" + cache = RegistryArtifactCache(temp_cache_dir) + oldest = write_image_entry(temp_cache_dir, "oldest", size=4096, mtime=100.0) + if oldest_has_tarball: + write_tarball_entry(temp_cache_dir, "oldest") + oldest_entry = cache._paths_for("oldest").entry_dir + os.utime(oldest_entry, (100.0, 100.0)) + older = write_image_entry(temp_cache_dir, "older", size=4096, mtime=200.0) + newest = write_image_entry(temp_cache_dir, "newest", size=4096, mtime=300.0) + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 9000), + ): + within_budget = await cache._enforce_cache_budget(protected_key="pending") + + assert within_budget is True + assert not oldest.exists() + assert not cache._paths_for("oldest").tarball_target_dir.exists() + assert older.exists() + assert newest.exists() + + @pytest.mark.anyio + async def test_final_lease_release_unmounts_and_retains_image(self, temp_cache_dir): + """An idle entry releases its loop device without deleting its image.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/site-packages.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + mounted = {paths.squashfs_mount_dir} + + process = AsyncMock() + process.communicate.return_value = (b"", b"") + process.returncode = 0 + + async def mock_umount(*args, **kwargs): + mounted.discard(paths.squashfs_mount_dir) + return process + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/sbin/umount", + ), + patch.object( + asyncio, + "create_subprocess_exec", + side_effect=mock_umount, + ), + ): + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [paths.squashfs_mount_dir] + assert paths.squashfs_mount_dir in mounted + + assert paths.squashfs_mount_dir not in mounted + + assert paths.squashfs_image_path.read_bytes() == b"squashfs" + assert paths.squashfs_mount_dir.is_dir() + + @pytest.mark.anyio + async def test_concurrent_budget_passes_only_evict_once(self, temp_cache_dir): + """A waiting budget pass must re-scan after the active pass evicts.""" + cache = RegistryArtifactCache(temp_cache_dir) + oldest = write_image_entry(temp_cache_dir, "oldest", size=16, mtime=100.0) + retained = write_image_entry(temp_cache_dir, "retained", size=16, mtime=200.0) + original_scan = cache._scan_cache_entries + first_scan_started = threading.Event() + second_scan_started = threading.Event() + release_first_scan = threading.Event() + release_second_scan = threading.Event() + scan_count_lock = threading.Lock() + scan_count = 0 + + def controlled_scan(): + nonlocal scan_count + entries = original_scan() + with scan_count_lock: + scan_index = scan_count + scan_count += 1 + if scan_index == 0: + first_scan_started.set() + release_first_scan.wait(timeout=5) + elif scan_index == 1: + second_scan_started.set() + release_second_scan.wait(timeout=5) + return entries + + eviction_started = asyncio.Event() + finish_eviction = asyncio.Event() + extra_eviction_finished = asyncio.Event() + evicted_keys: list[str] = [] + + async def controlled_evict( + cache_key: str, + ) -> RegistryArtifactEviction: + if cache_key == "oldest": + if eviction_started.is_set(): + return RegistryArtifactEviction( + retired=False, + reclaimed=False, + ) + eviction_started.set() + await finish_eviction.wait() + _delete_cache_path(cache._paths_for("oldest").entry_dir) + else: + _delete_cache_path(cache._paths_for("retained").entry_dir) + extra_eviction_finished.set() + evicted_keys.append(cache_key) + return RegistryArtifactEviction(retired=True, reclaimed=True) + + with ( + patch.object(cache, "_scan_cache_entries", side_effect=controlled_scan), + patch.object(cache, "_evict_entry", side_effect=controlled_evict), + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + first_pass = asyncio.create_task(cache._enforce_cache_budget()) + assert await asyncio.to_thread(first_scan_started.wait, 1) + second_pass = asyncio.create_task(cache._enforce_cache_budget()) + second_scan_overlapped = await asyncio.to_thread( + second_scan_started.wait, 0.5 + ) + + release_first_scan.set() + await asyncio.wait_for(eviction_started.wait(), timeout=1) + release_second_scan.set() + try: + await asyncio.wait_for(extra_eviction_finished.wait(), timeout=0.2) + except TimeoutError: + pass + finish_eviction.set() + await asyncio.gather(first_pass, second_pass) + + assert second_scan_overlapped is False + assert scan_count == 2 + assert evicted_keys == ["oldest"] + assert not oldest.exists() + assert retained.exists() + + @pytest.mark.anyio + async def test_enforce_budget_ignores_missing_protected_key(self, temp_cache_dir): + """Only successfully published entries count against the entry budget.""" + cache = RegistryArtifactCache(temp_cache_dir) + existing = write_image_entry(temp_cache_dir, "existing", size=16, mtime=100.0) + + with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): + within_budget = await cache._enforce_cache_budget(protected_key="missing") + + assert within_budget is True + assert existing.exists() + + @pytest.mark.anyio + async def test_enforce_budget_never_evicts_the_protected_key(self, temp_cache_dir): + """The newly materialized key is exempt even when it is the LRU entry.""" + cache = RegistryArtifactCache(temp_cache_dir) + protected = write_image_entry( + temp_cache_dir, "protected", size=4096, mtime=100.0 + ) + other = write_image_entry(temp_cache_dir, "other", size=4096, mtime=300.0) + + with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): + await cache._enforce_cache_budget(protected_key="protected") + + assert protected.exists() + assert not other.exists() + + @pytest.mark.anyio + async def test_enforce_budget_proceeds_over_budget_when_everything_is_leased( + self, temp_cache_dir + ): + """An over-budget cache must degrade, not fail the action.""" + cache = RegistryArtifactCache(temp_cache_dir) + leased = write_image_entry(temp_cache_dir, "leased", size=4096, mtime=100.0) + cache._acquire_lease("leased") + + with patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 1): + within_budget = await cache._enforce_cache_budget(protected_key="missing") + + assert within_budget is False + assert leased.exists() diff --git a/tests/unit/executor/test_registry_artifact_eviction.py b/tests/unit/executor/test_registry_artifact_eviction.py new file mode 100644 index 0000000000..d4d2108ad7 --- /dev/null +++ b/tests/unit/executor/test_registry_artifact_eviction.py @@ -0,0 +1,293 @@ +"""Registry artifact retirement and unmount tests.""" + +from __future__ import annotations + +import asyncio +import os +import threading +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import pytest + +from tracecat.executor.registry_artifacts import ( + RegistryArtifactCache, + RegistryArtifactEviction, + TarballArtifact, + _delete_cache_path, + compute_registry_artifact_cache_key, +) + +from .registry_artifact_test_helpers import ( + MAX_BYTES_CONFIG, + MAX_ENTRIES_CONFIG, + BlockingSubprocess, + tarball_payload, + write_image_entry, + write_tarball_entry, +) + + +class TestRegistryArtifactCacheEviction: + """Retire idle cache entries without disrupting live leases.""" + + @pytest.mark.anyio + async def test_eviction_unmounts_before_deleting_the_image(self, temp_cache_dir): + """Unlinking a mounted image would strand an open-file zombie.""" + cache = RegistryArtifactCache(temp_cache_dir) + paths = cache._paths_for("mounted") + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + mounted = {paths.squashfs_mount_dir} + image_present_at_umount: list[bool] = [] + + process = AsyncMock() + process.communicate.return_value = (b"", b"") + process.returncode = 0 + + async def mock_umount(*args, **kwargs): + image_present_at_umount.append(paths.squashfs_image_path.exists()) + mounted.discard(paths.squashfs_mount_dir) + return process + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/sbin/umount", + ), + patch( + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", + side_effect=mock_umount, + ) as create_subprocess_exec, + ): + evicted = await cache._evict_entry("mounted") + + assert evicted == RegistryArtifactEviction(retired=True, reclaimed=True) + assert image_present_at_umount == [True] + assert not paths.squashfs_image_path.exists() + assert not paths.squashfs_mount_dir.exists() + create_subprocess_exec.assert_called_once() + assert create_subprocess_exec.call_args.args == ( + "/sbin/umount", + str(paths.squashfs_mount_dir), + ) + + @pytest.mark.anyio + async def test_cancelled_unmount_kills_and_reaps_before_releasing_key_lock( + self, temp_cache_dir + ): + """Cancellation leaves a consistent entry for the next admission.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/cancelled-unmount.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") + mounted = {paths.squashfs_mount_dir} + blocked_process = BlockingSubprocess() + released_process = AsyncMock() + released_process.communicate.return_value = (b"", b"") + released_process.returncode = 0 + unmount_attempts = 0 + + async def mock_umount(*args, **kwargs): + nonlocal unmount_attempts + unmount_attempts += 1 + if unmount_attempts == 1: + return blocked_process + mounted.discard(paths.squashfs_mount_dir) + return released_process + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/sbin/umount", + ), + patch( + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", + side_effect=mock_umount, + ), + ): + eviction = asyncio.create_task(cache._evict_entry(cache_key)) + await blocked_process.communicate_started.wait() + eviction.cancel() + + with pytest.raises(asyncio.CancelledError): + await eviction + + assert blocked_process.cleanup_calls == ["kill", "wait"] + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [paths.squashfs_mount_dir] + assert registry_paths[0].is_dir() + assert (registry_paths[0] / "module.py").read_text() == "VALUE = 1" + + assert paths.squashfs_image_path.is_file() + assert paths.squashfs_mount_dir.is_dir() + assert paths.squashfs_mount_dir not in mounted + + @pytest.mark.anyio + async def test_cancelled_background_deletion_leaves_a_clean_miss( + self, temp_cache_dir + ): + """Cancellation cannot expose live paths that deletion still owns.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/cancelled-eviction.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + original_target = write_tarball_entry(temp_cache_dir, cache_key) + delete_started = threading.Event() + finish_delete = threading.Event() + delete_finished = threading.Event() + doomed: list[Path] = [] + + def blocked_delete(path: Path) -> bool: + doomed.append(path) + delete_started.set() + finish_delete.wait(timeout=5) + deleted = _delete_cache_path(path) + delete_finished.set() + return deleted + + async def mock_download(self, ctx, path): + path.write_bytes(tarball_payload(size=1)) + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch( + "tracecat.executor.registry_artifact_storage._delete_cache_path", + side_effect=blocked_delete, + ), + patch( + "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", + return_value=9, + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + eviction = asyncio.create_task(cache._evict_entry(cache_key)) + assert await asyncio.to_thread(delete_started.wait, 1) + eviction.cancel() + try: + await asyncio.sleep(0) + eviction.cancel() + await asyncio.sleep(0) + assert not eviction.done() + assert not original_target.exists() + assert (doomed[0] / "tarball").is_dir() + assert cache._discover_cache_keys() == set() + finally: + finish_delete.set() + + with pytest.raises(asyncio.CancelledError): + await eviction + assert await asyncio.to_thread(delete_finished.wait, 1) + assert not doomed[0].exists() + + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [original_target] + assert (original_target / "module.py").read_text() == "VALUE = 2" + + assert original_target.is_dir() + + @pytest.mark.anyio + async def test_doomed_eviction_names_are_startup_scratch(self, temp_cache_dir): + """Every renamed entry root is invisible and reclaimed on startup.""" + cache = RegistryArtifactCache(temp_cache_dir) + cache_key = "squashfs-doomed" + paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + paths.squashfs_extract_dir.mkdir() + paths.tarball_target_dir.mkdir() + + with patch( + "tracecat.executor.registry_artifact_storage._delete_cache_path", + return_value=True, + ) as delete_cache_path: + assert await cache._evict_entry(cache_key) == RegistryArtifactEviction( + retired=True, reclaimed=True + ) + + delete_cache_path.assert_called_once() + trash_path = delete_cache_path.call_args.args[0] + assert trash_path.parent == cache.trash_dir + assert trash_path.exists() + assert cache._discover_cache_keys() == set() + + cache._sweep_startup_state() + + assert not trash_path.exists() + + @pytest.mark.anyio + async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): + """A failed unmount skips the key instead of forcing a lazy detach.""" + cache = RegistryArtifactCache(temp_cache_dir) + stuck = cache._paths_for("stuck") + stuck.entry_dir.mkdir(parents=True) + stuck.squashfs_image_path.write_bytes(b"squashfs") + stuck.squashfs_mount_dir.mkdir() + os.utime(stuck.squashfs_image_path, (100.0, 100.0)) + idle = write_image_entry(temp_cache_dir, "idle", size=16, mtime=300.0) + mounted = {stuck.squashfs_mount_dir} + + process = AsyncMock() + process.communicate.return_value = (b"", b"target is busy") + process.returncode = 32 + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/sbin/umount", + ), + patch( + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ), + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + await cache._enforce_cache_budget(protected_key="pending") + + assert stuck.squashfs_image_path.exists() + assert stuck.squashfs_mount_dir.exists() + assert not idle.exists() + + @pytest.mark.anyio + async def test_eviction_keeps_stable_runtime_state(self, temp_cache_dir): + """A key keeps one lock and zeroed lease state for the process lifetime.""" + cache = RegistryArtifactCache(temp_cache_dir) + write_tarball_entry(temp_cache_dir, "bookkeeping") + cache._acquire_lease("bookkeeping") + cache._release_lease("bookkeeping") + lock = cache._runtime_for("bookkeeping").lock + + assert await cache._evict_entry("bookkeeping") == RegistryArtifactEviction( + retired=True, reclaimed=True + ) + runtime = cache._runtime["bookkeeping"] + assert runtime.lock is lock + assert runtime.refcount == 0 + + @pytest.mark.anyio + async def test_eviction_skips_busy_key(self, temp_cache_dir): + """A key another task is materializing is never evicted underneath it.""" + cache = RegistryArtifactCache(temp_cache_dir) + target_dir = write_tarball_entry(temp_cache_dir, "busy") + lock = cache._runtime_for("busy").lock + + async with lock: + assert await cache._evict_entry("busy") == RegistryArtifactEviction( + retired=False, reclaimed=False + ) + + assert target_dir.is_dir() diff --git a/tests/unit/executor/test_registry_artifact_leases.py b/tests/unit/executor/test_registry_artifact_leases.py new file mode 100644 index 0000000000..4c5c188470 --- /dev/null +++ b/tests/unit/executor/test_registry_artifact_leases.py @@ -0,0 +1,785 @@ +"""Registry artifact lease lifetime and concurrency tests.""" + +from __future__ import annotations + +import asyncio +import os +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import pytest +import tracecat_registry + +from tracecat.executor.registry_artifacts import ( + RegistryArtifactCache, + RegistryArtifactCacheLoopError, + RegistryArtifactEviction, + RegistryArtifactMaterializationContext, + SquashfsArtifact, + TarballArtifact, + bundled_builtin_registry_uri, + compute_registry_artifact_cache_key, +) + +from .registry_artifact_test_helpers import ( + MAX_BYTES_CONFIG, + MAX_ENTRIES_CONFIG, + SQUASHFS_ENABLED_CONFIG, + SquashfsMountHarness, + write_image_entry, + write_tarball_entry, +) + + +class TestRegistryArtifactCacheLease: + """Tests for lease-based pinning of registry artifact cache entries.""" + + @pytest.mark.anyio + async def test_cache_rejects_use_from_a_second_event_loop( + self, temp_cache_dir: Path + ) -> None: + """Protect the process-wide cache from thread-local Temporal loops. + + The cache owns asyncio locks, tasks, and refcounts as one unit. A typed + ownership failure prevents a future synchronous activity from silently + sharing only part of that state across thread-local event loops, which + previously caused intermittent cross-loop failures in storage caches. + """ + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + + def use_cache_from_another_thread() -> None: + asyncio.run(cache.ensure_swept()) + + with pytest.raises(RegistryArtifactCacheLoopError): + await asyncio.to_thread(use_cache_from_another_thread) + + # A rejected caller must not poison the owning loop's cache. + async with cache.lease(None) as registry_paths: + assert registry_paths == [temp_cache_dir / "base"] + + def test_touch_entry_refreshes_tarball_root_mtime(self, temp_cache_dir): + """Touching a tarball-only entry persists its restart-safe recency.""" + cache = RegistryArtifactCache(temp_cache_dir) + cache_key = "tarball-only" + write_tarball_entry(temp_cache_dir, cache_key) + entry_dir = cache._paths_for(cache_key).entry_dir + os.utime(entry_dir, (100.0, 100.0)) + + cache._touch_entry(cache_key) + + assert entry_dir.stat().st_mtime > 100.0 + + @pytest.mark.anyio + async def test_lease_refcounts_and_touches_image_mtime(self, temp_cache_dir): + """A lease pins its entry and refreshes the restart-safe LRU timestamp.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/leased.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + target_dir = write_tarball_entry(temp_cache_dir, cache_key) + image_path = write_image_entry(temp_cache_dir, cache_key, size=16, mtime=100.0) + entry_dir = cache._paths_for(cache_key).entry_dir + + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [target_dir] + assert cache._refcount(cache_key) == 1 + assert entry_dir.stat().st_mtime > 100.0 + + assert cache._refcount(cache_key) == 0 + assert image_path.is_file() + + @pytest.mark.anyio + async def test_lease_releases_refcount_when_materialization_fails( + self, temp_cache_dir + ): + """A failed materialization must not leak a pin or empty cache entry.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/broken.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def mock_download(self, ctx, path): + raise RuntimeError("download failed") + + with patch.object(TarballArtifact, "download", mock_download): + with pytest.raises(RuntimeError, match="download failed"): + async with cache.lease([artifact_uri]): + pass + + assert cache._refcount(cache_key) == 0 + assert not cache._paths_for(cache_key).entry_dir.exists() + assert cache_key not in cache._discover_cache_keys() + + @pytest.mark.anyio + async def test_failed_first_admission_converges_deposited_image( + self, temp_cache_dir: Path + ) -> None: + """A failed first artifact must not strand an over-budget image.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/failed-first.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) + + async def fail_after_deposit( + artifact: SquashfsArtifact, + ctx: RegistryArtifactMaterializationContext, + ) -> list[Path]: + del artifact + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + ctx.paths.squashfs_image_path.write_bytes(b"reusable image") + raise RuntimeError("mount failed") + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 1), + patch.object( + cache, + "_artifact_candidates", + new_callable=AsyncMock, + return_value=[artifact], + ), + patch.object(SquashfsArtifact, "materialize", fail_after_deposit), + ): + with pytest.raises(RuntimeError, match="mount failed"): + async with cache.lease([artifact_uri]): + pass + + assert cache._refcount(cache_key) == 0 + assert not cache._paths_for(cache_key).entry_dir.exists() + assert cache._discover_cache_keys() == set() + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_cancelled_lease_admission_releases_refcount(self, temp_cache_dir): + """Cancellation during candidate lookup must not leak a permanent pin.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/site-packages.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + lookup_started = asyncio.Event() + finish_lookup = asyncio.Event() + + async def blocked_sidecar_lookup(**kwargs): + lookup_started.set() + await finish_lookup.wait() + return False + + async def take_lease() -> None: + async with cache.lease([artifact_uri]): + pass + + with ( + patch(SQUASHFS_ENABLED_CONFIG, True), + patch.object(cache, "_sidecar_exists", blocked_sidecar_lookup), + ): + acquisition = asyncio.create_task(take_lease()) + await lookup_started.wait() + assert cache._refcount(cache_key) == 1 + acquisition.cancel() + with pytest.raises(asyncio.CancelledError): + await acquisition + + assert cache._refcount(cache_key) == 0 + assert not cache._paths_for(cache_key).entry_dir.exists() + + @pytest.mark.anyio + async def test_cancelled_waiter_preserves_existing_same_key_lease( + self, temp_cache_dir: Path + ) -> None: + """A waiter cancelled before admission cannot release another holder.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/path/shared.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + lock = cache._runtime_for(cache_key).lock + + async def take_lease() -> None: + async with cache.lease([artifact_uri]): + pytest.fail("cancelled waiter must not enter the lease context") + + unmount_idle_entry = AsyncMock() + with patch.object(cache, "_unmount_idle_entry", unmount_idle_entry): + async with lock: + cache._acquire_lease(cache_key) + try: + waiter = asyncio.create_task(take_lease()) + await asyncio.sleep(0) + assert not waiter.done() + + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + + assert cache._refcount(cache_key) == 1 + unmount_idle_entry.assert_not_awaited() + finally: + cache._release_lease(cache_key) + + @pytest.mark.anyio + async def test_lease_without_uris_returns_base_pythonpath_dir(self, temp_cache_dir): + """No artifact URIs still yields the base PYTHONPATH directory.""" + cache = RegistryArtifactCache(temp_cache_dir) + + async with cache.lease(None) as registry_paths: + assert registry_paths == [temp_cache_dir / "base"] + assert registry_paths[0].is_dir() + + assert cache._runtime == {} + + @pytest.mark.anyio + async def test_lease_preserves_uri_order(self, temp_cache_dir): + """Multiple artifacts keep their deterministic PYTHONPATH order.""" + cache = RegistryArtifactCache(temp_cache_dir) + uris = ["s3://bucket/first.tar.gz", "s3://bucket/second.tar.gz"] + expected = [ + write_tarball_entry( + temp_cache_dir, compute_registry_artifact_cache_key(uri) + ) + for uri in uris + ] + + async with cache.lease(uris) as registry_paths: + assert registry_paths == expected + + @pytest.mark.anyio + async def test_overlapping_same_key_leases_share_one_mount_until_final_release( + self, temp_cache_dir: Path + ) -> None: + """Protect refcount and loop-device lifetime under heavy fan-in. + + Every holder must observe one published mount, intermediate releases + must leave that mount usable, and exactly the final release may reclaim + its loop device while retaining the downloaded image for a remount. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/shared.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + harness = SquashfsMountHarness(cache) + holder_count = 32 + entered = 0 + all_entered = asyncio.Event() + releases = [asyncio.Event() for _ in range(holder_count)] + + async def hold_lease(index: int) -> None: + nonlocal entered + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [ + cache._paths_for(cache_key).squashfs_mount_dir + ] + entered += 1 + if entered == holder_count: + all_entered.set() + await releases[index].wait() + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", harness.mount), + patch.object(SquashfsArtifact, "extract", harness.extract), + patch.object(cache, "_unmount", harness.unmount), + ): + holders = [ + asyncio.create_task(hold_lease(index)) for index in range(holder_count) + ] + await asyncio.wait_for(all_entered.wait(), timeout=5) + + mount_dir = cache._paths_for(cache_key).squashfs_mount_dir + assert cache._refcount(cache_key) == holder_count + assert harness.mount_attempts == [cache_key] + assert harness.extraction_attempts == [] + assert mount_dir in harness.mounted + + for release in releases[:-1]: + release.set() + await asyncio.gather(*holders[:-1]) + + assert cache._refcount(cache_key) == 1 + assert mount_dir in harness.mounted + assert harness.unmounts == [] + + releases[-1].set() + await holders[-1] + + assert cache._refcount(cache_key) == 0 + assert harness.unmounts == [mount_dir] + assert mount_dir not in harness.mounted + assert cache._paths_for(cache_key).squashfs_image_path.is_file() + assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) + assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) + + @pytest.mark.anyio + async def test_cancelling_one_holder_preserves_sibling_leases( + self, temp_cache_dir: Path + ) -> None: + """Protect sibling actions from another holder's cancellation. + + Cancellation must release exactly one pin without unmounting the shared + artifact beneath surviving actions; the last surviving holder remains + solely responsible for loop-device reclamation. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/cancel-one.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + harness = SquashfsMountHarness(cache) + entered = [asyncio.Event() for _ in range(3)] + releases = [asyncio.Event() for _ in range(3)] + + async def hold_lease(index: int) -> None: + async with cache.lease([artifact_uri]): + entered[index].set() + await releases[index].wait() + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", harness.mount), + patch.object(SquashfsArtifact, "extract", harness.extract), + patch.object(cache, "_unmount", harness.unmount), + ): + holders = [asyncio.create_task(hold_lease(index)) for index in range(3)] + await asyncio.wait_for( + asyncio.gather(*(event.wait() for event in entered)), + timeout=5, + ) + + holders[0].cancel() + with pytest.raises(asyncio.CancelledError): + await holders[0] + + mount_dir = cache._paths_for(cache_key).squashfs_mount_dir + assert cache._refcount(cache_key) == 2 + assert mount_dir in harness.mounted + assert harness.unmounts == [] + + releases[1].set() + await holders[1] + assert cache._refcount(cache_key) == 1 + assert mount_dir in harness.mounted + assert harness.unmounts == [] + + releases[2].set() + await holders[2] + + assert cache._refcount(cache_key) == 0 + assert harness.mount_attempts == [cache_key] + assert harness.unmounts == [mount_dir] + + @pytest.mark.anyio + async def test_repeated_cancellation_finishes_all_lease_cleanup( + self, temp_cache_dir: Path + ) -> None: + """Every unmount and budget pass finishes before cancellation propagates.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uris = [ + "s3://bucket/first.tar.gz", + "s3://bucket/second.tar.gz", + ] + cache_keys = [compute_registry_artifact_cache_key(uri) for uri in artifact_uris] + for cache_key in cache_keys: + write_tarball_entry(temp_cache_dir, cache_key) + + lease_entered = asyncio.Event() + first_unmount_started = asyncio.Event() + finish_first_unmount = asyncio.Event() + cleanup_calls: list[str] = [] + + async def mock_unmount_idle_entry(cache_key: str) -> None: + cleanup_calls.append(cache_key) + if cache_key == cache_keys[0]: + first_unmount_started.set() + await finish_first_unmount.wait() + + async def mock_converge_cache_budget() -> None: + cleanup_calls.append("converge") + + async def hold_lease() -> None: + async with cache.lease(artifact_uris): + lease_entered.set() + await asyncio.Event().wait() + + with ( + patch.object(cache, "_unmount_idle_entry", mock_unmount_idle_entry), + patch.object(cache, "_converge_cache_budget", mock_converge_cache_budget), + ): + holder = asyncio.create_task(hold_lease()) + await lease_entered.wait() + holder.cancel() + await first_unmount_started.wait() + + holder.cancel() + await asyncio.sleep(0) + assert not holder.done() + + finish_first_unmount.set() + with pytest.raises(asyncio.CancelledError): + await holder + + assert cleanup_calls == [*cache_keys, "converge"] + assert all(cache._refcount(cache_key) == 0 for cache_key in cache_keys) + + @pytest.mark.anyio + async def test_new_lease_racing_final_release_prevents_stale_unmount( + self, temp_cache_dir: Path + ) -> None: + """Protect the zero-refcount-to-unmount handoff from a new acquisition. + + A lease can arrive after the old holder decrements to zero but before + unmount takes the key lock. The lock-time refcount recheck must preserve + the mount for that newcomer instead of returning a path being torn down. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/release-race.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + harness = SquashfsMountHarness(cache) + mount_dir = harness.seed_mount(cache_key) + first_entered = asyncio.Event() + release_first = asyncio.Event() + release_reached_unmount = asyncio.Event() + allow_unmount_recheck = asyncio.Event() + newcomer_entered = asyncio.Event() + release_newcomer = asyncio.Event() + original_unmount_idle_entry = cache._unmount_idle_entry + unmount_requests = 0 + + async def pause_first_unmount_request(requested_key: str) -> None: + nonlocal unmount_requests + unmount_requests += 1 + if unmount_requests == 1: + release_reached_unmount.set() + await allow_unmount_recheck.wait() + await original_unmount_idle_entry(requested_key) + + async def old_holder() -> None: + async with cache.lease([artifact_uri]): + first_entered.set() + await release_first.wait() + + async def new_holder() -> None: + async with cache.lease([artifact_uri]): + newcomer_entered.set() + await release_newcomer.wait() + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch.object(cache, "_unmount", harness.unmount), + patch.object( + cache, + "_unmount_idle_entry", + pause_first_unmount_request, + ), + ): + old_task = asyncio.create_task(old_holder()) + await first_entered.wait() + release_first.set() + await release_reached_unmount.wait() + assert cache._refcount(cache_key) == 0 + + new_task = asyncio.create_task(new_holder()) + await newcomer_entered.wait() + assert cache._refcount(cache_key) == 1 + + allow_unmount_recheck.set() + await old_task + + assert mount_dir in harness.mounted + assert harness.unmounts == [] + + release_newcomer.set() + await new_task + + assert cache._refcount(cache_key) == 0 + assert harness.unmounts == [mount_dir] + + @pytest.mark.anyio + async def test_partial_multi_artifact_failure_rolls_back_prior_leases( + self, temp_cache_dir: Path + ) -> None: + """Protect sequential multi-artifact admission from partial pin leaks. + + If a middle artifact fails, every earlier acquisition must be released + and unmounted, the failed key must leave no shell, later artifacts must + never be acquired, and the release path must still request convergence. + """ + cache = RegistryArtifactCache(temp_cache_dir) + first_uri = "s3://bucket/first.squashfs" + failed_uri = "s3://bucket/failed.tar.gz" + untouched_uri = "s3://bucket/untouched.tar.gz" + first_key = compute_registry_artifact_cache_key(first_uri) + failed_key = compute_registry_artifact_cache_key(failed_uri) + untouched_key = compute_registry_artifact_cache_key(untouched_uri) + harness = SquashfsMountHarness(cache) + first_mount = harness.seed_mount(first_key) + untouched_path = write_tarball_entry(temp_cache_dir, untouched_key) + + async def fail_download( + artifact: TarballArtifact, + ctx: RegistryArtifactMaterializationContext, + output_path: Path, + ) -> None: + del artifact, ctx, output_path + raise RuntimeError("download failed") + + tracked_lease_artifact = AsyncMock(wraps=cache._lease_artifact) + converge_cache_budget = AsyncMock() + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "download", fail_download), + patch.object(cache, "_unmount", harness.unmount), + patch.object(cache, "_lease_artifact", tracked_lease_artifact), + patch.object( + cache, + "_converge_cache_budget", + converge_cache_budget, + ), + ): + with pytest.raises(RuntimeError, match="download failed"): + async with cache.lease([first_uri, failed_uri, untouched_uri]): + pass + + requested_uris = [ + await_call.args[0] for await_call in tracked_lease_artifact.await_args_list + ] + assert requested_uris == [first_uri, failed_uri] + assert cache._refcount(first_key) == 0 + assert cache._refcount(failed_key) == 0 + assert cache._refcount(untouched_key) == 0 + assert harness.unmounts == [first_mount] + assert cache._paths_for(first_key).squashfs_image_path.is_file() + assert not cache._paths_for(failed_key).entry_dir.exists() + assert untouched_path.is_dir() + converge_cache_budget.assert_awaited_once_with() + assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) + assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) + + @pytest.mark.anyio + async def test_waiter_retries_after_first_same_key_materializer_fails( + self, temp_cache_dir: Path + ) -> None: + """Protect same-key waiters from a failed cold-cache publisher. + + The first materializer may fail after writing scratch. Its waiter must + retry against a clean staging area, publish the sole valid entry, and + leave neither a leaked pin nor an empty LRU-visible cache shell. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/retry-after-failure.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + first_download_started = asyncio.Event() + fail_first_download = asyncio.Event() + retry_download_started = asyncio.Event() + allow_retry_download = asyncio.Event() + download_attempts = 0 + retry_saw_clean_staging = False + + async def controlled_download( + artifact: TarballArtifact, + ctx: RegistryArtifactMaterializationContext, + output_path: Path, + ) -> None: + del artifact, ctx + nonlocal download_attempts, retry_saw_clean_staging + download_attempts += 1 + if download_attempts == 1: + output_path.write_bytes(b"partial") + first_download_started.set() + await fail_first_download.wait() + raise RuntimeError("first publisher failed") + + retry_saw_clean_staging = not any(cache.staging_dir.iterdir()) + retry_download_started.set() + await allow_retry_download.wait() + output_path.write_bytes(b"complete") + + async def mock_extract( + artifact: TarballArtifact, + tarball_path: Path, + target_dir: Path, + ) -> None: + del artifact + assert tarball_path.read_bytes() == b"complete" + (target_dir / "module.py").write_text("VALUE = 1") + + async def take_lease() -> list[Path]: + async with cache.lease([artifact_uri]) as registry_paths: + return registry_paths + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch( + "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", + return_value=1, + ), + patch.object(TarballArtifact, "download", controlled_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + first = asyncio.create_task(take_lease()) + await first_download_started.wait() + waiter = asyncio.create_task(take_lease()) + await asyncio.sleep(0) + assert download_attempts == 1 + + fail_first_download.set() + await retry_download_started.wait() + allow_retry_download.set() + registry_paths = await waiter + with pytest.raises(RuntimeError, match="first publisher failed"): + await first + + target_dir = cache._paths_for(cache_key).tarball_target_dir + assert registry_paths == [target_dir] + assert (target_dir / "module.py").read_text() == "VALUE = 1" + assert download_attempts == 2 + assert retry_saw_clean_staging is True + assert cache._discover_cache_keys() == {cache_key} + assert cache._refcount(cache_key) == 0 + assert not any(cache.staging_dir.iterdir()) + assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) + + @pytest.mark.anyio + async def test_duplicate_uri_balances_each_acquisition_and_release( + self, temp_cache_dir: Path + ) -> None: + """Protect list bookkeeping when one lease requests the same URI twice. + + Duplicate PYTHONPATH entries intentionally acquire two pins. Both must + be released, while materialization and final unmount still occur once; + accidental deduplication on only one side would leak or underflow pins. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/duplicate.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + harness = SquashfsMountHarness(cache) + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", harness.mount), + patch.object(SquashfsArtifact, "extract", harness.extract), + patch.object(cache, "_unmount", harness.unmount), + ): + async with cache.lease([artifact_uri, artifact_uri]) as registry_paths: + mount_dir = cache._paths_for(cache_key).squashfs_mount_dir + assert registry_paths == [mount_dir, mount_dir] + assert cache._refcount(cache_key) == 2 + assert harness.mount_attempts == [cache_key] + + assert cache._refcount(cache_key) == 0 + assert harness.unmounts == [mount_dir] + assert cache._paths_for(cache_key).squashfs_image_path.is_file() + + @pytest.mark.anyio + async def test_lease_is_never_admitted_across_an_in_flight_eviction( + self, temp_cache_dir + ): + """A lease must not return a mount an in-flight eviction is deleting.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/path/site-packages.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + mounted = {paths.squashfs_mount_dir} + umount_started = asyncio.Event() + finish_umount = asyncio.Event() + remounts: list[str] = [] + + umount_process = AsyncMock() + umount_process.communicate.return_value = (b"", b"") + umount_process.returncode = 0 + + async def mock_umount(*args, **kwargs): + umount_started.set() + await finish_umount.wait() + mounted.discard(paths.squashfs_mount_dir) + return umount_process + + async def mock_mount(self, ctx, image_path): + remounts.append(ctx.cache_key) + target_dir = ctx.paths.squashfs_mount_dir + target_dir.mkdir(parents=True, exist_ok=True) + (target_dir / "module.py").write_text("VALUE = 1") + mounted.add(target_dir) + return target_dir + + leased_paths: list[Path] = [] + leased_path_exists: list[bool] = [] + + async def take_lease() -> None: + async with cache.lease([artifact_uri]) as registry_paths: + leased_paths.extend(registry_paths) + leased_path_exists.append(registry_paths[0].is_dir()) + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/sbin/umount", + ), + patch( + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", + side_effect=mock_umount, + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + ): + eviction = asyncio.create_task(cache._evict_entry(cache_key)) + await umount_started.wait() + lease = asyncio.create_task(take_lease()) + # Let the lease block on the per-key lock the eviction holds. + await asyncio.sleep(0) + finish_umount.set() + evicted, _ = await asyncio.gather(eviction, lease) + + assert evicted == RegistryArtifactEviction(retired=True, reclaimed=True) + # The lease waited for the eviction and re-materialized the entry. + assert remounts == [cache_key] + assert leased_paths == [paths.squashfs_mount_dir] + assert leased_path_exists == [True] + assert (paths.squashfs_mount_dir / "module.py").read_text() == "VALUE = 1" + + @pytest.mark.anyio + async def test_builtin_artifact_is_exempt_from_cache_accounting( + self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch + ): + """The bundled builtin registry is never a cache entry.""" + version = "1.2.3" + site_packages = temp_cache_dir / "venv" / "site-packages" + package_dir = site_packages / "tracecat_registry" + package_dir.mkdir(parents=True) + package_file = package_dir / "__init__.py" + package_file.write_text("") + + monkeypatch.setattr(tracecat_registry, "__version__", version) + monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) + monkeypatch.setattr( + "tracecat.executor.registry_artifact_materialization.sysconfig.get_path", + lambda name: str(site_packages) if name == "purelib" else None, + ) + + cache = RegistryArtifactCache(temp_cache_dir) + + with patch.object( + cache, + "_enforce_cache_budget", + new_callable=AsyncMock, + ) as enforce_cache_budget: + async with cache.lease([bundled_builtin_registry_uri(version)]) as paths: + assert paths == [site_packages.resolve()] + assert cache._runtime == {} + + enforce_cache_budget.assert_not_awaited() diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py new file mode 100644 index 0000000000..702cdc9f4e --- /dev/null +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -0,0 +1,698 @@ +"""Artifact selection, download, and materialization tests.""" + +from __future__ import annotations + +import asyncio +import tarfile +import threading +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import pytest + +from tracecat.executor.registry_artifacts import ( + SQUASHFS_MOUNT_OPTIONS, + RegistryArtifactCache, + RegistryArtifactMaterializationContext, + SquashfsArtifact, + SquashfsMountCommandError, + TarballArtifact, + _squashfs_listing_size, + compute_registry_artifact_cache_key, +) + +from .registry_artifact_test_helpers import ( + SQUASHFS_ENABLED_CONFIG, + BlockingSubprocess, + CapturedSubprocess, + SquashfsMountHarness, + lease_paths, + tarball_payload, +) + + +class TestRegistryArtifactMaterialization: + """Materialize and reuse executor-local artifact formats.""" + + @pytest.mark.anyio + async def test_same_key_cold_fan_in_materializes_and_enforces_once( + self, temp_cache_dir + ): + """Same-key waiters share candidate lookup, materialization, and enforcement.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/path/site-packages.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + materialization_started = asyncio.Event() + finish_materialization = asyncio.Event() + + async def mock_materialize( + self: TarballArtifact, + ctx: RegistryArtifactMaterializationContext, + ) -> list[Path]: + materialization_started.set() + await finish_materialization.wait() + ctx.paths.tarball_target_dir.mkdir(parents=True) + return [ctx.paths.tarball_target_dir] + + async def take_lease() -> list[Path]: + async with cache.lease([artifact_uri]) as paths: + return paths + + with ( + patch.object( + cache, + "_sidecar_exists", + new_callable=AsyncMock, + return_value=False, + ) as sidecar_exists, + patch.object( + cache, + "_enforce_cache_budget", + new_callable=AsyncMock, + return_value=True, + ) as enforce_cache_budget, + patch.object(cache, "_converge_cache_budget", new_callable=AsyncMock), + patch.object(TarballArtifact, "materialize", mock_materialize), + ): + leases = [asyncio.create_task(take_lease()) for _ in range(5)] + await materialization_started.wait() + await asyncio.sleep(0) + finish_materialization.set() + results = await asyncio.gather(*leases) + + expected_paths = [cache._paths_for(cache_key).tarball_target_dir] + assert results == [expected_paths] * 5 + sidecar_exists.assert_awaited_once() + enforce_cache_budget.assert_awaited_once_with(protected_key=cache_key) + + @pytest.mark.anyio + async def test_different_cold_keys_serialize_materialization( + self, temp_cache_dir: Path + ) -> None: + """Distinct cold keys cannot multiply staging and extraction peaks.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + uris = ["s3://bucket/first.tar.gz", "s3://bucket/second.tar.gz"] + first_started = asyncio.Event() + release_first = asyncio.Event() + second_started = asyncio.Event() + started_keys: list[str] = [] + + async def controlled_materialize( + self: TarballArtifact, + ctx: RegistryArtifactMaterializationContext, + ) -> list[Path]: + del self + started_keys.append(ctx.cache_key) + if len(started_keys) == 1: + first_started.set() + await release_first.wait() + else: + second_started.set() + ctx.paths.tarball_target_dir.mkdir(parents=True) + return [ctx.paths.tarball_target_dir] + + async def take_lease(uri: str) -> None: + async with cache.lease([uri]): + pass + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "materialize", controlled_materialize), + patch.object(cache, "_enforce_cache_budget", new_callable=AsyncMock), + patch.object(cache, "_converge_cache_budget", new_callable=AsyncMock), + ): + first = asyncio.create_task(take_lease(uris[0])) + await first_started.wait() + second = asyncio.create_task(take_lease(uris[1])) + await asyncio.sleep(0) + assert not second_started.is_set() + + release_first.set() + await second_started.wait() + await asyncio.gather(first, second) + + assert started_keys == [ + compute_registry_artifact_cache_key(uri) for uri in uris + ] + + def test_squashfs_listing_size_sums_files_and_symlinks(self) -> None: + listing = b"\n".join( + [ + b"drwxr-xr-x 0/0 64 2026-01-01 00:00 squashfs-root", + b"-rw-r--r-- 0/0 123 2026-01-01 00:00 squashfs-root/module.py", + b"lrwxrwxrwx 0/0 9 2026-01-01 00:00 squashfs-root/current -> module.py", + ] + ) + + assert _squashfs_listing_size(listing) == 132 + + def test_squashfs_listing_size_rejects_unparseable_files(self) -> None: + with pytest.raises(ValueError, match="Could not parse SquashFS listing"): + _squashfs_listing_size(b"-rw-r--r-- malformed") + + @pytest.mark.anyio + async def test_materialize_mounts_squashfs_sidecar(self, temp_cache_dir): + """Test that a SquashFS sidecar is mounted instead of extracting tarballs.""" + cache = RegistryArtifactCache(temp_cache_dir) + + async def mock_mount(self, ctx, image_path): + assert image_path.name.endswith(".squashfs") + target_dir = ctx.paths.squashfs_mount_dir + target_dir.mkdir(parents=True, exist_ok=True) + (target_dir / "module.py").write_text("VALUE = 1") + return target_dir + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + return_value=True, + ), + patch( + "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", + True, + ), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + patch.object( + TarballArtifact, + "materialize", + new_callable=AsyncMock, + ) as tarball_materialize, + ): + result = await lease_paths( + cache, + "s3://bucket/path/site-packages.tar.gz", + ) + + assert len(result) == 1 + assert (result[0] / "module.py").read_text() == "VALUE = 1" + tarball_materialize.assert_not_awaited() + + @pytest.mark.anyio + async def test_mount_squashfs_uses_hardened_read_only_options( + self, + temp_cache_dir, + ): + """Test that SquashFS images are mounted read-only without device/setuid bits.""" + cache_key = "cache-key" + cache = RegistryArtifactCache(temp_cache_dir) + ctx = cache._context_for(cache_key) + artifact = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key=cache_key, + ) + image_path = ctx.paths.squashfs_image_path + target_dir = ctx.paths.squashfs_mount_dir + ctx.paths.entry_dir.mkdir(parents=True) + image_path.write_bytes(b"squashfs") + target_dir.mkdir() + process = AsyncMock() + process.communicate.return_value = (b"", b"") + process.returncode = 0 + + with patch( + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ) as create_subprocess_exec: + await artifact.mount(ctx, image_path) + + create_subprocess_exec.assert_awaited_once_with( + "mount", + "-t", + "squashfs", + "-o", + SQUASHFS_MOUNT_OPTIONS, + str(image_path), + str(target_dir), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + + @pytest.mark.anyio + async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): + """Cancellation cannot leave an orphan mount process after lock release.""" + cache_key = "cancelled-mount" + cache = RegistryArtifactCache(temp_cache_dir) + ctx = cache._context_for(cache_key) + artifact = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key=cache_key, + ) + image_path = ctx.paths.squashfs_image_path + target_dir = ctx.paths.squashfs_mount_dir + ctx.paths.entry_dir.mkdir(parents=True) + image_path.write_bytes(b"squashfs") + target_dir.mkdir() + process = BlockingSubprocess() + + with patch( + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ): + mounting = asyncio.create_task( + artifact._mount_image(image_path, target_dir) + ) + await process.communicate_started.wait() + mounting.cancel() + + with pytest.raises(asyncio.CancelledError): + await mounting + + assert process.cleanup_calls == ["kill", "wait"] + assert target_dir.is_dir() + assert not target_dir.is_mount() + + @pytest.mark.anyio + async def test_cancelled_squashfs_extract_kills_and_reaps_subprocess( + self, temp_cache_dir + ): + """Cancellation cannot leave unsquashfs writing into scratch.""" + cache_key = "cancelled-extract" + cache = RegistryArtifactCache(temp_cache_dir) + ctx = cache._context_for(cache_key) + artifact = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key=cache_key, + ) + image_path = ctx.paths.squashfs_image_path + target_dir = ctx.paths.squashfs_extract_dir + ctx.paths.entry_dir.mkdir(parents=True) + image_path.write_bytes(b"squashfs") + target_dir.mkdir() + real_create_subprocess_exec = asyncio.create_subprocess_exec + process_started = asyncio.Event() + captured_processes: list[CapturedSubprocess] = [] + + async def create_sleep_subprocess( + *args: object, **kwargs: object + ) -> CapturedSubprocess: + del args, kwargs + process = await real_create_subprocess_exec( + "/bin/sleep", + "30", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + captured = CapturedSubprocess(process) + captured_processes.append(captured) + process_started.set() + return captured + + with ( + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/usr/bin/unsquashfs", + ), + patch( + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", + side_effect=create_sleep_subprocess, + ), + ): + extracting = asyncio.create_task( + artifact._extract_image(image_path, target_dir) + ) + await process_started.wait() + captured = captured_processes[0] + extracting.cancel() + + try: + with pytest.raises(asyncio.CancelledError): + await extracting + finally: + if captured.returncode is None: + captured.kill() + await captured.wait() + + assert captured.killed is True + assert captured.reaped is True + assert captured.returncode is not None + + @pytest.mark.anyio + async def test_repeatedly_cancelled_tarball_extract_rejoins_thread( + self, temp_cache_dir + ): + """Repeated cancellation waits until the tar extractor stops writing.""" + artifact = TarballArtifact( + uri="s3://bucket/path/site-packages.tar.gz", + cache_key="cancelled-tarball-extract", + ) + tarball_path = temp_cache_dir / "artifact.tar.gz" + target_dir = temp_cache_dir / "target" + target_dir.mkdir() + with tarfile.open(tarball_path, "w:gz"): + pass + + extraction_started = threading.Event() + extraction_release = threading.Event() + extraction_finished = threading.Event() + + def blocking_extractall(*args: object, **kwargs: object) -> None: + del args, kwargs + extraction_started.set() + extraction_release.wait() + extraction_finished.set() + + with patch.object( + tarfile.TarFile, + "extractall", + new=blocking_extractall, + ): + extracting = asyncio.create_task(artifact.extract(tarball_path, target_dir)) + assert await asyncio.to_thread(extraction_started.wait, 1) + extracting.cancel() + done, _ = await asyncio.wait({extracting}, timeout=0.05) + first_cancellation_propagated_early = bool(done) + + extracting.cancel() + done, _ = await asyncio.wait({extracting}, timeout=0.05) + second_cancellation_propagated_early = bool(done) + extraction_release.set() + + with pytest.raises(asyncio.CancelledError): + await extracting + + assert extraction_finished.is_set() + assert first_cancellation_propagated_early is False + assert second_cancellation_propagated_early is False + + @pytest.mark.anyio + async def test_materialize_extracts_squashfs_when_mount_fails(self, temp_cache_dir): + """Test that SquashFS mount failures fall back to unsquashfs extraction.""" + cache = RegistryArtifactCache(temp_cache_dir) + + async def mock_mount(self, ctx, image_path): + raise SquashfsMountCommandError("operation not permitted") + + async def mock_extract(self, ctx, image_path): + target_dir = ctx.paths.squashfs_extract_dir + target_dir.mkdir(parents=True, exist_ok=True) + (target_dir / "module.py").write_text("VALUE = 1") + return target_dir + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + return_value=True, + ), + patch( + "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", + True, + ), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + patch.object(SquashfsArtifact, "extract", mock_extract), + patch.object( + TarballArtifact, + "materialize", + new_callable=AsyncMock, + ) as tarball_materialize, + ): + result = await lease_paths( + cache, + "s3://bucket/path/site-packages.tar.gz", + ) + + assert len(result) == 1 + assert (result[0] / "module.py").read_text() == "VALUE = 1" + assert result[0].name == "extracted" + tarball_materialize.assert_not_awaited() + + @pytest.mark.anyio + async def test_materialize_extracts_squashfs_without_mount_binary( + self, temp_cache_dir + ): + """Test that SquashFS is still preferred when only unsquashfs is available.""" + cache = RegistryArtifactCache(temp_cache_dir) + + async def mock_extract(self, ctx, image_path): + target_dir = ctx.paths.squashfs_extract_dir + target_dir.mkdir(parents=True, exist_ok=True) + (target_dir / "module.py").write_text("VALUE = 1") + return target_dir + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + return_value=True, + ), + patch( + "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", + True, + ), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value=None, + ), + patch.object(SquashfsArtifact, "extract", mock_extract), + ): + result = await lease_paths( + cache, + "s3://bucket/path/site-packages.tar.gz", + ) + + assert len(result) == 1 + assert (result[0] / "module.py").read_text() == "VALUE = 1" + assert result[0].name == "extracted" + + @pytest.mark.anyio + async def test_materialize_falls_back_to_gzip_when_squashfs_extract_fails( + self, temp_cache_dir + ): + """Test that legacy gzip remains the final compatibility fallback.""" + cache = RegistryArtifactCache(temp_cache_dir) + source = temp_cache_dir / "source" + source.mkdir() + (source / "module.py").write_text("VALUE = 1") + + async def mock_tarball_download(self, ctx, path): + with tarfile.open(path, "w:gz") as tar: + tar.add(source / "module.py", arcname="module.py") + + async def mock_mount(self, ctx, image_path): + raise SquashfsMountCommandError("operation not permitted") + + async def mock_extract(self, ctx, image_path): + raise RuntimeError("unsquashfs unavailable") + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + side_effect=[True, False], + ), + patch( + "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", + True, + ), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + patch.object(SquashfsArtifact, "extract", mock_extract), + patch.object(TarballArtifact, "download", mock_tarball_download), + ): + result = await lease_paths( + cache, + "s3://bucket/path/site-packages.tar.gz", + ) + + assert len(result) == 1 + assert (result[0] / "module.py").read_text() == "VALUE = 1" + assert result[0].name == "tarball" + + @pytest.mark.anyio + async def test_materialize_treats_unknown_suffix_as_gzip(self, temp_cache_dir): + """Test that existing gzip artifacts can use arbitrary S3 key suffixes.""" + cache = RegistryArtifactCache(temp_cache_dir) + source = temp_cache_dir / "source" + source.mkdir() + (source / "module.py").write_text("VALUE = 1") + + async def mock_download(self, ctx, path): + assert path.name.endswith(".tar.gz") + with tarfile.open(path, "w:gz") as tar: + tar.add(source / "module.py", arcname="module.py") + + with patch.object(TarballArtifact, "download", mock_download): + result = await lease_paths( + cache, + "s3://bucket/path/custom-key", + ) + + assert len(result) == 1 + assert (result[0] / "module.py").read_text() == "VALUE = 1" + + @pytest.mark.anyio + async def test_materialize_caches_result(self, temp_cache_dir): + """Test that tarball extraction is cached.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/test.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + target_dir = cache._paths_for(cache_key).tarball_target_dir + target_dir.mkdir(parents=True) + + result = await lease_paths(cache, artifact_uri) + + assert result == [target_dir] + + @pytest.mark.anyio + async def test_materialize_concurrent_requests(self, temp_cache_dir): + """Test that concurrent requests for same artifact do not race.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/test.tar.gz" + download_count = 0 + + async def mock_download(self, ctx, path): + nonlocal download_count + download_count += 1 + await asyncio.sleep(0.1) + path.write_bytes(tarball_payload(size=1)) + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "extracted.txt").write_text("extracted") + + with ( + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + results = await asyncio.gather( + lease_paths(cache, artifact_uri), + lease_paths(cache, artifact_uri), + lease_paths(cache, artifact_uri), + ) + + assert all(r == results[0] for r in results) + assert download_count == 1 + + +class TestSquashfsMountPolicy: + """Tests for per-artifact SquashFS fallback.""" + + @pytest.mark.anyio + async def test_loop_device_exhaustion_isolated_sticky_extraction_fallback( + self, temp_cache_dir: Path + ) -> None: + """Protect the cache's fail-open policy when loop devices are saturated. + + A mount-command failure must extract only the affected cold artifact, + preserve already-leased mounts, reuse that extraction on later leases, + and still let unrelated artifacts attempt mounting. This models loop + exhaustion deterministically without consuming host-global devices. + """ + cache = RegistryArtifactCache(temp_cache_dir) + held_uri = "s3://bucket/already-mounted.squashfs" + saturated_uri = "s3://bucket/no-loop-available.squashfs" + later_uri = "s3://bucket/later-artifact.squashfs" + held_key = compute_registry_artifact_cache_key(held_uri) + saturated_key = compute_registry_artifact_cache_key(saturated_uri) + later_key = compute_registry_artifact_cache_key(later_uri) + harness = SquashfsMountHarness( + cache, + failed_mount_keys={saturated_key}, + ) + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", harness.mount), + patch.object(SquashfsArtifact, "extract", harness.extract), + patch.object(cache, "_unmount", harness.unmount), + ): + async with cache.lease([held_uri]) as held_paths: + held_mount = cache._paths_for(held_key).squashfs_mount_dir + assert held_paths == [held_mount] + + async with cache.lease([saturated_uri]) as saturated_paths: + extracted = cache._paths_for(saturated_key).squashfs_extract_dir + assert saturated_paths == [extracted] + assert held_mount in harness.mounted + assert cache._refcount(held_key) == 1 + assert harness.unmounts == [] + + # The extracted directory is the cache entry's sticky path; + # freeing a loop elsewhere does not trigger a remount attempt. + async with cache.lease([saturated_uri]) as cached_paths: + assert cached_paths == [extracted] + + async with cache.lease([later_uri]) as later_paths: + later_mount = cache._paths_for(later_key).squashfs_mount_dir + assert later_paths == [later_mount] + assert held_mount in harness.mounted + + assert held_mount in harness.mounted + + assert harness.mount_attempts == [held_key, saturated_key, later_key] + assert harness.extraction_attempts == [saturated_key] + assert harness.unmounts == [later_mount, held_mount] + assert cache._paths_for(saturated_key).squashfs_image_path.is_file() + assert cache._paths_for(saturated_key).squashfs_extract_dir.is_dir() + assert cache._refcount(held_key) == 0 + assert cache._refcount(saturated_key) == 0 + assert cache._refcount(later_key) == 0 + + @pytest.mark.anyio + async def test_mount_failure_does_not_disable_later_artifacts( + self, temp_cache_dir + ) -> None: + """One failed mount falls back without changing later mount attempts.""" + cache = RegistryArtifactCache(temp_cache_dir) + first_ctx = cache._context_for("first") + second_ctx = cache._context_for("second") + first = SquashfsArtifact( + uri="s3://bucket/first.squashfs", + cache_key="first", + ) + second = SquashfsArtifact( + uri="s3://bucket/second.squashfs", + cache_key="second", + ) + mount_attempts: list[str] = [] + + async def mock_mount(self, ctx, image_path): + del image_path + mount_attempts.append(ctx.cache_key) + if ctx.cache_key == "first": + raise SquashfsMountCommandError("operation not permitted") + ctx.paths.squashfs_mount_dir.mkdir(parents=True) + return ctx.paths.squashfs_mount_dir + + async def mock_extract(self, ctx, image_path): + del image_path + ctx.paths.squashfs_extract_dir.mkdir(parents=True) + return ctx.paths.squashfs_extract_dir + + with ( + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + patch.object(SquashfsArtifact, "extract", mock_extract), + ): + assert await first.materialize(first_ctx) == [ + first_ctx.paths.squashfs_extract_dir + ] + assert await second.materialize(second_ctx) == [ + second_ctx.paths.squashfs_mount_dir + ] + + assert mount_attempts == ["first", "second"] diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py new file mode 100644 index 0000000000..b5e7f6acb6 --- /dev/null +++ b/tests/unit/executor/test_registry_artifact_resolution.py @@ -0,0 +1,345 @@ +"""Registry artifact URI, key, and candidate resolution tests.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import httpx +import pytest +import tracecat_registry + +from tracecat.executor.registry_artifacts import ( + RegistryArtifactCache, + RegistryArtifactFormat, + SquashfsArtifact, + TarballArtifact, + bundled_builtin_registry_uri, + compute_registry_artifact_cache_key, +) +from tracecat.registry.artifact_keys import parse_s3_uri + + +class TestParseS3Uri: + """Tests for parse_s3_uri function.""" + + def test_valid_uri(self): + """Test parsing a valid S3 URI.""" + bucket, key = parse_s3_uri("s3://my-bucket/path/to/file.tar.gz") + assert bucket == "my-bucket" + assert key == "path/to/file.tar.gz" + + def test_uri_with_nested_path(self): + """Test parsing URI with deeply nested path.""" + bucket, key = parse_s3_uri("s3://bucket/a/b/c/d/e/file.tar.gz") + assert bucket == "bucket" + assert key == "a/b/c/d/e/file.tar.gz" + + def test_invalid_uri_no_prefix(self): + """Test that non-S3 URIs raise ValueError.""" + with pytest.raises(ValueError, match="Invalid S3 URI"): + parse_s3_uri("https://bucket/key") + + def test_invalid_uri_no_key(self): + """Test that URIs without keys raise ValueError.""" + with pytest.raises(ValueError, match="Invalid S3 URI"): + parse_s3_uri("s3://bucket") + + def test_invalid_uri_empty_bucket(self): + """Test that URIs with empty bucket raise ValueError.""" + with pytest.raises(ValueError, match="Invalid S3 URI"): + parse_s3_uri("s3:///key") + + +class TestRegistryArtifactResolution: + """Resolve artifact identities and preferred formats.""" + + def test_compute_registry_artifact_cache_key_deterministic(self): + """Test that cache key computation is deterministic.""" + uri = "s3://bucket/path/to/registry-v1.2.3.tar.gz" + + key1 = compute_registry_artifact_cache_key(uri) + key2 = compute_registry_artifact_cache_key(uri) + + assert key1 == key2 + assert len(key1) == 16 + + def test_compute_registry_artifact_cache_key_case_sensitive(self): + """Test that cache key is case-sensitive because S3 keys are case-sensitive.""" + key1 = compute_registry_artifact_cache_key("s3://BUCKET/PATH/FILE.tar.gz") + key2 = compute_registry_artifact_cache_key("s3://bucket/path/file.tar.gz") + + assert key1 != key2 + + def test_compute_registry_artifact_cache_key_empty(self): + """Test that empty URI returns the base cache key.""" + assert compute_registry_artifact_cache_key("") == "base" + + @pytest.mark.anyio + async def test_download_artifact_uses_blob_download_file_to_path( + self, temp_cache_dir + ): + """Test that artifact downloads stay behind the blob storage helper.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key="download-test", + ) + ctx = cache._context_for(artifact.cache_key) + output_path = temp_cache_dir / "artifact.squashfs" + + async def mock_download_file_to_path( + *, + key: str, + bucket: str, + output_path: Path, + ) -> None: + output_path.write_bytes(b"squashfs") + + with patch( + "tracecat.executor.registry_artifacts.blob.download_file_to_path", + new_callable=AsyncMock, + side_effect=mock_download_file_to_path, + ) as download_file_to_path: + await artifact.download(ctx, output_path) + + download_file_to_path.assert_awaited_once() + await_args = download_file_to_path.await_args + assert await_args is not None + assert await_args.kwargs["key"] == "path/site-packages.squashfs" + assert await_args.kwargs["bucket"] == "bucket" + assert output_path.read_bytes() == b"squashfs" + + @pytest.mark.anyio + async def test_lease_uses_bundled_current_builtin( + self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch + ): + """In-tree builtin registry returns only the installed site-packages.""" + version = "1.2.3" + site_packages = temp_cache_dir / "venv" / "site-packages" + package_dir = site_packages / "tracecat_registry" + package_dir.mkdir(parents=True) + package_file = package_dir / "__init__.py" + package_file.write_text("") + + monkeypatch.setattr(tracecat_registry, "__version__", version) + monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) + monkeypatch.setattr( + "tracecat.executor.registry_artifact_materialization.sysconfig.get_path", + lambda name: str(site_packages) if name == "purelib" else None, + ) + + cache = RegistryArtifactCache(temp_cache_dir) + async with cache.lease([bundled_builtin_registry_uri(version)]) as result: + assert result == [site_packages.resolve()] + + @pytest.mark.anyio + async def test_lease_exposes_editable_builtin_parent( + self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch + ): + """Editable builtin registry exposes the package wrapper + site-packages.""" + version = "1.2.3" + site_packages = temp_cache_dir / "venv" / "site-packages" + dependency_dir = site_packages / "orjson" + dependency_dir.mkdir(parents=True) + (dependency_dir / "__init__.py").write_text("VALUE = 1\n") + source_root = temp_cache_dir / "src" / "tracecat-registry" + package_dir = source_root / "tracecat_registry" + package_dir.mkdir(parents=True) + package_file = package_dir / "__init__.py" + package_file.write_text("__version__ = '1.2.3'\n") + + monkeypatch.setattr(tracecat_registry, "__version__", version) + monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) + monkeypatch.setattr( + "tracecat.executor.registry_artifact_materialization.sysconfig.get_path", + lambda name: str(site_packages) if name == "purelib" else None, + ) + + cache = RegistryArtifactCache(temp_cache_dir) + async with cache.lease([bundled_builtin_registry_uri(version)]) as result: + assert result == [source_root.resolve(), site_packages.resolve()] + + @pytest.mark.anyio + async def test_lease_rejects_stale_bundled_builtin( + self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch + ): + """Bundled pseudo-URIs must match this executor's installed package.""" + monkeypatch.setattr(tracecat_registry, "__version__", "1.2.3") + + cache = RegistryArtifactCache(temp_cache_dir) + with pytest.raises(RuntimeError, match="does not match installed version"): + async with cache.lease([bundled_builtin_registry_uri("1.2.4")]): + pass + + @pytest.mark.anyio + async def test_download_artifact_normalizes_missing_objects_to_http_404( + self, temp_cache_dir + ): + """Preserve the missing-artifact error contract from presigned downloads.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact = TarballArtifact( + uri="s3://bucket/path/site-packages.tar.gz", + cache_key="missing-test", + ) + ctx = cache._context_for(artifact.cache_key) + output_path = temp_cache_dir / "artifact.tar.gz" + + with patch( + "tracecat.executor.registry_artifacts.blob.download_file_to_path", + new_callable=AsyncMock, + side_effect=FileNotFoundError, + ): + with pytest.raises(httpx.HTTPStatusError) as exc_info: + await artifact.download(ctx, output_path) + + assert exc_info.value.response.status_code == 404 + assert isinstance(exc_info.value.__cause__, FileNotFoundError) + + @pytest.mark.anyio + async def test_artifact_candidates_prefer_squashfs_sidecar(self, temp_cache_dir): + """Test that gzip tarballs prefer a sibling SquashFS sidecar.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + return_value=True, + ) as file_exists, + patch.object(cache, "_can_try_squashfs", return_value=True), + ): + cache_key = compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.tar.gz" + ) + ctx = cache._context_for(cache_key) + candidates = await cache._artifact_candidates( + ctx, "s3://bucket/path/site-packages.tar.gz" + ) + + artifact = candidates[0] + assert len(candidates) == 2 + assert isinstance(artifact, SquashfsArtifact) + assert isinstance(candidates[1], TarballArtifact) + assert artifact.uri == "s3://bucket/path/site-packages.squashfs" + assert artifact.format == RegistryArtifactFormat.SQUASHFS + file_exists.assert_awaited_once_with( + key="path/site-packages.squashfs", + bucket="bucket", + ) + + @pytest.mark.anyio + async def test_artifact_candidates_direct_squashfs_include_gzip_fallback( + self, temp_cache_dir + ): + """Test direct SquashFS URIs fall back to sibling gzip tarballs.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with patch.object(cache, "_can_try_squashfs") as can_try_squashfs: + cache_key = compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.squashfs" + ) + ctx = cache._context_for(cache_key) + candidates = await cache._artifact_candidates( + ctx, + "s3://bucket/path/site-packages.squashfs", + ) + + assert isinstance(candidates[0], SquashfsArtifact) + assert isinstance(candidates[1], TarballArtifact) + assert [artifact.uri for artifact in candidates] == [ + "s3://bucket/path/site-packages.squashfs", + "s3://bucket/path/site-packages.tar.gz", + ] + assert [artifact.format for artifact in candidates] == [ + RegistryArtifactFormat.SQUASHFS, + RegistryArtifactFormat.TAR_GZ, + ] + can_try_squashfs.assert_not_called() + + @pytest.mark.anyio + async def test_artifact_candidates_fall_back_to_gzip(self, temp_cache_dir): + """Test that gzip tarballs are used when no sidecar exists.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + return_value=False, + ), + patch.object(cache, "_can_try_squashfs", return_value=True), + ): + cache_key = compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.tar.gz" + ) + ctx = cache._context_for(cache_key) + candidates = await cache._artifact_candidates( + ctx, "s3://bucket/path/site-packages.tar.gz" + ) + + artifact = candidates[0] + assert len(candidates) == 1 + assert isinstance(artifact, TarballArtifact) + assert artifact.uri == "s3://bucket/path/site-packages.tar.gz" + assert artifact.format == RegistryArtifactFormat.TAR_GZ + + def test_can_try_squashfs_does_not_require_mount_binary(self, temp_cache_dir): + """Prefer SquashFS whenever enabled; extraction may work without mounts.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value=None, + ), + patch( + "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", + True, + ), + ): + ctx = cache._context_for("squashfs-test") + assert cache._can_try_squashfs() is True + assert ctx.can_mount_squashfs() is False + + @pytest.mark.anyio + async def test_artifact_candidates_skip_non_registry_tarballs(self, temp_cache_dir): + """Test that arbitrary gzip tarballs do not trigger sidecar lookups.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + ) as file_exists: + cache_key = compute_registry_artifact_cache_key( + "s3://bucket/path/custom.tar.gz" + ) + ctx = cache._context_for(cache_key) + candidates = await cache._artifact_candidates( + ctx, "s3://bucket/path/custom.tar.gz" + ) + + artifact = candidates[0] + assert len(candidates) == 1 + assert isinstance(artifact, TarballArtifact) + assert artifact.uri == "s3://bucket/path/custom.tar.gz" + assert artifact.format == RegistryArtifactFormat.TAR_GZ + file_exists.assert_not_awaited() + + def test_runtime_for_same_key(self, temp_cache_dir): + """The same cache key returns the same runtime state.""" + cache = RegistryArtifactCache(temp_cache_dir) + + runtime1 = cache._runtime_for("key1") + runtime2 = cache._runtime_for("key1") + + assert runtime1 is runtime2 + + def test_runtime_for_different_keys(self, temp_cache_dir): + """Different cache keys return different runtime state.""" + cache = RegistryArtifactCache(temp_cache_dir) + + runtime1 = cache._runtime_for("key1") + runtime2 = cache._runtime_for("key2") + + assert runtime1 is not runtime2 diff --git a/tests/unit/executor/test_registry_artifact_startup.py b/tests/unit/executor/test_registry_artifact_startup.py new file mode 100644 index 0000000000..bc2d6d83a9 --- /dev/null +++ b/tests/unit/executor/test_registry_artifact_startup.py @@ -0,0 +1,395 @@ +"""Registry artifact startup recovery tests.""" + +from __future__ import annotations + +import asyncio +import os +import threading +from pathlib import Path +from unittest.mock import patch + +import pytest + +from tracecat.executor.registry_artifacts import ( + RegistryArtifactCache, + _delete_cache_path, +) + +from .registry_artifact_test_helpers import ( + MAX_BYTES_CONFIG, + MAX_ENTRIES_CONFIG, + write_image_entry, + write_tarball_entry, +) + + +class TestRegistryArtifactCacheStartupSweep: + """Tests for the startup sweep that reclaims state from a dead process.""" + + @pytest.mark.anyio + async def test_sweep_uses_entry_root_mtime_for_restart_safe_lru( + self, temp_cache_dir + ): + """Startup trimming preserves a touched older tarball-only entry.""" + previous_cache = RegistryArtifactCache(temp_cache_dir) + old = write_tarball_entry(temp_cache_dir, "old") + previous_cache._touch_entry("old") + + old_mtime = previous_cache._paths_for("old").entry_dir.stat().st_mtime + new = write_tarball_entry(temp_cache_dir, "new") + new_entry_dir = previous_cache._paths_for("new").entry_dir + os.utime(new_entry_dir, (old_mtime - 1, old_mtime - 1)) + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + + assert old.is_dir() + assert not new.exists() + + @pytest.mark.anyio + async def test_sweep_tolerates_missing_cache_dir(self, temp_cache_dir): + """A cache directory that does not exist yet is a no-op.""" + cache_dir = temp_cache_dir / "missing" + + cache = RegistryArtifactCache(cache_dir) + await cache.ensure_swept() + + assert cache.cache_dir == cache_dir + assert not cache_dir.exists() + + @pytest.mark.anyio + async def test_sweep_removes_orphaned_work(self, temp_cache_dir): + """Interrupted work is reclaimed without touching unrelated paths.""" + cache = RegistryArtifactCache(temp_cache_dir) + orphaned = cache.staging_dir / "abc123.999999.4321.squashfs" + orphaned.parent.mkdir() + orphaned.write_bytes(b"partial") + orphaned_dir = cache.trash_dir / "abc123.999999.4321" + orphaned_dir.mkdir(parents=True) + unrelated = temp_cache_dir / "unrelated" + unrelated.mkdir() + unrelated_file = unrelated / "keep.txt" + unrelated_file.write_text("keep") + entry_dir = write_tarball_entry(temp_cache_dir, "abc123") + + await cache.ensure_swept() + + assert not orphaned.exists() + assert not orphaned_dir.exists() + assert unrelated_file.read_text() == "keep" + assert entry_dir.is_dir() + + @pytest.mark.anyio + async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): + """A live mountpoint belongs to a running process and must survive.""" + cache = RegistryArtifactCache(temp_cache_dir) + paths = cache._paths_for("abc123") + paths.entry_dir.mkdir(parents=True) + mount_dir = paths.squashfs_mount_dir + mount_dir.mkdir() + + with patch.object(Path, "is_mount", lambda self: self == mount_dir): + await cache.ensure_swept() + + assert mount_dir.is_dir() + + @pytest.mark.anyio + async def test_ensure_swept_runs_once(self, temp_cache_dir): + """Repeated startup-sweep requests invoke the sweep once.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with patch.object( + cache, + "_sweep_startup_state", + wraps=cache._sweep_startup_state, + ) as sweep: + await cache.ensure_swept() + await cache.ensure_swept() + + assert sweep.call_count == 1 + + @pytest.mark.anyio + async def test_concurrent_ensure_swept_runs_once(self, temp_cache_dir): + """Concurrent startup-sweep requests invoke the sweep once.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with patch.object( + cache, + "_sweep_startup_state", + wraps=cache._sweep_startup_state, + ) as sweep: + await asyncio.gather(*(cache.ensure_swept() for _ in range(4))) + + assert sweep.call_count == 1 + + @pytest.mark.anyio + async def test_cancelled_waiter_joins_same_startup_sweep(self, temp_cache_dir): + """A cancelled waiter cannot start a second concurrent startup sweep.""" + cache = RegistryArtifactCache(temp_cache_dir) + sweep_started = threading.Event() + sweep_release = threading.Event() + invocation_count = 0 + + def blocking_sweep() -> None: + nonlocal invocation_count + invocation_count += 1 + sweep_started.set() + sweep_release.wait() + + with patch.object( + cache, + "_sweep_startup_state", + side_effect=blocking_sweep, + ): + first_waiter = asyncio.create_task(cache.ensure_swept()) + assert await asyncio.to_thread(sweep_started.wait, 1) + first_waiter.cancel() + + with pytest.raises(asyncio.CancelledError): + await first_waiter + + second_waiter = asyncio.create_task(cache.ensure_swept()) + await asyncio.sleep(0) + sweep_release.set() + await second_waiter + + assert invocation_count == 1 + assert cache._swept is True + + @pytest.mark.anyio + async def test_lease_triggers_startup_sweep(self, temp_cache_dir): + """Lease admission reclaims startup scratch before yielding paths.""" + cache = RegistryArtifactCache(temp_cache_dir) + orphaned_dir = cache.staging_dir / "abc123.999999.4321" + orphaned_dir.mkdir(parents=True) + + async with cache.lease(None): + assert not orphaned_dir.exists() + + @pytest.mark.anyio + async def test_failed_ensure_swept_retries(self, temp_cache_dir): + """A failed startup sweep is retried by the next caller.""" + cache = RegistryArtifactCache(temp_cache_dir) + orphaned_dir = cache.staging_dir / "abc123.999999.4321" + orphaned_dir.mkdir(parents=True) + + with ( + patch.object( + cache, + "_clear_work_dir", + side_effect=OSError("simulated sweep failure"), + ), + pytest.raises(OSError), + ): + await cache.ensure_swept() + + assert orphaned_dir.is_dir() + + await cache.ensure_swept() + + assert not orphaned_dir.exists() + + @pytest.mark.parametrize("work_dir_name", ["staging", "trash"]) + def test_work_directory_inspection_errors_propagate( + self, + temp_cache_dir, + work_dir_name: str, + ): + """Unreadable work directories must trigger a later cleanup retry.""" + cache = RegistryArtifactCache(temp_cache_dir) + work_dir = getattr(cache, f"{work_dir_name}_dir") + work_dir.mkdir() + + with ( + patch.object(Path, "iterdir", side_effect=PermissionError("denied")), + pytest.raises(PermissionError, match="denied"), + ): + cache._clear_work_dir(work_dir) + + @pytest.mark.anyio + async def test_active_entry_scan_errors_propagate(self, temp_cache_dir): + """Unreadable active entries cannot be treated as an empty cache.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch( + "tracecat.executor.registry_artifact_storage.os.scandir", + side_effect=PermissionError("denied"), + ), + pytest.raises(PermissionError, match="denied"), + ): + await cache._enforce_cache_budget() + + @pytest.mark.anyio + async def test_failed_trash_cleanup_skips_startup_trimming(self, temp_cache_dir): + """A failed orphan deletion must not cascade into active retirements.""" + cache = RegistryArtifactCache(temp_cache_dir) + orphan = cache.trash_dir / "orphan" + orphan.mkdir(parents=True) + oldest = write_image_entry( + temp_cache_dir, + "oldest", + size=16, + mtime=100.0, + ) + newest = write_image_entry( + temp_cache_dir, + "newest", + size=16, + mtime=200.0, + ) + real_delete = _delete_cache_path + + def fail_orphan(path: Path) -> bool: + if path == orphan: + return False + return real_delete(path) + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifact_storage._delete_cache_path", + side_effect=fail_orphan, + ), + ): + await cache.ensure_swept() + + assert orphan.is_dir() + assert oldest.is_file() + assert newest.is_file() + assert cache._budget_dirty is True + + @pytest.mark.anyio + async def test_failed_startup_cleanup_retries_exact_path(self, temp_cache_dir): + """Failed startup scratch cleanup retries without sweeping live staging.""" + cache = RegistryArtifactCache(temp_cache_dir) + orphaned = cache.staging_dir / "orphaned.123.456.tmp" + orphaned.parent.mkdir(parents=True) + orphaned.write_bytes(b"partial") + real_delete = _delete_cache_path + failed_once = False + + def fail_once(path: Path) -> bool: + nonlocal failed_once + if path == orphaned and not failed_once: + failed_once = True + return False + return real_delete(path) + + with patch( + "tracecat.executor.registry_artifact_storage._delete_cache_path", + side_effect=fail_once, + ): + await cache.ensure_swept() + assert orphaned.is_file() + assert cache._failed_startup_cleanup == {orphaned} + assert cache._budget_dirty is True + + assert await cache._enforce_cache_budget() is True + + assert not orphaned.exists() + assert cache._failed_startup_cleanup == set() + + @pytest.mark.anyio + async def test_failed_startup_retirement_stays_dirty_and_retries( + self, temp_cache_dir + ): + """A startup rename failure preserves entries for later enforcement.""" + oldest = write_image_entry( + temp_cache_dir, + "oldest", + size=16, + mtime=100.0, + ) + newest = write_image_entry( + temp_cache_dir, + "newest", + size=16, + mtime=200.0, + ) + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifact_storage._move_entry_to_trash", + side_effect=OSError("rename failed"), + ), + ): + await cache.ensure_swept() + + assert oldest.is_file() + assert newest.is_file() + assert cache._budget_dirty is True + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + await cache._converge_cache_budget() + + assert not oldest.exists() + assert newest.is_file() + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_failed_startup_physical_delete_retries_exact_path( + self, temp_cache_dir + ): + """Startup stops retiring entries when bytes were not reclaimed.""" + oldest = write_image_entry( + temp_cache_dir, + "oldest", + size=16, + mtime=100.0, + ) + older = write_image_entry( + temp_cache_dir, + "older", + size=16, + mtime=200.0, + ) + newest = write_image_entry( + temp_cache_dir, + "newest", + size=16, + mtime=300.0, + ) + cache = RegistryArtifactCache(temp_cache_dir) + real_delete = _delete_cache_path + failed_once = False + + def fail_once(path: Path) -> bool: + nonlocal failed_once + if path.parent == cache.trash_dir and not failed_once: + failed_once = True + return False + return real_delete(path) + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 16), + patch( + "tracecat.executor.registry_artifact_storage._delete_cache_path", + side_effect=fail_once, + ), + ): + await cache.ensure_swept() + assert not oldest.exists() + assert older.is_file() + assert newest.is_file() + assert len(tuple(cache.trash_dir.iterdir())) == 1 + assert cache._budget_dirty is True + + await cache._converge_cache_budget() + + assert not older.exists() + assert newest.is_file() + assert not any(cache.trash_dir.iterdir()) + assert cache._budget_dirty is False diff --git a/tests/unit/test_multitenant_registry.py b/tests/unit/executor/test_registry_artifact_tarball.py similarity index 83% rename from tests/unit/test_multitenant_registry.py rename to tests/unit/executor/test_registry_artifact_tarball.py index 7857dd6326..c25937e17d 100644 --- a/tests/unit/test_multitenant_registry.py +++ b/tests/unit/executor/test_registry_artifact_tarball.py @@ -1,14 +1,8 @@ -"""Tests for tarball cache behavior in registry action runner. - -These tests verify: -1. Tarball cache behavior (concurrent downloads, cache keys) -2. Cache key isolation per tarball URI -""" +"""Tarball cache behavior through the public lease API.""" from __future__ import annotations import asyncio -import tempfile from pathlib import Path from unittest.mock import patch @@ -20,29 +14,10 @@ compute_registry_artifact_cache_key, ) -# ============================================================================= -# Fixtures -# ============================================================================= - - -@pytest.fixture -def temp_cache_dir(): - """Create a temporary cache directory for each test.""" - with tempfile.TemporaryDirectory() as tmpdir: - yield Path(tmpdir) - - -async def _lease_paths( - cache: RegistryArtifactCache, - artifact_uri: str, -) -> list[Path]: - async with cache.lease([artifact_uri]) as paths: - return paths - - -# ============================================================================= -# Test Class: Tarball Cache Behavior -# ============================================================================= +from .registry_artifact_test_helpers import ( + lease_paths, + tarball_payload, +) class TestTarballCacheBehavior: @@ -77,7 +52,7 @@ async def test_concurrent_same_uri_single_download(self, temp_cache_dir: Path): async def mock_download(self, ctx, path: Path): download_count[0] += 1 await asyncio.sleep(0.1) # Simulate network delay - path.write_bytes(b"fake tarball") + path.write_bytes(tarball_payload(size=1)) async def mock_extract(self, tarball_path: Path, target_dir: Path): (target_dir / "extracted.txt").write_text("content") @@ -88,10 +63,10 @@ async def mock_extract(self, tarball_path: Path, target_dir: Path): ): # Launch multiple concurrent requests results = await asyncio.gather( - _lease_paths(cache, tarball_uri), - _lease_paths(cache, tarball_uri), - _lease_paths(cache, tarball_uri), - _lease_paths(cache, tarball_uri), + lease_paths(cache, tarball_uri), + lease_paths(cache, tarball_uri), + lease_paths(cache, tarball_uri), + lease_paths(cache, tarball_uri), ) # All should return same path @@ -124,7 +99,7 @@ async def test_different_uris_separate_cache_entries(self, temp_cache_dir: Path) async def mock_download(self, ctx, path: Path): download_calls.append(self.uri) - path.write_bytes(b"fake tarball") + path.write_bytes(tarball_payload(size=1)) async def mock_extract(self, tarball_path: Path, target_dir: Path): (target_dir / "extracted.txt").write_text("content") @@ -135,7 +110,7 @@ async def mock_extract(self, tarball_path: Path, target_dir: Path): ): results = [] for uri in uris: - result = await _lease_paths(cache, uri) + result = await lease_paths(cache, uri) results.append(result) # All results should be different paths @@ -165,7 +140,7 @@ async def test_failed_extraction_cleans_up_temp_files(self, temp_cache_dir: Path cache_key = compute_registry_artifact_cache_key(tarball_uri) async def mock_download(self, ctx, path: Path): - path.write_bytes(b"corrupt tarball") + path.write_bytes(tarball_payload(size=1)) async def mock_extract(self, tarball_path: Path, target_dir: Path): raise RuntimeError("Extraction failed - corrupt tarball") @@ -175,7 +150,7 @@ async def mock_extract(self, tarball_path: Path, target_dir: Path): patch.object(TarballArtifact, "extract", mock_extract), ): with pytest.raises(RuntimeError, match="Extraction failed"): - await _lease_paths(cache, tarball_uri) + await lease_paths(cache, tarball_uri) # Verify no temp files remain temp_files = ( @@ -206,7 +181,7 @@ async def test_cache_reused_on_second_request(self, temp_cache_dir: Path): async def mock_download(self, ctx, path: Path): download_count[0] += 1 - path.write_bytes(b"tarball") + path.write_bytes(tarball_payload(size=1)) async def mock_extract(self, tarball_path: Path, target_dir: Path): (target_dir / "file.txt").write_text("content") @@ -216,11 +191,11 @@ async def mock_extract(self, tarball_path: Path, target_dir: Path): patch.object(TarballArtifact, "extract", mock_extract), ): # First request - result1 = await _lease_paths(cache, tarball_uri) + result1 = await lease_paths(cache, tarball_uri) assert download_count[0] == 1 # Second request (should use cache) - result2 = await _lease_paths(cache, tarball_uri) + result2 = await lease_paths(cache, tarball_uri) assert download_count[0] == 1 # No additional download assert result1 == result2 diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py deleted file mode 100644 index f0b8ea5e68..0000000000 --- a/tests/unit/test_registry_artifacts.py +++ /dev/null @@ -1,3423 +0,0 @@ -"""Tests for executor registry artifact materialization.""" - -from __future__ import annotations - -import asyncio -import io -import os -import tarfile -import tempfile -import threading -from collections.abc import Awaitable, Callable -from dataclasses import dataclass, field -from pathlib import Path -from unittest.mock import ANY, AsyncMock, patch - -import httpx -import pytest -import tracecat_registry - -from tracecat.executor.registry_artifacts import ( - SQUASHFS_MOUNT_OPTIONS, - RegistryArtifactCache, - RegistryArtifactCacheCapacityError, - RegistryArtifactCacheLoopError, - RegistryArtifactEviction, - RegistryArtifactFormat, - RegistryArtifactMaterializationContext, - SquashfsArtifact, - SquashfsMountCommandError, - TarballArtifact, - _delete_cache_path, - _squashfs_listing_size, - bundled_builtin_registry_uri, - compute_registry_artifact_cache_key, -) -from tracecat.registry.artifact_keys import parse_s3_uri - -MAX_ENTRIES_CONFIG = ( - "tracecat.executor.registry_artifacts.config" - ".TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES" -) -MAX_BYTES_CONFIG = ( - "tracecat.executor.registry_artifacts.config" - ".TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES" -) -SQUASHFS_ENABLED_CONFIG = ( - "tracecat.executor.registry_artifacts.config" - ".TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED" -) - - -def _write_tarball_entry(cache_dir: Path, cache_key: str) -> Path: - """Create a materialized tarball cache entry on disk.""" - target_dir = cache_dir / "entries" / cache_key / "tarball" - target_dir.mkdir(parents=True) - (target_dir / "module.py").write_text("VALUE = 1") - return target_dir - - -def _write_image_entry( - cache_dir: Path, cache_key: str, *, size: int, mtime: float -) -> Path: - """Create a downloaded SquashFS image cache entry with a fixed mtime.""" - entry_dir = cache_dir / "entries" / cache_key - entry_dir.mkdir(parents=True, exist_ok=True) - image_path = entry_dir / "image.squashfs" - image_path.write_bytes(b"x" * size) - os.utime(image_path, (mtime, mtime)) - os.utime(entry_dir, (mtime, mtime)) - return image_path - - -def _tarball_payload(*, size: int) -> bytes: - """Return a gzip tarball containing one synthetic regular file.""" - payload = b"x" * size - output = io.BytesIO() - with tarfile.open(fileobj=output, mode="w:gz") as tar: - member = tarfile.TarInfo("module.py") - member.size = len(payload) - tar.addfile(member, io.BytesIO(payload)) - return output.getvalue() - - -@dataclass(slots=True) -class _SquashfsMountHarness: - """Model mount ownership without consuming host loop devices.""" - - cache: RegistryArtifactCache - failed_mount_keys: set[str] = field(default_factory=set) - mounted: set[Path] = field(default_factory=set) - mount_attempts: list[str] = field(default_factory=list) - extraction_attempts: list[str] = field(default_factory=list) - unmounts: list[Path] = field(default_factory=list) - - def seed_mount(self, cache_key: str) -> Path: - """Create one reusable image and mark its mount directory active.""" - paths = self.cache._paths_for(cache_key) - paths.entry_dir.mkdir(parents=True, exist_ok=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir(exist_ok=True) - (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") - self.mounted.add(paths.squashfs_mount_dir) - return paths.squashfs_mount_dir - - async def mount( - self, - ctx: RegistryArtifactMaterializationContext, - image_path: Path, - ) -> Path: - """Record a mount attempt, failing selected keys like loop exhaustion.""" - self.mount_attempts.append(ctx.cache_key) - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - image_path.write_bytes(b"squashfs") - if ctx.cache_key in self.failed_mount_keys: - raise SquashfsMountCommandError("no free loop device") - ctx.paths.squashfs_mount_dir.mkdir(exist_ok=True) - (ctx.paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") - self.mounted.add(ctx.paths.squashfs_mount_dir) - return ctx.paths.squashfs_mount_dir - - async def extract( - self, - ctx: RegistryArtifactMaterializationContext, - image_path: Path, - ) -> Path: - """Publish an extracted fallback from the image retained by mount.""" - assert image_path.is_file() - self.extraction_attempts.append(ctx.cache_key) - ctx.paths.squashfs_extract_dir.mkdir(parents=True, exist_ok=True) - (ctx.paths.squashfs_extract_dir / "module.py").write_text("VALUE = 1") - return ctx.paths.squashfs_extract_dir - - async def unmount(self, mount_dir: Path) -> bool: - """Record final-release unmounts and release the modeled loop device.""" - self.unmounts.append(mount_dir) - self.mounted.discard(mount_dir) - return True - - -async def _materialize( - cache: RegistryArtifactCache, - cache_key: str, - artifact_uri: str, -) -> list[Path]: - """Exercise internal materialization while releasing its test-only lease.""" - await cache.ensure_swept() - ctx = cache._context_for(cache_key) - lock = cache._runtime_for(cache_key).lock - lease_acquired = False - try: - async with lock: - cache._acquire_lease(cache_key) - lease_acquired = True - candidates = await cache._artifact_candidates(ctx, artifact_uri) - if cached_paths := cache._first_cached_path(candidates, ctx): - return cached_paths - paths = await cache._materialize_candidates(ctx, candidates) - cache._touch_entry(cache_key) - await cache._enforce_cache_budget(protected_key=cache_key) - return paths - finally: - if lease_acquired: - cache._release_lease(cache_key) - - -class _BlockingSubprocess: - """Fake subprocess that blocks in communicate until it is cancelled.""" - - def __init__(self) -> None: - self.communicate_started = asyncio.Event() - self.cleanup_calls: list[str] = [] - self.returncode: int | None = None - - async def communicate(self) -> tuple[bytes, bytes]: - """Block until the task awaiting subprocess completion is cancelled.""" - self.communicate_started.set() - await asyncio.Event().wait() - return b"", b"" - - def kill(self) -> None: - """Record that the subprocess was killed.""" - self.cleanup_calls.append("kill") - self.returncode = -9 - - async def wait(self) -> int: - """Record that the killed subprocess was reaped.""" - self.cleanup_calls.append("wait") - return -9 - - -class _CapturedSubprocess: - """Capture cleanup of a real subprocess used by cancellation tests.""" - - def __init__(self, process: asyncio.subprocess.Process) -> None: - self.process = process - self.killed = False - self.reaped = False - - @property - def returncode(self) -> int | None: - """Return the wrapped subprocess exit status.""" - return self.process.returncode - - async def communicate(self) -> tuple[bytes, bytes]: - """Wait for the wrapped subprocess and collect its output.""" - return await self.process.communicate() - - def kill(self) -> None: - """Kill the wrapped subprocess and record the signal.""" - self.killed = True - self.process.kill() - - async def wait(self) -> int: - """Reap the wrapped subprocess and record completion.""" - returncode = await self.process.wait() - self.reaped = True - return returncode - - -@pytest.fixture -def temp_cache_dir(): - """Create a temporary cache directory.""" - with tempfile.TemporaryDirectory() as tmpdir: - yield Path(tmpdir) - - -class TestParseS3Uri: - """Tests for parse_s3_uri function.""" - - def test_valid_uri(self): - """Test parsing a valid S3 URI.""" - bucket, key = parse_s3_uri("s3://my-bucket/path/to/file.tar.gz") - assert bucket == "my-bucket" - assert key == "path/to/file.tar.gz" - - def test_uri_with_nested_path(self): - """Test parsing URI with deeply nested path.""" - bucket, key = parse_s3_uri("s3://bucket/a/b/c/d/e/file.tar.gz") - assert bucket == "bucket" - assert key == "a/b/c/d/e/file.tar.gz" - - def test_invalid_uri_no_prefix(self): - """Test that non-S3 URIs raise ValueError.""" - with pytest.raises(ValueError, match="Invalid S3 URI"): - parse_s3_uri("https://bucket/key") - - def test_invalid_uri_no_key(self): - """Test that URIs without keys raise ValueError.""" - with pytest.raises(ValueError, match="Invalid S3 URI"): - parse_s3_uri("s3://bucket") - - def test_invalid_uri_empty_bucket(self): - """Test that URIs with empty bucket raise ValueError.""" - with pytest.raises(ValueError, match="Invalid S3 URI"): - parse_s3_uri("s3:///key") - - -class TestRegistryArtifactCache: - """Tests for registry artifact cache behavior.""" - - def test_compute_registry_artifact_cache_key_deterministic(self): - """Test that cache key computation is deterministic.""" - uri = "s3://bucket/path/to/registry-v1.2.3.tar.gz" - - key1 = compute_registry_artifact_cache_key(uri) - key2 = compute_registry_artifact_cache_key(uri) - - assert key1 == key2 - assert len(key1) == 16 - - def test_compute_registry_artifact_cache_key_case_sensitive(self): - """Test that cache key is case-sensitive because S3 keys are case-sensitive.""" - key1 = compute_registry_artifact_cache_key("s3://BUCKET/PATH/FILE.tar.gz") - key2 = compute_registry_artifact_cache_key("s3://bucket/path/file.tar.gz") - - assert key1 != key2 - - def test_compute_registry_artifact_cache_key_empty(self): - """Test that empty URI returns the base cache key.""" - assert compute_registry_artifact_cache_key("") == "base" - - @pytest.mark.anyio - async def test_download_artifact_uses_blob_download_file_to_path( - self, temp_cache_dir - ): - """Test that artifact downloads stay behind the blob storage helper.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact = SquashfsArtifact( - uri="s3://bucket/path/site-packages.squashfs", - cache_key="download-test", - ) - ctx = cache._context_for(artifact.cache_key) - output_path = temp_cache_dir / "artifact.squashfs" - - async def mock_download_file_to_path( - *, - key: str, - bucket: str, - output_path: Path, - ) -> None: - output_path.write_bytes(b"squashfs") - - with patch( - "tracecat.executor.registry_artifacts.blob.download_file_to_path", - new_callable=AsyncMock, - side_effect=mock_download_file_to_path, - ) as download_file_to_path: - await artifact.download(ctx, output_path) - - download_file_to_path.assert_awaited_once() - await_args = download_file_to_path.await_args - assert await_args is not None - assert await_args.kwargs["key"] == "path/site-packages.squashfs" - assert await_args.kwargs["bucket"] == "bucket" - assert output_path.read_bytes() == b"squashfs" - - @pytest.mark.anyio - async def test_lease_uses_bundled_current_builtin( - self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch - ): - """In-tree builtin registry returns only the installed site-packages.""" - version = "1.2.3" - site_packages = temp_cache_dir / "venv" / "site-packages" - package_dir = site_packages / "tracecat_registry" - package_dir.mkdir(parents=True) - package_file = package_dir / "__init__.py" - package_file.write_text("") - - monkeypatch.setattr(tracecat_registry, "__version__", version) - monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) - monkeypatch.setattr( - "tracecat.executor.registry_artifact_materialization.sysconfig.get_path", - lambda name: str(site_packages) if name == "purelib" else None, - ) - - cache = RegistryArtifactCache(temp_cache_dir) - async with cache.lease([bundled_builtin_registry_uri(version)]) as result: - assert result == [site_packages.resolve()] - - @pytest.mark.anyio - async def test_lease_exposes_editable_builtin_parent( - self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch - ): - """Editable builtin registry exposes the package wrapper + site-packages.""" - version = "1.2.3" - site_packages = temp_cache_dir / "venv" / "site-packages" - dependency_dir = site_packages / "orjson" - dependency_dir.mkdir(parents=True) - (dependency_dir / "__init__.py").write_text("VALUE = 1\n") - source_root = temp_cache_dir / "src" / "tracecat-registry" - package_dir = source_root / "tracecat_registry" - package_dir.mkdir(parents=True) - package_file = package_dir / "__init__.py" - package_file.write_text("__version__ = '1.2.3'\n") - - monkeypatch.setattr(tracecat_registry, "__version__", version) - monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) - monkeypatch.setattr( - "tracecat.executor.registry_artifact_materialization.sysconfig.get_path", - lambda name: str(site_packages) if name == "purelib" else None, - ) - - cache = RegistryArtifactCache(temp_cache_dir) - async with cache.lease([bundled_builtin_registry_uri(version)]) as result: - assert result == [source_root.resolve(), site_packages.resolve()] - - @pytest.mark.anyio - async def test_lease_rejects_stale_bundled_builtin( - self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch - ): - """Bundled pseudo-URIs must match this executor's installed package.""" - monkeypatch.setattr(tracecat_registry, "__version__", "1.2.3") - - cache = RegistryArtifactCache(temp_cache_dir) - with pytest.raises(RuntimeError, match="does not match installed version"): - async with cache.lease([bundled_builtin_registry_uri("1.2.4")]): - pass - - @pytest.mark.anyio - async def test_download_artifact_normalizes_missing_objects_to_http_404( - self, temp_cache_dir - ): - """Preserve the missing-artifact error contract from presigned downloads.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact = TarballArtifact( - uri="s3://bucket/path/site-packages.tar.gz", - cache_key="missing-test", - ) - ctx = cache._context_for(artifact.cache_key) - output_path = temp_cache_dir / "artifact.tar.gz" - - with patch( - "tracecat.executor.registry_artifacts.blob.download_file_to_path", - new_callable=AsyncMock, - side_effect=FileNotFoundError, - ): - with pytest.raises(httpx.HTTPStatusError) as exc_info: - await artifact.download(ctx, output_path) - - assert exc_info.value.response.status_code == 404 - assert isinstance(exc_info.value.__cause__, FileNotFoundError) - - @pytest.mark.anyio - async def test_artifact_candidates_prefer_squashfs_sidecar(self, temp_cache_dir): - """Test that gzip tarballs prefer a sibling SquashFS sidecar.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - return_value=True, - ) as file_exists, - patch.object(cache, "_can_try_squashfs", return_value=True), - ): - cache_key = compute_registry_artifact_cache_key( - "s3://bucket/path/site-packages.tar.gz" - ) - ctx = cache._context_for(cache_key) - candidates = await cache._artifact_candidates( - ctx, "s3://bucket/path/site-packages.tar.gz" - ) - - artifact = candidates[0] - assert len(candidates) == 2 - assert isinstance(artifact, SquashfsArtifact) - assert isinstance(candidates[1], TarballArtifact) - assert artifact.uri == "s3://bucket/path/site-packages.squashfs" - assert artifact.format == RegistryArtifactFormat.SQUASHFS - file_exists.assert_awaited_once_with( - key="path/site-packages.squashfs", - bucket="bucket", - ) - - @pytest.mark.anyio - async def test_same_key_cold_fan_in_materializes_and_enforces_once( - self, temp_cache_dir - ): - """Same-key waiters share candidate lookup, materialization, and enforcement.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/path/site-packages.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - materialization_started = asyncio.Event() - finish_materialization = asyncio.Event() - - async def mock_materialize( - self: TarballArtifact, - ctx: RegistryArtifactMaterializationContext, - ) -> list[Path]: - materialization_started.set() - await finish_materialization.wait() - ctx.paths.tarball_target_dir.mkdir(parents=True) - return [ctx.paths.tarball_target_dir] - - async def take_lease() -> list[Path]: - async with cache.lease([artifact_uri]) as paths: - return paths - - with ( - patch.object( - cache, - "_sidecar_exists", - new_callable=AsyncMock, - return_value=False, - ) as sidecar_exists, - patch.object( - cache, - "_enforce_cache_budget", - new_callable=AsyncMock, - return_value=True, - ) as enforce_cache_budget, - patch.object(cache, "_converge_cache_budget", new_callable=AsyncMock), - patch.object(TarballArtifact, "materialize", mock_materialize), - ): - leases = [asyncio.create_task(take_lease()) for _ in range(5)] - await materialization_started.wait() - await asyncio.sleep(0) - finish_materialization.set() - results = await asyncio.gather(*leases) - - expected_paths = [cache._paths_for(cache_key).tarball_target_dir] - assert results == [expected_paths] * 5 - sidecar_exists.assert_awaited_once() - enforce_cache_budget.assert_awaited_once_with(protected_key=cache_key) - - @pytest.mark.anyio - async def test_different_cold_keys_serialize_materialization( - self, temp_cache_dir: Path - ) -> None: - """Distinct cold keys cannot multiply staging and extraction peaks.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - uris = ["s3://bucket/first.tar.gz", "s3://bucket/second.tar.gz"] - first_started = asyncio.Event() - release_first = asyncio.Event() - second_started = asyncio.Event() - started_keys: list[str] = [] - - async def controlled_materialize( - self: TarballArtifact, - ctx: RegistryArtifactMaterializationContext, - ) -> list[Path]: - del self - started_keys.append(ctx.cache_key) - if len(started_keys) == 1: - first_started.set() - await release_first.wait() - else: - second_started.set() - ctx.paths.tarball_target_dir.mkdir(parents=True) - return [ctx.paths.tarball_target_dir] - - async def take_lease(uri: str) -> None: - async with cache.lease([uri]): - pass - - with ( - patch(SQUASHFS_ENABLED_CONFIG, False), - patch.object(TarballArtifact, "materialize", controlled_materialize), - patch.object(cache, "_enforce_cache_budget", new_callable=AsyncMock), - patch.object(cache, "_converge_cache_budget", new_callable=AsyncMock), - ): - first = asyncio.create_task(take_lease(uris[0])) - await first_started.wait() - second = asyncio.create_task(take_lease(uris[1])) - await asyncio.sleep(0) - assert not second_started.is_set() - - release_first.set() - await second_started.wait() - await asyncio.gather(first, second) - - assert started_keys == [ - compute_registry_artifact_cache_key(uri) for uri in uris - ] - - def test_squashfs_listing_size_sums_files_and_symlinks(self) -> None: - listing = b"\n".join( - [ - b"drwxr-xr-x 0/0 64 2026-01-01 00:00 squashfs-root", - b"-rw-r--r-- 0/0 123 2026-01-01 00:00 squashfs-root/module.py", - b"lrwxrwxrwx 0/0 9 2026-01-01 00:00 squashfs-root/current -> module.py", - ] - ) - - assert _squashfs_listing_size(listing) == 132 - - def test_squashfs_listing_size_rejects_unparseable_files(self) -> None: - with pytest.raises(ValueError, match="Could not parse SquashFS listing"): - _squashfs_listing_size(b"-rw-r--r-- malformed") - - @pytest.mark.anyio - async def test_artifact_candidates_direct_squashfs_include_gzip_fallback( - self, temp_cache_dir - ): - """Test direct SquashFS URIs fall back to sibling gzip tarballs.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with patch.object(cache, "_can_try_squashfs") as can_try_squashfs: - cache_key = compute_registry_artifact_cache_key( - "s3://bucket/path/site-packages.squashfs" - ) - ctx = cache._context_for(cache_key) - candidates = await cache._artifact_candidates( - ctx, - "s3://bucket/path/site-packages.squashfs", - ) - - assert isinstance(candidates[0], SquashfsArtifact) - assert isinstance(candidates[1], TarballArtifact) - assert [artifact.uri for artifact in candidates] == [ - "s3://bucket/path/site-packages.squashfs", - "s3://bucket/path/site-packages.tar.gz", - ] - assert [artifact.format for artifact in candidates] == [ - RegistryArtifactFormat.SQUASHFS, - RegistryArtifactFormat.TAR_GZ, - ] - can_try_squashfs.assert_not_called() - - @pytest.mark.anyio - async def test_artifact_candidates_fall_back_to_gzip(self, temp_cache_dir): - """Test that gzip tarballs are used when no sidecar exists.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - return_value=False, - ), - patch.object(cache, "_can_try_squashfs", return_value=True), - ): - cache_key = compute_registry_artifact_cache_key( - "s3://bucket/path/site-packages.tar.gz" - ) - ctx = cache._context_for(cache_key) - candidates = await cache._artifact_candidates( - ctx, "s3://bucket/path/site-packages.tar.gz" - ) - - artifact = candidates[0] - assert len(candidates) == 1 - assert isinstance(artifact, TarballArtifact) - assert artifact.uri == "s3://bucket/path/site-packages.tar.gz" - assert artifact.format == RegistryArtifactFormat.TAR_GZ - - def test_can_try_squashfs_does_not_require_mount_binary(self, temp_cache_dir): - """Prefer SquashFS whenever enabled; extraction may work without mounts.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value=None, - ), - patch( - "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", - True, - ), - ): - ctx = cache._context_for("squashfs-test") - assert cache._can_try_squashfs() is True - assert ctx.can_mount_squashfs() is False - - @pytest.mark.anyio - async def test_artifact_candidates_skip_non_registry_tarballs(self, temp_cache_dir): - """Test that arbitrary gzip tarballs do not trigger sidecar lookups.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - ) as file_exists: - cache_key = compute_registry_artifact_cache_key( - "s3://bucket/path/custom.tar.gz" - ) - ctx = cache._context_for(cache_key) - candidates = await cache._artifact_candidates( - ctx, "s3://bucket/path/custom.tar.gz" - ) - - artifact = candidates[0] - assert len(candidates) == 1 - assert isinstance(artifact, TarballArtifact) - assert artifact.uri == "s3://bucket/path/custom.tar.gz" - assert artifact.format == RegistryArtifactFormat.TAR_GZ - file_exists.assert_not_awaited() - - @pytest.mark.anyio - async def test_materialize_mounts_squashfs_sidecar(self, temp_cache_dir): - """Test that a SquashFS sidecar is mounted instead of extracting tarballs.""" - cache = RegistryArtifactCache(temp_cache_dir) - - async def mock_mount(self, ctx, image_path): - assert image_path.name.endswith(".squashfs") - target_dir = ctx.paths.squashfs_mount_dir - target_dir.mkdir(parents=True, exist_ok=True) - (target_dir / "module.py").write_text("VALUE = 1") - return target_dir - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - return_value=True, - ), - patch( - "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", - True, - ), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", mock_mount), - patch.object( - TarballArtifact, - "materialize", - new_callable=AsyncMock, - ) as tarball_materialize, - ): - result = await _materialize( - cache, - "squashfs-key", - "s3://bucket/path/site-packages.tar.gz", - ) - - assert len(result) == 1 - assert (result[0] / "module.py").read_text() == "VALUE = 1" - tarball_materialize.assert_not_awaited() - - @pytest.mark.anyio - async def test_mount_squashfs_uses_hardened_read_only_options( - self, - temp_cache_dir, - ): - """Test that SquashFS images are mounted read-only without device/setuid bits.""" - cache_key = "cache-key" - cache = RegistryArtifactCache(temp_cache_dir) - ctx = cache._context_for(cache_key) - artifact = SquashfsArtifact( - uri="s3://bucket/path/site-packages.squashfs", - cache_key=cache_key, - ) - image_path = ctx.paths.squashfs_image_path - target_dir = ctx.paths.squashfs_mount_dir - ctx.paths.entry_dir.mkdir(parents=True) - image_path.write_bytes(b"squashfs") - target_dir.mkdir() - process = AsyncMock() - process.communicate.return_value = (b"", b"") - process.returncode = 0 - - with patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - new_callable=AsyncMock, - return_value=process, - ) as create_subprocess_exec: - await artifact.mount(ctx, image_path) - - create_subprocess_exec.assert_awaited_once_with( - "mount", - "-t", - "squashfs", - "-o", - SQUASHFS_MOUNT_OPTIONS, - str(image_path), - str(target_dir), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - - @pytest.mark.anyio - async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): - """Cancellation cannot leave an orphan mount process after lock release.""" - cache_key = "cancelled-mount" - cache = RegistryArtifactCache(temp_cache_dir) - ctx = cache._context_for(cache_key) - artifact = SquashfsArtifact( - uri="s3://bucket/path/site-packages.squashfs", - cache_key=cache_key, - ) - image_path = ctx.paths.squashfs_image_path - target_dir = ctx.paths.squashfs_mount_dir - ctx.paths.entry_dir.mkdir(parents=True) - image_path.write_bytes(b"squashfs") - target_dir.mkdir() - process = _BlockingSubprocess() - - with patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - new_callable=AsyncMock, - return_value=process, - ): - mounting = asyncio.create_task( - artifact._mount_image(image_path, target_dir) - ) - await process.communicate_started.wait() - mounting.cancel() - - with pytest.raises(asyncio.CancelledError): - await mounting - - assert process.cleanup_calls == ["kill", "wait"] - assert target_dir.is_dir() - assert not target_dir.is_mount() - - @pytest.mark.anyio - async def test_cancelled_squashfs_extract_kills_and_reaps_subprocess( - self, temp_cache_dir - ): - """Cancellation cannot leave unsquashfs writing into scratch.""" - cache_key = "cancelled-extract" - cache = RegistryArtifactCache(temp_cache_dir) - ctx = cache._context_for(cache_key) - artifact = SquashfsArtifact( - uri="s3://bucket/path/site-packages.squashfs", - cache_key=cache_key, - ) - image_path = ctx.paths.squashfs_image_path - target_dir = ctx.paths.squashfs_extract_dir - ctx.paths.entry_dir.mkdir(parents=True) - image_path.write_bytes(b"squashfs") - target_dir.mkdir() - real_create_subprocess_exec = asyncio.create_subprocess_exec - process_started = asyncio.Event() - captured_processes: list[_CapturedSubprocess] = [] - - async def create_sleep_subprocess( - *args: object, **kwargs: object - ) -> _CapturedSubprocess: - del args, kwargs - process = await real_create_subprocess_exec( - "/bin/sleep", - "30", - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - captured = _CapturedSubprocess(process) - captured_processes.append(captured) - process_started.set() - return captured - - with ( - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/usr/bin/unsquashfs", - ), - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - side_effect=create_sleep_subprocess, - ), - ): - extracting = asyncio.create_task( - artifact._extract_image(image_path, target_dir) - ) - await process_started.wait() - captured = captured_processes[0] - extracting.cancel() - - try: - with pytest.raises(asyncio.CancelledError): - await extracting - finally: - if captured.returncode is None: - captured.kill() - await captured.wait() - - assert captured.killed is True - assert captured.reaped is True - assert captured.returncode is not None - - @pytest.mark.anyio - async def test_repeatedly_cancelled_tarball_extract_rejoins_thread( - self, temp_cache_dir - ): - """Repeated cancellation waits until the tar extractor stops writing.""" - artifact = TarballArtifact( - uri="s3://bucket/path/site-packages.tar.gz", - cache_key="cancelled-tarball-extract", - ) - tarball_path = temp_cache_dir / "artifact.tar.gz" - target_dir = temp_cache_dir / "target" - target_dir.mkdir() - with tarfile.open(tarball_path, "w:gz"): - pass - - extraction_started = threading.Event() - extraction_release = threading.Event() - extraction_finished = threading.Event() - - def blocking_extractall(*args: object, **kwargs: object) -> None: - del args, kwargs - extraction_started.set() - extraction_release.wait() - extraction_finished.set() - - with patch.object( - tarfile.TarFile, - "extractall", - new=blocking_extractall, - ): - extracting = asyncio.create_task(artifact.extract(tarball_path, target_dir)) - assert await asyncio.to_thread(extraction_started.wait, 1) - extracting.cancel() - done, _ = await asyncio.wait({extracting}, timeout=0.05) - first_cancellation_propagated_early = bool(done) - - extracting.cancel() - done, _ = await asyncio.wait({extracting}, timeout=0.05) - second_cancellation_propagated_early = bool(done) - extraction_release.set() - - with pytest.raises(asyncio.CancelledError): - await extracting - - assert extraction_finished.is_set() - assert first_cancellation_propagated_early is False - assert second_cancellation_propagated_early is False - - @pytest.mark.anyio - async def test_materialize_extracts_squashfs_when_mount_fails(self, temp_cache_dir): - """Test that SquashFS mount failures fall back to unsquashfs extraction.""" - cache = RegistryArtifactCache(temp_cache_dir) - - async def mock_mount(self, ctx, image_path): - raise SquashfsMountCommandError("operation not permitted") - - async def mock_extract(self, ctx, image_path): - target_dir = ctx.paths.squashfs_extract_dir - target_dir.mkdir(parents=True, exist_ok=True) - (target_dir / "module.py").write_text("VALUE = 1") - return target_dir - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - return_value=True, - ), - patch( - "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", - True, - ), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", mock_mount), - patch.object(SquashfsArtifact, "extract", mock_extract), - patch.object( - TarballArtifact, - "materialize", - new_callable=AsyncMock, - ) as tarball_materialize, - ): - result = await _materialize( - cache, - "fallback-key", - "s3://bucket/path/site-packages.tar.gz", - ) - - assert len(result) == 1 - assert (result[0] / "module.py").read_text() == "VALUE = 1" - assert result[0].name == "extracted" - tarball_materialize.assert_not_awaited() - - @pytest.mark.anyio - async def test_materialize_extracts_squashfs_without_mount_binary( - self, temp_cache_dir - ): - """Test that SquashFS is still preferred when only unsquashfs is available.""" - cache = RegistryArtifactCache(temp_cache_dir) - - async def mock_extract(self, ctx, image_path): - target_dir = ctx.paths.squashfs_extract_dir - target_dir.mkdir(parents=True, exist_ok=True) - (target_dir / "module.py").write_text("VALUE = 1") - return target_dir - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - return_value=True, - ), - patch( - "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", - True, - ), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value=None, - ), - patch.object(SquashfsArtifact, "extract", mock_extract), - ): - result = await _materialize( - cache, - "extract-key", - "s3://bucket/path/site-packages.tar.gz", - ) - - assert len(result) == 1 - assert (result[0] / "module.py").read_text() == "VALUE = 1" - assert result[0].name == "extracted" - - @pytest.mark.anyio - async def test_materialize_falls_back_to_gzip_when_squashfs_extract_fails( - self, temp_cache_dir - ): - """Test that legacy gzip remains the final compatibility fallback.""" - cache = RegistryArtifactCache(temp_cache_dir) - source = temp_cache_dir / "source" - source.mkdir() - (source / "module.py").write_text("VALUE = 1") - - async def mock_tarball_download(self, ctx, path): - with tarfile.open(path, "w:gz") as tar: - tar.add(source / "module.py", arcname="module.py") - - async def mock_mount(self, ctx, image_path): - raise SquashfsMountCommandError("operation not permitted") - - async def mock_extract(self, ctx, image_path): - raise RuntimeError("unsquashfs unavailable") - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - side_effect=[True, False], - ), - patch( - "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", - True, - ), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", mock_mount), - patch.object(SquashfsArtifact, "extract", mock_extract), - patch.object(TarballArtifact, "download", mock_tarball_download), - ): - result = await _materialize( - cache, - "gzip-fallback-key", - "s3://bucket/path/site-packages.tar.gz", - ) - - assert len(result) == 1 - assert (result[0] / "module.py").read_text() == "VALUE = 1" - assert result[0].name == "tarball" - - @pytest.mark.anyio - async def test_materialize_treats_unknown_suffix_as_gzip(self, temp_cache_dir): - """Test that existing gzip artifacts can use arbitrary S3 key suffixes.""" - cache = RegistryArtifactCache(temp_cache_dir) - source = temp_cache_dir / "source" - source.mkdir() - (source / "module.py").write_text("VALUE = 1") - - async def mock_download(self, ctx, path): - assert path.name.endswith(".tar.gz") - with tarfile.open(path, "w:gz") as tar: - tar.add(source / "module.py", arcname="module.py") - - with patch.object(TarballArtifact, "download", mock_download): - result = await _materialize( - cache, - "custom-key-test", - "s3://bucket/path/custom-key", - ) - - assert len(result) == 1 - assert (result[0] / "module.py").read_text() == "VALUE = 1" - - @pytest.mark.anyio - async def test_materialize_caches_result(self, temp_cache_dir): - """Test that tarball extraction is cached.""" - cache = RegistryArtifactCache(temp_cache_dir) - cache_key = "test-cache-key" - target_dir = cache._paths_for(cache_key).tarball_target_dir - target_dir.mkdir(parents=True) - - result = await _materialize( - cache, - cache_key, - "s3://bucket/test.tar.gz", - ) - - assert result == [target_dir] - - @pytest.mark.anyio - async def test_materialize_concurrent_requests(self, temp_cache_dir): - """Test that concurrent requests for same artifact do not race.""" - cache = RegistryArtifactCache(temp_cache_dir) - cache_key = "concurrent-test" - download_count = 0 - - async def mock_download(self, ctx, path): - nonlocal download_count - download_count += 1 - await asyncio.sleep(0.1) - path.write_bytes(b"fake tarball content") - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "extracted.txt").write_text("extracted") - - with ( - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - results = await asyncio.gather( - _materialize(cache, cache_key, "s3://bucket/test.tar.gz"), - _materialize(cache, cache_key, "s3://bucket/test.tar.gz"), - _materialize(cache, cache_key, "s3://bucket/test.tar.gz"), - ) - - assert all(r == results[0] for r in results) - assert download_count == 1 - - def test_runtime_for_same_key(self, temp_cache_dir): - """The same cache key returns the same runtime state.""" - cache = RegistryArtifactCache(temp_cache_dir) - - runtime1 = cache._runtime_for("key1") - runtime2 = cache._runtime_for("key1") - - assert runtime1 is runtime2 - - def test_runtime_for_different_keys(self, temp_cache_dir): - """Different cache keys return different runtime state.""" - cache = RegistryArtifactCache(temp_cache_dir) - - runtime1 = cache._runtime_for("key1") - runtime2 = cache._runtime_for("key2") - - assert runtime1 is not runtime2 - - -class TestRegistryArtifactCacheLease: - """Tests for lease-based pinning of registry artifact cache entries.""" - - @pytest.mark.anyio - async def test_cache_rejects_use_from_a_second_event_loop( - self, temp_cache_dir: Path - ) -> None: - """Protect the process-wide cache from thread-local Temporal loops. - - The cache owns asyncio locks, tasks, and refcounts as one unit. A typed - ownership failure prevents a future synchronous activity from silently - sharing only part of that state across thread-local event loops, which - previously caused intermittent cross-loop failures in storage caches. - """ - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - - def use_cache_from_another_thread() -> None: - asyncio.run(cache.ensure_swept()) - - with pytest.raises(RegistryArtifactCacheLoopError): - await asyncio.to_thread(use_cache_from_another_thread) - - # A rejected caller must not poison the owning loop's cache. - async with cache.lease(None) as registry_paths: - assert registry_paths == [temp_cache_dir / "base"] - - def test_touch_entry_refreshes_tarball_root_mtime(self, temp_cache_dir): - """Touching a tarball-only entry persists its restart-safe recency.""" - cache = RegistryArtifactCache(temp_cache_dir) - cache_key = "tarball-only" - _write_tarball_entry(temp_cache_dir, cache_key) - entry_dir = cache._paths_for(cache_key).entry_dir - os.utime(entry_dir, (100.0, 100.0)) - - cache._touch_entry(cache_key) - - assert entry_dir.stat().st_mtime > 100.0 - - @pytest.mark.anyio - async def test_lease_refcounts_and_touches_image_mtime(self, temp_cache_dir): - """A lease pins its entry and refreshes the restart-safe LRU timestamp.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/leased.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - target_dir = _write_tarball_entry(temp_cache_dir, cache_key) - image_path = _write_image_entry(temp_cache_dir, cache_key, size=16, mtime=100.0) - entry_dir = cache._paths_for(cache_key).entry_dir - - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [target_dir] - assert cache._refcount(cache_key) == 1 - assert entry_dir.stat().st_mtime > 100.0 - - assert cache._refcount(cache_key) == 0 - assert image_path.is_file() - - @pytest.mark.anyio - async def test_lease_releases_refcount_when_materialization_fails( - self, temp_cache_dir - ): - """A failed materialization must not leak a pin or empty cache entry.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/broken.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - - async def mock_download(self, ctx, path): - raise RuntimeError("download failed") - - with patch.object(TarballArtifact, "download", mock_download): - with pytest.raises(RuntimeError, match="download failed"): - async with cache.lease([artifact_uri]): - pass - - assert cache._refcount(cache_key) == 0 - assert not cache._paths_for(cache_key).entry_dir.exists() - assert cache_key not in cache._discover_cache_keys() - - @pytest.mark.anyio - async def test_failed_first_admission_converges_deposited_image( - self, temp_cache_dir: Path - ) -> None: - """A failed first artifact must not strand an over-budget image.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/failed-first.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) - - async def fail_after_deposit( - artifact: SquashfsArtifact, - ctx: RegistryArtifactMaterializationContext, - ) -> list[Path]: - del artifact - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - ctx.paths.squashfs_image_path.write_bytes(b"reusable image") - raise RuntimeError("mount failed") - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 1), - patch.object( - cache, - "_artifact_candidates", - new_callable=AsyncMock, - return_value=[artifact], - ), - patch.object(SquashfsArtifact, "materialize", fail_after_deposit), - ): - with pytest.raises(RuntimeError, match="mount failed"): - async with cache.lease([artifact_uri]): - pass - - assert cache._refcount(cache_key) == 0 - assert not cache._paths_for(cache_key).entry_dir.exists() - assert cache._discover_cache_keys() == set() - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_cancelled_lease_admission_releases_refcount(self, temp_cache_dir): - """Cancellation during candidate lookup must not leak a permanent pin.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/path/site-packages.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - lookup_started = asyncio.Event() - finish_lookup = asyncio.Event() - - async def blocked_sidecar_lookup(**kwargs): - lookup_started.set() - await finish_lookup.wait() - return False - - async def take_lease() -> None: - async with cache.lease([artifact_uri]): - pass - - with ( - patch(SQUASHFS_ENABLED_CONFIG, True), - patch.object(cache, "_sidecar_exists", blocked_sidecar_lookup), - ): - acquisition = asyncio.create_task(take_lease()) - await lookup_started.wait() - assert cache._refcount(cache_key) == 1 - acquisition.cancel() - with pytest.raises(asyncio.CancelledError): - await acquisition - - assert cache._refcount(cache_key) == 0 - assert not cache._paths_for(cache_key).entry_dir.exists() - - @pytest.mark.anyio - async def test_cancelled_waiter_preserves_existing_same_key_lease( - self, temp_cache_dir: Path - ) -> None: - """A waiter cancelled before admission cannot release another holder.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/path/shared.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - lock = cache._runtime_for(cache_key).lock - - async def take_lease() -> None: - async with cache.lease([artifact_uri]): - pytest.fail("cancelled waiter must not enter the lease context") - - unmount_idle_entry = AsyncMock() - with patch.object(cache, "_unmount_idle_entry", unmount_idle_entry): - async with lock: - cache._acquire_lease(cache_key) - try: - waiter = asyncio.create_task(take_lease()) - await asyncio.sleep(0) - assert not waiter.done() - - waiter.cancel() - with pytest.raises(asyncio.CancelledError): - await waiter - - assert cache._refcount(cache_key) == 1 - unmount_idle_entry.assert_not_awaited() - finally: - cache._release_lease(cache_key) - - @pytest.mark.anyio - async def test_lease_without_uris_returns_base_pythonpath_dir(self, temp_cache_dir): - """No artifact URIs still yields the base PYTHONPATH directory.""" - cache = RegistryArtifactCache(temp_cache_dir) - - async with cache.lease(None) as registry_paths: - assert registry_paths == [temp_cache_dir / "base"] - assert registry_paths[0].is_dir() - - assert cache._runtime == {} - - @pytest.mark.anyio - async def test_lease_preserves_uri_order(self, temp_cache_dir): - """Multiple artifacts keep their deterministic PYTHONPATH order.""" - cache = RegistryArtifactCache(temp_cache_dir) - uris = ["s3://bucket/first.tar.gz", "s3://bucket/second.tar.gz"] - expected = [ - _write_tarball_entry( - temp_cache_dir, compute_registry_artifact_cache_key(uri) - ) - for uri in uris - ] - - async with cache.lease(uris) as registry_paths: - assert registry_paths == expected - - @pytest.mark.anyio - async def test_overlapping_same_key_leases_share_one_mount_until_final_release( - self, temp_cache_dir: Path - ) -> None: - """Protect refcount and loop-device lifetime under heavy fan-in. - - Every holder must observe one published mount, intermediate releases - must leave that mount usable, and exactly the final release may reclaim - its loop device while retaining the downloaded image for a remount. - """ - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/shared.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - harness = _SquashfsMountHarness(cache) - holder_count = 32 - entered = 0 - all_entered = asyncio.Event() - releases = [asyncio.Event() for _ in range(holder_count)] - - async def hold_lease(index: int) -> None: - nonlocal entered - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [ - cache._paths_for(cache_key).squashfs_mount_dir - ] - entered += 1 - if entered == holder_count: - all_entered.set() - await releases[index].wait() - - with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", harness.mount), - patch.object(SquashfsArtifact, "extract", harness.extract), - patch.object(cache, "_unmount", harness.unmount), - ): - holders = [ - asyncio.create_task(hold_lease(index)) for index in range(holder_count) - ] - await asyncio.wait_for(all_entered.wait(), timeout=5) - - mount_dir = cache._paths_for(cache_key).squashfs_mount_dir - assert cache._refcount(cache_key) == holder_count - assert harness.mount_attempts == [cache_key] - assert harness.extraction_attempts == [] - assert mount_dir in harness.mounted - - for release in releases[:-1]: - release.set() - await asyncio.gather(*holders[:-1]) - - assert cache._refcount(cache_key) == 1 - assert mount_dir in harness.mounted - assert harness.unmounts == [] - - releases[-1].set() - await holders[-1] - - assert cache._refcount(cache_key) == 0 - assert harness.unmounts == [mount_dir] - assert mount_dir not in harness.mounted - assert cache._paths_for(cache_key).squashfs_image_path.is_file() - assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) - assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) - - @pytest.mark.anyio - async def test_cancelling_one_holder_preserves_sibling_leases( - self, temp_cache_dir: Path - ) -> None: - """Protect sibling actions from another holder's cancellation. - - Cancellation must release exactly one pin without unmounting the shared - artifact beneath surviving actions; the last surviving holder remains - solely responsible for loop-device reclamation. - """ - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/cancel-one.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - harness = _SquashfsMountHarness(cache) - entered = [asyncio.Event() for _ in range(3)] - releases = [asyncio.Event() for _ in range(3)] - - async def hold_lease(index: int) -> None: - async with cache.lease([artifact_uri]): - entered[index].set() - await releases[index].wait() - - with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", harness.mount), - patch.object(SquashfsArtifact, "extract", harness.extract), - patch.object(cache, "_unmount", harness.unmount), - ): - holders = [asyncio.create_task(hold_lease(index)) for index in range(3)] - await asyncio.wait_for( - asyncio.gather(*(event.wait() for event in entered)), - timeout=5, - ) - - holders[0].cancel() - with pytest.raises(asyncio.CancelledError): - await holders[0] - - mount_dir = cache._paths_for(cache_key).squashfs_mount_dir - assert cache._refcount(cache_key) == 2 - assert mount_dir in harness.mounted - assert harness.unmounts == [] - - releases[1].set() - await holders[1] - assert cache._refcount(cache_key) == 1 - assert mount_dir in harness.mounted - assert harness.unmounts == [] - - releases[2].set() - await holders[2] - - assert cache._refcount(cache_key) == 0 - assert harness.mount_attempts == [cache_key] - assert harness.unmounts == [mount_dir] - - @pytest.mark.anyio - async def test_repeated_cancellation_finishes_all_lease_cleanup( - self, temp_cache_dir: Path - ) -> None: - """Every unmount and budget pass finishes before cancellation propagates.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uris = [ - "s3://bucket/first.tar.gz", - "s3://bucket/second.tar.gz", - ] - cache_keys = [compute_registry_artifact_cache_key(uri) for uri in artifact_uris] - for cache_key in cache_keys: - _write_tarball_entry(temp_cache_dir, cache_key) - - lease_entered = asyncio.Event() - first_unmount_started = asyncio.Event() - finish_first_unmount = asyncio.Event() - cleanup_calls: list[str] = [] - - async def mock_unmount_idle_entry(cache_key: str) -> None: - cleanup_calls.append(cache_key) - if cache_key == cache_keys[0]: - first_unmount_started.set() - await finish_first_unmount.wait() - - async def mock_converge_cache_budget() -> None: - cleanup_calls.append("converge") - - async def hold_lease() -> None: - async with cache.lease(artifact_uris): - lease_entered.set() - await asyncio.Event().wait() - - with ( - patch.object(cache, "_unmount_idle_entry", mock_unmount_idle_entry), - patch.object(cache, "_converge_cache_budget", mock_converge_cache_budget), - ): - holder = asyncio.create_task(hold_lease()) - await lease_entered.wait() - holder.cancel() - await first_unmount_started.wait() - - holder.cancel() - await asyncio.sleep(0) - assert not holder.done() - - finish_first_unmount.set() - with pytest.raises(asyncio.CancelledError): - await holder - - assert cleanup_calls == [*cache_keys, "converge"] - assert all(cache._refcount(cache_key) == 0 for cache_key in cache_keys) - - @pytest.mark.anyio - async def test_new_lease_racing_final_release_prevents_stale_unmount( - self, temp_cache_dir: Path - ) -> None: - """Protect the zero-refcount-to-unmount handoff from a new acquisition. - - A lease can arrive after the old holder decrements to zero but before - unmount takes the key lock. The lock-time refcount recheck must preserve - the mount for that newcomer instead of returning a path being torn down. - """ - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/release-race.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - harness = _SquashfsMountHarness(cache) - mount_dir = harness.seed_mount(cache_key) - first_entered = asyncio.Event() - release_first = asyncio.Event() - release_reached_unmount = asyncio.Event() - allow_unmount_recheck = asyncio.Event() - newcomer_entered = asyncio.Event() - release_newcomer = asyncio.Event() - original_unmount_idle_entry = cache._unmount_idle_entry - unmount_requests = 0 - - async def pause_first_unmount_request(requested_key: str) -> None: - nonlocal unmount_requests - unmount_requests += 1 - if unmount_requests == 1: - release_reached_unmount.set() - await allow_unmount_recheck.wait() - await original_unmount_idle_entry(requested_key) - - async def old_holder() -> None: - async with cache.lease([artifact_uri]): - first_entered.set() - await release_first.wait() - - async def new_holder() -> None: - async with cache.lease([artifact_uri]): - newcomer_entered.set() - await release_newcomer.wait() - - with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), - patch.object(cache, "_unmount", harness.unmount), - patch.object( - cache, - "_unmount_idle_entry", - pause_first_unmount_request, - ), - ): - old_task = asyncio.create_task(old_holder()) - await first_entered.wait() - release_first.set() - await release_reached_unmount.wait() - assert cache._refcount(cache_key) == 0 - - new_task = asyncio.create_task(new_holder()) - await newcomer_entered.wait() - assert cache._refcount(cache_key) == 1 - - allow_unmount_recheck.set() - await old_task - - assert mount_dir in harness.mounted - assert harness.unmounts == [] - - release_newcomer.set() - await new_task - - assert cache._refcount(cache_key) == 0 - assert harness.unmounts == [mount_dir] - - @pytest.mark.anyio - async def test_partial_multi_artifact_failure_rolls_back_prior_leases( - self, temp_cache_dir: Path - ) -> None: - """Protect sequential multi-artifact admission from partial pin leaks. - - If a middle artifact fails, every earlier acquisition must be released - and unmounted, the failed key must leave no shell, later artifacts must - never be acquired, and the release path must still request convergence. - """ - cache = RegistryArtifactCache(temp_cache_dir) - first_uri = "s3://bucket/first.squashfs" - failed_uri = "s3://bucket/failed.tar.gz" - untouched_uri = "s3://bucket/untouched.tar.gz" - first_key = compute_registry_artifact_cache_key(first_uri) - failed_key = compute_registry_artifact_cache_key(failed_uri) - untouched_key = compute_registry_artifact_cache_key(untouched_uri) - harness = _SquashfsMountHarness(cache) - first_mount = harness.seed_mount(first_key) - untouched_path = _write_tarball_entry(temp_cache_dir, untouched_key) - - async def fail_download( - artifact: TarballArtifact, - ctx: RegistryArtifactMaterializationContext, - output_path: Path, - ) -> None: - del artifact, ctx, output_path - raise RuntimeError("download failed") - - tracked_lease_artifact = AsyncMock(wraps=cache._lease_artifact) - converge_cache_budget = AsyncMock() - - with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), - patch(SQUASHFS_ENABLED_CONFIG, False), - patch.object(TarballArtifact, "download", fail_download), - patch.object(cache, "_unmount", harness.unmount), - patch.object(cache, "_lease_artifact", tracked_lease_artifact), - patch.object( - cache, - "_converge_cache_budget", - converge_cache_budget, - ), - ): - with pytest.raises(RuntimeError, match="download failed"): - async with cache.lease([first_uri, failed_uri, untouched_uri]): - pass - - requested_uris = [ - await_call.args[0] for await_call in tracked_lease_artifact.await_args_list - ] - assert requested_uris == [first_uri, failed_uri] - assert cache._refcount(first_key) == 0 - assert cache._refcount(failed_key) == 0 - assert cache._refcount(untouched_key) == 0 - assert harness.unmounts == [first_mount] - assert cache._paths_for(first_key).squashfs_image_path.is_file() - assert not cache._paths_for(failed_key).entry_dir.exists() - assert untouched_path.is_dir() - converge_cache_budget.assert_awaited_once_with() - assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) - assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) - - @pytest.mark.anyio - async def test_waiter_retries_after_first_same_key_materializer_fails( - self, temp_cache_dir: Path - ) -> None: - """Protect same-key waiters from a failed cold-cache publisher. - - The first materializer may fail after writing scratch. Its waiter must - retry against a clean staging area, publish the sole valid entry, and - leave neither a leaked pin nor an empty LRU-visible cache shell. - """ - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/retry-after-failure.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - first_download_started = asyncio.Event() - fail_first_download = asyncio.Event() - retry_download_started = asyncio.Event() - allow_retry_download = asyncio.Event() - download_attempts = 0 - retry_saw_clean_staging = False - - async def controlled_download( - artifact: TarballArtifact, - ctx: RegistryArtifactMaterializationContext, - output_path: Path, - ) -> None: - del artifact, ctx - nonlocal download_attempts, retry_saw_clean_staging - download_attempts += 1 - if download_attempts == 1: - output_path.write_bytes(b"partial") - first_download_started.set() - await fail_first_download.wait() - raise RuntimeError("first publisher failed") - - retry_saw_clean_staging = not any(cache.staging_dir.iterdir()) - retry_download_started.set() - await allow_retry_download.wait() - output_path.write_bytes(b"complete") - - async def mock_extract( - artifact: TarballArtifact, - tarball_path: Path, - target_dir: Path, - ) -> None: - del artifact - assert tarball_path.read_bytes() == b"complete" - (target_dir / "module.py").write_text("VALUE = 1") - - async def take_lease() -> list[Path]: - async with cache.lease([artifact_uri]) as registry_paths: - return registry_paths - - with ( - patch(SQUASHFS_ENABLED_CONFIG, False), - patch( - "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", - return_value=1, - ), - patch.object(TarballArtifact, "download", controlled_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - first = asyncio.create_task(take_lease()) - await first_download_started.wait() - waiter = asyncio.create_task(take_lease()) - await asyncio.sleep(0) - assert download_attempts == 1 - - fail_first_download.set() - await retry_download_started.wait() - allow_retry_download.set() - registry_paths = await waiter - with pytest.raises(RuntimeError, match="first publisher failed"): - await first - - target_dir = cache._paths_for(cache_key).tarball_target_dir - assert registry_paths == [target_dir] - assert (target_dir / "module.py").read_text() == "VALUE = 1" - assert download_attempts == 2 - assert retry_saw_clean_staging is True - assert cache._discover_cache_keys() == {cache_key} - assert cache._refcount(cache_key) == 0 - assert not any(cache.staging_dir.iterdir()) - assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) - - @pytest.mark.anyio - async def test_duplicate_uri_balances_each_acquisition_and_release( - self, temp_cache_dir: Path - ) -> None: - """Protect list bookkeeping when one lease requests the same URI twice. - - Duplicate PYTHONPATH entries intentionally acquire two pins. Both must - be released, while materialization and final unmount still occur once; - accidental deduplication on only one side would leak or underflow pins. - """ - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/duplicate.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - harness = _SquashfsMountHarness(cache) - - with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", harness.mount), - patch.object(SquashfsArtifact, "extract", harness.extract), - patch.object(cache, "_unmount", harness.unmount), - ): - async with cache.lease([artifact_uri, artifact_uri]) as registry_paths: - mount_dir = cache._paths_for(cache_key).squashfs_mount_dir - assert registry_paths == [mount_dir, mount_dir] - assert cache._refcount(cache_key) == 2 - assert harness.mount_attempts == [cache_key] - - assert cache._refcount(cache_key) == 0 - assert harness.unmounts == [mount_dir] - assert cache._paths_for(cache_key).squashfs_image_path.is_file() - - @pytest.mark.anyio - async def test_lease_is_never_admitted_across_an_in_flight_eviction( - self, temp_cache_dir - ): - """A lease must not return a mount an in-flight eviction is deleting.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/path/site-packages.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - paths = cache._paths_for(cache_key) - paths.entry_dir.mkdir(parents=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir() - mounted = {paths.squashfs_mount_dir} - umount_started = asyncio.Event() - finish_umount = asyncio.Event() - remounts: list[str] = [] - - umount_process = AsyncMock() - umount_process.communicate.return_value = (b"", b"") - umount_process.returncode = 0 - - async def mock_umount(*args, **kwargs): - umount_started.set() - await finish_umount.wait() - mounted.discard(paths.squashfs_mount_dir) - return umount_process - - async def mock_mount(self, ctx, image_path): - remounts.append(ctx.cache_key) - target_dir = ctx.paths.squashfs_mount_dir - target_dir.mkdir(parents=True, exist_ok=True) - (target_dir / "module.py").write_text("VALUE = 1") - mounted.add(target_dir) - return target_dir - - leased_paths: list[Path] = [] - leased_path_exists: list[bool] = [] - - async def take_lease() -> None: - async with cache.lease([artifact_uri]) as registry_paths: - leased_paths.extend(registry_paths) - leased_path_exists.append(registry_paths[0].is_dir()) - - with ( - patch.object(Path, "is_mount", lambda self: self in mounted), - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/umount", - ), - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - side_effect=mock_umount, - ), - patch.object(SquashfsArtifact, "mount", mock_mount), - ): - eviction = asyncio.create_task(cache._evict_entry(cache_key)) - await umount_started.wait() - lease = asyncio.create_task(take_lease()) - # Let the lease block on the per-key lock the eviction holds. - await asyncio.sleep(0) - finish_umount.set() - evicted, _ = await asyncio.gather(eviction, lease) - - assert evicted == RegistryArtifactEviction(retired=True, reclaimed=True) - # The lease waited for the eviction and re-materialized the entry. - assert remounts == [cache_key] - assert leased_paths == [paths.squashfs_mount_dir] - assert leased_path_exists == [True] - assert (paths.squashfs_mount_dir / "module.py").read_text() == "VALUE = 1" - - @pytest.mark.anyio - async def test_builtin_artifact_is_exempt_from_cache_accounting( - self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch - ): - """The bundled builtin registry is never a cache entry.""" - version = "1.2.3" - site_packages = temp_cache_dir / "venv" / "site-packages" - package_dir = site_packages / "tracecat_registry" - package_dir.mkdir(parents=True) - package_file = package_dir / "__init__.py" - package_file.write_text("") - - monkeypatch.setattr(tracecat_registry, "__version__", version) - monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) - monkeypatch.setattr( - "tracecat.executor.registry_artifact_materialization.sysconfig.get_path", - lambda name: str(site_packages) if name == "purelib" else None, - ) - - cache = RegistryArtifactCache(temp_cache_dir) - - with patch.object( - cache, - "_enforce_cache_budget", - new_callable=AsyncMock, - ) as enforce_cache_budget: - async with cache.lease([bundled_builtin_registry_uri(version)]) as paths: - assert paths == [site_packages.resolve()] - assert cache._runtime == {} - - enforce_cache_budget.assert_not_awaited() - - -class TestRegistryArtifactCacheEviction: - """Tests for bounded eviction of registry artifact cache entries.""" - - def test_delete_cache_path_reports_directory_failure(self, temp_cache_dir): - """Physical deletion failures are observable instead of suppressed.""" - entry_dir = temp_cache_dir / "entry" - entry_dir.mkdir() - - with ( - patch( - "tracecat.executor.registry_artifact_storage.shutil.rmtree", - side_effect=OSError("permission denied"), - ), - patch( - "tracecat.executor.registry_artifact_storage.logger.warning" - ) as warning, - ): - deleted = _delete_cache_path(entry_dir) - - assert deleted is False - assert entry_dir.is_dir() - warning.assert_called_once_with( - "Failed to delete registry artifact cache path", - path=str(entry_dir), - error="permission denied", - ) - - @pytest.mark.anyio - async def test_leased_entry_survives_eviction_of_idle_entry(self, temp_cache_dir): - """Eviction must never remove an entry a live action is importing from.""" - cache = RegistryArtifactCache(temp_cache_dir) - leased_uri = "s3://bucket/leased.tar.gz" - idle_uri = "s3://bucket/idle.tar.gz" - new_uri = "s3://bucket/new.tar.gz" - leased_dir = _write_tarball_entry( - temp_cache_dir, compute_registry_artifact_cache_key(leased_uri) - ) - idle_dir = _write_tarball_entry( - temp_cache_dir, compute_registry_artifact_cache_key(idle_uri) - ) - - async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_text("VALUE = 2") - - with ( - patch(MAX_ENTRIES_CONFIG, 2), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - async with cache.lease([leased_uri]): - await _materialize( - cache, compute_registry_artifact_cache_key(new_uri), new_uri - ) - - assert leased_dir.is_dir() - assert not idle_dir.exists() - - @pytest.mark.anyio - async def test_cold_download_reserves_space_before_writing( - self, temp_cache_dir: Path - ) -> None: - """Admission evicts idle bytes before a new download enters staging.""" - cache = RegistryArtifactCache(temp_cache_dir) - idle = _write_image_entry(temp_cache_dir, "idle", size=80, mtime=100.0) - artifact_uri = "s3://bucket/new.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - payload = _tarball_payload(size=32) - max_bytes = len(payload) + 32 - capacity_checked = False - - async def download_file_to_path( - *, - key: str, - bucket: str, - output_path: Path, - max_bytes: int, - ensure_capacity: Callable[[int], Awaitable[None]], - ) -> int: - del key, bucket - nonlocal capacity_checked - assert max_bytes == len(payload) + 32 - await ensure_capacity(len(payload)) - capacity_checked = True - assert not idle.exists() - output_path.write_bytes(payload) - return len(payload) - - with ( - patch(SQUASHFS_ENABLED_CONFIG, False), - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, max_bytes), - patch( - "tracecat.executor.registry_artifacts.blob.download_file_to_path", - side_effect=download_file_to_path, - ), - ): - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [ - cache._paths_for(cache_key).tarball_target_dir - ] - - assert capacity_checked is True - assert not idle.exists() - assert (registry_paths[0] / "module.py").read_bytes() == b"x" * 32 - - @pytest.mark.anyio - async def test_compression_heavy_tarball_is_rejected_before_extraction( - self, temp_cache_dir: Path - ) -> None: - """Compressed bytes plus declared extraction cannot exceed the cache cap.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/compression-heavy.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - payload = _tarball_payload(size=4096) - max_bytes = len(payload) + 256 - - async def download_file_to_path( - *, - key: str, - bucket: str, - output_path: Path, - max_bytes: int, - ensure_capacity: Callable[[int], Awaitable[None]], - ) -> int: - del key, bucket - assert max_bytes == len(payload) + 256 - await ensure_capacity(len(payload)) - output_path.write_bytes(payload) - return len(payload) - - with ( - patch(SQUASHFS_ENABLED_CONFIG, False), - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, max_bytes), - patch( - "tracecat.executor.registry_artifacts.blob.download_file_to_path", - side_effect=download_file_to_path, - ), - patch.object( - TarballArtifact, - "extract", - new_callable=AsyncMock, - ) as extract, - ): - with pytest.raises(RegistryArtifactCacheCapacityError) as raised: - async with cache.lease([artifact_uri]): - pass - - assert raised.value.additional_bytes == 4096 - assert raised.value.max_bytes == max_bytes - extract.assert_not_awaited() - assert not cache._paths_for(cache_key).entry_dir.exists() - assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) - - @pytest.mark.anyio - async def test_squashfs_expansion_is_rejected_before_extraction( - self, temp_cache_dir: Path - ) -> None: - """SquashFS metadata is accounted before unsquashfs writes scratch.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/oversized.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - - async def download( - self: SquashfsArtifact, - ctx: RegistryArtifactMaterializationContext, - image_path: Path, - ) -> float: - del self, ctx - image_path.parent.mkdir(parents=True, exist_ok=True) - image_path.write_bytes(b"image") - return 0.0 - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 100), - patch.object( - RegistryArtifactMaterializationContext, - "can_mount_squashfs", - return_value=False, - ), - patch.object(SquashfsArtifact, "download", download), - patch.object( - SquashfsArtifact, - "_squashfs_extracted_size", - new_callable=AsyncMock, - return_value=101, - ), - patch.object( - SquashfsArtifact, - "_extract_image", - new_callable=AsyncMock, - ) as extract_image, - ): - with pytest.raises(RegistryArtifactCacheCapacityError) as raised: - async with cache.lease([artifact_uri]): - pass - - assert raised.value.additional_bytes == 101 - extract_image.assert_not_awaited() - paths = cache._paths_for(cache_key) - assert paths.squashfs_image_path.read_bytes() == b"image" - assert not paths.squashfs_extract_dir.exists() - - @pytest.mark.anyio - async def test_successful_admission_enforces_actual_size_before_yield( - self, temp_cache_dir - ): - """Post-publication enforcement sees the new entry's actual size.""" - cache = RegistryArtifactCache(temp_cache_dir) - idle = _write_image_entry(temp_cache_dir, "idle", size=4096, mtime=100.0) - new_uri = "s3://bucket/new.tar.gz" - new_key = compute_registry_artifact_cache_key(new_uri) - - async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_bytes(b"x" * 4096) - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 6000), - patch( - "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", - return_value=4096, - ), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - async with cache.lease([new_uri]) as registry_paths: - assert registry_paths == [cache._paths_for(new_key).tarball_target_dir] - assert not idle.exists() - - assert not idle.exists() - assert cache._paths_for(new_key).tarball_target_dir.is_dir() - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_deletion_failure_does_not_block_materialization( - self, temp_cache_dir - ): - """Failed cleanup is reported while cache admission remains fail-open.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - _write_image_entry(temp_cache_dir, "idle", size=4096, mtime=100.0) - artifact_uri = "s3://bucket/new.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - - async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_text("VALUE = 2") - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - return_value=False, - ), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - patch( - "tracecat.executor.registry_artifact_storage.logger.warning" - ) as warning, - ): - registry_paths = await _materialize(cache, cache_key, artifact_uri) - - assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] - assert registry_paths[0].is_dir() - assert cache._budget_dirty is True - warning.assert_any_call( - "Registry artifact eviction remains pending physical deletion", - cache_key="idle", - trash_path=ANY, - ) - - @pytest.mark.anyio - async def test_failed_physical_delete_retries_without_extra_eviction( - self, temp_cache_dir - ): - """Failed byte reclamation stops eviction until trash cleanup succeeds.""" - cache = RegistryArtifactCache(temp_cache_dir) - oldest = _write_image_entry( - temp_cache_dir, - "oldest", - size=16, - mtime=100.0, - ) - older = _write_image_entry( - temp_cache_dir, - "older", - size=16, - mtime=200.0, - ) - newest = _write_image_entry( - temp_cache_dir, - "newest", - size=16, - mtime=300.0, - ) - real_delete = _delete_cache_path - failed_once = False - - def fail_once(path: Path) -> bool: - nonlocal failed_once - if not failed_once: - failed_once = True - return False - return real_delete(path) - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 16), - patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - side_effect=fail_once, - ), - ): - assert await cache._enforce_cache_budget() is False - assert not oldest.exists() - assert older.exists() - assert newest.exists() - assert cache._discover_cache_keys() == {"older", "newest"} - assert len(tuple(cache.trash_dir.iterdir())) == 1 - - assert await cache._enforce_cache_budget() is True - - assert not older.exists() - assert newest.exists() - assert not any(cache.trash_dir.iterdir()) - - @pytest.mark.anyio - async def test_rename_failure_does_not_block_materialization(self, temp_cache_dir): - """A failed atomic retirement keeps the old entry and admits new work.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - idle = _write_image_entry(temp_cache_dir, "idle", size=16, mtime=100.0) - artifact_uri = "s3://bucket/new-after-rename-failure.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - - async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_text("VALUE = 2") - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - patch( - "tracecat.executor.registry_artifact_storage._move_entry_to_trash", - side_effect=OSError("rename failed"), - ), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - registry_paths = await _materialize(cache, cache_key, artifact_uri) - - assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] - assert registry_paths[0].is_dir() - assert idle.is_file() - assert cache._budget_dirty is True - - @pytest.mark.anyio - async def test_cache_scan_failure_does_not_block_materialization( - self, temp_cache_dir - ): - """Maintenance errors are observable while artifact admission stays open.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/new-after-scan-failure.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - - async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_text("VALUE = 2") - - with ( - patch.object( - cache, - "_enforce_cache_budget", - new_callable=AsyncMock, - side_effect=PermissionError("denied"), - ), - patch( - "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", - return_value=9, - ), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [ - cache._paths_for(cache_key).tarball_target_dir - ] - - @pytest.mark.anyio - async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( - self, temp_cache_dir - ): - """Steady-state cache hits must not pay for a full cache scan.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/cached.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - _write_tarball_entry(temp_cache_dir, cache_key) - cache._budget_dirty = False - - with patch.object( - cache, - "_scan_cache_entries", - side_effect=AssertionError("cache hits must not scan the cache dir"), - ): - async with cache.lease([artifact_uri]): - pass - - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_failed_cold_admission_keeps_warm_lru(self, temp_cache_dir): - """A missing cold artifact must not evict an existing warm entry.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - warm_uri = "s3://bucket/warm.tar.gz" - warm_key = compute_registry_artifact_cache_key(warm_uri) - warm_dir = _write_tarball_entry(temp_cache_dir, warm_key) - missing_uri = "s3://bucket/missing.tar.gz" - missing_key = compute_registry_artifact_cache_key(missing_uri) - - async def mock_download(self, ctx, path): - raise FileNotFoundError("missing artifact") - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - patch(SQUASHFS_ENABLED_CONFIG, False), - patch.object(TarballArtifact, "download", mock_download), - ): - with pytest.raises(FileNotFoundError, match="missing artifact"): - async with cache.lease([missing_uri]): - pytest.fail("failed admission must not yield a lease") - - assert warm_dir.is_dir() - assert not cache._paths_for(missing_key).entry_dir.exists() - - @pytest.mark.anyio - async def test_failed_materialization_rearms_budget_dirty(self, temp_cache_dir): - """A failed materialization may leave a canonical image to evict.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - assert cache._budget_dirty is False - - artifact_uri = "s3://bucket/path/site-packages.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) - - async def mock_materialize(self, ctx): - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - ctx.paths.squashfs_image_path.write_bytes(b"orphaned image") - raise RuntimeError("mount failed") - - with ( - patch.object( - cache, - "_artifact_candidates", - new_callable=AsyncMock, - return_value=[artifact], - ), - patch.object(SquashfsArtifact, "materialize", mock_materialize), - ): - with pytest.raises(RuntimeError, match="mount failed"): - await _materialize(cache, cache_key, artifact_uri) - - assert cache._paths_for(cache_key).squashfs_image_path.is_file() - assert cache._budget_dirty is True - - @pytest.mark.anyio - async def test_materialization_rearms_budget_dirty_consumed_mid_flight( - self, temp_cache_dir - ): - """A convergence pass consuming the signal mid-download cannot unarm it.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - assert cache._budget_dirty is False - - artifact_uri = "s3://bucket/path/site-packages.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) - - async def mock_materialize(self, ctx): - # A concurrent lease release consumes the dirty signal and finishes - # its scan before this attempt lands its canonical bytes. - cache._budget_dirty = False - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - ctx.paths.squashfs_image_path.write_bytes(b"late image") - return [ctx.paths.squashfs_mount_dir] - - with ( - patch.object( - cache, - "_artifact_candidates", - new_callable=AsyncMock, - return_value=[artifact], - ), - patch.object(SquashfsArtifact, "materialize", mock_materialize), - ): - await _materialize(cache, cache_key, artifact_uri) - - assert cache._budget_dirty is True - - @pytest.mark.anyio - async def test_release_keeps_retrying_while_the_cache_stays_over_budget( - self, temp_cache_dir - ): - """A cache that cannot shrink yet must stay marked for re-enforcement.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/pinned.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - _write_image_entry(temp_cache_dir, cache_key, size=4096, mtime=100.0) - _write_tarball_entry(temp_cache_dir, cache_key) - cache._budget_dirty = True - # A second holder keeps the entry pinned past the inner lease. - cache._acquire_lease(cache_key) - - with patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 1): - async with cache.lease([artifact_uri]): - pass - - assert cache._budget_dirty is True - - @pytest.mark.anyio - async def test_convergence_rescans_after_concurrent_materialization( - self, temp_cache_dir - ): - """A materialization during a budget scan must schedule another scan.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/concurrent.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - scan_started = asyncio.Event() - materialized = asyncio.Event() - convergence_scans = 0 - - async def mock_enforce_cache_budget( - *, protected_key: str | None = None - ) -> bool: - nonlocal convergence_scans - if protected_key is not None: - return True - - convergence_scans += 1 - if convergence_scans == 1: - scan_started.set() - await materialized.wait() - return True - - async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_text("VALUE = 1") - - cache._budget_dirty = True - with ( - patch.object( - cache, - "_enforce_cache_budget", - side_effect=mock_enforce_cache_budget, - ), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - convergence = asyncio.create_task(cache._converge_cache_budget()) - await scan_started.wait() - registry_paths = await _materialize(cache, cache_key, artifact_uri) - materialized.set() - await convergence - - assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] - assert convergence_scans == 2 - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_cancelled_convergence_rearms_budget_dirty( - self, temp_cache_dir - ) -> None: - """Cancellation must preserve the need for a later budget pass.""" - cache = RegistryArtifactCache(temp_cache_dir) - enforcement_started = asyncio.Event() - finish_enforcement = asyncio.Event() - - async def mock_enforce_cache_budget( - *, protected_key: str | None = None - ) -> bool: - del protected_key - enforcement_started.set() - await finish_enforcement.wait() - return True - - cache._budget_dirty = True - with patch.object( - cache, - "_enforce_cache_budget", - side_effect=mock_enforce_cache_budget, - ): - convergence = asyncio.create_task(cache._converge_cache_budget()) - await enforcement_started.wait() - convergence.cancel() - - with pytest.raises(asyncio.CancelledError): - await convergence - - assert cache._budget_dirty is True - - @pytest.mark.parametrize( - "oldest_has_tarball", - [False, True], - ids=["squashfs-only", "tarball-bearing"], - ) - @pytest.mark.anyio - async def test_enforce_budget_evicts_least_recently_used_until_under_max_bytes( - self, temp_cache_dir, oldest_has_tarball: bool - ): - """Size eviction stops once the cache is within budget.""" - cache = RegistryArtifactCache(temp_cache_dir) - oldest = _write_image_entry(temp_cache_dir, "oldest", size=4096, mtime=100.0) - if oldest_has_tarball: - _write_tarball_entry(temp_cache_dir, "oldest") - oldest_entry = cache._paths_for("oldest").entry_dir - os.utime(oldest_entry, (100.0, 100.0)) - older = _write_image_entry(temp_cache_dir, "older", size=4096, mtime=200.0) - newest = _write_image_entry(temp_cache_dir, "newest", size=4096, mtime=300.0) - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 9000), - ): - within_budget = await cache._enforce_cache_budget(protected_key="pending") - - assert within_budget is True - assert not oldest.exists() - assert not cache._paths_for("oldest").tarball_target_dir.exists() - assert older.exists() - assert newest.exists() - - @pytest.mark.anyio - async def test_final_lease_release_unmounts_and_retains_image(self, temp_cache_dir): - """An idle entry releases its loop device without deleting its image.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/path/site-packages.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - paths = cache._paths_for(cache_key) - paths.entry_dir.mkdir(parents=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir() - mounted = {paths.squashfs_mount_dir} - - process = AsyncMock() - process.communicate.return_value = (b"", b"") - process.returncode = 0 - - async def mock_umount(*args, **kwargs): - mounted.discard(paths.squashfs_mount_dir) - return process - - with ( - patch.object(Path, "is_mount", lambda self: self in mounted), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/umount", - ), - patch.object( - asyncio, - "create_subprocess_exec", - side_effect=mock_umount, - ), - ): - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [paths.squashfs_mount_dir] - assert paths.squashfs_mount_dir in mounted - - assert paths.squashfs_mount_dir not in mounted - - assert paths.squashfs_image_path.read_bytes() == b"squashfs" - assert paths.squashfs_mount_dir.is_dir() - - @pytest.mark.anyio - async def test_concurrent_budget_passes_only_evict_once(self, temp_cache_dir): - """A waiting budget pass must re-scan after the active pass evicts.""" - cache = RegistryArtifactCache(temp_cache_dir) - oldest = _write_image_entry(temp_cache_dir, "oldest", size=16, mtime=100.0) - retained = _write_image_entry(temp_cache_dir, "retained", size=16, mtime=200.0) - original_scan = cache._scan_cache_entries - first_scan_started = threading.Event() - second_scan_started = threading.Event() - release_first_scan = threading.Event() - release_second_scan = threading.Event() - scan_count_lock = threading.Lock() - scan_count = 0 - - def controlled_scan(): - nonlocal scan_count - entries = original_scan() - with scan_count_lock: - scan_index = scan_count - scan_count += 1 - if scan_index == 0: - first_scan_started.set() - release_first_scan.wait(timeout=5) - elif scan_index == 1: - second_scan_started.set() - release_second_scan.wait(timeout=5) - return entries - - eviction_started = asyncio.Event() - finish_eviction = asyncio.Event() - extra_eviction_finished = asyncio.Event() - evicted_keys: list[str] = [] - - async def controlled_evict( - cache_key: str, - ) -> RegistryArtifactEviction: - if cache_key == "oldest": - if eviction_started.is_set(): - return RegistryArtifactEviction( - retired=False, - reclaimed=False, - ) - eviction_started.set() - await finish_eviction.wait() - _delete_cache_path(cache._paths_for("oldest").entry_dir) - else: - _delete_cache_path(cache._paths_for("retained").entry_dir) - extra_eviction_finished.set() - evicted_keys.append(cache_key) - return RegistryArtifactEviction(retired=True, reclaimed=True) - - with ( - patch.object(cache, "_scan_cache_entries", side_effect=controlled_scan), - patch.object(cache, "_evict_entry", side_effect=controlled_evict), - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - ): - first_pass = asyncio.create_task(cache._enforce_cache_budget()) - assert await asyncio.to_thread(first_scan_started.wait, 1) - second_pass = asyncio.create_task(cache._enforce_cache_budget()) - second_scan_overlapped = await asyncio.to_thread( - second_scan_started.wait, 0.5 - ) - - release_first_scan.set() - await asyncio.wait_for(eviction_started.wait(), timeout=1) - release_second_scan.set() - try: - await asyncio.wait_for(extra_eviction_finished.wait(), timeout=0.2) - except TimeoutError: - pass - finish_eviction.set() - await asyncio.gather(first_pass, second_pass) - - assert second_scan_overlapped is False - assert scan_count == 2 - assert evicted_keys == ["oldest"] - assert not oldest.exists() - assert retained.exists() - - @pytest.mark.anyio - async def test_enforce_budget_ignores_missing_protected_key(self, temp_cache_dir): - """Only successfully published entries count against the entry budget.""" - cache = RegistryArtifactCache(temp_cache_dir) - existing = _write_image_entry(temp_cache_dir, "existing", size=16, mtime=100.0) - - with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): - within_budget = await cache._enforce_cache_budget(protected_key="missing") - - assert within_budget is True - assert existing.exists() - - @pytest.mark.anyio - async def test_enforce_budget_never_evicts_the_protected_key(self, temp_cache_dir): - """The newly materialized key is exempt even when it is the LRU entry.""" - cache = RegistryArtifactCache(temp_cache_dir) - protected = _write_image_entry( - temp_cache_dir, "protected", size=4096, mtime=100.0 - ) - other = _write_image_entry(temp_cache_dir, "other", size=4096, mtime=300.0) - - with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): - await cache._enforce_cache_budget(protected_key="protected") - - assert protected.exists() - assert not other.exists() - - @pytest.mark.anyio - async def test_enforce_budget_proceeds_over_budget_when_everything_is_leased( - self, temp_cache_dir - ): - """An over-budget cache must degrade, not fail the action.""" - cache = RegistryArtifactCache(temp_cache_dir) - leased = _write_image_entry(temp_cache_dir, "leased", size=4096, mtime=100.0) - cache._acquire_lease("leased") - - with patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 1): - within_budget = await cache._enforce_cache_budget(protected_key="missing") - - assert within_budget is False - assert leased.exists() - - @pytest.mark.anyio - async def test_eviction_unmounts_before_deleting_the_image(self, temp_cache_dir): - """Unlinking a mounted image would strand an open-file zombie.""" - cache = RegistryArtifactCache(temp_cache_dir) - paths = cache._paths_for("mounted") - paths.entry_dir.mkdir(parents=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir() - mounted = {paths.squashfs_mount_dir} - image_present_at_umount: list[bool] = [] - - process = AsyncMock() - process.communicate.return_value = (b"", b"") - process.returncode = 0 - - async def mock_umount(*args, **kwargs): - image_present_at_umount.append(paths.squashfs_image_path.exists()) - mounted.discard(paths.squashfs_mount_dir) - return process - - with ( - patch.object(Path, "is_mount", lambda self: self in mounted), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/umount", - ), - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - side_effect=mock_umount, - ) as create_subprocess_exec, - ): - evicted = await cache._evict_entry("mounted") - - assert evicted == RegistryArtifactEviction(retired=True, reclaimed=True) - assert image_present_at_umount == [True] - assert not paths.squashfs_image_path.exists() - assert not paths.squashfs_mount_dir.exists() - create_subprocess_exec.assert_called_once() - assert create_subprocess_exec.call_args.args == ( - "/sbin/umount", - str(paths.squashfs_mount_dir), - ) - - @pytest.mark.anyio - async def test_cancelled_unmount_kills_and_reaps_before_releasing_key_lock( - self, temp_cache_dir - ): - """Cancellation leaves a consistent entry for the next admission.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/path/cancelled-unmount.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - paths = cache._paths_for(cache_key) - paths.entry_dir.mkdir(parents=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir() - (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") - mounted = {paths.squashfs_mount_dir} - blocked_process = _BlockingSubprocess() - released_process = AsyncMock() - released_process.communicate.return_value = (b"", b"") - released_process.returncode = 0 - unmount_attempts = 0 - - async def mock_umount(*args, **kwargs): - nonlocal unmount_attempts - unmount_attempts += 1 - if unmount_attempts == 1: - return blocked_process - mounted.discard(paths.squashfs_mount_dir) - return released_process - - with ( - patch.object(Path, "is_mount", lambda self: self in mounted), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/umount", - ), - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - side_effect=mock_umount, - ), - ): - eviction = asyncio.create_task(cache._evict_entry(cache_key)) - await blocked_process.communicate_started.wait() - eviction.cancel() - - with pytest.raises(asyncio.CancelledError): - await eviction - - assert blocked_process.cleanup_calls == ["kill", "wait"] - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [paths.squashfs_mount_dir] - assert registry_paths[0].is_dir() - assert (registry_paths[0] / "module.py").read_text() == "VALUE = 1" - - assert paths.squashfs_image_path.is_file() - assert paths.squashfs_mount_dir.is_dir() - assert paths.squashfs_mount_dir not in mounted - - @pytest.mark.anyio - async def test_cancelled_background_deletion_leaves_a_clean_miss( - self, temp_cache_dir - ): - """Cancellation cannot expose live paths that deletion still owns.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/cancelled-eviction.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - original_target = _write_tarball_entry(temp_cache_dir, cache_key) - delete_started = threading.Event() - finish_delete = threading.Event() - delete_finished = threading.Event() - doomed: list[Path] = [] - - def blocked_delete(path: Path) -> bool: - doomed.append(path) - delete_started.set() - finish_delete.wait(timeout=5) - deleted = _delete_cache_path(path) - delete_finished.set() - return deleted - - async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_text("VALUE = 2") - - with ( - patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - side_effect=blocked_delete, - ), - patch( - "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", - return_value=9, - ), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - eviction = asyncio.create_task(cache._evict_entry(cache_key)) - assert await asyncio.to_thread(delete_started.wait, 1) - eviction.cancel() - try: - await asyncio.sleep(0) - eviction.cancel() - await asyncio.sleep(0) - assert not eviction.done() - assert not original_target.exists() - assert (doomed[0] / "tarball").is_dir() - assert cache._discover_cache_keys() == set() - finally: - finish_delete.set() - - with pytest.raises(asyncio.CancelledError): - await eviction - assert await asyncio.to_thread(delete_finished.wait, 1) - assert not doomed[0].exists() - - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [original_target] - assert (original_target / "module.py").read_text() == "VALUE = 2" - - assert original_target.is_dir() - - @pytest.mark.anyio - async def test_doomed_eviction_names_are_startup_scratch(self, temp_cache_dir): - """Every renamed entry root is invisible and reclaimed on startup.""" - cache = RegistryArtifactCache(temp_cache_dir) - cache_key = "squashfs-doomed" - paths = cache._paths_for(cache_key) - paths.entry_dir.mkdir(parents=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir() - paths.squashfs_extract_dir.mkdir() - paths.tarball_target_dir.mkdir() - - with patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - return_value=True, - ) as delete_cache_path: - assert await cache._evict_entry(cache_key) == RegistryArtifactEviction( - retired=True, reclaimed=True - ) - - delete_cache_path.assert_called_once() - trash_path = delete_cache_path.call_args.args[0] - assert trash_path.parent == cache.trash_dir - assert trash_path.exists() - assert cache._discover_cache_keys() == set() - - cache._sweep_startup_state() - - assert not trash_path.exists() - - @pytest.mark.anyio - async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): - """A failed unmount skips the key instead of forcing a lazy detach.""" - cache = RegistryArtifactCache(temp_cache_dir) - stuck = cache._paths_for("stuck") - stuck.entry_dir.mkdir(parents=True) - stuck.squashfs_image_path.write_bytes(b"squashfs") - stuck.squashfs_mount_dir.mkdir() - os.utime(stuck.squashfs_image_path, (100.0, 100.0)) - idle = _write_image_entry(temp_cache_dir, "idle", size=16, mtime=300.0) - mounted = {stuck.squashfs_mount_dir} - - process = AsyncMock() - process.communicate.return_value = (b"", b"target is busy") - process.returncode = 32 - - with ( - patch.object(Path, "is_mount", lambda self: self in mounted), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/umount", - ), - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - new_callable=AsyncMock, - return_value=process, - ), - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - ): - await cache._enforce_cache_budget(protected_key="pending") - - assert stuck.squashfs_image_path.exists() - assert stuck.squashfs_mount_dir.exists() - assert not idle.exists() - - @pytest.mark.anyio - async def test_eviction_keeps_stable_runtime_state(self, temp_cache_dir): - """A key keeps one lock and zeroed lease state for the process lifetime.""" - cache = RegistryArtifactCache(temp_cache_dir) - _write_tarball_entry(temp_cache_dir, "bookkeeping") - cache._acquire_lease("bookkeeping") - cache._release_lease("bookkeeping") - lock = cache._runtime_for("bookkeeping").lock - - assert await cache._evict_entry("bookkeeping") == RegistryArtifactEviction( - retired=True, reclaimed=True - ) - runtime = cache._runtime["bookkeeping"] - assert runtime.lock is lock - assert runtime.refcount == 0 - - @pytest.mark.anyio - async def test_eviction_skips_busy_key(self, temp_cache_dir): - """A key another task is materializing is never evicted underneath it.""" - cache = RegistryArtifactCache(temp_cache_dir) - target_dir = _write_tarball_entry(temp_cache_dir, "busy") - lock = cache._runtime_for("busy").lock - - async with lock: - assert await cache._evict_entry("busy") == RegistryArtifactEviction( - retired=False, reclaimed=False - ) - - assert target_dir.is_dir() - - -class TestRegistryArtifactCacheStartupSweep: - """Tests for the startup sweep that reclaims state from a dead process.""" - - @pytest.mark.anyio - async def test_sweep_uses_entry_root_mtime_for_restart_safe_lru( - self, temp_cache_dir - ): - """Startup trimming preserves a touched older tarball-only entry.""" - previous_cache = RegistryArtifactCache(temp_cache_dir) - old = _write_tarball_entry(temp_cache_dir, "old") - previous_cache._touch_entry("old") - - old_mtime = previous_cache._paths_for("old").entry_dir.stat().st_mtime - new = _write_tarball_entry(temp_cache_dir, "new") - new_entry_dir = previous_cache._paths_for("new").entry_dir - os.utime(new_entry_dir, (old_mtime - 1, old_mtime - 1)) - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - ): - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - - assert old.is_dir() - assert not new.exists() - - @pytest.mark.anyio - async def test_sweep_tolerates_missing_cache_dir(self, temp_cache_dir): - """A cache directory that does not exist yet is a no-op.""" - cache_dir = temp_cache_dir / "missing" - - cache = RegistryArtifactCache(cache_dir) - await cache.ensure_swept() - - assert cache.cache_dir == cache_dir - assert not cache_dir.exists() - - @pytest.mark.anyio - async def test_sweep_removes_orphaned_work(self, temp_cache_dir): - """Interrupted work is reclaimed without touching unrelated paths.""" - cache = RegistryArtifactCache(temp_cache_dir) - orphaned = cache.staging_dir / "abc123.999999.4321.squashfs" - orphaned.parent.mkdir() - orphaned.write_bytes(b"partial") - orphaned_dir = cache.trash_dir / "abc123.999999.4321" - orphaned_dir.mkdir(parents=True) - unrelated = temp_cache_dir / "unrelated" - unrelated.mkdir() - unrelated_file = unrelated / "keep.txt" - unrelated_file.write_text("keep") - entry_dir = _write_tarball_entry(temp_cache_dir, "abc123") - - await cache.ensure_swept() - - assert not orphaned.exists() - assert not orphaned_dir.exists() - assert unrelated_file.read_text() == "keep" - assert entry_dir.is_dir() - - @pytest.mark.anyio - async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): - """A live mountpoint belongs to a running process and must survive.""" - cache = RegistryArtifactCache(temp_cache_dir) - paths = cache._paths_for("abc123") - paths.entry_dir.mkdir(parents=True) - mount_dir = paths.squashfs_mount_dir - mount_dir.mkdir() - - with patch.object(Path, "is_mount", lambda self: self == mount_dir): - await cache.ensure_swept() - - assert mount_dir.is_dir() - - @pytest.mark.anyio - async def test_ensure_swept_runs_once(self, temp_cache_dir): - """Repeated startup-sweep requests invoke the sweep once.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with patch.object( - cache, - "_sweep_startup_state", - wraps=cache._sweep_startup_state, - ) as sweep: - await cache.ensure_swept() - await cache.ensure_swept() - - assert sweep.call_count == 1 - - @pytest.mark.anyio - async def test_concurrent_ensure_swept_runs_once(self, temp_cache_dir): - """Concurrent startup-sweep requests invoke the sweep once.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with patch.object( - cache, - "_sweep_startup_state", - wraps=cache._sweep_startup_state, - ) as sweep: - await asyncio.gather(*(cache.ensure_swept() for _ in range(4))) - - assert sweep.call_count == 1 - - @pytest.mark.anyio - async def test_cancelled_waiter_joins_same_startup_sweep(self, temp_cache_dir): - """A cancelled waiter cannot start a second concurrent startup sweep.""" - cache = RegistryArtifactCache(temp_cache_dir) - sweep_started = threading.Event() - sweep_release = threading.Event() - invocation_count = 0 - - def blocking_sweep() -> None: - nonlocal invocation_count - invocation_count += 1 - sweep_started.set() - sweep_release.wait() - - with patch.object( - cache, - "_sweep_startup_state", - side_effect=blocking_sweep, - ): - first_waiter = asyncio.create_task(cache.ensure_swept()) - assert await asyncio.to_thread(sweep_started.wait, 1) - first_waiter.cancel() - - with pytest.raises(asyncio.CancelledError): - await first_waiter - - second_waiter = asyncio.create_task(cache.ensure_swept()) - await asyncio.sleep(0) - sweep_release.set() - await second_waiter - - assert invocation_count == 1 - assert cache._swept is True - - @pytest.mark.anyio - async def test_lease_triggers_startup_sweep(self, temp_cache_dir): - """Lease admission reclaims startup scratch before yielding paths.""" - cache = RegistryArtifactCache(temp_cache_dir) - orphaned_dir = cache.staging_dir / "abc123.999999.4321" - orphaned_dir.mkdir(parents=True) - - async with cache.lease(None): - assert not orphaned_dir.exists() - - @pytest.mark.anyio - async def test_failed_ensure_swept_retries(self, temp_cache_dir): - """A failed startup sweep is retried by the next caller.""" - cache = RegistryArtifactCache(temp_cache_dir) - orphaned_dir = cache.staging_dir / "abc123.999999.4321" - orphaned_dir.mkdir(parents=True) - - with ( - patch.object( - cache, - "_clear_work_dir", - side_effect=OSError("simulated sweep failure"), - ), - pytest.raises(OSError), - ): - await cache.ensure_swept() - - assert orphaned_dir.is_dir() - - await cache.ensure_swept() - - assert not orphaned_dir.exists() - - @pytest.mark.parametrize("work_dir_name", ["staging", "trash"]) - def test_work_directory_inspection_errors_propagate( - self, - temp_cache_dir, - work_dir_name: str, - ): - """Unreadable work directories must trigger a later cleanup retry.""" - cache = RegistryArtifactCache(temp_cache_dir) - work_dir = getattr(cache, f"{work_dir_name}_dir") - work_dir.mkdir() - - with ( - patch.object(Path, "iterdir", side_effect=PermissionError("denied")), - pytest.raises(PermissionError, match="denied"), - ): - cache._clear_work_dir(work_dir) - - @pytest.mark.anyio - async def test_active_entry_scan_errors_propagate(self, temp_cache_dir): - """Unreadable active entries cannot be treated as an empty cache.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch( - "tracecat.executor.registry_artifact_storage.os.scandir", - side_effect=PermissionError("denied"), - ), - pytest.raises(PermissionError, match="denied"), - ): - await cache._enforce_cache_budget() - - @pytest.mark.anyio - async def test_failed_trash_cleanup_skips_startup_trimming(self, temp_cache_dir): - """A failed orphan deletion must not cascade into active retirements.""" - cache = RegistryArtifactCache(temp_cache_dir) - orphan = cache.trash_dir / "orphan" - orphan.mkdir(parents=True) - oldest = _write_image_entry( - temp_cache_dir, - "oldest", - size=16, - mtime=100.0, - ) - newest = _write_image_entry( - temp_cache_dir, - "newest", - size=16, - mtime=200.0, - ) - real_delete = _delete_cache_path - - def fail_orphan(path: Path) -> bool: - if path == orphan: - return False - return real_delete(path) - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - side_effect=fail_orphan, - ), - ): - await cache.ensure_swept() - - assert orphan.is_dir() - assert oldest.is_file() - assert newest.is_file() - assert cache._budget_dirty is True - - @pytest.mark.anyio - async def test_failed_startup_cleanup_retries_exact_path(self, temp_cache_dir): - """Failed startup scratch cleanup retries without sweeping live staging.""" - cache = RegistryArtifactCache(temp_cache_dir) - orphaned = cache.staging_dir / "orphaned.123.456.tmp" - orphaned.parent.mkdir(parents=True) - orphaned.write_bytes(b"partial") - real_delete = _delete_cache_path - failed_once = False - - def fail_once(path: Path) -> bool: - nonlocal failed_once - if path == orphaned and not failed_once: - failed_once = True - return False - return real_delete(path) - - with patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - side_effect=fail_once, - ): - await cache.ensure_swept() - assert orphaned.is_file() - assert cache._failed_startup_cleanup == {orphaned} - assert cache._budget_dirty is True - - assert await cache._enforce_cache_budget() is True - - assert not orphaned.exists() - assert cache._failed_startup_cleanup == set() - - @pytest.mark.anyio - async def test_failed_startup_retirement_stays_dirty_and_retries( - self, temp_cache_dir - ): - """A startup rename failure preserves entries for later enforcement.""" - oldest = _write_image_entry( - temp_cache_dir, - "oldest", - size=16, - mtime=100.0, - ) - newest = _write_image_entry( - temp_cache_dir, - "newest", - size=16, - mtime=200.0, - ) - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - patch( - "tracecat.executor.registry_artifact_storage._move_entry_to_trash", - side_effect=OSError("rename failed"), - ), - ): - await cache.ensure_swept() - - assert oldest.is_file() - assert newest.is_file() - assert cache._budget_dirty is True - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - ): - await cache._converge_cache_budget() - - assert not oldest.exists() - assert newest.is_file() - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_failed_startup_physical_delete_retries_exact_path( - self, temp_cache_dir - ): - """Startup stops retiring entries when bytes were not reclaimed.""" - oldest = _write_image_entry( - temp_cache_dir, - "oldest", - size=16, - mtime=100.0, - ) - older = _write_image_entry( - temp_cache_dir, - "older", - size=16, - mtime=200.0, - ) - newest = _write_image_entry( - temp_cache_dir, - "newest", - size=16, - mtime=300.0, - ) - cache = RegistryArtifactCache(temp_cache_dir) - real_delete = _delete_cache_path - failed_once = False - - def fail_once(path: Path) -> bool: - nonlocal failed_once - if path.parent == cache.trash_dir and not failed_once: - failed_once = True - return False - return real_delete(path) - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 16), - patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - side_effect=fail_once, - ), - ): - await cache.ensure_swept() - assert not oldest.exists() - assert older.is_file() - assert newest.is_file() - assert len(tuple(cache.trash_dir.iterdir())) == 1 - assert cache._budget_dirty is True - - await cache._converge_cache_budget() - - assert not older.exists() - assert newest.is_file() - assert not any(cache.trash_dir.iterdir()) - assert cache._budget_dirty is False - - -class TestSquashfsMountPolicy: - """Tests for per-artifact SquashFS fallback.""" - - @pytest.mark.anyio - async def test_loop_device_exhaustion_isolated_sticky_extraction_fallback( - self, temp_cache_dir: Path - ) -> None: - """Protect the cache's fail-open policy when loop devices are saturated. - - A mount-command failure must extract only the affected cold artifact, - preserve already-leased mounts, reuse that extraction on later leases, - and still let unrelated artifacts attempt mounting. This models loop - exhaustion deterministically without consuming host-global devices. - """ - cache = RegistryArtifactCache(temp_cache_dir) - held_uri = "s3://bucket/already-mounted.squashfs" - saturated_uri = "s3://bucket/no-loop-available.squashfs" - later_uri = "s3://bucket/later-artifact.squashfs" - held_key = compute_registry_artifact_cache_key(held_uri) - saturated_key = compute_registry_artifact_cache_key(saturated_uri) - later_key = compute_registry_artifact_cache_key(later_uri) - harness = _SquashfsMountHarness( - cache, - failed_mount_keys={saturated_key}, - ) - - with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", harness.mount), - patch.object(SquashfsArtifact, "extract", harness.extract), - patch.object(cache, "_unmount", harness.unmount), - ): - async with cache.lease([held_uri]) as held_paths: - held_mount = cache._paths_for(held_key).squashfs_mount_dir - assert held_paths == [held_mount] - - async with cache.lease([saturated_uri]) as saturated_paths: - extracted = cache._paths_for(saturated_key).squashfs_extract_dir - assert saturated_paths == [extracted] - assert held_mount in harness.mounted - assert cache._refcount(held_key) == 1 - assert harness.unmounts == [] - - # The extracted directory is the cache entry's sticky path; - # freeing a loop elsewhere does not trigger a remount attempt. - async with cache.lease([saturated_uri]) as cached_paths: - assert cached_paths == [extracted] - - async with cache.lease([later_uri]) as later_paths: - later_mount = cache._paths_for(later_key).squashfs_mount_dir - assert later_paths == [later_mount] - assert held_mount in harness.mounted - - assert held_mount in harness.mounted - - assert harness.mount_attempts == [held_key, saturated_key, later_key] - assert harness.extraction_attempts == [saturated_key] - assert harness.unmounts == [later_mount, held_mount] - assert cache._paths_for(saturated_key).squashfs_image_path.is_file() - assert cache._paths_for(saturated_key).squashfs_extract_dir.is_dir() - assert cache._refcount(held_key) == 0 - assert cache._refcount(saturated_key) == 0 - assert cache._refcount(later_key) == 0 - - @pytest.mark.anyio - async def test_mount_failure_does_not_disable_later_artifacts( - self, temp_cache_dir - ) -> None: - """One failed mount falls back without changing later mount attempts.""" - cache = RegistryArtifactCache(temp_cache_dir) - first_ctx = cache._context_for("first") - second_ctx = cache._context_for("second") - first = SquashfsArtifact( - uri="s3://bucket/first.squashfs", - cache_key="first", - ) - second = SquashfsArtifact( - uri="s3://bucket/second.squashfs", - cache_key="second", - ) - mount_attempts: list[str] = [] - - async def mock_mount(self, ctx, image_path): - del image_path - mount_attempts.append(ctx.cache_key) - if ctx.cache_key == "first": - raise SquashfsMountCommandError("operation not permitted") - ctx.paths.squashfs_mount_dir.mkdir(parents=True) - return ctx.paths.squashfs_mount_dir - - async def mock_extract(self, ctx, image_path): - del image_path - ctx.paths.squashfs_extract_dir.mkdir(parents=True) - return ctx.paths.squashfs_extract_dir - - with ( - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", mock_mount), - patch.object(SquashfsArtifact, "extract", mock_extract), - ): - assert await first.materialize(first_ctx) == [ - first_ctx.paths.squashfs_extract_dir - ] - assert await second.materialize(second_ctx) == [ - second_ctx.paths.squashfs_mount_dir - ] - - assert mount_attempts == ["first", "second"] From e1308e8616acd837f9faa33625b003b4067ef7a8 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 15:57:42 -0400 Subject: [PATCH 042/161] fix(executor): rejoin tarball sizing on cancellation --- .../test_registry_artifact_materialization.py | 60 +++++++++++++++++++ .../registry_artifact_materialization.py | 47 ++++++++------- 2 files changed, 85 insertions(+), 22 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 702cdc9f4e..b5c741c3de 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -383,6 +383,66 @@ def blocking_extractall(*args: object, **kwargs: object) -> None: assert first_cancellation_propagated_early is False assert second_cancellation_propagated_early is False + @pytest.mark.anyio + async def test_repeatedly_cancelled_tarball_size_scan_rejoins_thread( + self, temp_cache_dir + ): + """Cancellation cannot unlink a tarball while its size scan is running.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/slow-size-scan.tar.gz" + downloaded_paths: list[Path] = [] + scan_started = threading.Event() + scan_release = threading.Event() + scan_finished = threading.Event() + input_present_at_finish: list[bool] = [] + + async def mock_download(self, ctx, path): + del self, ctx + path.write_bytes(tarball_payload(size=1)) + downloaded_paths.append(path) + + def blocking_size_scan(path: Path) -> int: + scan_started.set() + scan_release.wait() + input_present_at_finish.append(path.exists()) + scan_finished.set() + return 1 + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "download", mock_download), + patch( + "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", + side_effect=blocking_size_scan, + ), + patch.object( + TarballArtifact, + "extract", + new_callable=AsyncMock, + ) as extract, + ): + materializing = asyncio.create_task(lease_paths(cache, artifact_uri)) + assert await asyncio.to_thread(scan_started.wait, 1) + materializing.cancel() + done, _ = await asyncio.wait({materializing}, timeout=0.05) + first_cancellation_propagated_early = bool(done) + + materializing.cancel() + done, _ = await asyncio.wait({materializing}, timeout=0.05) + second_cancellation_propagated_early = bool(done) + assert downloaded_paths[0].exists() + scan_release.set() + + with pytest.raises(asyncio.CancelledError): + await materializing + + assert scan_finished.is_set() + assert input_present_at_finish == [True] + assert not downloaded_paths[0].exists() + assert first_cancellation_propagated_early is False + assert second_cancellation_propagated_early is False + extract.assert_not_awaited() + @pytest.mark.anyio async def test_materialize_extracts_squashfs_when_mount_fails(self, temp_cache_dir): """Test that SquashFS mount failures fall back to unsquashfs extraction.""" diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index f2f1f68bf1..c50baa60d3 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -129,6 +129,28 @@ def _temp_path( return ctx.staging_dir / f"{self.cache_key}.{os.getpid()}.{unique_id}{suffix}" +async def _run_blocking_rejoin_on_cancel[T](operation: Callable[[], T]) -> T: + """Run blocking work without abandoning its thread on cancellation.""" + worker = asyncio.ensure_future(asyncio.to_thread(operation)) + try: + return await asyncio.shield(worker) + except asyncio.CancelledError: + # A thread cannot be killed. Rejoin it before callers remove its input + # or output paths. Repeated cancellation can interrupt shield without + # stopping the thread, so keep waiting for a terminal state. + while not worker.done(): + try: + await asyncio.shield(worker) + except asyncio.CancelledError: + continue + except Exception: + break + if not worker.cancelled(): + with contextlib.suppress(Exception): + worker.result() + raise + + @dataclass(frozen=True, slots=True) class BuiltinArtifact(RegistryArtifact): """Current builtin registry package already installed in the executor image.""" @@ -481,9 +503,8 @@ async def materialize( download_elapsed = (time.monotonic() - download_start) * 1000 if ctx.admission is not None: - extracted_size = await asyncio.to_thread( - _tarball_extracted_size, - temp_tarball, + extracted_size = await _run_blocking_rejoin_on_cancel( + lambda: _tarball_extracted_size(temp_tarball) ) await ctx.admission.ensure_capacity(extracted_size) @@ -549,25 +570,7 @@ def _do_extract() -> None: raise ValueError(f"Unsupported tarball format: {tarball_path}") - extraction = asyncio.ensure_future(asyncio.to_thread(_do_extract)) - try: - await asyncio.shield(extraction) - except asyncio.CancelledError: - # A thread cannot be killed. Rejoin it before materialize removes - # scratch. Each cancellation can interrupt shield without stopping - # the thread, so keep waiting until extraction reaches a terminal - # state before propagating the original cancellation. - while not extraction.done(): - try: - await asyncio.shield(extraction) - except asyncio.CancelledError: - continue - except Exception: - break - if not extraction.cancelled(): - with contextlib.suppress(Exception): - extraction.result() - raise + await _run_blocking_rejoin_on_cancel(_do_extract) logger.debug( "Tarball extracted", From b5fbf2264f341214cea14e784e770d2022d5030e Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:00:27 -0400 Subject: [PATCH 043/161] fix(executor): discard failed artifact candidates --- .../executor/test_registry_artifact_budget.py | 51 +++++++++++++++++++ .../registry_artifact_materialization.py | 40 +++++++++++++++ tracecat/executor/registry_artifacts.py | 1 + 3 files changed, 92 insertions(+) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index e3438a4a90..5463692a82 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -188,6 +188,57 @@ async def download_file_to_path( assert not cache._paths_for(cache_key).entry_dir.exists() assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) + @pytest.mark.anyio + async def test_failed_squashfs_bytes_do_not_block_tarball_fallback( + self, temp_cache_dir: Path + ) -> None: + """An unusable image cannot consume the tarball fallback's budget.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/site-packages.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + payload = tarball_payload(size=32) + max_bytes = len(payload) + 32 + + async def fail_after_squashfs_download( + self: SquashfsArtifact, + ctx: RegistryArtifactMaterializationContext, + ) -> list[Path]: + del self + assert ctx.admission is not None + await ctx.admission.ensure_capacity(max_bytes) + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + ctx.paths.squashfs_image_path.write_bytes(b"x" * max_bytes) + raise RuntimeError("unusable SquashFS image") + + async def download_tarball( + self: TarballArtifact, + ctx: RegistryArtifactMaterializationContext, + path: Path, + ) -> None: + del self + assert ctx.admission is not None + await ctx.admission.ensure_capacity(len(payload)) + path.write_bytes(payload) + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, max_bytes), + patch.object( + SquashfsArtifact, + "materialize", + fail_after_squashfs_download, + ), + patch.object(TarballArtifact, "download", download_tarball), + ): + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [ + cache._paths_for(cache_key).tarball_target_dir + ] + + paths = cache._paths_for(cache_key) + assert not paths.squashfs_image_path.exists() + assert (paths.tarball_target_dir / "module.py").read_bytes() == b"x" * 32 + @pytest.mark.anyio async def test_squashfs_expansion_is_rejected_before_extraction( self, temp_cache_dir: Path diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index c50baa60d3..dd230b5be9 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -119,6 +119,12 @@ async def materialize( ) -> list[Path]: """Return importable Python paths, materializing the artifact if needed.""" + def discard_failed_materialization( + self, ctx: RegistryArtifactMaterializationContext + ) -> None: + """Discard canonical bytes that cannot serve a fallback candidate.""" + del ctx + def _temp_path( self, ctx: RegistryArtifactMaterializationContext, @@ -204,6 +210,40 @@ def cached_path( return [ctx.paths.squashfs_extract_dir] return None + def discard_failed_materialization( + self, ctx: RegistryArtifactMaterializationContext + ) -> None: + """Remove an unusable image before admitting a tarball fallback.""" + try: + if self.cached_path(ctx) is not None: + return + except OSError as e: + logger.warning( + "Cannot determine whether failed SquashFS candidate is reusable", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + error=str(e), + ) + return + + try: + ctx.paths.squashfs_image_path.unlink(missing_ok=True) + except OSError as e: + logger.warning( + "Failed to discard unusable SquashFS candidate", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + error=str(e), + ) + + for directory in ( + ctx.paths.squashfs_extract_dir, + ctx.paths.squashfs_mount_dir, + ctx.paths.entry_dir, + ): + with contextlib.suppress(OSError): + directory.rmdir() + async def materialize( self, ctx: RegistryArtifactMaterializationContext ) -> list[Path]: diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index f038418bd8..cbce7c922a 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -246,6 +246,7 @@ async def _materialize_candidates( except Exception as e: if index == len(candidates) - 1: raise + artifact.discard_failed_materialization(ctx) logger.warning( "Failed to materialize registry artifact candidate, trying fallback", cache_key=cache_key, From b7bd0a87179de7446318f02a6cb842731d56f154 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:03:32 -0400 Subject: [PATCH 044/161] fix(sandbox): contain dependency setup processes --- tests/unit/test_unsafe_pid_executor.py | 72 +++++++++++++++++++++++++ tracecat/sandbox/unsafe_pid_executor.py | 23 ++------ 2 files changed, 77 insertions(+), 18 deletions(-) diff --git a/tests/unit/test_unsafe_pid_executor.py b/tests/unit/test_unsafe_pid_executor.py index 6667d749b0..3772623314 100644 --- a/tests/unit/test_unsafe_pid_executor.py +++ b/tests/unit/test_unsafe_pid_executor.py @@ -5,6 +5,7 @@ import logging import os import signal +import sys from pathlib import Path import pytest @@ -189,6 +190,77 @@ async def fake_wait_for(awaitable, *args, **kwargs): unsafe_pid_executor.pid_namespace_probe_error() == "unshare probe timed out" ) + @pytest.mark.parametrize("operation", ["create-venv", "install-packages"]) + @pytest.mark.anyio + async def test_cancelled_dependency_setup_kills_process_group( + self, + executor: UnsafePidExecutor, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + operation: str, + ) -> None: + """Cancellation cannot leave package-build descendants running.""" + pid_file = tmp_path / f"{operation}-child.pid" + real_create_subprocess_exec = asyncio.create_subprocess_exec + created_processes: list[asyncio.subprocess.Process] = [] + setup_script = """ +import subprocess +import sys +import time +from pathlib import Path + +child = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(30)"], + stdin=subprocess.DEVNULL, +) +Path(sys.argv[1]).write_text(str(child.pid)) +time.sleep(30) +""" + + async def create_dependency_process(*args, **kwargs): + del args + assert kwargs["start_new_session"] is True + process = await real_create_subprocess_exec( + sys.executable, + "-c", + setup_script, + str(pid_file), + stdout=kwargs["stdout"], + stderr=kwargs["stderr"], + start_new_session=True, + ) + created_processes.append(process) + return process + + monkeypatch.setattr( + asyncio, + "create_subprocess_exec", + create_dependency_process, + ) + + if operation == "create-venv": + setup = executor._create_venv(tmp_path / "venv") + else: + setup = executor._install_packages( + tmp_path / "venv", + ["synthetic-package"], + ) + + task = asyncio.create_task(setup) + await _wait_for_file(pid_file) + child_pid = int(pid_file.read_text()) + try: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert len(created_processes) == 1 + assert created_processes[0].returncode is not None + await _wait_for_process_exit(child_pid) + finally: + with contextlib.suppress(ProcessLookupError): + os.kill(child_pid, signal.SIGKILL) + @pytest.mark.anyio async def test_execute_basic_script(self, executor: UnsafePidExecutor) -> None: script = """ diff --git a/tracecat/sandbox/unsafe_pid_executor.py b/tracecat/sandbox/unsafe_pid_executor.py index 7c8ebdc400..2be02133cd 100644 --- a/tracecat/sandbox/unsafe_pid_executor.py +++ b/tracecat/sandbox/unsafe_pid_executor.py @@ -6,7 +6,6 @@ """ import asyncio -import contextlib import hashlib import json import logging @@ -331,17 +330,11 @@ async def _create_venv(self, venv_path: Path) -> None: "HOME": os.environ.get("HOME", "/tmp"), "UV_CACHE_DIR": str(self.uv_cache), }, + start_new_session=True, ) try: - _, stderr = await asyncio.wait_for(process.communicate(), timeout=60) - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - raise + _, stderr = await communicate_process_group(process, timeout=60) except TimeoutError as e: - process.kill() - await process.wait() raise PackageInstallError("Virtual environment creation timed out") from e if process.returncode != 0: raise PackageInstallError( @@ -377,21 +370,15 @@ async def _install_packages( "HOME": os.environ.get("HOME", "/tmp"), "UV_CACHE_DIR": str(self.uv_cache), }, + start_new_session=True, ) try: - _, stderr = await asyncio.wait_for( - process.communicate(), + _, stderr = await communicate_process_group( + process, timeout=timeout_seconds, ) - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - raise except TimeoutError as e: - process.kill() - await process.wait() raise PackageInstallError( f"Package installation timed out after {timeout_seconds}s" ) from e From b2c3fe7c47e15660f3f00273757ff5f4373b4890 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:05:44 -0400 Subject: [PATCH 045/161] fix(executor): clean partial artifacts off loop --- .../test_registry_artifact_materialization.py | 100 ++++++++++++++++++ .../registry_artifact_materialization.py | 15 ++- 2 files changed, 111 insertions(+), 4 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index b5c741c3de..7538271f4f 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import shutil import tarfile import threading from pathlib import Path @@ -443,6 +444,105 @@ def blocking_size_scan(path: Path) -> int: assert second_cancellation_propagated_early is False extract.assert_not_awaited() + @pytest.mark.parametrize("artifact_format", ["squashfs", "tarball"]) + @pytest.mark.anyio + async def test_repeatedly_cancelled_partial_cleanup_rejoins_thread( + self, + temp_cache_dir, + artifact_format: str, + ): + """Large partial environments are deleted off-loop before cancellation.""" + cache = RegistryArtifactCache(temp_cache_dir) + cleanup_started = threading.Event() + cleanup_release = threading.Event() + cleanup_finished = threading.Event() + real_rmtree = shutil.rmtree + + def blocking_rmtree(path: Path, *, ignore_errors: bool = False) -> None: + cleanup_started.set() + cleanup_release.wait() + real_rmtree(path, ignore_errors=ignore_errors) + cleanup_finished.set() + + async def assert_cleanup(task: asyncio.Task[list[Path]]) -> None: + assert await asyncio.to_thread(cleanup_started.wait, 1) + task.cancel() + done, _ = await asyncio.wait({task}, timeout=0.05) + first_cancellation_propagated_early = bool(done) + + task.cancel() + done, _ = await asyncio.wait({task}, timeout=0.05) + second_cancellation_propagated_early = bool(done) + cleanup_release.set() + + with pytest.raises(asyncio.CancelledError): + await task + + assert cleanup_finished.is_set() + assert first_cancellation_propagated_early is False + assert second_cancellation_propagated_early is False + + async def fail_tarball_extract(self, tarball_path, target_dir): + del self, tarball_path + target_dir.mkdir(parents=True) + (target_dir / "partial.py").write_text("partial") + raise RuntimeError("tarball extraction failed") + + async def download_tarball(self, ctx, path): + del self, ctx + path.write_bytes(tarball_payload(size=1)) + + async def download_squashfs(self, ctx, image_path): + del self, ctx + image_path.parent.mkdir(parents=True, exist_ok=True) + image_path.write_bytes(b"image") + return 0.0 + + async def fail_squashfs_extract(self, image_path, target_dir): + del self, image_path + target_dir.mkdir(parents=True) + (target_dir / "partial.py").write_text("partial") + raise RuntimeError("SquashFS extraction failed") + + with patch( + "tracecat.executor.registry_artifact_materialization.shutil.rmtree", + side_effect=blocking_rmtree, + ): + if artifact_format == "tarball": + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "download", download_tarball), + patch.object(TarballArtifact, "extract", fail_tarball_extract), + ): + task = asyncio.create_task( + lease_paths(cache, "s3://bucket/partial.tar.gz") + ) + await assert_cleanup(task) + else: + with ( + patch.object( + RegistryArtifactMaterializationContext, + "can_mount_squashfs", + return_value=False, + ), + patch.object(SquashfsArtifact, "download", download_squashfs), + patch.object( + SquashfsArtifact, + "_squashfs_extracted_size", + new_callable=AsyncMock, + return_value=1, + ), + patch.object( + SquashfsArtifact, + "_extract_image", + fail_squashfs_extract, + ), + ): + task = asyncio.create_task( + lease_paths(cache, "s3://bucket/partial.squashfs") + ) + await assert_cleanup(task) + @pytest.mark.anyio async def test_materialize_extracts_squashfs_when_mount_fails(self, temp_cache_dir): """Test that SquashFS mount failures fall back to unsquashfs extraction.""" diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index dd230b5be9..26e8ef45e6 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -157,6 +157,15 @@ async def _run_blocking_rejoin_on_cancel[T](operation: Callable[[], T]) -> T: raise +async def _remove_tree_rejoin_on_cancel(path: Path) -> None: + """Remove a directory off-loop without abandoning cleanup on cancellation.""" + if not path.exists(): + return + await _run_blocking_rejoin_on_cancel( + lambda: shutil.rmtree(path, ignore_errors=True) + ) + + @dataclass(frozen=True, slots=True) class BuiltinArtifact(RegistryArtifact): """Current builtin registry package already installed in the executor image.""" @@ -392,8 +401,7 @@ async def extract( else: raise finally: - if temp_dir.exists(): - shutil.rmtree(temp_dir, ignore_errors=True) + await _remove_tree_rejoin_on_cancel(temp_dir) return target_dir @@ -576,8 +584,7 @@ async def materialize( else: raise finally: - if temp_dir.exists(): - shutil.rmtree(temp_dir, ignore_errors=True) + await _remove_tree_rejoin_on_cancel(temp_dir) if temp_tarball.exists(): temp_tarball.unlink(missing_ok=True) From a301b66427528692a1a67739ae2734a8ed31bada Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:06:57 -0400 Subject: [PATCH 046/161] fix(executor): persist lease release recency --- tests/unit/executor/test_registry_artifact_leases.py | 4 +++- tracecat/executor/registry_artifact_cache_state.py | 5 ++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_leases.py b/tests/unit/executor/test_registry_artifact_leases.py index 4c5c188470..6a3c91a3c0 100644 --- a/tests/unit/executor/test_registry_artifact_leases.py +++ b/tests/unit/executor/test_registry_artifact_leases.py @@ -72,7 +72,7 @@ def test_touch_entry_refreshes_tarball_root_mtime(self, temp_cache_dir): @pytest.mark.anyio async def test_lease_refcounts_and_touches_image_mtime(self, temp_cache_dir): - """A lease pins its entry and refreshes the restart-safe LRU timestamp.""" + """Acquire and final release persist the restart-safe LRU timestamp.""" cache = RegistryArtifactCache(temp_cache_dir) artifact_uri = "s3://bucket/leased.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) @@ -84,8 +84,10 @@ async def test_lease_refcounts_and_touches_image_mtime(self, temp_cache_dir): assert registry_paths == [target_dir] assert cache._refcount(cache_key) == 1 assert entry_dir.stat().st_mtime > 100.0 + os.utime(entry_dir, (100.0, 100.0)) assert cache._refcount(cache_key) == 0 + assert entry_dir.stat().st_mtime > 100.0 assert image_path.is_file() @pytest.mark.anyio diff --git a/tracecat/executor/registry_artifact_cache_state.py b/tracecat/executor/registry_artifact_cache_state.py index 5f921bc1c5..aabe995f13 100644 --- a/tracecat/executor/registry_artifact_cache_state.py +++ b/tracecat/executor/registry_artifact_cache_state.py @@ -145,7 +145,10 @@ def _release_lease(self, cache_key: str) -> bool: return False runtime.refcount -= 1 runtime.last_used = time.time() - return runtime.refcount == 0 + became_idle = runtime.refcount == 0 + if became_idle: + self._touch_entry(cache_key) + return became_idle def _refcount(self, cache_key: str) -> int: """Return the number of live leases on a cache entry.""" From d6653eb126e0b62616ad9229d011926ea4a64c7c Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:08:32 -0400 Subject: [PATCH 047/161] perf(executor): skip redundant budget rescans --- .../executor/test_registry_artifact_budget.py | 28 +++++++++++++++++++ .../executor/registry_artifact_storage.py | 7 ++++- 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 5463692a82..0f09eaa25d 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -507,6 +507,34 @@ async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( assert cache._budget_dirty is False + @pytest.mark.anyio + async def test_successful_cold_admission_skips_release_rescan(self, temp_cache_dir): + """A successful protected pass consumes the materialization dirty signal.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/new-with-one-budget-pass.tar.gz" + + async def mock_download(self, ctx, path): + del self, ctx + path.write_bytes(tarball_payload(size=1)) + + with ( + patch(MAX_ENTRIES_CONFIG, 10), + patch(MAX_BYTES_CONFIG, 0), + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "download", mock_download), + patch.object( + cache, + "_scan_cache_entries", + wraps=cache._scan_cache_entries, + ) as scan_cache_entries, + ): + async with cache.lease([artifact_uri]): + assert cache._budget_dirty is False + + assert scan_cache_entries.call_count == 1 + assert cache._budget_dirty is False + @pytest.mark.anyio async def test_failed_cold_admission_keeps_warm_lru(self, temp_cache_dir): """A missing cold artifact must not evict an existing warm entry.""" diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 35c930fab4..f1c7f06d6f 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -268,9 +268,14 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo """ async with self._admission_lock: async with self._budget_lock: - return await self._enforce_cache_budget_locked( + within_budget = await self._enforce_cache_budget_locked( protected_key=protected_key ) + if within_budget: + # Clear while cold writers remain excluded so a later + # materialization cannot have its dirty signal erased. + self._budget_dirty = False + return within_budget async def _enforce_cache_budget_locked( self, From b6d3a53d11713aecd58309e31e0b3f20ce111773 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:09:30 -0400 Subject: [PATCH 048/161] test(executor): exercise mounted startup protection --- tests/unit/executor/test_registry_artifact_startup.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_startup.py b/tests/unit/executor/test_registry_artifact_startup.py index bc2d6d83a9..cac94a85bb 100644 --- a/tests/unit/executor/test_registry_artifact_startup.py +++ b/tests/unit/executor/test_registry_artifact_startup.py @@ -87,15 +87,22 @@ async def test_sweep_removes_orphaned_work(self, temp_cache_dir): async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): """A live mountpoint belongs to a running process and must survive.""" cache = RegistryArtifactCache(temp_cache_dir) - paths = cache._paths_for("abc123") + paths = cache._paths_for("mounted") paths.entry_dir.mkdir(parents=True) mount_dir = paths.squashfs_mount_dir mount_dir.mkdir() + idle_dir = write_tarball_entry(temp_cache_dir, "idle") - with patch.object(Path, "is_mount", lambda self: self == mount_dir): + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch.object(Path, "is_mount", lambda self: self == mount_dir), + ): await cache.ensure_swept() assert mount_dir.is_dir() + assert not idle_dir.exists() + assert cache._budget_dirty is False @pytest.mark.anyio async def test_ensure_swept_runs_once(self, temp_cache_dir): From 66fe8fbb58ff7007b98dd77d7d626aadeb84648e Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:11:37 -0400 Subject: [PATCH 049/161] test(executor): align budget dirty expectations --- .../executor/test_registry_artifact_budget.py | 35 ------------------- 1 file changed, 35 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 0f09eaa25d..a6ecdccb23 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -562,41 +562,6 @@ async def mock_download(self, ctx, path): assert warm_dir.is_dir() assert not cache._paths_for(missing_key).entry_dir.exists() - @pytest.mark.anyio - async def test_materialization_rearms_budget_dirty_consumed_mid_flight( - self, temp_cache_dir - ): - """A convergence pass consuming the signal mid-download cannot unarm it.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - assert cache._budget_dirty is False - - artifact_uri = "s3://bucket/path/site-packages.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) - - async def mock_materialize(self, ctx): - # A concurrent lease release consumes the dirty signal and finishes - # its scan before this attempt lands its canonical bytes. - cache._budget_dirty = False - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - ctx.paths.squashfs_image_path.write_bytes(b"late image") - return [ctx.paths.squashfs_mount_dir] - - with ( - patch.object( - cache, - "_artifact_candidates", - new_callable=AsyncMock, - return_value=[artifact], - ), - patch.object(SquashfsArtifact, "materialize", mock_materialize), - ): - async with cache.lease([artifact_uri]): - assert cache._budget_dirty is True - - assert cache._budget_dirty is False - @pytest.mark.anyio async def test_release_keeps_retrying_while_the_cache_stays_over_budget( self, temp_cache_dir From 2906ad602493ae3599c3a9d8cb10ff220678c969 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:20:46 -0400 Subject: [PATCH 050/161] fix(executor): always unlink cancelled tarballs --- .../unit/executor/test_registry_artifact_materialization.py | 3 +++ tracecat/executor/registry_artifact_materialization.py | 5 +++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 7538271f4f..7521873fc1 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -481,6 +481,9 @@ async def assert_cleanup(task: asyncio.Task[list[Path]]) -> None: assert cleanup_finished.is_set() assert first_cancellation_propagated_early is False assert second_cancellation_propagated_early is False + assert not cache.staging_dir.exists() or not any( + cache.staging_dir.iterdir() + ) async def fail_tarball_extract(self, tarball_path, target_dir): del self, tarball_path diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 26e8ef45e6..2bbe88f379 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -584,8 +584,9 @@ async def materialize( else: raise finally: - await _remove_tree_rejoin_on_cancel(temp_dir) - if temp_tarball.exists(): + try: + await _remove_tree_rejoin_on_cancel(temp_dir) + finally: temp_tarball.unlink(missing_ok=True) return [target_dir] From e40d54fdcf0371f94b692b99549921794a83bcd3 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:26:33 -0400 Subject: [PATCH 051/161] fix(executor): preserve warm cache on trash failure --- .../executor/test_registry_artifact_budget.py | 33 +++++++++++++++++++ .../executor/registry_artifact_storage.py | 11 ++++++- 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index a6ecdccb23..72360419dc 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -419,6 +419,39 @@ def fail_once(path: Path) -> bool: assert newest.exists() assert not any(cache.trash_dir.iterdir()) + @pytest.mark.anyio + async def test_undeletable_trash_does_not_evict_warm_entries_for_admission( + self, temp_cache_dir + ): + """Unreclaimed trash blocks an oversized write without extra eviction.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + warm = write_image_entry(temp_cache_dir, "warm", size=32, mtime=100.0) + stale_trash = cache.trash_dir / "stale" + stale_trash.mkdir(parents=True) + (stale_trash / "image.squashfs").write_bytes(b"x" * 32) + real_delete = _delete_cache_path + + def fail_stale_trash(path: Path) -> bool: + if path == stale_trash: + return False + return real_delete(path) + + with patch( + "tracecat.executor.registry_artifact_storage._delete_cache_path", + side_effect=fail_stale_trash, + ): + with pytest.raises(RegistryArtifactCacheCapacityError) as raised: + await cache._ensure_cache_capacity( + additional_bytes=16, + protected_key="new", + max_bytes=64, + ) + + assert raised.value.current_bytes == 64 + assert warm.exists() + assert stale_trash.exists() + @pytest.mark.anyio async def test_rename_failure_does_not_block_materialization(self, temp_cache_dir): """A failed atomic retirement keeps the old entry and admits new work.""" diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index f1c7f06d6f..694db3182b 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -347,7 +347,7 @@ async def _ensure_cache_capacity( raise ValueError("additional_bytes must be non-negative") async with self._budget_lock: - await asyncio.gather( + trash_clean, startup_clean = await asyncio.gather( asyncio.to_thread(self._clear_work_dir, self.trash_dir), asyncio.to_thread(self._retry_failed_startup_cleanup), ) @@ -361,6 +361,15 @@ async def _ensure_cache_capacity( + staging_bytes + trash_bytes ) + if ( + not (trash_clean and startup_clean) + and total_bytes + additional_bytes > max_bytes + ): + raise RegistryArtifactCacheCapacityError( + current_bytes=total_bytes, + additional_bytes=additional_bytes, + max_bytes=max_bytes, + ) skipped = {protected_key} while total_bytes + additional_bytes > max_bytes: From 6166abeba9d5c13726d33c40b170d4409e038e68 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:27:20 -0400 Subject: [PATCH 052/161] test(executor): reject mixed broken artifact batches --- tests/unit/executor/test_test_backend_no_registry_action.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/executor/test_test_backend_no_registry_action.py b/tests/unit/executor/test_test_backend_no_registry_action.py index 50363e03fc..f954eb6ae9 100644 --- a/tests/unit/executor/test_test_backend_no_registry_action.py +++ b/tests/unit/executor/test_test_backend_no_registry_action.py @@ -275,7 +275,7 @@ def __init__(self) -> None: async def lease( self, artifact_uris: list[str] | None = None ) -> AsyncIterator[list[Path]]: - if artifact_uris == [broken_uri]: + if broken_uri in (artifact_uris or []): raise RuntimeError("artifact unavailable") self.active += 1 try: From ebe7245acfa13782e6970021ce2a6fe6d0385b33 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:29:38 -0400 Subject: [PATCH 053/161] test(executor): remove vacuous mount sweep check --- .../test_registry_artifact_cache_mount_lifecycle.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py index 6982282480..639539d632 100644 --- a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py +++ b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py @@ -184,9 +184,8 @@ def test_registry_artifact_cache_mount_lifecycle() -> None: assert payload["converged_loop_device_released"] is True assert payload["converged_entries_remaining"] == 2 - # The startup sweep trims to budget and removes stale mount directories. + # The startup sweep trims retained images to budget. assert payload["startup_sweep_trimmed"] is True - assert payload["startup_sweep_removed_stale_mount_dir"] is True def _squashfs_mounts(cache_dir: Path) -> dict[str, str]: @@ -401,7 +400,7 @@ async def hold_concurrent_lease(index: int) -> None: ) payload["converged_entries_remaining"] = len(cache._discover_cache_keys()) - # (f) The startup sweep trims to budget and drops stale mount directories. + # (f) The startup sweep trims retained images to budget. sweep_dir = root / "sweep-cache" sweep_dir.mkdir() sweep_cache = RegistryArtifactCache(sweep_dir) @@ -412,12 +411,6 @@ async def hold_concurrent_lease(index: int) -> None: image_path = sweep_paths.squashfs_image_path image_path.write_bytes(b"x" * 4096) os.utime(sweep_paths.entry_dir, (100.0 + index, 100.0 + index)) - stale_mount_dir = sweep_cache._paths_for(sweep_keys[0]).squashfs_mount_dir - stale_mount_dir.mkdir() - os.utime( - sweep_cache._paths_for(sweep_keys[0]).entry_dir, - (100.0, 100.0), - ) config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES = 1 await sweep_cache.ensure_swept() @@ -427,7 +420,6 @@ async def hold_concurrent_lease(index: int) -> None: not evicted_paths.squashfs_image_path.exists() and retained_paths.squashfs_image_path.exists() ) - payload["startup_sweep_removed_stale_mount_dir"] = not stale_mount_dir.exists() # Leave no mounts behind for the container teardown. for cache_key in sorted(cache._discover_cache_keys()): From b548ba94e275125ff4c769624b6a0028e816dbdb Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:32:24 -0400 Subject: [PATCH 054/161] perf(executor): skip non-final release rescans --- .../executor/test_registry_artifact_budget.py | 14 ++++++++++--- tracecat/executor/registry_artifacts.py | 21 +++++++++++++++---- 2 files changed, 28 insertions(+), 7 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 72360419dc..14f7218eb3 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -596,10 +596,10 @@ async def mock_download(self, ctx, path): assert not cache._paths_for(missing_key).entry_dir.exists() @pytest.mark.anyio - async def test_release_keeps_retrying_while_the_cache_stays_over_budget( + async def test_non_final_release_defers_retry_while_cache_stays_over_budget( self, temp_cache_dir ): - """A cache that cannot shrink yet must stay marked for re-enforcement.""" + """A non-final release defers rescanning until an entry becomes idle.""" cache = RegistryArtifactCache(temp_cache_dir) await cache.ensure_swept() artifact_uri = "s3://bucket/pinned.tar.gz" @@ -610,7 +610,15 @@ async def test_release_keeps_retrying_while_the_cache_stays_over_budget( # A second holder keeps the entry pinned past the inner lease. cache._acquire_lease(cache_key) - with patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 1): + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 1), + patch.object( + cache, + "_scan_cache_entries", + side_effect=AssertionError("non-final releases must not rescan"), + ), + ): async with cache.lease([artifact_uri]): pass diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index cbce7c922a..27cbc18177 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -122,6 +122,7 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat return leased_keys: list[str] = [] + lease_setup_complete = False try: registry_paths: list[Path] = [] for artifact_uri in artifact_uris: @@ -133,12 +134,18 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat "Using registry artifact environments", count=len(registry_paths), ) + lease_setup_complete = True yield registry_paths finally: idle_keys = [ cache_key for cache_key in leased_keys if self._release_lease(cache_key) ] - cleanup_task = asyncio.ensure_future(self._finish_lease_cleanup(idle_keys)) + cleanup_task = asyncio.ensure_future( + self._finish_lease_cleanup( + idle_keys, + converge=not lease_setup_complete or bool(idle_keys), + ) + ) pending_cancellation: asyncio.CancelledError | None = None while True: try: @@ -152,11 +159,17 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat if pending_cancellation is not None: raise pending_cancellation - async def _finish_lease_cleanup(self, idle_keys: list[str]) -> None: - """Unmount every newly idle entry and converge the cache budget.""" + async def _finish_lease_cleanup( + self, + idle_keys: list[str], + *, + converge: bool, + ) -> None: + """Unmount newly idle entries and converge after meaningful changes.""" for cache_key in idle_keys: await self._unmount_idle_entry(cache_key) - await self._converge_cache_budget() + if converge: + await self._converge_cache_budget() async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Path]]: """Pin and materialize one artifact, returning its releasable cache key.""" From 6da311be34bb5cf650916ec028c329b2bf54947b Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:41:40 -0400 Subject: [PATCH 055/161] test(executor): isolate eviction handoff race --- tests/unit/executor/test_registry_artifact_leases.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/unit/executor/test_registry_artifact_leases.py b/tests/unit/executor/test_registry_artifact_leases.py index 6a3c91a3c0..1227fb89d2 100644 --- a/tests/unit/executor/test_registry_artifact_leases.py +++ b/tests/unit/executor/test_registry_artifact_leases.py @@ -738,6 +738,15 @@ async def take_lease() -> None: side_effect=mock_umount, ), patch.object(SquashfsArtifact, "mount", mock_mount), + # This test targets the per-key eviction/lease handoff. Keep the + # lease's budget pass from concurrently sweeping the same trash + # path and turning physical reclamation into a two-deleter race. + patch.object( + cache, + "_enforce_cache_budget", + new_callable=AsyncMock, + return_value=True, + ), ): eviction = asyncio.create_task(cache._evict_entry(cache_key)) await umount_started.wait() From d3acc6bf45687ee7ae15d5dfaccdcf7dac519c39 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:49:17 -0400 Subject: [PATCH 056/161] fix(executor): retry failed idle unmounts --- .../executor/test_registry_artifact_budget.py | 43 +++++++++++++++++++ .../executor/registry_artifact_cache_state.py | 3 ++ .../executor/registry_artifact_storage.py | 30 ++++++++++++- tracecat/executor/registry_artifacts.py | 1 + 4 files changed, 76 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 14f7218eb3..596c60d09d 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -779,6 +779,49 @@ async def mock_umount(*args, **kwargs): assert paths.squashfs_image_path.read_bytes() == b"squashfs" assert paths.squashfs_mount_dir.is_dir() + @pytest.mark.anyio + async def test_failed_final_release_unmount_retries_on_later_cleanup( + self, temp_cache_dir + ): + """A transient unmount failure is retried after another lease release.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/path/retry-unmount.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + mounted = {paths.squashfs_mount_dir} + retry_uri = "s3://bucket/path/retry-trigger.tar.gz" + retry_key = compute_registry_artifact_cache_key(retry_uri) + write_tarball_entry(temp_cache_dir, retry_key) + attempts: list[Path] = [] + + async def flaky_unmount(mount_dir: Path) -> bool: + attempts.append(mount_dir) + if len(attempts) == 1: + return False + mounted.discard(mount_dir) + return True + + with ( + patch.object(Path, "is_mount", lambda path: path in mounted), + patch.object(cache, "_unmount", side_effect=flaky_unmount), + ): + async with cache.lease([artifact_uri]): + pass + + assert cache._failed_unmounts == {cache_key} + assert paths.squashfs_mount_dir in mounted + + async with cache.lease([retry_uri]): + pass + + assert attempts == [paths.squashfs_mount_dir, paths.squashfs_mount_dir] + assert cache._failed_unmounts == set() + assert paths.squashfs_mount_dir not in mounted + @pytest.mark.anyio async def test_concurrent_budget_passes_only_evict_once(self, temp_cache_dir): """A waiting budget pass must re-scan after the active pass evicts.""" diff --git a/tracecat/executor/registry_artifact_cache_state.py b/tracecat/executor/registry_artifact_cache_state.py index aabe995f13..875a795733 100644 --- a/tracecat/executor/registry_artifact_cache_state.py +++ b/tracecat/executor/registry_artifact_cache_state.py @@ -71,6 +71,9 @@ def __init__(self, cache_dir: Path): # Startup is the only time the whole staging directory is swept. Exact # paths that could not be removed are safe to retry later. self._failed_startup_cleanup: set[Path] = set() + # Final-release unmount failures are retried by later lease cleanup so + # transient errors cannot accumulate idle loop devices indefinitely. + self._failed_unmounts: set[str] = set() # Whether the on-disk cache may exceed its budget. Set when a new entry # is materialized and cleared once enforcement measures a cache that # fits, so steady-state cache hits never pay for a disk scan. diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 694db3182b..8b29d5c4e4 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -205,13 +205,41 @@ async def ensure_capacity(additional_bytes: int) -> None: async def _unmount_idle_entry(self, cache_key: str) -> None: """Best-effort unmount an entry after its final lease is released.""" try: - await self._unmount_entry(cache_key) + unmounted = await self._unmount_entry(cache_key) except OSError as e: + self._failed_unmounts.add(cache_key) logger.warning( "Failed to release idle registry artifact mount", cache_key=cache_key, error=str(e), ) + return + + if unmounted: + self._failed_unmounts.discard(cache_key) + return + + mount_dir = self._paths_for(cache_key).squashfs_mount_dir + try: + retry = self._refcount(cache_key) == 0 and mount_dir.is_mount() + except OSError as e: + retry = True + logger.warning( + "Failed to inspect idle registry artifact mount", + cache_key=cache_key, + mount_dir=str(mount_dir), + error=str(e), + ) + + if retry: + self._failed_unmounts.add(cache_key) + else: + self._failed_unmounts.discard(cache_key) + + async def _retry_failed_unmounts(self, *, excluded: set[str]) -> None: + """Retry prior unmount failures once on a later lease cleanup.""" + for cache_key in sorted(self._failed_unmounts - excluded): + await self._unmount_idle_entry(cache_key) async def _converge_cache_budget(self) -> None: """Bring an idle cache back under budget after a lease is released. diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 27cbc18177..525c19ee26 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -168,6 +168,7 @@ async def _finish_lease_cleanup( """Unmount newly idle entries and converge after meaningful changes.""" for cache_key in idle_keys: await self._unmount_idle_entry(cache_key) + await self._retry_failed_unmounts(excluded=set(idle_keys)) if converge: await self._converge_cache_budget() From d5109202ae1bfaded6d8db03f06c2c7c2ef56766 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:51:25 -0400 Subject: [PATCH 057/161] fix(executor): retry failed staging cleanup --- .../test_registry_artifact_materialization.py | 45 +++++++++++++++++++ .../test_registry_artifact_startup.py | 4 +- .../executor/registry_artifact_cache_state.py | 5 ++- .../registry_artifact_materialization.py | 41 ++++++++++++++--- .../executor/registry_artifact_storage.py | 18 ++++---- 5 files changed, 93 insertions(+), 20 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 7521873fc1..8962319e26 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -444,6 +444,51 @@ def blocking_size_scan(path: Path) -> int: assert second_cancellation_propagated_early is False extract.assert_not_awaited() + @pytest.mark.anyio + async def test_failed_partial_cleanup_is_deferred_for_capacity_retry( + self, temp_cache_dir + ): + """A failed staging-tree deletion remains discoverable and retryable.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/failed-cleanup.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + artifact = TarballArtifact(uri=artifact_uri, cache_key=cache_key) + ctx = cache._context_for(cache_key) + cleanup_attempts: list[Path] = [] + + async def download(self, ctx, path): + del self, ctx + path.write_bytes(b"archive") + + async def fail_extract(self, tarball_path, target_dir): + del self, tarball_path + (target_dir / "partial.py").write_text("partial") + raise RuntimeError("extraction failed") + + def fail_cleanup(path: Path) -> None: + cleanup_attempts.append(path) + raise PermissionError("cleanup denied") + + with ( + patch.object(TarballArtifact, "download", download), + patch.object(TarballArtifact, "extract", fail_extract), + patch( + "tracecat.executor.registry_artifact_materialization.shutil.rmtree", + side_effect=fail_cleanup, + ), + ): + with pytest.raises(RuntimeError, match="extraction failed"): + await artifact.materialize(ctx) + + assert len(cleanup_attempts) == 1 + assert set(cleanup_attempts) == cache._deferred_staging_cleanup + deferred_path = cleanup_attempts[0] + assert deferred_path.is_dir() + + assert cache._retry_deferred_staging_cleanup() is True + assert cache._deferred_staging_cleanup == set() + assert not deferred_path.exists() + @pytest.mark.parametrize("artifact_format", ["squashfs", "tarball"]) @pytest.mark.anyio async def test_repeatedly_cancelled_partial_cleanup_rejoins_thread( diff --git a/tests/unit/executor/test_registry_artifact_startup.py b/tests/unit/executor/test_registry_artifact_startup.py index cac94a85bb..657a1d34d4 100644 --- a/tests/unit/executor/test_registry_artifact_startup.py +++ b/tests/unit/executor/test_registry_artifact_startup.py @@ -294,13 +294,13 @@ def fail_once(path: Path) -> bool: ): await cache.ensure_swept() assert orphaned.is_file() - assert cache._failed_startup_cleanup == {orphaned} + assert cache._deferred_staging_cleanup == {orphaned} assert cache._budget_dirty is True assert await cache._enforce_cache_budget() is True assert not orphaned.exists() - assert cache._failed_startup_cleanup == set() + assert cache._deferred_staging_cleanup == set() @pytest.mark.anyio async def test_failed_startup_retirement_stays_dirty_and_retries( diff --git a/tracecat/executor/registry_artifact_cache_state.py b/tracecat/executor/registry_artifact_cache_state.py index 875a795733..4c1ae9bee3 100644 --- a/tracecat/executor/registry_artifact_cache_state.py +++ b/tracecat/executor/registry_artifact_cache_state.py @@ -69,8 +69,8 @@ def __init__(self, cache_dir: Path): self._sweep_task: asyncio.Task[None] | None = None self._sweep_lock = asyncio.Lock() # Startup is the only time the whole staging directory is swept. Exact - # paths that could not be removed are safe to retry later. - self._failed_startup_cleanup: set[Path] = set() + # startup or runtime paths that could not be removed are safe to retry. + self._deferred_staging_cleanup: set[Path] = set() # Final-release unmount failures are retried by later lease cleanup so # transient errors cannot accumulate idle loop devices indefinitely. self._failed_unmounts: set[str] = set() @@ -121,6 +121,7 @@ def _context_for( cache_key=cache_key, staging_dir=self.staging_dir, paths=self._paths_for(cache_key), + defer_cleanup=self._deferred_staging_cleanup.add, admission=admission, ) diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 2bbe88f379..4b7a7d1992 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -87,6 +87,7 @@ class RegistryArtifactMaterializationContext: cache_key: str staging_dir: Path paths: RegistryArtifactPaths + defer_cleanup: Callable[[Path], None] admission: RegistryArtifactAdmission | None = None def can_mount_squashfs(self) -> bool: @@ -157,13 +158,33 @@ async def _run_blocking_rejoin_on_cancel[T](operation: Callable[[], T]) -> T: raise -async def _remove_tree_rejoin_on_cancel(path: Path) -> None: - """Remove a directory off-loop without abandoning cleanup on cancellation.""" +async def _remove_tree_rejoin_on_cancel( + path: Path, + *, + defer_cleanup: Callable[[Path], None], +) -> None: + """Remove a tree off-loop and retain failed paths for a later retry.""" if not path.exists(): return - await _run_blocking_rejoin_on_cancel( - lambda: shutil.rmtree(path, ignore_errors=True) - ) + try: + await _run_blocking_rejoin_on_cancel(lambda: shutil.rmtree(path)) + except FileNotFoundError: + return + except asyncio.CancelledError: + if path.exists(): + defer_cleanup(path) + logger.warning( + "Deferred cancelled registry artifact staging cleanup", + path=str(path), + ) + raise + except OSError as e: + defer_cleanup(path) + logger.warning( + "Deferred failed registry artifact staging cleanup", + path=str(path), + error=str(e), + ) @dataclass(frozen=True, slots=True) @@ -401,7 +422,10 @@ async def extract( else: raise finally: - await _remove_tree_rejoin_on_cancel(temp_dir) + await _remove_tree_rejoin_on_cancel( + temp_dir, + defer_cleanup=ctx.defer_cleanup, + ) return target_dir @@ -585,7 +609,10 @@ async def materialize( raise finally: try: - await _remove_tree_rejoin_on_cancel(temp_dir) + await _remove_tree_rejoin_on_cancel( + temp_dir, + defer_cleanup=ctx.defer_cleanup, + ) finally: temp_tarball.unlink(missing_ok=True) diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 8b29d5c4e4..d524d07f88 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -313,7 +313,7 @@ async def _enforce_cache_budget_locked( """Enforce entry and byte limits while both cache-wide locks are held.""" trash_clean, startup_clean = await asyncio.gather( asyncio.to_thread(self._clear_work_dir, self.trash_dir), - asyncio.to_thread(self._retry_failed_startup_cleanup), + asyncio.to_thread(self._retry_deferred_staging_cleanup), ) cleanup_complete = trash_clean and startup_clean if not cleanup_complete: @@ -377,7 +377,7 @@ async def _ensure_cache_capacity( async with self._budget_lock: trash_clean, startup_clean = await asyncio.gather( asyncio.to_thread(self._clear_work_dir, self.trash_dir), - asyncio.to_thread(self._retry_failed_startup_cleanup), + asyncio.to_thread(self._retry_deferred_staging_cleanup), ) entries = await asyncio.to_thread(self._scan_cache_entries) staging_bytes, trash_bytes = await asyncio.gather( @@ -708,7 +708,7 @@ def _clear_work_dir( for path in paths: if _delete_cache_path(path): if remember_failures: - self._failed_startup_cleanup.discard(path) + self._deferred_staging_cleanup.discard(path) logger.info( "Removed registry artifact work path", path=str(path), @@ -716,15 +716,15 @@ def _clear_work_dir( else: deleted = False if remember_failures: - self._failed_startup_cleanup.add(path) + self._deferred_staging_cleanup.add(path) return deleted - def _retry_failed_startup_cleanup(self) -> bool: - """Retry exact startup paths without sweeping live staging work.""" - for path in tuple(self._failed_startup_cleanup): + def _retry_deferred_staging_cleanup(self) -> bool: + """Retry exact failed paths without sweeping live staging work.""" + for path in tuple(self._deferred_staging_cleanup): if _delete_cache_path(path): - self._failed_startup_cleanup.discard(path) - return not self._failed_startup_cleanup + self._deferred_staging_cleanup.discard(path) + return not self._deferred_staging_cleanup def _trim_startup_cache(self) -> bool: """Trim the cache to budget before any artifact is leased. From f85157d689d7f10cd552e10d05149b3118e465c9 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:58:07 -0400 Subject: [PATCH 058/161] fix(executor): rejoin cancelled artifact subprocesses --- .../registry_artifact_test_helpers.py | 8 ++- .../test_registry_artifact_materialization.py | 57 +++++++++++++++++++ .../registry_artifact_materialization.py | 56 +++++++++++------- 3 files changed, 99 insertions(+), 22 deletions(-) diff --git a/tests/unit/executor/registry_artifact_test_helpers.py b/tests/unit/executor/registry_artifact_test_helpers.py index 062a0d0c71..37e00f4248 100644 --- a/tests/unit/executor/registry_artifact_test_helpers.py +++ b/tests/unit/executor/registry_artifact_test_helpers.py @@ -129,10 +129,13 @@ async def lease_paths( class BlockingSubprocess: """Fake subprocess that blocks in communicate until it is cancelled.""" - def __init__(self) -> None: + def __init__(self, *, block_wait: bool = False) -> None: self.communicate_started = asyncio.Event() + self.wait_started = asyncio.Event() + self.release_wait = asyncio.Event() self.cleanup_calls: list[str] = [] self.returncode: int | None = None + self._block_wait = block_wait async def communicate(self) -> tuple[bytes, bytes]: """Block until the task awaiting subprocess completion is cancelled.""" @@ -148,6 +151,9 @@ def kill(self) -> None: async def wait(self) -> int: """Record that the killed subprocess was reaped.""" self.cleanup_calls.append("wait") + self.wait_started.set() + if self._block_wait: + await self.release_wait.wait() return -9 diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 8962319e26..cf244d0931 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -336,6 +336,63 @@ async def create_sleep_subprocess( assert captured.reaped is True assert captured.returncode is not None + @pytest.mark.parametrize("operation", ["mount", "extract", "size"]) + @pytest.mark.anyio + async def test_repeated_cancellation_reaps_squashfs_subprocess( + self, + temp_cache_dir: Path, + operation: str, + ) -> None: + """A second cancellation cannot abandon a killed SquashFS child.""" + artifact = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key="repeated-subprocess-cancellation", + ) + image_path = temp_cache_dir / "image.squashfs" + image_path.write_bytes(b"squashfs") + target_dir = temp_cache_dir / "target" + target_dir.mkdir() + process = BlockingSubprocess(block_wait=True) + + with ( + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/usr/bin/unsquashfs", + ), + patch( + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ), + ): + if operation == "mount": + running = asyncio.create_task( + artifact._mount_image(image_path, target_dir) + ) + elif operation == "extract": + running = asyncio.create_task( + artifact._extract_image(image_path, target_dir) + ) + else: + running = asyncio.create_task( + artifact._squashfs_extracted_size(image_path) + ) + + await process.communicate_started.wait() + running.cancel() + await process.wait_started.wait() + + running.cancel() + done, _ = await asyncio.wait({running}, timeout=0.05) + second_cancellation_propagated_early = bool(done) + process.release_wait.set() + + with pytest.raises(asyncio.CancelledError): + await running + + assert second_cancellation_propagated_early is False + assert process.cleanup_calls == ["kill", "wait"] + @pytest.mark.anyio async def test_repeatedly_cancelled_tarball_extract_rejoins_thread( self, temp_cache_dir diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 4b7a7d1992..ef97c7dac0 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -158,6 +158,38 @@ async def _run_blocking_rejoin_on_cancel[T](operation: Callable[[], T]) -> T: raise +async def _kill_and_reap_subprocess(process: asyncio.subprocess.Process) -> None: + """Kill a subprocess and wait until its child state is reaped.""" + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + + +async def _communicate_rejoin_on_cancel( + process: asyncio.subprocess.Process, +) -> tuple[bytes, bytes]: + """Communicate while keeping cancellation from abandoning child cleanup.""" + try: + stdout, stderr = await process.communicate() + except asyncio.CancelledError: + reaper = asyncio.ensure_future(_kill_and_reap_subprocess(process)) + while not reaper.done(): + try: + await asyncio.shield(reaper) + except asyncio.CancelledError: + continue + except Exception: + break + if not reaper.cancelled(): + with contextlib.suppress(Exception): + reaper.result() + raise + + if stdout is None or stderr is None: + raise RuntimeError("Captured subprocess output is required") + return stdout, stderr + + async def _remove_tree_rejoin_on_cancel( path: Path, *, @@ -458,13 +490,7 @@ async def _mount_image(self, image_path: Path, target_dir: Path) -> None: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - try: - stdout, stderr = await proc.communicate() - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - proc.kill() - await proc.wait() - raise + stdout, stderr = await _communicate_rejoin_on_cancel(proc) if proc.returncode == 0 or target_dir.is_mount(): return @@ -492,13 +518,7 @@ async def _extract_image(self, image_path: Path, target_dir: Path) -> None: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - try: - stdout, stderr = await proc.communicate() - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - proc.kill() - await proc.wait() - raise + stdout, stderr = await _communicate_rejoin_on_cancel(proc) if proc.returncode == 0: return @@ -519,13 +539,7 @@ async def _squashfs_extracted_size(self, image_path: Path) -> int: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - try: - stdout, stderr = await proc.communicate() - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - proc.kill() - await proc.wait() - raise + stdout, stderr = await _communicate_rejoin_on_cancel(proc) if proc.returncode != 0: output = (stderr or stdout).decode(errors="replace").strip() From 88eef13ddd0e68d742b9a226911e1df2f3bc3702 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:59:01 -0400 Subject: [PATCH 059/161] fix(executor): rejoin cancelled artifact unmounts --- .../unit/executor/test_registry_artifact_eviction.py | 11 ++++++++--- tracecat/executor/registry_artifact_storage.py | 9 ++------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_eviction.py b/tests/unit/executor/test_registry_artifact_eviction.py index d4d2108ad7..1a994a9bdf 100644 --- a/tests/unit/executor/test_registry_artifact_eviction.py +++ b/tests/unit/executor/test_registry_artifact_eviction.py @@ -75,10 +75,10 @@ async def mock_umount(*args, **kwargs): ) @pytest.mark.anyio - async def test_cancelled_unmount_kills_and_reaps_before_releasing_key_lock( + async def test_repeatedly_cancelled_unmount_reaps_before_releasing_key_lock( self, temp_cache_dir ): - """Cancellation leaves a consistent entry for the next admission.""" + """Repeated cancellation leaves a consistent entry for next admission.""" cache = RegistryArtifactCache(temp_cache_dir) artifact_uri = "s3://bucket/path/cancelled-unmount.squashfs" cache_key = compute_registry_artifact_cache_key(artifact_uri) @@ -88,7 +88,7 @@ async def test_cancelled_unmount_kills_and_reaps_before_releasing_key_lock( paths.squashfs_mount_dir.mkdir() (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") mounted = {paths.squashfs_mount_dir} - blocked_process = BlockingSubprocess() + blocked_process = BlockingSubprocess(block_wait=True) released_process = AsyncMock() released_process.communicate.return_value = (b"", b"") released_process.returncode = 0 @@ -116,7 +116,12 @@ async def mock_umount(*args, **kwargs): eviction = asyncio.create_task(cache._evict_entry(cache_key)) await blocked_process.communicate_started.wait() eviction.cancel() + await blocked_process.wait_started.wait() + eviction.cancel() + done, _ = await asyncio.wait({eviction}, timeout=0.05) + assert not done + blocked_process.release_wait.set() with pytest.raises(asyncio.CancelledError): await eviction diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index d524d07f88..ccaa3623fb 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -17,6 +17,7 @@ ) from tracecat.executor.registry_artifact_materialization import ( RegistryArtifactAdmission, + _communicate_rejoin_on_cancel, ) from tracecat.logger import logger @@ -587,13 +588,7 @@ async def _unmount(self, mount_dir: Path) -> bool: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - try: - stdout, stderr = await proc.communicate() - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - proc.kill() - await proc.wait() - raise + stdout, stderr = await _communicate_rejoin_on_cancel(proc) if proc.returncode == 0 or not mount_dir.is_mount(): return True From 092bd2d08fa5304130087c66ceb8f91102be4da4 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:59:54 -0400 Subject: [PATCH 060/161] test(sandbox): wait for complete descendant pid --- tests/unit/test_executor_sandbox_nsjail.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_executor_sandbox_nsjail.py b/tests/unit/test_executor_sandbox_nsjail.py index 36389a6e0f..c4eb6e130a 100644 --- a/tests/unit/test_executor_sandbox_nsjail.py +++ b/tests/unit/test_executor_sandbox_nsjail.py @@ -759,11 +759,13 @@ async def capture_subprocess(*args, **kwargs): try: await process_started.wait() for _ in range(100): - if descendant_pid_path.exists(): + try: + descendant_pid = int(descendant_pid_path.read_text()) + except (FileNotFoundError, ValueError): + await asyncio.sleep(0.01) + else: break - await asyncio.sleep(0.01) - assert descendant_pid_path.is_file() - descendant_pid = int(descendant_pid_path.read_text()) + assert descendant_pid is not None execution.cancel() From 4f3baf074f38b250ba09d413cc46e4a1f8da6dc2 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 17:13:50 -0400 Subject: [PATCH 061/161] fix(executor): discard idle cache runtime states --- .../test_registry_artifact_eviction.py | 9 ++--- .../executor/test_registry_artifact_leases.py | 20 ++++++++++ .../executor/registry_artifact_cache_state.py | 40 ++++++++++++++++++- .../executor/registry_artifact_storage.py | 12 +++--- tracecat/executor/registry_artifacts.py | 5 +-- 5 files changed, 69 insertions(+), 17 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_eviction.py b/tests/unit/executor/test_registry_artifact_eviction.py index 1a994a9bdf..aefc71b15d 100644 --- a/tests/unit/executor/test_registry_artifact_eviction.py +++ b/tests/unit/executor/test_registry_artifact_eviction.py @@ -268,20 +268,17 @@ async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): assert not idle.exists() @pytest.mark.anyio - async def test_eviction_keeps_stable_runtime_state(self, temp_cache_dir): - """A key keeps one lock and zeroed lease state for the process lifetime.""" + async def test_eviction_discards_idle_runtime_state(self, temp_cache_dir): + """An evicted key releases runtime state after every lock user exits.""" cache = RegistryArtifactCache(temp_cache_dir) write_tarball_entry(temp_cache_dir, "bookkeeping") cache._acquire_lease("bookkeeping") cache._release_lease("bookkeeping") - lock = cache._runtime_for("bookkeeping").lock assert await cache._evict_entry("bookkeeping") == RegistryArtifactEviction( retired=True, reclaimed=True ) - runtime = cache._runtime["bookkeeping"] - assert runtime.lock is lock - assert runtime.refcount == 0 + assert "bookkeeping" not in cache._runtime @pytest.mark.anyio async def test_eviction_skips_busy_key(self, temp_cache_dir): diff --git a/tests/unit/executor/test_registry_artifact_leases.py b/tests/unit/executor/test_registry_artifact_leases.py index 1227fb89d2..34d7cc18ce 100644 --- a/tests/unit/executor/test_registry_artifact_leases.py +++ b/tests/unit/executor/test_registry_artifact_leases.py @@ -110,6 +110,26 @@ async def mock_download(self, ctx, path): assert cache._refcount(cache_key) == 0 assert not cache._paths_for(cache_key).entry_dir.exists() assert cache_key not in cache._discover_cache_keys() + assert cache_key not in cache._runtime + + @pytest.mark.anyio + async def test_distinct_failed_keys_do_not_accumulate_runtime_states( + self, temp_cache_dir: Path + ) -> None: + """Failed cold keys release lock state after their last waiter exits.""" + cache = RegistryArtifactCache(temp_cache_dir) + + async def fail_download(self, ctx, path): + del self, ctx, path + raise RuntimeError("download failed") + + with patch.object(TarballArtifact, "download", fail_download): + for index in range(100): + with pytest.raises(RuntimeError, match="download failed"): + async with cache.lease([f"s3://bucket/broken-{index}.tar.gz"]): + pass + + assert cache._runtime == {} @pytest.mark.anyio async def test_failed_first_admission_converges_deposited_image( diff --git a/tracecat/executor/registry_artifact_cache_state.py b/tracecat/executor/registry_artifact_cache_state.py index 4c1ae9bee3..b1daac941f 100644 --- a/tracecat/executor/registry_artifact_cache_state.py +++ b/tracecat/executor/registry_artifact_cache_state.py @@ -6,6 +6,8 @@ import os import threading import time +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from dataclasses import dataclass, field from pathlib import Path @@ -40,6 +42,7 @@ class RegistryArtifactRuntimeState: lock: asyncio.Lock = field(default_factory=asyncio.Lock) refcount: int = 0 last_used: float = 0.0 + users: int = 0 class _RegistryArtifactCacheState: @@ -50,8 +53,9 @@ def __init__(self, cache_dir: Path): self.entries_dir = cache_dir / CACHE_ENTRIES_DIR_NAME self.staging_dir = cache_dir / CACHE_STAGING_DIR_NAME self.trash_dir = cache_dir / CACHE_TRASH_DIR_NAME - # Runtime states live for the process lifetime so every operation for a - # key always serializes on the same lock. + # Runtime states remain stable while their entry, leases, or lock users + # exist. Failed and evicted keys are discarded once no waiter can still + # hold the old lock identity. self._runtime: dict[str, RegistryArtifactRuntimeState] = {} # The cache contains asyncio locks, tasks, and multi-step lease state. # Bind the public API to one loop/thread so a future synchronous @@ -110,6 +114,38 @@ def _runtime_for(self, cache_key: str) -> RegistryArtifactRuntimeState: self._runtime[cache_key] = runtime return runtime + @asynccontextmanager + async def _runtime_lock( + self, cache_key: str + ) -> AsyncIterator[RegistryArtifactRuntimeState]: + """Lock one key while keeping its runtime identity pinned for waiters.""" + runtime = self._runtime_for(cache_key) + runtime.users += 1 + try: + async with runtime.lock: + yield runtime + finally: + runtime.users -= 1 + self._discard_idle_runtime(cache_key, runtime) + + def _discard_idle_runtime( + self, + cache_key: str, + runtime: RegistryArtifactRuntimeState, + ) -> None: + """Discard state only after its entry and every possible user are gone.""" + if runtime.users > 0 or runtime.refcount > 0 or runtime.lock.locked(): + return + try: + if self._paths_for(cache_key).entry_dir.exists(): + return + except OSError: + # Retaining a small state object is safer than splitting lock identity + # when the entry cannot be inspected. + return + if self._runtime.get(cache_key) is runtime: + del self._runtime[cache_key] + def _context_for( self, cache_key: str, diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index ccaa3623fb..b1a05f5bdf 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -464,15 +464,15 @@ async def _unmount_entry(self, cache_key: str) -> bool: Returns: Whether a mounted entry was unmounted. """ - lock = self._runtime_for(cache_key).lock - if lock.locked(): + runtime = self._runtime.get(cache_key) + if runtime is not None and runtime.lock.locked(): logger.debug( "Skipping unmount of busy registry artifact", cache_key=cache_key, ) return False - async with lock: + async with self._runtime_lock(cache_key): if self._refcount(cache_key) > 0: return False @@ -511,15 +511,15 @@ async def _evict_entry(self, cache_key: str) -> RegistryArtifactEviction: Returns: Whether the entry was retired and its bytes were reclaimed. """ - lock = self._runtime_for(cache_key).lock - if lock.locked(): + runtime = self._runtime.get(cache_key) + if runtime is not None and runtime.lock.locked(): logger.debug( "Skipping eviction of busy registry artifact", cache_key=cache_key, ) return RegistryArtifactEviction(retired=False, reclaimed=False) - async with lock: + async with self._runtime_lock(cache_key): if self._refcount(cache_key) > 0: return RegistryArtifactEviction(retired=False, reclaimed=False) diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 525c19ee26..284976a96d 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -180,17 +180,16 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat candidates = await self._artifact_candidates(ctx, artifact_uri) return None, await self._materialize_candidates(ctx, candidates) - lock = self._runtime_for(cache_key).lock lease_acquired = False try: - async with lock: + async with self._runtime_lock(cache_key): self._acquire_lease(cache_key) lease_acquired = True if cached_paths := self._locally_cached_path(ctx, artifact_uri): return cache_key, cached_paths async with self._admission_lock: - async with lock: + async with self._runtime_lock(cache_key): if cached_paths := self._locally_cached_path(ctx, artifact_uri): return cache_key, cached_paths ctx = self._context_for( From 7abcb544259fb8fec3492b7f7962b27a06adf03e Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 17:14:51 -0400 Subject: [PATCH 062/161] fix(executor): preserve results on cleanup failure --- .../executor/test_registry_artifact_leases.py | 35 +++++++++++++++++++ tracecat/executor/registry_artifacts.py | 7 ++++ 2 files changed, 42 insertions(+) diff --git a/tests/unit/executor/test_registry_artifact_leases.py b/tests/unit/executor/test_registry_artifact_leases.py index 34d7cc18ce..52b5ef0748 100644 --- a/tests/unit/executor/test_registry_artifact_leases.py +++ b/tests/unit/executor/test_registry_artifact_leases.py @@ -446,6 +446,41 @@ async def hold_lease() -> None: assert cleanup_calls == [*cache_keys, "converge"] assert all(cache._refcount(cache_key) == 0 for cache_key in cache_keys) + @pytest.mark.anyio + async def test_cleanup_failure_preserves_successful_lease_outcome( + self, temp_cache_dir: Path + ) -> None: + """Post-lease maintenance cannot replace a successful caller result.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/cleanup-failure.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + target_dir = write_tarball_entry(temp_cache_dir, cache_key) + + async def fail_cleanup( + idle_keys: list[str], + *, + converge: bool, + ) -> None: + del idle_keys, converge + raise RuntimeError("cleanup failed") + + with ( + patch.object(cache, "_finish_lease_cleanup", fail_cleanup), + patch( + "tracecat.executor.registry_artifacts.logger.exception" + ) as log_exception, + ): + async with cache.lease([artifact_uri]) as registry_paths: + result = registry_paths + + assert result == [target_dir] + assert cache._refcount(cache_key) == 0 + log_exception.assert_called_once_with( + "Registry artifact lease cleanup failed; preserving caller outcome", + cache_dir=str(temp_cache_dir), + error="cleanup failed", + ) + @pytest.mark.anyio async def test_new_lease_racing_final_release_prevents_stale_unmount( self, temp_cache_dir: Path diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 284976a96d..76ff51babe 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -155,6 +155,13 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat if cleanup_task.cancelled(): raise pending_cancellation = e + except Exception as e: + logger.exception( + "Registry artifact lease cleanup failed; preserving caller outcome", + cache_dir=str(self.cache_dir), + error=str(e), + ) + break if pending_cancellation is not None: raise pending_cancellation From 232de0a5738cb6524ec029cf465f43636329cf07 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 17:15:50 -0400 Subject: [PATCH 063/161] test(sandbox): bound nsjail startup wait --- tests/unit/test_executor_sandbox_nsjail.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_executor_sandbox_nsjail.py b/tests/unit/test_executor_sandbox_nsjail.py index c4eb6e130a..d8619ed137 100644 --- a/tests/unit/test_executor_sandbox_nsjail.py +++ b/tests/unit/test_executor_sandbox_nsjail.py @@ -757,7 +757,7 @@ async def capture_subprocess(*args, **kwargs): ) try: - await process_started.wait() + await asyncio.wait_for(process_started.wait(), timeout=5) for _ in range(100): try: descendant_pid = int(descendant_pid_path.read_text()) @@ -784,6 +784,10 @@ async def capture_subprocess(*args, **kwargs): else: pytest.fail("nsjail descendant survived process-group cleanup") finally: + if not execution.done(): + execution.cancel() + with contextlib.suppress(asyncio.CancelledError): + await execution if process is not None and process.returncode is None: process.kill() await process.wait() From 577e4003c725545b649408afb14271da243ff1f8 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 17:30:11 -0400 Subject: [PATCH 064/161] fix(executor): account allocated cache bytes --- tests/unit/executor/conftest.py | 25 ++++++ .../executor/test_registry_artifact_budget.py | 57 +++++++++++++- .../test_registry_artifact_materialization.py | 22 +++++- .../registry_artifact_materialization.py | 62 +++++++++++---- .../executor/registry_artifact_storage.py | 76 +++++++++++++++---- tracecat/executor/registry_artifacts.py | 2 + 6 files changed, 211 insertions(+), 33 deletions(-) diff --git a/tests/unit/executor/conftest.py b/tests/unit/executor/conftest.py index 5a218e3f97..26cd619eab 100644 --- a/tests/unit/executor/conftest.py +++ b/tests/unit/executor/conftest.py @@ -2,12 +2,37 @@ from __future__ import annotations +import os +import stat import tempfile from collections.abc import Iterator from pathlib import Path import pytest +from tracecat.executor import registry_artifact_storage + + +@pytest.fixture(autouse=True) +def logical_cache_sizes(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep synthetic byte budgets independent from host filesystem block sizes.""" + + def logical_stat_size(file_stat: os.stat_result) -> int: + if stat.S_ISDIR(file_stat.st_mode): + return 0 + return file_stat.st_size + + monkeypatch.setattr( + registry_artifact_storage, + "_allocated_stat_size", + logical_stat_size, + ) + monkeypatch.setattr( + registry_artifact_storage, + "_filesystem_allocation_unit", + lambda _path: 1, + ) + @pytest.fixture def temp_cache_dir() -> Iterator[Path]: diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 596c60d09d..84d39ed3f5 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -7,7 +7,7 @@ import threading from collections.abc import Awaitable, Callable from pathlib import Path -from unittest.mock import ANY, AsyncMock, patch +from unittest.mock import ANY, AsyncMock, MagicMock, patch import pytest @@ -18,7 +18,9 @@ RegistryArtifactMaterializationContext, SquashfsArtifact, TarballArtifact, + _allocated_stat_size, _delete_cache_path, + _directory_footprint, compute_registry_artifact_cache_key, ) @@ -36,6 +38,59 @@ class TestRegistryArtifactCacheBudget: """Enforce peak and steady-state cache capacity.""" + def test_allocated_stat_size_uses_filesystem_blocks(self) -> None: + file_stat = MagicMock(spec=os.stat_result) + file_stat.st_blocks = 7 + file_stat.st_size = 1 + + assert _allocated_stat_size(file_stat) == 7 * 512 + + def test_directory_footprint_includes_directory_inodes( + self, + temp_cache_dir: Path, + ) -> None: + nested = temp_cache_dir / "nested" + nested.mkdir() + (nested / "module.py").write_text("x") + + with patch( + "tracecat.executor.registry_artifact_storage._allocated_stat_size", + return_value=4096, + ) as allocated_size: + assert _directory_footprint(temp_cache_dir) == 3 * 4096 + + assert allocated_size.call_count == 3 + + @pytest.mark.anyio + async def test_admission_rounds_download_reservation_to_allocation_unit( + self, + temp_cache_dir: Path, + ) -> None: + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch(MAX_BYTES_CONFIG, 8192), + patch( + "tracecat.executor.registry_artifact_storage." + "_filesystem_allocation_unit", + return_value=4096, + ), + patch.object( + cache, + "_ensure_cache_capacity", + new_callable=AsyncMock, + ) as ensure_capacity, + ): + admission = cache._admission_for("new") + assert admission is not None + await admission.ensure_capacity(1) + + ensure_capacity.assert_awaited_once_with( + additional_bytes=4096, + protected_key="new", + max_bytes=8192, + ) + def test_delete_cache_path_reports_directory_failure(self, temp_cache_dir): """Physical deletion failures are observable instead of suppressed.""" entry_dir = temp_cache_dir / "entry" diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index cf244d0931..59338e0883 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import io import shutil import tarfile import threading @@ -19,6 +20,7 @@ SquashfsMountCommandError, TarballArtifact, _squashfs_listing_size, + _tarball_extracted_size, compute_registry_artifact_cache_key, ) @@ -138,7 +140,7 @@ async def take_lease(uri: str) -> None: compute_registry_artifact_cache_key(uri) for uri in uris ] - def test_squashfs_listing_size_sums_files_and_symlinks(self) -> None: + def test_squashfs_listing_size_bounds_each_inode_allocation(self) -> None: listing = b"\n".join( [ b"drwxr-xr-x 0/0 64 2026-01-01 00:00 squashfs-root", @@ -147,7 +149,20 @@ def test_squashfs_listing_size_sums_files_and_symlinks(self) -> None: ] ) - assert _squashfs_listing_size(listing) == 132 + assert _squashfs_listing_size(listing, allocation_unit=4096) == 12_288 + + def test_tarball_size_bounds_each_member_allocation( + self, + temp_cache_dir: Path, + ) -> None: + tarball_path = temp_cache_dir / "many-small-files.tar.gz" + with tarfile.open(tarball_path, "w:gz") as tar: + for index in range(3): + member = tarfile.TarInfo(f"module-{index}.py") + member.size = 1 + tar.addfile(member, io.BytesIO(b"x")) + + assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 12_288 def test_squashfs_listing_size_rejects_unparseable_files(self) -> None: with pytest.raises(ValueError, match="Could not parse SquashFS listing"): @@ -459,7 +474,8 @@ async def mock_download(self, ctx, path): path.write_bytes(tarball_payload(size=1)) downloaded_paths.append(path) - def blocking_size_scan(path: Path) -> int: + def blocking_size_scan(path: Path, *, allocation_unit: int) -> int: + assert allocation_unit == 1 scan_started.set() scan_release.wait() input_present_at_finish.append(path.exists()) diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index ef97c7dac0..d280ea71be 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -77,6 +77,7 @@ class RegistryArtifactAdmission: """Byte-bound admission hook shared by one cold materialization.""" max_bytes: int + allocation_unit: int ensure_capacity: Callable[[int], Awaitable[None]] @@ -424,7 +425,10 @@ async def extract( temp_dir = self._temp_path(ctx, ".unsquashfs") try: if ctx.admission is not None: - extracted_size = await self._squashfs_extracted_size(image_path) + extracted_size = await self._squashfs_extracted_size( + image_path, + allocation_unit=ctx.admission.allocation_unit, + ) await ctx.admission.ensure_capacity(extracted_size) extract_start = time.monotonic() temp_dir.mkdir(parents=True, exist_ok=True) @@ -526,8 +530,13 @@ async def _extract_image(self, image_path: Path, target_dir: Path) -> None: output = (stderr or stdout).decode(errors="replace").strip() raise RuntimeError(output or "unsquashfs command failed") - async def _squashfs_extracted_size(self, image_path: Path) -> int: - """Return a conservative logical size for an extracted SquashFS image.""" + async def _squashfs_extracted_size( + self, + image_path: Path, + *, + allocation_unit: int = 1, + ) -> int: + """Return a conservative allocated size for a SquashFS extraction.""" unsquashfs = shutil.which("unsquashfs") if unsquashfs is None: raise RuntimeError("unsquashfs command is not installed") @@ -544,7 +553,7 @@ async def _squashfs_extracted_size(self, image_path: Path) -> int: if proc.returncode != 0: output = (stderr or stdout).decode(errors="replace").strip() raise RuntimeError(output or "unsquashfs listing failed") - return _squashfs_listing_size(stdout) + return _squashfs_listing_size(stdout, allocation_unit=allocation_unit) @dataclass(frozen=True, slots=True) @@ -588,11 +597,15 @@ async def materialize( await self.download(ctx, temp_tarball) download_elapsed = (time.monotonic() - download_start) * 1000 - if ctx.admission is not None: + admission = ctx.admission + if admission is not None: extracted_size = await _run_blocking_rejoin_on_cancel( - lambda: _tarball_extracted_size(temp_tarball) + lambda: _tarball_extracted_size( + temp_tarball, + allocation_unit=admission.allocation_unit, + ) ) - await ctx.admission.ensure_capacity(extracted_size) + await admission.ensure_capacity(extracted_size) extract_start = time.monotonic() temp_dir.mkdir(parents=True, exist_ok=True) @@ -789,8 +802,23 @@ def _is_cache_entry_uri(artifact_uri: str) -> bool: return _bundled_builtin_registry_version(artifact_uri) is None -def _tarball_extracted_size(tarball_path: Path) -> int: - """Return a conservative logical size for all tarball members.""" +def _allocated_size_bound(size_bytes: int, *, allocation_unit: int) -> int: + """Round one filesystem object up to its minimum allocated footprint.""" + if size_bytes < 0: + raise ValueError("size_bytes must be non-negative") + if allocation_unit <= 0: + raise ValueError("allocation_unit must be positive") + return ( + max(1, (size_bytes + allocation_unit - 1) // allocation_unit) * allocation_unit + ) + + +def _tarball_extracted_size( + tarball_path: Path, + *, + allocation_unit: int = 1, +) -> int: + """Return a conservative allocated size bound for all tarball members.""" total_bytes = 0 with tarfile.open(tarball_path, "r:gz") as tar: for member in tar: @@ -798,12 +826,15 @@ def _tarball_extracted_size(tarball_path: Path) -> int: raise ValueError( f"Registry tarball member has a negative size: {member.name}" ) - total_bytes += member.size + total_bytes += _allocated_size_bound( + member.size, + allocation_unit=allocation_unit, + ) return total_bytes -def _squashfs_listing_size(output: bytes) -> int: - """Sum file sizes from ``unsquashfs -lln`` output, failing closed.""" +def _squashfs_listing_size(output: bytes, *, allocation_unit: int = 1) -> int: + """Bound allocated bytes from ``unsquashfs -lln`` output, failing closed.""" total_bytes = 0 for raw_line in output.decode(errors="strict").splitlines(): line = raw_line.strip() @@ -813,9 +844,10 @@ def _squashfs_listing_size(output: bytes) -> int: mode = fields[0] if len(mode) != 10 or mode[0] not in "bcdlps-": continue - if mode[0] not in "-l": - continue if len(fields) < 5 or "/" not in fields[1] or not fields[2].isdigit(): raise ValueError(f"Could not parse SquashFS listing line: {line}") - total_bytes += int(fields[2]) + total_bytes += _allocated_size_bound( + int(fields[2]), + allocation_unit=allocation_unit, + ) return total_bytes diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index b1a05f5bdf..fca4d5af2c 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -6,6 +6,7 @@ import contextlib import os import shutil +import stat import time from collections.abc import Iterable from dataclasses import dataclass @@ -17,6 +18,7 @@ ) from tracecat.executor.registry_artifact_materialization import ( RegistryArtifactAdmission, + _allocated_size_bound, _communicate_rejoin_on_cancel, ) from tracecat.logger import logger @@ -59,14 +61,38 @@ class RegistryArtifactCacheEntry: last_used: float +def _allocated_stat_size(file_stat: os.stat_result) -> int: + """Return allocated bytes for one inode, falling back to logical size.""" + blocks = getattr(file_stat, "st_blocks", None) + if blocks is None: + return file_stat.st_size + return blocks * 512 + + +def _filesystem_allocation_unit(path: Path) -> int: + """Return the allocation unit for a path or its nearest existing parent.""" + candidate = path + while True: + try: + filesystem = os.statvfs(candidate) + break + except FileNotFoundError: + parent = candidate.parent + if parent == candidate: + raise + candidate = parent + + return filesystem.f_frsize or filesystem.f_bsize or 1 + + def _directory_footprint(directory: Path) -> int: - """Return the total file size of a cache directory. + """Return the allocated footprint of a cache directory tree. Args: directory: Cache directory to measure. Returns: - Total byte size of contained files, or zero when the directory is + Total allocated bytes of contained inodes, or zero when the directory is missing. """ @@ -76,12 +102,25 @@ def raise_walk_error(error: OSError) -> None: total_bytes = 0 try: walker = os.walk(directory, onerror=raise_walk_error) - for root, _dirs, files in walker: + for root, dirs, files in walker: + try: + total_bytes += _allocated_stat_size(os.lstat(root)) + except FileNotFoundError: + continue for file_name in files: try: - total_bytes += os.lstat(os.path.join(root, file_name)).st_size + total_bytes += _allocated_stat_size( + os.lstat(os.path.join(root, file_name)) + ) + except FileNotFoundError: + continue + for directory_name in dirs: + try: + directory_stat = os.lstat(os.path.join(root, directory_name)) except FileNotFoundError: continue + if stat.S_ISLNK(directory_stat.st_mode): + total_bytes += _allocated_stat_size(directory_stat) except FileNotFoundError: return 0 return total_bytes @@ -191,15 +230,21 @@ def _admission_for(self, cache_key: str) -> RegistryArtifactAdmission | None: if max_bytes <= 0: return None + allocation_unit = _filesystem_allocation_unit(self.cache_dir) + async def ensure_capacity(additional_bytes: int) -> None: await self._ensure_cache_capacity( - additional_bytes=additional_bytes, + additional_bytes=_allocated_size_bound( + additional_bytes, + allocation_unit=allocation_unit, + ), protected_key=cache_key, max_bytes=max_bytes, ) return RegistryArtifactAdmission( max_bytes=max_bytes, + allocation_unit=allocation_unit, ensure_capacity=ensure_capacity, ) @@ -622,19 +667,22 @@ def _discover_cache_keys(self) -> set[str]: def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: """Measure the on-disk footprint and recency of one cache entry. - The mount directory is excluded because a mounted view only costs the - image file that backs it. The image is measured with a single ``stat`` - so a concurrent eviction deleting it cannot fail the scan. + Mounted contents are excluded because the image already accounts for + their backing bytes. The entry root, mount-point inode, and image are + measured individually so a concurrent eviction cannot fail the scan. """ paths = self._paths_for(cache_key) size_bytes = 0 - try: - image_stat = paths.squashfs_image_path.stat() - except FileNotFoundError: - pass - else: - size_bytes += image_stat.st_size + for path in ( + paths.entry_dir, + paths.squashfs_image_path, + paths.squashfs_mount_dir, + ): + try: + size_bytes += _allocated_stat_size(path.lstat()) + except FileNotFoundError: + continue for directory in (paths.squashfs_extract_dir, paths.tarball_target_dir): size_bytes += _directory_footprint(directory) diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 76ff51babe..791dfb2b58 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -45,6 +45,7 @@ RegistryArtifactCacheCapacityError, RegistryArtifactCacheEntry, RegistryArtifactEviction, + _allocated_stat_size, _delete_cache_path, _delete_cache_path_off_loop, _directory_footprint, @@ -79,6 +80,7 @@ "SquashfsMountCommandError", "TarballArtifact", "_artifact_format", + "_allocated_stat_size", "_bundled_builtin_registry_import_paths", "_bundled_builtin_registry_version", "_delete_cache_path", From 6bcd2f3acb7dc9c036165b2e5531c0d55be45529 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 17:36:53 -0400 Subject: [PATCH 065/161] fix(executor): rejoin failed lease rollback --- .../executor/test_registry_artifact_leases.py | 44 +++++++++++++++++++ tracecat/executor/registry_artifacts.py | 14 +++++- 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_leases.py b/tests/unit/executor/test_registry_artifact_leases.py index 52b5ef0748..854659145d 100644 --- a/tests/unit/executor/test_registry_artifact_leases.py +++ b/tests/unit/executor/test_registry_artifact_leases.py @@ -203,6 +203,50 @@ async def take_lease() -> None: assert cache._refcount(cache_key) == 0 assert not cache._paths_for(cache_key).entry_dir.exists() + @pytest.mark.anyio + async def test_repeated_cancellation_finishes_acquisition_rollback( + self, + temp_cache_dir: Path, + ) -> None: + """A second cancellation cannot abandon the final rollback unmount.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/path/site-packages.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + lookup_started = asyncio.Event() + rollback_started = asyncio.Event() + finish_rollback = asyncio.Event() + + async def blocked_sidecar_lookup(**kwargs): + del kwargs + lookup_started.set() + await asyncio.Event().wait() + + async def blocked_unmount(requested_key: str) -> None: + assert requested_key == cache_key + rollback_started.set() + await finish_rollback.wait() + + with ( + patch(SQUASHFS_ENABLED_CONFIG, True), + patch.object(cache, "_sidecar_exists", blocked_sidecar_lookup), + patch.object(cache, "_unmount_idle_entry", blocked_unmount), + ): + acquisition = asyncio.create_task(cache._lease_artifact(artifact_uri)) + await lookup_started.wait() + acquisition.cancel() + await rollback_started.wait() + + acquisition.cancel() + await asyncio.sleep(0) + assert not acquisition.done() + + finish_rollback.set() + with pytest.raises(asyncio.CancelledError): + await acquisition + + assert cache._refcount(cache_key) == 0 + @pytest.mark.anyio async def test_cancelled_waiter_preserves_existing_same_key_lease( self, temp_cache_dir: Path diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 791dfb2b58..983642aa78 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -227,7 +227,19 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat return cache_key, paths except BaseException: if lease_acquired and self._release_lease(cache_key): - await self._unmount_idle_entry(cache_key) + rollback_task = asyncio.ensure_future( + self._unmount_idle_entry(cache_key) + ) + while not rollback_task.done(): + try: + await asyncio.shield(rollback_task) + except asyncio.CancelledError: + continue + except Exception: + break + if not rollback_task.cancelled(): + with contextlib.suppress(Exception): + rollback_task.result() raise async def _materialize_candidates( From d432fde40a7203b15e5ce2e3eca73335e865668f Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 17:39:49 -0400 Subject: [PATCH 066/161] fix(executor): contain artifact utility processes --- .../registry_artifact_test_helpers.py | 6 +++++ .../executor/test_registry_artifact_budget.py | 1 + .../test_registry_artifact_eviction.py | 8 +++++++ .../test_registry_artifact_materialization.py | 22 ++++++++++++++----- .../registry_artifact_materialization.py | 10 +++++---- .../executor/registry_artifact_storage.py | 1 + 6 files changed, 39 insertions(+), 9 deletions(-) diff --git a/tests/unit/executor/registry_artifact_test_helpers.py b/tests/unit/executor/registry_artifact_test_helpers.py index 37e00f4248..670521cf83 100644 --- a/tests/unit/executor/registry_artifact_test_helpers.py +++ b/tests/unit/executor/registry_artifact_test_helpers.py @@ -130,6 +130,7 @@ class BlockingSubprocess: """Fake subprocess that blocks in communicate until it is cancelled.""" def __init__(self, *, block_wait: bool = False) -> None: + self.pid = 999_999_999 self.communicate_started = asyncio.Event() self.wait_started = asyncio.Event() self.release_wait = asyncio.Event() @@ -170,6 +171,11 @@ def returncode(self) -> int | None: """Return the wrapped subprocess exit status.""" return self.process.returncode + @property + def pid(self) -> int: + """Return the wrapped subprocess process-group identifier.""" + return self.process.pid + async def communicate(self) -> tuple[bytes, bytes]: """Wait for the wrapped subprocess and collect its output.""" return await self.process.communicate() diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 84d39ed3f5..30189e938a 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -810,6 +810,7 @@ async def test_final_lease_release_unmounts_and_retains_image(self, temp_cache_d process.returncode = 0 async def mock_umount(*args, **kwargs): + assert kwargs["start_new_session"] is True mounted.discard(paths.squashfs_mount_dir) return process diff --git a/tests/unit/executor/test_registry_artifact_eviction.py b/tests/unit/executor/test_registry_artifact_eviction.py index aefc71b15d..80348e0880 100644 --- a/tests/unit/executor/test_registry_artifact_eviction.py +++ b/tests/unit/executor/test_registry_artifact_eviction.py @@ -4,6 +4,7 @@ import asyncio import os +import signal import threading from pathlib import Path from unittest.mock import AsyncMock, patch @@ -47,6 +48,7 @@ async def test_eviction_unmounts_before_deleting_the_image(self, temp_cache_dir) process.returncode = 0 async def mock_umount(*args, **kwargs): + assert kwargs["start_new_session"] is True image_present_at_umount.append(paths.squashfs_image_path.exists()) mounted.discard(paths.squashfs_mount_dir) return process @@ -96,6 +98,7 @@ async def test_repeatedly_cancelled_unmount_reaps_before_releasing_key_lock( async def mock_umount(*args, **kwargs): nonlocal unmount_attempts + assert kwargs["start_new_session"] is True unmount_attempts += 1 if unmount_attempts == 1: return blocked_process @@ -112,6 +115,7 @@ async def mock_umount(*args, **kwargs): "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", side_effect=mock_umount, ), + patch("tracecat.sandbox.utils.os.killpg") as kill_group, ): eviction = asyncio.create_task(cache._evict_entry(cache_key)) await blocked_process.communicate_started.wait() @@ -126,6 +130,10 @@ async def mock_umount(*args, **kwargs): await eviction assert blocked_process.cleanup_calls == ["kill", "wait"] + kill_group.assert_called_once_with( + blocked_process.pid, + signal.SIGKILL, + ) async with cache.lease([artifact_uri]) as registry_paths: assert registry_paths == [paths.squashfs_mount_dir] assert registry_paths[0].is_dir() diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 59338e0883..2a57ddb875 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -5,6 +5,7 @@ import asyncio import io import shutil +import signal import tarfile import threading from pathlib import Path @@ -249,6 +250,7 @@ async def test_mount_squashfs_uses_hardened_read_only_options( str(target_dir), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + start_new_session=True, ) @pytest.mark.anyio @@ -268,10 +270,13 @@ async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): target_dir.mkdir() process = BlockingSubprocess() - with patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - new_callable=AsyncMock, - return_value=process, + with ( + patch( + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ), + patch("tracecat.sandbox.utils.os.killpg") as kill_group, ): mounting = asyncio.create_task( artifact._mount_image(image_path, target_dir) @@ -283,6 +288,7 @@ async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): await mounting assert process.cleanup_calls == ["kill", "wait"] + kill_group.assert_called_once_with(process.pid, signal.SIGKILL) assert target_dir.is_dir() assert not target_dir.is_mount() @@ -316,6 +322,7 @@ async def create_sleep_subprocess( "30", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + start_new_session=True, ) captured = CapturedSubprocess(process) captured_processes.append(captured) @@ -378,7 +385,8 @@ async def test_repeated_cancellation_reaps_squashfs_subprocess( "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process, - ), + ) as create_subprocess_exec, + patch("tracecat.sandbox.utils.os.killpg") as kill_group, ): if operation == "mount": running = asyncio.create_task( @@ -407,6 +415,10 @@ async def test_repeated_cancellation_reaps_squashfs_subprocess( assert second_cancellation_propagated_early is False assert process.cleanup_calls == ["kill", "wait"] + await_args = create_subprocess_exec.await_args + assert await_args is not None + assert await_args.kwargs["start_new_session"] is True + kill_group.assert_called_once_with(process.pid, signal.SIGKILL) @pytest.mark.anyio async def test_repeatedly_cancelled_tarball_extract_rejoins_thread( diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index d280ea71be..c2686a7dde 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -23,6 +23,7 @@ from tracecat.logger import logger from tracecat.registry.artifact_keys import parse_s3_uri from tracecat.registry.constants import DEFAULT_REGISTRY_ORIGIN +from tracecat.sandbox.utils import terminate_process_group from tracecat.storage import blob __all__ = [ @@ -160,10 +161,8 @@ async def _run_blocking_rejoin_on_cancel[T](operation: Callable[[], T]) -> T: async def _kill_and_reap_subprocess(process: asyncio.subprocess.Process) -> None: - """Kill a subprocess and wait until its child state is reaped.""" - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() + """Kill a subprocess group and wait until its leader is reaped.""" + await terminate_process_group(process) async def _communicate_rejoin_on_cancel( @@ -493,6 +492,7 @@ async def _mount_image(self, image_path: Path, target_dir: Path) -> None: str(target_dir), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + start_new_session=True, ) stdout, stderr = await _communicate_rejoin_on_cancel(proc) @@ -521,6 +521,7 @@ async def _extract_image(self, image_path: Path, target_dir: Path) -> None: str(image_path), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + start_new_session=True, ) stdout, stderr = await _communicate_rejoin_on_cancel(proc) @@ -547,6 +548,7 @@ async def _squashfs_extracted_size( str(image_path), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + start_new_session=True, ) stdout, stderr = await _communicate_rejoin_on_cancel(proc) diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index fca4d5af2c..97373540ca 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -632,6 +632,7 @@ async def _unmount(self, mount_dir: Path) -> bool: str(mount_dir), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + start_new_session=True, ) stdout, stderr = await _communicate_rejoin_on_cancel(proc) if proc.returncode == 0 or not mount_dir.is_mount(): From 084e50047f7d07950b4252354adbdb4726cb4935 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 17:41:14 -0400 Subject: [PATCH 067/161] test(executor): synchronize concurrent budget passes --- .../executor/test_registry_artifact_budget.py | 32 ++++++++++++------- 1 file changed, 21 insertions(+), 11 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 30189e938a..9867f233a3 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -7,6 +7,7 @@ import threading from collections.abc import Awaitable, Callable from pathlib import Path +from typing import Literal from unittest.mock import ANY, AsyncMock, MagicMock, patch import pytest @@ -882,6 +883,21 @@ async def flaky_unmount(mount_dir: Path) -> bool: async def test_concurrent_budget_passes_only_evict_once(self, temp_cache_dir): """A waiting budget pass must re-scan after the active pass evicts.""" cache = RegistryArtifactCache(temp_cache_dir) + + class ObservedLock(asyncio.Lock): + def __init__(self) -> None: + super().__init__() + self.second_acquire_started = asyncio.Event() + self._acquire_attempts = 0 + + async def acquire(self) -> Literal[True]: + self._acquire_attempts += 1 + if self._acquire_attempts == 2: + self.second_acquire_started.set() + return await super().acquire() + + admission_lock = ObservedLock() + cache._admission_lock = admission_lock oldest = write_image_entry(temp_cache_dir, "oldest", size=16, mtime=100.0) retained = write_image_entry(temp_cache_dir, "retained", size=16, mtime=200.0) original_scan = cache._scan_cache_entries @@ -908,7 +924,6 @@ def controlled_scan(): eviction_started = asyncio.Event() finish_eviction = asyncio.Event() - extra_eviction_finished = asyncio.Event() evicted_keys: list[str] = [] async def controlled_evict( @@ -925,7 +940,6 @@ async def controlled_evict( _delete_cache_path(cache._paths_for("oldest").entry_dir) else: _delete_cache_path(cache._paths_for("retained").entry_dir) - extra_eviction_finished.set() evicted_keys.append(cache_key) return RegistryArtifactEviction(retired=True, reclaimed=True) @@ -938,21 +952,17 @@ async def controlled_evict( first_pass = asyncio.create_task(cache._enforce_cache_budget()) assert await asyncio.to_thread(first_scan_started.wait, 1) second_pass = asyncio.create_task(cache._enforce_cache_budget()) - second_scan_overlapped = await asyncio.to_thread( - second_scan_started.wait, 0.5 - ) + await admission_lock.second_acquire_started.wait() + assert not second_scan_started.is_set() release_first_scan.set() await asyncio.wait_for(eviction_started.wait(), timeout=1) - release_second_scan.set() - try: - await asyncio.wait_for(extra_eviction_finished.wait(), timeout=0.2) - except TimeoutError: - pass + assert not second_scan_started.is_set() finish_eviction.set() + assert await asyncio.to_thread(second_scan_started.wait, 1) + release_second_scan.set() await asyncio.gather(first_pass, second_pass) - assert second_scan_overlapped is False assert scan_count == 2 assert evicted_keys == ["oldest"] assert not oldest.exists() From b02a3c72d976610ac20ef5a795d4e8753f6a9317 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 17:42:39 -0400 Subject: [PATCH 068/161] test(sandbox): always clean dependency setup processes --- tests/unit/test_unsafe_pid_executor.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_unsafe_pid_executor.py b/tests/unit/test_unsafe_pid_executor.py index 3772623314..1ecef7b382 100644 --- a/tests/unit/test_unsafe_pid_executor.py +++ b/tests/unit/test_unsafe_pid_executor.py @@ -247,9 +247,10 @@ async def create_dependency_process(*args, **kwargs): ) task = asyncio.create_task(setup) - await _wait_for_file(pid_file) - child_pid = int(pid_file.read_text()) + child_pid: int | None = None try: + await _wait_for_file(pid_file) + child_pid = int(pid_file.read_text()) task.cancel() with pytest.raises(asyncio.CancelledError): await task @@ -258,8 +259,20 @@ async def create_dependency_process(*args, **kwargs): assert created_processes[0].returncode is not None await _wait_for_process_exit(child_pid) finally: - with contextlib.suppress(ProcessLookupError): - os.kill(child_pid, signal.SIGKILL) + if not task.done(): + task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + for process in created_processes: + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + if process.returncode is None: + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + if child_pid is not None: + with contextlib.suppress(ProcessLookupError): + os.kill(child_pid, signal.SIGKILL) @pytest.mark.anyio async def test_execute_basic_script(self, executor: UnsafePidExecutor) -> None: From 3e114d0091af234697fddc01bf1a5e0e6de31e3f Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 17:43:39 -0400 Subject: [PATCH 069/161] test(executor): synchronize eviction lease handoff --- .../unit/executor/test_registry_artifact_leases.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_leases.py b/tests/unit/executor/test_registry_artifact_leases.py index 854659145d..957a2985a0 100644 --- a/tests/unit/executor/test_registry_artifact_leases.py +++ b/tests/unit/executor/test_registry_artifact_leases.py @@ -797,7 +797,15 @@ async def test_lease_is_never_admitted_across_an_in_flight_eviction( mounted = {paths.squashfs_mount_dir} umount_started = asyncio.Event() finish_umount = asyncio.Event() + lease_waiting_for_key = asyncio.Event() remounts: list[str] = [] + original_runtime_for = cache._runtime_for + + def observed_runtime_for(requested_key: str): + runtime = original_runtime_for(requested_key) + if requested_key == cache_key and umount_started.is_set(): + lease_waiting_for_key.set() + return runtime umount_process = AsyncMock() umount_process.communicate.return_value = (b"", b"") @@ -836,6 +844,7 @@ async def take_lease() -> None: "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", side_effect=mock_umount, ), + patch.object(cache, "_runtime_for", side_effect=observed_runtime_for), patch.object(SquashfsArtifact, "mount", mock_mount), # This test targets the per-key eviction/lease handoff. Keep the # lease's budget pass from concurrently sweeping the same trash @@ -850,8 +859,8 @@ async def take_lease() -> None: eviction = asyncio.create_task(cache._evict_entry(cache_key)) await umount_started.wait() lease = asyncio.create_task(take_lease()) - # Let the lease block on the per-key lock the eviction holds. - await asyncio.sleep(0) + await lease_waiting_for_key.wait() + assert cache._runtime[cache_key].users == 2 finish_umount.set() evicted, _ = await asyncio.gather(eviction, lease) From 64a7a5262d2476fc9b80bfc2a6d72374151459dd Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 17:44:27 -0400 Subject: [PATCH 070/161] test(executor): narrow temporal cache assertions --- .../test_registry_artifact_cache_temporal_worker.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/integration/test_registry_artifact_cache_temporal_worker.py b/tests/integration/test_registry_artifact_cache_temporal_worker.py index 47fae9f691..aa122cdc71 100644 --- a/tests/integration/test_registry_artifact_cache_temporal_worker.py +++ b/tests/integration/test_registry_artifact_cache_temporal_worker.py @@ -35,9 +35,8 @@ async def temporal_env() -> AsyncGenerator[WorkflowEnvironment, None]: @dataclass(frozen=True, slots=True) class _ProbeResult: - """Runtime identity observed by one Temporal activity.""" + """Runtime observations from one Temporal activity.""" - cache_instance_id: int event_loop_id: int thread_id: int registry_path: str @@ -73,7 +72,6 @@ async def run(self, index: int) -> _ProbeResult: try: await self.release.wait() return _ProbeResult( - cache_instance_id=id(self.cache), event_loop_id=id(asyncio.get_running_loop()), thread_id=threading.get_ident(), registry_path=str(registry_paths[0]), @@ -112,10 +110,9 @@ async def test_one_temporal_worker_uses_one_cache_loop_and_thread( """Protect the production async-activity ownership contract. A real Temporal worker must schedule overlapping cache users on one event - loop and thread while sharing one process-wide cache instance. The explicit - production-activity assertion also prevents action execution from quietly - moving into Temporal's synchronous thread pool, the historical failure mode - for process-wide async storage state. + loop and thread and balance their leases. The explicit production-activity + assertion also prevents action execution from quietly moving into Temporal's + synchronous thread pool, the historical failure mode for async storage state. """ assert inspect.iscoroutinefunction(ExecutorActivities.execute_action_activity) @@ -166,7 +163,6 @@ async def test_one_temporal_worker_uses_one_cache_loop_and_thread( assert probe.peak_holders == activity_count assert probe.active_holders == 0 - assert {result.cache_instance_id for result in results} == {id(cache)} assert len({result.event_loop_id for result in results}) == 1 assert len({result.thread_id for result in results}) == 1 assert {result.registry_path for result in results} == {str(registry_path)} From 58bcae6a70ac5f947e83a3ab82edda095f34e485 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 17:45:19 -0400 Subject: [PATCH 071/161] test(executor): preserve subprocess cleanup evidence --- .../test_registry_artifact_materialization.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 2a57ddb875..16b11f2495 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -346,16 +346,20 @@ async def create_sleep_subprocess( captured = captured_processes[0] extracting.cancel() + production_killed = False + production_reaped = False try: with pytest.raises(asyncio.CancelledError): await extracting + production_killed = captured.killed + production_reaped = captured.reaped finally: if captured.returncode is None: - captured.kill() - await captured.wait() + captured.process.kill() + await captured.process.wait() - assert captured.killed is True - assert captured.reaped is True + assert production_killed is True + assert production_reaped is True assert captured.returncode is not None @pytest.mark.parametrize("operation", ["mount", "extract", "size"]) From 8055be9b4cceffe8d269aa43d9553401d174a4f2 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:01:39 -0400 Subject: [PATCH 072/161] fix(sandbox): rejoin process cleanup on cancellation --- tests/unit/test_action_runner.py | 32 ++++++++++++++- tests/unit/test_sandbox_utils.py | 70 ++++++++++++++++++++++++++++++++ tracecat/sandbox/utils.py | 60 ++++++++++++++++++++++----- 3 files changed, 151 insertions(+), 11 deletions(-) create mode 100644 tests/unit/test_sandbox_utils.py diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index ac749d35e8..536ceadb56 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio +import contextlib import tempfile import uuid from datetime import UTC, datetime @@ -31,6 +32,7 @@ from tracecat.executor.secret_preprocessors import SecretEnvProjection from tracecat.identifiers.workflow import WorkflowUUID from tracecat.registry.lock.types import RegistryLock +from tracecat.sandbox import utils as sandbox_utils from tracecat.sandbox.types import SandboxResult @@ -747,7 +749,10 @@ async def test_cancelled_action_reaps_child_before_releasing_mounted_artifact( ) real_create_subprocess_exec = asyncio.create_subprocess_exec + real_terminate_process_group = sandbox_utils.terminate_process_group process_started = asyncio.Event() + termination_started = asyncio.Event() + finish_termination = asyncio.Event() process: asyncio.subprocess.Process | None = None reaped_before_unmount: list[bool] = [] @@ -757,6 +762,13 @@ async def capture_subprocess(*args, **kwargs): process_started.set() return process + async def controlled_termination( + requested_process: asyncio.subprocess.Process, + ) -> None: + termination_started.set() + await finish_termination.wait() + await real_terminate_process_group(requested_process) + async def release_mount(mount_dir: Path) -> bool: reaped_before_unmount.append( process is not None and process.returncode is not None @@ -775,6 +787,11 @@ async def release_mount(mount_dir: Path) -> bool: "tracecat.executor.action_runner.asyncio.create_subprocess_exec", side_effect=capture_subprocess, ), + patch.object( + sandbox_utils, + "terminate_process_group", + side_effect=controlled_termination, + ), patch.object(cache, "_unmount", side_effect=release_mount), ): execution = asyncio.create_task( @@ -787,13 +804,26 @@ async def release_mount(mount_dir: Path) -> bool: ) ) try: - await process_started.wait() + await asyncio.wait_for(process_started.wait(), timeout=5) assert cache._refcount(cache_key) == 1 execution.cancel() + await termination_started.wait() + execution.cancel() + await asyncio.sleep(0) + assert not execution.done() + assert cache._refcount(cache_key) == 1 + assert paths.squashfs_mount_dir in mounted + + finish_termination.set() with pytest.raises(asyncio.CancelledError): await execution finally: + finish_termination.set() + if not execution.done(): + execution.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await execution if process is not None and process.returncode is None: process.kill() await process.wait() diff --git a/tests/unit/test_sandbox_utils.py b/tests/unit/test_sandbox_utils.py new file mode 100644 index 0000000000..676ceb180c --- /dev/null +++ b/tests/unit/test_sandbox_utils.py @@ -0,0 +1,70 @@ +"""Tests for shared sandbox process utilities.""" + +from __future__ import annotations + +import asyncio +from typing import cast +from unittest.mock import patch + +import pytest + +from tracecat.sandbox.utils import communicate_process_group + + +class _BlockingProcess: + """Minimal process double whose communication blocks until cancellation.""" + + def __init__(self) -> None: + self.returncode: int | None = None + self.communicate_started = asyncio.Event() + self.communicate_finished = asyncio.Event() + + async def communicate( + self, + input: bytes | None = None, # noqa: A002 + ) -> tuple[bytes, bytes]: + del input + self.communicate_started.set() + try: + await asyncio.Event().wait() + finally: + self.communicate_finished.set() + return b"", b"" + + +@pytest.mark.anyio +async def test_repeated_cancellation_rejoins_process_group_cleanup() -> None: + """A second cancellation cannot return while group termination is live.""" + fake_process = _BlockingProcess() + process = cast(asyncio.subprocess.Process, fake_process) + termination_started = asyncio.Event() + finish_termination = asyncio.Event() + termination_finished = asyncio.Event() + + async def blocking_termination( + requested_process: asyncio.subprocess.Process, + ) -> None: + assert requested_process is process + termination_started.set() + await finish_termination.wait() + termination_finished.set() + + with patch( + "tracecat.sandbox.utils.terminate_process_group", + side_effect=blocking_termination, + ): + communication = asyncio.create_task(communicate_process_group(process)) + await fake_process.communicate_started.wait() + communication.cancel() + await termination_started.wait() + + communication.cancel() + await asyncio.sleep(0) + assert not communication.done() + + finish_termination.set() + with pytest.raises(asyncio.CancelledError): + await communication + + assert termination_finished.is_set() + assert fake_process.communicate_finished.is_set() diff --git a/tracecat/sandbox/utils.py b/tracecat/sandbox/utils.py index 360b1349e3..8fdc6b39ae 100644 --- a/tracecat/sandbox/utils.py +++ b/tracecat/sandbox/utils.py @@ -41,6 +41,42 @@ async def terminate_process_group(process: asyncio.subprocess.Process) -> None: await process.wait() +async def _finish_process_group_cleanup( + process: asyncio.subprocess.Process, + communicate_task: asyncio.Task[tuple[bytes | None, bytes | None]], + termination_task: asyncio.Task[None] | None, +) -> None: + """Finish process termination and consume the communication task.""" + if termination_task is None: + termination_task = asyncio.create_task(terminate_process_group(process)) + try: + await termination_task + finally: + if not communicate_task.done(): + communicate_task.cancel() + with suppress(asyncio.CancelledError): + await communicate_task + + +async def _rejoin_cleanup_through_cancellation( + cleanup_task: asyncio.Task[None], +) -> None: + """Wait for cleanup despite repeated caller cancellation.""" + pending_cancellation: asyncio.CancelledError | None = None + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError as e: + if cleanup_task.cancelled(): + raise + pending_cancellation = e + + if not cleanup_task.cancelled(): + cleanup_task.result() + if pending_cancellation is not None: + raise pending_cancellation + + async def communicate_process_group( process: asyncio.subprocess.Process, *, @@ -53,24 +89,28 @@ async def communicate_process_group( ``communicate()`` and asyncio's ``wait()`` to wait after the leader exits. Polling ``returncode`` observes the leader exit independently, letting us terminate the group immediately and close those pipes. Cancellation also - terminates the group before it propagates. + terminates the group before it propagates. Cleanup runs in an independent + task and is rejoined through repeated cancellation so callers cannot release + resources while the process group is still alive. """ communicate_task = asyncio.create_task(process.communicate(input=input)) - group_terminated = False + termination_task: asyncio.Task[None] | None = None try: async with asyncio.timeout(timeout): while process.returncode is None: await asyncio.sleep(_PROCESS_EXIT_POLL_INTERVAL_SECONDS) - await terminate_process_group(process) - group_terminated = True + termination_task = asyncio.create_task(terminate_process_group(process)) + await asyncio.shield(termination_task) stdout, stderr = await communicate_task finally: - if not group_terminated: - await terminate_process_group(process) - if not communicate_task.done(): - communicate_task.cancel() - with suppress(asyncio.CancelledError): - await communicate_task + cleanup_task = asyncio.create_task( + _finish_process_group_cleanup( + process, + communicate_task, + termination_task, + ) + ) + await _rejoin_cleanup_through_cancellation(cleanup_task) if stdout is None or stderr is None: raise RuntimeError("Captured stdout and stderr are required") From 9a63597a4dfdd1149c9084b0df01e6074ce2d1c1 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:07:21 -0400 Subject: [PATCH 073/161] fix(storage): rejoin cancelled file writes --- tests/unit/test_storage_blob.py | 76 +++++++++++++++++++++++++++++++++ tracecat/storage/blob.py | 25 ++++++++++- 2 files changed, 100 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_storage_blob.py b/tests/unit/test_storage_blob.py index c960a2cf5b..7d551f72de 100644 --- a/tests/unit/test_storage_blob.py +++ b/tests/unit/test_storage_blob.py @@ -2,6 +2,7 @@ import asyncio import hashlib +import threading from contextlib import asynccontextmanager from pathlib import Path from unittest.mock import AsyncMock, patch @@ -918,6 +919,81 @@ async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 assert not out.exists() assert not (tmp_path / "out.bin.part").exists() + @pytest.mark.anyio + async def test_download_file_to_path_cancellation_rejoins_active_write( + self, tmp_path: Path, monkeypatch + ): + """Cancellation cannot unlink a partial file under an active writer.""" + write_started = threading.Event() + write_release = threading.Event() + write_finished = threading.Event() + temp_path = tmp_path / "out.bin.part" + + class DummyStream: + async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 + yield b"partial" + + class BlockingFile: + async def __aenter__(self): + temp_path.touch() + return self + + async def __aexit__(self, exc_type, exc, traceback): + del exc_type, exc, traceback + + async def write(self, chunk: bytes) -> int: + def blocking_write() -> int: + write_started.set() + write_release.wait() + temp_path.write_bytes(chunk) + write_finished.set() + return len(chunk) + + return await asyncio.to_thread(blocking_write) + + @asynccontextmanager + async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 + yield DummyStream(), len(b"partial") + + monkeypatch.setattr( + "tracecat.storage.blob.open_download_stream", + _fake_open_download_stream, + ) + monkeypatch.setattr( + "tracecat.storage.blob.aiofiles.open", + lambda *args, **kwargs: BlockingFile(), + ) + + out = tmp_path / "out.bin" + download = asyncio.create_task( + download_file_to_path( + key="k", + bucket="b", + output_path=out, + ) + ) + try: + assert await asyncio.to_thread(write_started.wait, 1.0) + + download.cancel() + await asyncio.sleep(0) + assert not download.done() + assert temp_path.exists() + + download.cancel() + await asyncio.sleep(0) + assert not download.done() + assert temp_path.exists() + finally: + write_release.set() + + with pytest.raises(asyncio.CancelledError): + await download + + assert write_finished.is_set() + assert not out.exists() + assert not temp_path.exists() + @pytest.mark.anyio @patch("tracecat.storage.blob.get_storage_client") async def test_ensure_bucket_exists_create_error_propagates(self, mock_get_client): diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index 79fb6e5e22..24167410d6 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -785,6 +785,29 @@ async def download_file_to_path( hasher = hashlib.sha256() if expected_sha256 is not None else None bytes_written = 0 + async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: + """Write one chunk without abandoning the aiofiles worker thread.""" + writer = asyncio.ensure_future(file.write(chunk)) + try: + await asyncio.shield(writer) + except asyncio.CancelledError: + # aiofiles delegates writes to a thread that cannot be killed. Keep + # the partial file live until the worker stops touching it, even if + # the caller is cancelled repeatedly while cleanup is in progress. + while not writer.done(): + try: + await asyncio.shield(writer) + except asyncio.CancelledError: + continue + except Exception: + break + if not writer.cancelled(): + try: + writer.result() + except Exception: + pass + raise + try: async with open_download_stream(key=key, bucket=bucket) as ( stream, @@ -830,7 +853,7 @@ async def download_file_to_path( ) if hasher is not None: hasher.update(chunk) - await f.write(chunk) + await write_chunk_rejoin_on_cancel(f, chunk) if hasher is not None: actual_sha256 = hasher.hexdigest() From bc59df6d4f5f587dc890b0b38f80d3c30f82f2bb Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:08:48 -0400 Subject: [PATCH 074/161] fix(executor): defer failed tarball cleanup --- .../test_registry_artifact_materialization.py | 48 +++++++++++++++++++ .../registry_artifact_materialization.py | 22 ++++++++- 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 16b11f2495..9c366547a7 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -578,6 +578,54 @@ def fail_cleanup(path: Path) -> None: assert cache._deferred_staging_cleanup == set() assert not deferred_path.exists() + @pytest.mark.anyio + async def test_failed_tarball_unlink_is_deferred_without_masking_success( + self, temp_cache_dir + ): + """A failed tarball unlink preserves success and remains retryable.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/failed-tarball-cleanup.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + artifact = TarballArtifact(uri=artifact_uri, cache_key=cache_key) + ctx = cache._context_for(cache_key) + downloaded_paths: list[Path] = [] + real_unlink = Path.unlink + + async def download(self, ctx, path): + del self, ctx + path.write_bytes(b"archive") + downloaded_paths.append(path) + + async def extract(self, tarball_path, target_dir): + del self, tarball_path + (target_dir / "module.py").write_text("VALUE = 1") + + def fail_download_unlink( + path: Path, + missing_ok: bool = False, + ) -> None: + if path in downloaded_paths: + raise PermissionError("cleanup denied") + real_unlink(path, missing_ok=missing_ok) + + with ( + patch.object(TarballArtifact, "download", download), + patch.object(TarballArtifact, "extract", extract), + patch.object(Path, "unlink", fail_download_unlink), + ): + result = await artifact.materialize(ctx) + + assert result == [ctx.paths.tarball_target_dir] + assert (result[0] / "module.py").read_text() == "VALUE = 1" + assert len(downloaded_paths) == 1 + deferred_path = downloaded_paths[0] + assert cache._deferred_staging_cleanup == {deferred_path} + assert deferred_path.exists() + + assert cache._retry_deferred_staging_cleanup() is True + assert cache._deferred_staging_cleanup == set() + assert not deferred_path.exists() + @pytest.mark.parametrize("artifact_format", ["squashfs", "tarball"]) @pytest.mark.anyio async def test_repeatedly_cancelled_partial_cleanup_rejoins_thread( diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index c2686a7dde..28d214ef38 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -219,6 +219,23 @@ async def _remove_tree_rejoin_on_cancel( ) +def _remove_file_or_defer( + path: Path, + *, + defer_cleanup: Callable[[Path], None], +) -> None: + """Remove one staging file without masking the materialization outcome.""" + try: + path.unlink(missing_ok=True) + except OSError as e: + defer_cleanup(path) + logger.warning( + "Deferred failed registry artifact staging cleanup", + path=str(path), + error=str(e), + ) + + @dataclass(frozen=True, slots=True) class BuiltinArtifact(RegistryArtifact): """Current builtin registry package already installed in the executor image.""" @@ -643,7 +660,10 @@ async def materialize( defer_cleanup=ctx.defer_cleanup, ) finally: - temp_tarball.unlink(missing_ok=True) + _remove_file_or_defer( + temp_tarball, + defer_cleanup=ctx.defer_cleanup, + ) return [target_dir] From 10648c6bfa58a26457386abb086bf247b598fc63 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:09:39 -0400 Subject: [PATCH 075/161] test(executor): exercise failed unmount eviction --- tests/unit/executor/test_registry_artifact_eviction.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_eviction.py b/tests/unit/executor/test_registry_artifact_eviction.py index 80348e0880..4fff9d6542 100644 --- a/tests/unit/executor/test_registry_artifact_eviction.py +++ b/tests/unit/executor/test_registry_artifact_eviction.py @@ -248,6 +248,7 @@ async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): stuck.squashfs_image_path.write_bytes(b"squashfs") stuck.squashfs_mount_dir.mkdir() os.utime(stuck.squashfs_image_path, (100.0, 100.0)) + os.utime(stuck.entry_dir, (100.0, 100.0)) idle = write_image_entry(temp_cache_dir, "idle", size=16, mtime=300.0) mounted = {stuck.squashfs_mount_dir} @@ -265,7 +266,7 @@ async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process, - ), + ) as create_subprocess_exec, patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0), ): @@ -274,6 +275,8 @@ async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): assert stuck.squashfs_image_path.exists() assert stuck.squashfs_mount_dir.exists() assert not idle.exists() + create_subprocess_exec.assert_awaited_once() + process.communicate.assert_awaited_once() @pytest.mark.anyio async def test_eviction_discards_idle_runtime_state(self, temp_cache_dir): From 6af7b079a20f14f4d570a1a82858d3cf2fe4136f Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:12:58 -0400 Subject: [PATCH 076/161] fix(storage): grow unknown download reservations --- tests/unit/test_storage_blob.py | 43 +++++++++++++++++++++++++++++++++ tracecat/storage/blob.py | 25 ++++++++++++------- 2 files changed, 59 insertions(+), 9 deletions(-) diff --git a/tests/unit/test_storage_blob.py b/tests/unit/test_storage_blob.py index 7d551f72de..4b017a6128 100644 --- a/tests/unit/test_storage_blob.py +++ b/tests/unit/test_storage_blob.py @@ -809,6 +809,49 @@ async def ensure_capacity(size_bytes: int) -> None: assert reservations == [7] assert out.read_bytes() == b"payload" + @pytest.mark.anyio + async def test_download_file_to_path_grows_unknown_length_reservation( + self, tmp_path: Path, monkeypatch + ): + """Unknown-length downloads reserve each chunk above current usage.""" + + class DummyStream: + async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 + yield b"abc" + yield b"defg" + + @asynccontextmanager + async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 + yield DummyStream(), None + + monkeypatch.setattr( + "tracecat.storage.blob.open_download_stream", + _fake_open_download_stream, + ) + + out = tmp_path / "out.bin" + temp_path = tmp_path / "out.bin.part" + reservations: list[int] = [] + partial_sizes: list[int] = [] + + async def ensure_capacity(size_bytes: int) -> None: + partial_size = temp_path.stat().st_size + assert 2 + partial_size + size_bytes <= 10 + reservations.append(size_bytes) + partial_sizes.append(partial_size) + + await download_file_to_path( + key="k", + bucket="b", + output_path=out, + max_bytes=10, + ensure_capacity=ensure_capacity, + ) + + assert reservations == [3, 4] + assert partial_sizes == [0, 3] + assert out.read_bytes() == b"abcdefg" + @pytest.mark.anyio async def test_download_file_to_path_never_exceeds_reserved_length( self, tmp_path: Path, monkeypatch diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index 24167410d6..9e44c8595b 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -774,7 +774,8 @@ async def download_file_to_path( expected_sha256: Optional integrity check; raise if computed SHA-256 differs. ensure_capacity: Optional callback invoked before the first disk write with the maximum number of bytes the download may occupy. When the server - omits ContentLength, max_bytes is required to provide that bound. + omits ContentLength, max_bytes is required and capacity is checked + incrementally before each chunk is written. Returns: Total bytes written. @@ -824,6 +825,7 @@ async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: ) download_limit = max_bytes + grow_reservation_by_chunk = False if ensure_capacity is not None: reserved_bytes = content_length if reserved_bytes is None: @@ -832,15 +834,18 @@ async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: "Cannot reserve disk capacity for a download without " f"ContentLength or max_bytes: {bucket}/{key}" ) - reserved_bytes = max_bytes - await ensure_capacity(reserved_bytes) - download_limit = ( - reserved_bytes - if download_limit is None - else min(download_limit, reserved_bytes) - ) + grow_reservation_by_chunk = True + else: + await ensure_capacity(reserved_bytes) + download_limit = ( + reserved_bytes + if download_limit is None + else min(download_limit, reserved_bytes) + ) - async with aiofiles.open(temp_path, "wb") as f: + # Unbuffered writes keep the partial file's allocated size visible + # to incremental capacity scans between unknown-length chunks. + async with aiofiles.open(temp_path, "wb", buffering=0) as f: async for chunk in stream.iter_chunks(chunk_size=chunk_size): if not chunk: continue @@ -851,6 +856,8 @@ async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: f"bytes_written={bytes_written} exceeds " f"max_bytes={download_limit}" ) + if grow_reservation_by_chunk and ensure_capacity is not None: + await ensure_capacity(len(chunk)) if hasher is not None: hasher.update(chunk) await write_chunk_rejoin_on_cancel(f, chunk) From dac9ec90bbe0a43ab11fe4b9d7098972b9e5ccb1 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:19:23 -0400 Subject: [PATCH 077/161] fix(sandbox): preserve cancellation on cleanup failure --- tests/unit/test_sandbox_utils.py | 33 ++++++++++++++++++++++++++++++++ tracecat/sandbox/utils.py | 17 ++++++++++++++-- 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_sandbox_utils.py b/tests/unit/test_sandbox_utils.py index 676ceb180c..d2f1fbdc05 100644 --- a/tests/unit/test_sandbox_utils.py +++ b/tests/unit/test_sandbox_utils.py @@ -68,3 +68,36 @@ async def blocking_termination( assert termination_finished.is_set() assert fake_process.communicate_finished.is_set() + + +@pytest.mark.anyio +async def test_cleanup_failure_preserves_caller_cancellation() -> None: + """A failed process cleanup remains context for the caller cancellation.""" + fake_process = _BlockingProcess() + process = cast(asyncio.subprocess.Process, fake_process) + termination_started = asyncio.Event() + finish_termination = asyncio.Event() + + async def failing_termination( + requested_process: asyncio.subprocess.Process, + ) -> None: + assert requested_process is process + termination_started.set() + await finish_termination.wait() + raise RuntimeError("process cleanup failed") + + with patch( + "tracecat.sandbox.utils.terminate_process_group", + side_effect=failing_termination, + ): + communication = asyncio.create_task(communicate_process_group(process)) + await fake_process.communicate_started.wait() + communication.cancel() + await termination_started.wait() + finish_termination.set() + + with pytest.raises(asyncio.CancelledError) as raised: + await communication + + assert isinstance(raised.value.__cause__, RuntimeError) + assert fake_process.communicate_finished.is_set() diff --git a/tracecat/sandbox/utils.py b/tracecat/sandbox/utils.py index 8fdc6b39ae..680592fae5 100644 --- a/tracecat/sandbox/utils.py +++ b/tracecat/sandbox/utils.py @@ -71,8 +71,12 @@ async def _rejoin_cleanup_through_cancellation( raise pending_cancellation = e - if not cleanup_task.cancelled(): + try: cleanup_task.result() + except BaseException as cleanup_error: + if pending_cancellation is not None: + raise pending_cancellation from cleanup_error + raise if pending_cancellation is not None: raise pending_cancellation @@ -95,6 +99,7 @@ async def communicate_process_group( """ communicate_task = asyncio.create_task(process.communicate(input=input)) termination_task: asyncio.Task[None] | None = None + operation_error: BaseException | None = None try: async with asyncio.timeout(timeout): while process.returncode is None: @@ -102,6 +107,9 @@ async def communicate_process_group( termination_task = asyncio.create_task(terminate_process_group(process)) await asyncio.shield(termination_task) stdout, stderr = await communicate_task + except BaseException as e: + operation_error = e + raise finally: cleanup_task = asyncio.create_task( _finish_process_group_cleanup( @@ -110,7 +118,12 @@ async def communicate_process_group( termination_task, ) ) - await _rejoin_cleanup_through_cancellation(cleanup_task) + try: + await _rejoin_cleanup_through_cancellation(cleanup_task) + except BaseException as cleanup_error: + if operation_error is not None: + raise operation_error from cleanup_error + raise if stdout is None or stderr is None: raise RuntimeError("Captured stdout and stderr are required") From 220a7229b0d0827cc53c99ea8c108786cb2e140a Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:24:00 -0400 Subject: [PATCH 078/161] fix(executor): contain artifact utility processes --- .../registry_artifact_test_helpers.py | 13 +++++-- .../test_registry_artifact_eviction.py | 6 +++ .../test_registry_artifact_materialization.py | 18 ++++++--- .../registry_artifact_materialization.py | 38 ++----------------- .../executor/registry_artifact_storage.py | 4 +- 5 files changed, 35 insertions(+), 44 deletions(-) diff --git a/tests/unit/executor/registry_artifact_test_helpers.py b/tests/unit/executor/registry_artifact_test_helpers.py index 670521cf83..b41c94b117 100644 --- a/tests/unit/executor/registry_artifact_test_helpers.py +++ b/tests/unit/executor/registry_artifact_test_helpers.py @@ -138,8 +138,12 @@ def __init__(self, *, block_wait: bool = False) -> None: self.returncode: int | None = None self._block_wait = block_wait - async def communicate(self) -> tuple[bytes, bytes]: + async def communicate( + self, + input: bytes | None = None, # noqa: A002 + ) -> tuple[bytes, bytes]: """Block until the task awaiting subprocess completion is cancelled.""" + del input self.communicate_started.set() await asyncio.Event().wait() return b"", b"" @@ -176,9 +180,12 @@ def pid(self) -> int: """Return the wrapped subprocess process-group identifier.""" return self.process.pid - async def communicate(self) -> tuple[bytes, bytes]: + async def communicate( + self, + input: bytes | None = None, # noqa: A002 + ) -> tuple[bytes, bytes]: """Wait for the wrapped subprocess and collect its output.""" - return await self.process.communicate() + return await self.process.communicate(input=input) def kill(self) -> None: """Kill the wrapped subprocess and record the signal.""" diff --git a/tests/unit/executor/test_registry_artifact_eviction.py b/tests/unit/executor/test_registry_artifact_eviction.py index 4fff9d6542..c51ebc2d01 100644 --- a/tests/unit/executor/test_registry_artifact_eviction.py +++ b/tests/unit/executor/test_registry_artifact_eviction.py @@ -44,6 +44,7 @@ async def test_eviction_unmounts_before_deleting_the_image(self, temp_cache_dir) image_present_at_umount: list[bool] = [] process = AsyncMock() + process.pid = 999_999_999 process.communicate.return_value = (b"", b"") process.returncode = 0 @@ -63,6 +64,7 @@ async def mock_umount(*args, **kwargs): "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", side_effect=mock_umount, ) as create_subprocess_exec, + patch("tracecat.sandbox.utils.os.killpg") as kill_group, ): evicted = await cache._evict_entry("mounted") @@ -75,6 +77,7 @@ async def mock_umount(*args, **kwargs): "/sbin/umount", str(paths.squashfs_mount_dir), ) + kill_group.assert_called_once_with(process.pid, signal.SIGKILL) @pytest.mark.anyio async def test_repeatedly_cancelled_unmount_reaps_before_releasing_key_lock( @@ -253,6 +256,7 @@ async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): mounted = {stuck.squashfs_mount_dir} process = AsyncMock() + process.pid = 999_999_999 process.communicate.return_value = (b"", b"target is busy") process.returncode = 32 @@ -267,6 +271,7 @@ async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): new_callable=AsyncMock, return_value=process, ) as create_subprocess_exec, + patch("tracecat.sandbox.utils.os.killpg") as kill_group, patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0), ): @@ -277,6 +282,7 @@ async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): assert not idle.exists() create_subprocess_exec.assert_awaited_once() process.communicate.assert_awaited_once() + kill_group.assert_called_once_with(process.pid, signal.SIGKILL) @pytest.mark.anyio async def test_eviction_discards_idle_runtime_state(self, temp_cache_dir): diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 9c366547a7..3e7224e043 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -230,14 +230,21 @@ async def test_mount_squashfs_uses_hardened_read_only_options( image_path.write_bytes(b"squashfs") target_dir.mkdir() process = AsyncMock() + process.pid = 1234 process.communicate.return_value = (b"", b"") process.returncode = 0 - with patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - new_callable=AsyncMock, - return_value=process, - ) as create_subprocess_exec: + with ( + patch( + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ) as create_subprocess_exec, + patch( + "tracecat.sandbox.utils.terminate_process_group", + new_callable=AsyncMock, + ) as terminate_process_group, + ): await artifact.mount(ctx, image_path) create_subprocess_exec.assert_awaited_once_with( @@ -252,6 +259,7 @@ async def test_mount_squashfs_uses_hardened_read_only_options( stderr=asyncio.subprocess.PIPE, start_new_session=True, ) + terminate_process_group.assert_awaited_once_with(process) @pytest.mark.anyio async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 28d214ef38..c1a9afa18f 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -23,7 +23,7 @@ from tracecat.logger import logger from tracecat.registry.artifact_keys import parse_s3_uri from tracecat.registry.constants import DEFAULT_REGISTRY_ORIGIN -from tracecat.sandbox.utils import terminate_process_group +from tracecat.sandbox.utils import communicate_process_group from tracecat.storage import blob __all__ = [ @@ -160,36 +160,6 @@ async def _run_blocking_rejoin_on_cancel[T](operation: Callable[[], T]) -> T: raise -async def _kill_and_reap_subprocess(process: asyncio.subprocess.Process) -> None: - """Kill a subprocess group and wait until its leader is reaped.""" - await terminate_process_group(process) - - -async def _communicate_rejoin_on_cancel( - process: asyncio.subprocess.Process, -) -> tuple[bytes, bytes]: - """Communicate while keeping cancellation from abandoning child cleanup.""" - try: - stdout, stderr = await process.communicate() - except asyncio.CancelledError: - reaper = asyncio.ensure_future(_kill_and_reap_subprocess(process)) - while not reaper.done(): - try: - await asyncio.shield(reaper) - except asyncio.CancelledError: - continue - except Exception: - break - if not reaper.cancelled(): - with contextlib.suppress(Exception): - reaper.result() - raise - - if stdout is None or stderr is None: - raise RuntimeError("Captured subprocess output is required") - return stdout, stderr - - async def _remove_tree_rejoin_on_cancel( path: Path, *, @@ -511,7 +481,7 @@ async def _mount_image(self, image_path: Path, target_dir: Path) -> None: stderr=asyncio.subprocess.PIPE, start_new_session=True, ) - stdout, stderr = await _communicate_rejoin_on_cancel(proc) + stdout, stderr = await communicate_process_group(proc) if proc.returncode == 0 or target_dir.is_mount(): return @@ -540,7 +510,7 @@ async def _extract_image(self, image_path: Path, target_dir: Path) -> None: stderr=asyncio.subprocess.PIPE, start_new_session=True, ) - stdout, stderr = await _communicate_rejoin_on_cancel(proc) + stdout, stderr = await communicate_process_group(proc) if proc.returncode == 0: return @@ -567,7 +537,7 @@ async def _squashfs_extracted_size( stderr=asyncio.subprocess.PIPE, start_new_session=True, ) - stdout, stderr = await _communicate_rejoin_on_cancel(proc) + stdout, stderr = await communicate_process_group(proc) if proc.returncode != 0: output = (stderr or stdout).decode(errors="replace").strip() diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 97373540ca..dd79282712 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -19,9 +19,9 @@ from tracecat.executor.registry_artifact_materialization import ( RegistryArtifactAdmission, _allocated_size_bound, - _communicate_rejoin_on_cancel, ) from tracecat.logger import logger +from tracecat.sandbox.utils import communicate_process_group class RegistryArtifactCacheCapacityError(RuntimeError): @@ -634,7 +634,7 @@ async def _unmount(self, mount_dir: Path) -> bool: stderr=asyncio.subprocess.PIPE, start_new_session=True, ) - stdout, stderr = await _communicate_rejoin_on_cancel(proc) + stdout, stderr = await communicate_process_group(proc) if proc.returncode == 0 or not mount_dir.is_mount(): return True From 65c5652390d01d8ac75441c2af5b0196e9e04e27 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:25:02 -0400 Subject: [PATCH 079/161] refactor(executor): reuse blocking rejoin helper --- .../executor/registry_artifact_storage.py | 21 ++----------------- 1 file changed, 2 insertions(+), 19 deletions(-) diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index dd79282712..39927a6c8a 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio -import contextlib import os import shutil import stat @@ -19,6 +18,7 @@ from tracecat.executor.registry_artifact_materialization import ( RegistryArtifactAdmission, _allocated_size_bound, + _run_blocking_rejoin_on_cancel, ) from tracecat.logger import logger from tracecat.sandbox.utils import communicate_process_group @@ -145,24 +145,7 @@ def _delete_cache_path(path: Path) -> bool: async def _delete_cache_path_off_loop(path: Path) -> bool: """Delete one path without abandoning its worker thread on cancellation.""" - deletion = asyncio.ensure_future(asyncio.to_thread(_delete_cache_path, path)) - try: - return await asyncio.shield(deletion) - except asyncio.CancelledError: - # A worker thread cannot be killed. Rejoin it so no live deletion can - # race a later trash-directory scan. Repeated cancellation can interrupt - # shield without stopping the thread, so keep waiting for termination. - while not deletion.done(): - try: - await asyncio.shield(deletion) - except asyncio.CancelledError: - continue - except Exception: - break - if not deletion.cancelled(): - with contextlib.suppress(Exception): - deletion.result() - raise + return await _run_blocking_rejoin_on_cancel(lambda: _delete_cache_path(path)) def _unique_work_path(root: Path, cache_key: str) -> Path: From 353a916cc5c67d5e8e077e490b5d05823c4aa443 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:27:06 -0400 Subject: [PATCH 080/161] fix(executor): rejoin cache cleanup workers --- .../executor/test_registry_artifact_budget.py | 59 +++++++++++++++++++ .../registry_artifact_materialization.py | 24 ++++---- .../executor/registry_artifact_storage.py | 7 ++- 3 files changed, 77 insertions(+), 13 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 9867f233a3..60eee04606 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -763,6 +763,65 @@ async def mock_enforce_cache_budget( assert cache._budget_dirty is True + @pytest.mark.parametrize("operation", ["budget", "admission"]) + @pytest.mark.anyio + async def test_cancelled_cleanup_rejoins_workers_before_releasing_locks( + self, + temp_cache_dir: Path, + operation: Literal["budget", "admission"], + ) -> None: + """Cache-wide locks outlive repeatedly cancelled cleanup workers.""" + cache = RegistryArtifactCache(temp_cache_dir) + cleanup_started = threading.Event() + cleanup_release = threading.Event() + cleanup_finished = threading.Event() + + def blocking_clear(work_dir: Path) -> bool: + assert work_dir == cache.trash_dir + cleanup_started.set() + cleanup_release.wait(timeout=5) + cleanup_finished.set() + return True + + async def run_operation() -> object: + if operation == "budget": + return await cache._enforce_cache_budget() + await cache._ensure_cache_capacity( + additional_bytes=0, + protected_key="pending", + max_bytes=1, + ) + return None + + with ( + patch.object(cache, "_clear_work_dir", side_effect=blocking_clear), + patch.object( + cache, + "_retry_deferred_staging_cleanup", + return_value=True, + ), + ): + running = asyncio.create_task(run_operation()) + try: + assert await asyncio.to_thread(cleanup_started.wait, 1) + running.cancel() + await asyncio.sleep(0) + running.cancel() + await asyncio.sleep(0) + + assert not running.done() + assert cache._budget_lock.locked() + assert cache._admission_lock.locked() is (operation == "budget") + finally: + cleanup_release.set() + + with pytest.raises(asyncio.CancelledError): + await running + + assert cleanup_finished.is_set() + assert not cache._budget_lock.locked() + assert not cache._admission_lock.locked() + @pytest.mark.parametrize( "oldest_has_tarball", [False, True], diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index c1a9afa18f..edfdb25a26 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -138,28 +138,30 @@ def _temp_path( return ctx.staging_dir / f"{self.cache_key}.{os.getpid()}.{unique_id}{suffix}" -async def _run_blocking_rejoin_on_cancel[T](operation: Callable[[], T]) -> T: - """Run blocking work without abandoning its thread on cancellation.""" - worker = asyncio.ensure_future(asyncio.to_thread(operation)) +async def _rejoin_future_on_cancel[T](future: asyncio.Future[T]) -> T: + """Shield a future and rejoin it through repeated caller cancellation.""" try: - return await asyncio.shield(worker) + return await asyncio.shield(future) except asyncio.CancelledError: - # A thread cannot be killed. Rejoin it before callers remove its input - # or output paths. Repeated cancellation can interrupt shield without - # stopping the thread, so keep waiting for a terminal state. - while not worker.done(): + while not future.done(): try: - await asyncio.shield(worker) + await asyncio.shield(future) except asyncio.CancelledError: continue except Exception: break - if not worker.cancelled(): + if not future.cancelled(): with contextlib.suppress(Exception): - worker.result() + future.result() raise +async def _run_blocking_rejoin_on_cancel[T](operation: Callable[[], T]) -> T: + """Run blocking work without abandoning its thread on cancellation.""" + worker = asyncio.ensure_future(asyncio.to_thread(operation)) + return await _rejoin_future_on_cancel(worker) + + async def _remove_tree_rejoin_on_cancel( path: Path, *, diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 39927a6c8a..23ceee15bd 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -18,6 +18,7 @@ from tracecat.executor.registry_artifact_materialization import ( RegistryArtifactAdmission, _allocated_size_bound, + _rejoin_future_on_cancel, _run_blocking_rejoin_on_cancel, ) from tracecat.logger import logger @@ -340,10 +341,11 @@ async def _enforce_cache_budget_locked( protected_key: str | None, ) -> bool: """Enforce entry and byte limits while both cache-wide locks are held.""" - trash_clean, startup_clean = await asyncio.gather( + cleanup = asyncio.gather( asyncio.to_thread(self._clear_work_dir, self.trash_dir), asyncio.to_thread(self._retry_deferred_staging_cleanup), ) + trash_clean, startup_clean = await _rejoin_future_on_cancel(cleanup) cleanup_complete = trash_clean and startup_clean if not cleanup_complete: return False @@ -404,10 +406,11 @@ async def _ensure_cache_capacity( raise ValueError("additional_bytes must be non-negative") async with self._budget_lock: - trash_clean, startup_clean = await asyncio.gather( + cleanup = asyncio.gather( asyncio.to_thread(self._clear_work_dir, self.trash_dir), asyncio.to_thread(self._retry_deferred_staging_cleanup), ) + trash_clean, startup_clean = await _rejoin_future_on_cancel(cleanup) entries = await asyncio.to_thread(self._scan_cache_entries) staging_bytes, trash_bytes = await asyncio.gather( asyncio.to_thread(_directory_footprint, self.staging_dir), From 176f3ef099e0cb09be3bfa112ba9dc2052b91e03 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:30:07 -0400 Subject: [PATCH 081/161] test(executor): model utility process groups --- tests/unit/executor/test_registry_artifact_leases.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unit/executor/test_registry_artifact_leases.py b/tests/unit/executor/test_registry_artifact_leases.py index 957a2985a0..f84c7efe50 100644 --- a/tests/unit/executor/test_registry_artifact_leases.py +++ b/tests/unit/executor/test_registry_artifact_leases.py @@ -4,6 +4,7 @@ import asyncio import os +import signal from pathlib import Path from unittest.mock import AsyncMock, patch @@ -808,6 +809,7 @@ def observed_runtime_for(requested_key: str): return runtime umount_process = AsyncMock() + umount_process.pid = 999_999_999 umount_process.communicate.return_value = (b"", b"") umount_process.returncode = 0 @@ -844,6 +846,7 @@ async def take_lease() -> None: "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", side_effect=mock_umount, ), + patch("tracecat.sandbox.utils.os.killpg") as kill_group, patch.object(cache, "_runtime_for", side_effect=observed_runtime_for), patch.object(SquashfsArtifact, "mount", mock_mount), # This test targets the per-key eviction/lease handoff. Keep the @@ -870,6 +873,8 @@ async def take_lease() -> None: assert leased_paths == [paths.squashfs_mount_dir] assert leased_path_exists == [True] assert (paths.squashfs_mount_dir / "module.py").read_text() == "VALUE = 1" + assert kill_group.call_count == 2 + kill_group.assert_called_with(umount_process.pid, signal.SIGKILL) @pytest.mark.anyio async def test_builtin_artifact_is_exempt_from_cache_accounting( From 302299c4ed1cf16bc97995cb113d25af7e03adab Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:34:46 -0400 Subject: [PATCH 082/161] fix(executor): reserve implicit tar directories --- .../test_registry_artifact_materialization.py | 13 +++++++++++++ .../executor/registry_artifact_materialization.py | 13 ++++++++++++- 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 3e7224e043..58a121051b 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -165,6 +165,19 @@ def test_tarball_size_bounds_each_member_allocation( assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 12_288 + def test_tarball_size_includes_implicit_parent_directories( + self, + temp_cache_dir: Path, + ) -> None: + """Extraction reserves directories omitted from the tar manifest.""" + tarball_path = temp_cache_dir / "implicit-directories.tar.gz" + with tarfile.open(tarball_path, "w:gz") as tar: + member = tarfile.TarInfo("one/two/three/module.py") + member.size = 0 + tar.addfile(member, io.BytesIO()) + + assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 16_384 + def test_squashfs_listing_size_rejects_unparseable_files(self) -> None: with pytest.raises(ValueError, match="Could not parse SquashFS listing"): _squashfs_listing_size(b"-rw-r--r-- malformed") diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index edfdb25a26..0b6ff43207 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -14,7 +14,7 @@ from collections.abc import Awaitable, Callable from dataclasses import dataclass from enum import StrEnum -from pathlib import Path +from pathlib import Path, PurePosixPath import httpx import tracecat_registry @@ -814,6 +814,8 @@ def _tarball_extracted_size( ) -> int: """Return a conservative allocated size bound for all tarball members.""" total_bytes = 0 + required_parent_dirs: set[PurePosixPath] = set() + explicit_dirs: set[PurePosixPath] = set() with tarfile.open(tarball_path, "r:gz") as tar: for member in tar: if member.size < 0: @@ -824,6 +826,15 @@ def _tarball_extracted_size( member.size, allocation_unit=allocation_unit, ) + member_path = PurePosixPath(member.name) + if member.isdir(): + explicit_dirs.add(member_path) + for parent in member_path.parents: + if parent == PurePosixPath("."): + break + required_parent_dirs.add(parent) + implicit_parent_dirs = required_parent_dirs - explicit_dirs + total_bytes += len(implicit_parent_dirs) * allocation_unit return total_bytes From 5b12ecb6206fdaf16e0d2f7eb7f5cd08c8c6d2c2 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:36:04 -0400 Subject: [PATCH 083/161] fix(executor): charge every cached inode --- tests/unit/executor/conftest.py | 7 ++- .../executor/test_registry_artifact_budget.py | 9 +++- .../executor/registry_artifact_storage.py | 49 +++++++++++++++---- 3 files changed, 53 insertions(+), 12 deletions(-) diff --git a/tests/unit/executor/conftest.py b/tests/unit/executor/conftest.py index 26cd619eab..0bfce17b4f 100644 --- a/tests/unit/executor/conftest.py +++ b/tests/unit/executor/conftest.py @@ -17,7 +17,12 @@ def logical_cache_sizes(monkeypatch: pytest.MonkeyPatch) -> None: """Keep synthetic byte budgets independent from host filesystem block sizes.""" - def logical_stat_size(file_stat: os.stat_result) -> int: + def logical_stat_size( + file_stat: os.stat_result, + *, + allocation_unit: int, + ) -> int: + del allocation_unit if stat.S_ISDIR(file_stat.st_mode): return 0 return file_stat.st_size diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 60eee04606..5587d4d726 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -44,7 +44,14 @@ def test_allocated_stat_size_uses_filesystem_blocks(self) -> None: file_stat.st_blocks = 7 file_stat.st_size = 1 - assert _allocated_stat_size(file_stat) == 7 * 512 + assert _allocated_stat_size(file_stat, allocation_unit=512) == 7 * 512 + + def test_allocated_stat_size_charges_zero_block_inode(self) -> None: + file_stat = MagicMock(spec=os.stat_result) + file_stat.st_blocks = 0 + file_stat.st_size = 0 + + assert _allocated_stat_size(file_stat, allocation_unit=4096) == 4096 def test_directory_footprint_includes_directory_inodes( self, diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 23ceee15bd..9cb682ff63 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -62,12 +62,20 @@ class RegistryArtifactCacheEntry: last_used: float -def _allocated_stat_size(file_stat: os.stat_result) -> int: - """Return allocated bytes for one inode, falling back to logical size.""" +def _allocated_stat_size( + file_stat: os.stat_result, + *, + allocation_unit: int, +) -> int: + """Return allocated bytes while charging at least one unit per inode.""" + if allocation_unit <= 0: + raise ValueError("allocation_unit must be positive") blocks = getattr(file_stat, "st_blocks", None) if blocks is None: - return file_stat.st_size - return blocks * 512 + allocated_bytes = file_stat.st_size + else: + allocated_bytes = blocks * 512 + return max(allocation_unit, allocated_bytes) def _filesystem_allocation_unit(path: Path) -> int: @@ -86,7 +94,11 @@ def _filesystem_allocation_unit(path: Path) -> int: return filesystem.f_frsize or filesystem.f_bsize or 1 -def _directory_footprint(directory: Path) -> int: +def _directory_footprint( + directory: Path, + *, + allocation_unit: int | None = None, +) -> int: """Return the allocated footprint of a cache directory tree. Args: @@ -100,18 +112,25 @@ def _directory_footprint(directory: Path) -> int: def raise_walk_error(error: OSError) -> None: raise error + if allocation_unit is None: + allocation_unit = _filesystem_allocation_unit(directory) + total_bytes = 0 try: walker = os.walk(directory, onerror=raise_walk_error) for root, dirs, files in walker: try: - total_bytes += _allocated_stat_size(os.lstat(root)) + total_bytes += _allocated_stat_size( + os.lstat(root), + allocation_unit=allocation_unit, + ) except FileNotFoundError: continue for file_name in files: try: total_bytes += _allocated_stat_size( - os.lstat(os.path.join(root, file_name)) + os.lstat(os.path.join(root, file_name)), + allocation_unit=allocation_unit, ) except FileNotFoundError: continue @@ -121,7 +140,10 @@ def raise_walk_error(error: OSError) -> None: except FileNotFoundError: continue if stat.S_ISLNK(directory_stat.st_mode): - total_bytes += _allocated_stat_size(directory_stat) + total_bytes += _allocated_stat_size( + directory_stat, + allocation_unit=allocation_unit, + ) except FileNotFoundError: return 0 return total_bytes @@ -659,6 +681,7 @@ def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: measured individually so a concurrent eviction cannot fail the scan. """ paths = self._paths_for(cache_key) + allocation_unit = _filesystem_allocation_unit(self.cache_dir) size_bytes = 0 for path in ( @@ -667,12 +690,18 @@ def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: paths.squashfs_mount_dir, ): try: - size_bytes += _allocated_stat_size(path.lstat()) + size_bytes += _allocated_stat_size( + path.lstat(), + allocation_unit=allocation_unit, + ) except FileNotFoundError: continue for directory in (paths.squashfs_extract_dir, paths.tarball_target_dir): - size_bytes += _directory_footprint(directory) + size_bytes += _directory_footprint( + directory, + allocation_unit=allocation_unit, + ) try: last_used = paths.entry_dir.stat().st_mtime From 8049a53ddf77ca69cd5de1ae3accb95ee22397b7 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:44:35 -0400 Subject: [PATCH 084/161] fix(executor): reserve tar extraction roots --- .../test_registry_artifact_materialization.py | 29 +++++++++++++++++-- .../registry_artifact_materialization.py | 10 +++++-- 2 files changed, 35 insertions(+), 4 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 58a121051b..5e02af5786 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -163,7 +163,32 @@ def test_tarball_size_bounds_each_member_allocation( member.size = 1 tar.addfile(member, io.BytesIO(b"x")) - assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 12_288 + assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 16_384 + + def test_tarball_size_includes_extraction_root( + self, + temp_cache_dir: Path, + ) -> None: + """Extraction reserves its root even when the manifest omits it.""" + tarball_path = temp_cache_dir / "implicit-root.tar.gz" + with tarfile.open(tarball_path, "w:gz") as tar: + member = tarfile.TarInfo("module.py") + member.size = 0 + tar.addfile(member, io.BytesIO()) + + assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 8192 + + def test_tarball_size_does_not_duplicate_explicit_root( + self, + temp_cache_dir: Path, + ) -> None: + tarball_path = temp_cache_dir / "explicit-root.tar.gz" + with tarfile.open(tarball_path, "w:gz") as tar: + root = tarfile.TarInfo(".") + root.type = tarfile.DIRTYPE + tar.addfile(root) + + assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 4096 def test_tarball_size_includes_implicit_parent_directories( self, @@ -176,7 +201,7 @@ def test_tarball_size_includes_implicit_parent_directories( member.size = 0 tar.addfile(member, io.BytesIO()) - assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 16_384 + assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 20_480 def test_squashfs_listing_size_rejects_unparseable_files(self) -> None: with pytest.raises(ValueError, match="Could not parse SquashFS listing"): diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 0b6ff43207..1944882204 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -812,8 +812,10 @@ def _tarball_extracted_size( *, allocation_unit: int = 1, ) -> int: - """Return a conservative allocated size bound for all tarball members.""" + """Return a conservative allocated size bound for a tarball extraction.""" total_bytes = 0 + root_path = PurePosixPath(".") + has_explicit_root_directory = False required_parent_dirs: set[PurePosixPath] = set() explicit_dirs: set[PurePosixPath] = set() with tarfile.open(tarball_path, "r:gz") as tar: @@ -829,10 +831,14 @@ def _tarball_extracted_size( member_path = PurePosixPath(member.name) if member.isdir(): explicit_dirs.add(member_path) + has_explicit_root_directory |= member_path == root_path for parent in member_path.parents: - if parent == PurePosixPath("."): + if parent == root_path: break required_parent_dirs.add(parent) + # Extraction creates a target root even when the tar manifest omits it. + if not has_explicit_root_directory: + total_bytes += _allocated_size_bound(0, allocation_unit=allocation_unit) implicit_parent_dirs = required_parent_dirs - explicit_dirs total_bytes += len(implicit_parent_dirs) * allocation_unit return total_bytes From 3cb6a05336267a4a720eac63560bc3a79c80ade6 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:54:58 -0400 Subject: [PATCH 085/161] fix(executor): deduplicate hard-linked cache inodes --- .../executor/test_registry_artifact_budget.py | 16 ++++++++++ .../executor/registry_artifact_storage.py | 31 +++++++++++-------- 2 files changed, 34 insertions(+), 13 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 5587d4d726..75a3e63034 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -69,6 +69,22 @@ def test_directory_footprint_includes_directory_inodes( assert allocated_size.call_count == 3 + def test_directory_footprint_counts_hard_linked_inode_once( + self, + temp_cache_dir: Path, + ) -> None: + payload = temp_cache_dir / "payload" + payload.write_text("x") + os.link(payload, temp_cache_dir / "payload-link") + + with patch( + "tracecat.executor.registry_artifact_storage._allocated_stat_size", + return_value=4096, + ) as allocated_size: + assert _directory_footprint(temp_cache_dir) == 2 * 4096 + + assert allocated_size.call_count == 2 + @pytest.mark.anyio async def test_admission_rounds_download_reservation_to_allocation_unit( self, diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 9cb682ff63..4e1e945d1c 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -105,8 +105,8 @@ def _directory_footprint( directory: Cache directory to measure. Returns: - Total allocated bytes of contained inodes, or zero when the directory is - missing. + Total allocated bytes of unique contained inodes, or zero when the + directory is missing. """ def raise_walk_error(error: OSError) -> None: @@ -116,21 +116,29 @@ def raise_walk_error(error: OSError) -> None: allocation_unit = _filesystem_allocation_unit(directory) total_bytes = 0 + seen_inodes: set[tuple[int, int]] = set() + + def allocated_inode_size(file_stat: os.stat_result) -> int: + inode_key = (file_stat.st_dev, file_stat.st_ino) + if inode_key in seen_inodes: + return 0 + seen_inodes.add(inode_key) + return _allocated_stat_size( + file_stat, + allocation_unit=allocation_unit, + ) + try: walker = os.walk(directory, onerror=raise_walk_error) for root, dirs, files in walker: try: - total_bytes += _allocated_stat_size( - os.lstat(root), - allocation_unit=allocation_unit, - ) + total_bytes += allocated_inode_size(os.lstat(root)) except FileNotFoundError: continue for file_name in files: try: - total_bytes += _allocated_stat_size( - os.lstat(os.path.join(root, file_name)), - allocation_unit=allocation_unit, + total_bytes += allocated_inode_size( + os.lstat(os.path.join(root, file_name)) ) except FileNotFoundError: continue @@ -140,10 +148,7 @@ def raise_walk_error(error: OSError) -> None: except FileNotFoundError: continue if stat.S_ISLNK(directory_stat.st_mode): - total_bytes += _allocated_stat_size( - directory_stat, - allocation_unit=allocation_unit, - ) + total_bytes += allocated_inode_size(directory_stat) except FileNotFoundError: return 0 return total_bytes From 9e8740b42a9524931bbfaecac031a6cd27912157 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:55:57 -0400 Subject: [PATCH 086/161] test(executor): account for tar extraction roots --- tests/unit/executor/test_registry_artifact_budget.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 75a3e63034..49ddfb6ec6 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -180,7 +180,7 @@ async def test_cold_download_reserves_space_before_writing( artifact_uri = "s3://bucket/new.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) payload = tarball_payload(size=32) - max_bytes = len(payload) + 32 + max_bytes = len(payload) + 33 capacity_checked = False async def download_file_to_path( @@ -193,7 +193,7 @@ async def download_file_to_path( ) -> int: del key, bucket nonlocal capacity_checked - assert max_bytes == len(payload) + 32 + assert max_bytes == len(payload) + 33 await ensure_capacity(len(payload)) capacity_checked = True assert not idle.exists() @@ -261,7 +261,7 @@ async def download_file_to_path( async with cache.lease([artifact_uri]): pass - assert raised.value.additional_bytes == 4096 + assert raised.value.additional_bytes == 4097 assert raised.value.max_bytes == max_bytes extract.assert_not_awaited() assert not cache._paths_for(cache_key).entry_dir.exists() @@ -276,7 +276,7 @@ async def test_failed_squashfs_bytes_do_not_block_tarball_fallback( artifact_uri = "s3://bucket/path/site-packages.squashfs" cache_key = compute_registry_artifact_cache_key(artifact_uri) payload = tarball_payload(size=32) - max_bytes = len(payload) + 32 + max_bytes = len(payload) + 33 async def fail_after_squashfs_download( self: SquashfsArtifact, From 8d2cf3b71870f877f1b5adfd1434944b770f87fa Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 19:05:24 -0400 Subject: [PATCH 087/161] fix(executor): bypass sweeps for cache-free leases --- .../test_registry_artifact_startup.py | 47 +++++++++++++++++-- tracecat/executor/registry_artifacts.py | 16 ++++++- 2 files changed, 57 insertions(+), 6 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_startup.py b/tests/unit/executor/test_registry_artifact_startup.py index 657a1d34d4..cd326b8e71 100644 --- a/tests/unit/executor/test_registry_artifact_startup.py +++ b/tests/unit/executor/test_registry_artifact_startup.py @@ -6,13 +6,14 @@ import os import threading from pathlib import Path -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest from tracecat.executor.registry_artifacts import ( RegistryArtifactCache, _delete_cache_path, + bundled_builtin_registry_uri, ) from .registry_artifact_test_helpers import ( @@ -168,14 +169,52 @@ def blocking_sweep() -> None: assert cache._swept is True @pytest.mark.anyio - async def test_lease_triggers_startup_sweep(self, temp_cache_dir): + async def test_cache_lease_triggers_startup_sweep(self, temp_cache_dir): """Lease admission reclaims startup scratch before yielding paths.""" cache = RegistryArtifactCache(temp_cache_dir) orphaned_dir = cache.staging_dir / "abc123.999999.4321" orphaned_dir.mkdir(parents=True) - async with cache.lease(None): - assert not orphaned_dir.exists() + with patch.object( + cache, + "_lease_artifact", + new_callable=AsyncMock, + return_value=(None, []), + ): + async with cache.lease(["s3://bucket/registry.tar.gz"]): + assert not orphaned_dir.exists() + + @pytest.mark.anyio + @pytest.mark.parametrize( + "artifact_uris", + [None, [bundled_builtin_registry_uri("1.2.3")]], + ) + async def test_cache_free_lease_skips_failed_startup_sweep( + self, + temp_cache_dir, + artifact_uris: list[str] | None, + ): + """Unrelated cache inspection failures cannot block cache-free actions.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch.object( + cache, + "ensure_swept", + new_callable=AsyncMock, + side_effect=OSError("simulated sweep failure"), + ) as ensure_swept, + patch.object( + cache, + "_lease_artifact", + new_callable=AsyncMock, + return_value=(None, [temp_cache_dir / "builtin"]), + ), + ): + async with cache.lease(artifact_uris): + pass + + ensure_swept.assert_not_awaited() @pytest.mark.anyio async def test_failed_ensure_swept_retries(self, temp_cache_dir): diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 983642aa78..c20c1fe669 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -116,13 +116,25 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat Yields: Importable Python paths for the requested artifacts. """ - await self.ensure_swept() - if not artifact_uris: logger.info("No registry artifact URIs provided, using base PYTHONPATH") yield [self._base_pythonpath_dir()] return + if not any(_is_cache_entry_uri(uri) for uri in artifact_uris): + cache_free_paths: list[Path] = [] + for artifact_uri in artifact_uris: + _, artifact_paths = await self._lease_artifact(artifact_uri) + cache_free_paths.extend(artifact_paths) + logger.info( + "Using cache-free registry artifact environments", + count=len(cache_free_paths), + ) + yield cache_free_paths + return + + await self.ensure_swept() + leased_keys: list[str] = [] lease_setup_complete = False try: From 02055ccda607a337df922df88d36b0d4eefc6937 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 19:20:56 -0400 Subject: [PATCH 088/161] fix(storage): defer failed partial download cleanup --- .../executor/test_registry_artifact_budget.py | 6 ++- .../test_registry_artifact_resolution.py | 3 ++ tests/unit/test_storage_blob.py | 40 +++++++++++++++++++ .../registry_artifact_materialization.py | 5 +++ tracecat/storage/blob.py | 5 +++ 5 files changed, 57 insertions(+), 2 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 49ddfb6ec6..0b3cc8a2c5 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -190,8 +190,9 @@ async def download_file_to_path( output_path: Path, max_bytes: int, ensure_capacity: Callable[[int], Awaitable[None]], + defer_cleanup: Callable[[Path], None], ) -> int: - del key, bucket + del key, bucket, defer_cleanup nonlocal capacity_checked assert max_bytes == len(payload) + 33 await ensure_capacity(len(payload)) @@ -236,8 +237,9 @@ async def download_file_to_path( output_path: Path, max_bytes: int, ensure_capacity: Callable[[int], Awaitable[None]], + defer_cleanup: Callable[[Path], None], ) -> int: - del key, bucket + del key, bucket, defer_cleanup assert max_bytes == len(payload) + 256 await ensure_capacity(len(payload)) output_path.write_bytes(payload) diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py index b5e7f6acb6..049e562ba2 100644 --- a/tests/unit/executor/test_registry_artifact_resolution.py +++ b/tests/unit/executor/test_registry_artifact_resolution.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Callable from pathlib import Path from unittest.mock import AsyncMock, patch @@ -93,7 +94,9 @@ async def mock_download_file_to_path( key: str, bucket: str, output_path: Path, + defer_cleanup: Callable[[Path], None], ) -> None: + assert defer_cleanup == ctx.defer_cleanup output_path.write_bytes(b"squashfs") with patch( diff --git a/tests/unit/test_storage_blob.py b/tests/unit/test_storage_blob.py index 4b017a6128..53abd9a8a1 100644 --- a/tests/unit/test_storage_blob.py +++ b/tests/unit/test_storage_blob.py @@ -922,6 +922,46 @@ async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 assert not out.exists() assert not (tmp_path / "out.bin.part").exists() + @pytest.mark.anyio + async def test_download_file_to_path_defers_failed_partial_cleanup( + self, + tmp_path: Path, + monkeypatch, + ): + """A failed partial unlink remains available for runtime cleanup.""" + + class DummyStream: + async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 + yield b"hello" + + @asynccontextmanager + async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 + yield DummyStream(), 5 + + monkeypatch.setattr( + "tracecat.storage.blob.open_download_stream", + _fake_open_download_stream, + ) + + out = tmp_path / "out.bin" + temp_path = tmp_path / "out.bin.part" + deferred_paths: list[Path] = [] + expected_sha256 = hashlib.sha256(b"hello").hexdigest() + with ( + patch.object(Path, "unlink", side_effect=PermissionError("busy")), + pytest.raises(ValueError, match="Integrity check failed"), + ): + await download_file_to_path( + key="k", + bucket="b", + output_path=out, + expected_sha256=expected_sha256 + "bad", + defer_cleanup=deferred_paths.append, + ) + + assert temp_path.exists() + assert deferred_paths == [temp_path] + @pytest.mark.anyio async def test_download_file_to_path_cancellation_cleans_partial( self, tmp_path: Path, monkeypatch diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 1944882204..bf09fb8775 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -330,6 +330,7 @@ async def download( self.uri, temp_image, admission=ctx.admission, + defer_cleanup=ctx.defer_cleanup, ) try: temp_image.rename(image_path) @@ -648,6 +649,7 @@ async def download( self.uri, output_path, admission=ctx.admission, + defer_cleanup=ctx.defer_cleanup, ) async def extract(self, tarball_path: Path, target_dir: Path) -> None: @@ -680,6 +682,7 @@ async def _download_s3_artifact( output_path: Path, *, admission: RegistryArtifactAdmission | None = None, + defer_cleanup: Callable[[Path], None], ) -> None: """Download an S3 registry artifact to a local path.""" bucket, key = parse_s3_uri(artifact_uri) @@ -689,6 +692,7 @@ async def _download_s3_artifact( key=key, bucket=bucket, output_path=output_path, + defer_cleanup=defer_cleanup, ) else: await blob.download_file_to_path( @@ -697,6 +701,7 @@ async def _download_s3_artifact( output_path=output_path, max_bytes=admission.max_bytes, ensure_capacity=admission.ensure_capacity, + defer_cleanup=defer_cleanup, ) except FileNotFoundError as e: request = httpx.Request("GET", artifact_uri) diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index 9e44c8595b..8aa43e32b1 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -759,6 +759,7 @@ async def download_file_to_path( max_bytes: int | None = None, expected_sha256: str | None = None, ensure_capacity: Callable[[int], Awaitable[None]] | None = None, + defer_cleanup: Callable[[Path], None] | None = None, ) -> int: """Stream an S3/MinIO object to a local file. @@ -776,6 +777,8 @@ async def download_file_to_path( the maximum number of bytes the download may occupy. When the server omits ContentLength, max_bytes is required and capacity is checked incrementally before each chunk is written. + defer_cleanup: Optional callback that retains a partial-file path for a + later cleanup retry when immediate deletion fails. Returns: Total bytes written. @@ -879,6 +882,8 @@ async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: "Failed to cleanup partial download", temp_path=str(temp_path), ) + if defer_cleanup is not None: + defer_cleanup(temp_path) raise logger.debug( From 39c634387a4644b628cdd3ddd80a0e3d4f917f48 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 19:46:15 -0400 Subject: [PATCH 089/161] test(executor): isolate mocked process group cleanup --- tests/unit/executor/test_registry_artifact_budget.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 0b3cc8a2c5..e526e45988 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -4,6 +4,7 @@ import asyncio import os +import signal import threading from collections.abc import Awaitable, Callable from pathlib import Path @@ -891,6 +892,7 @@ async def test_final_lease_release_unmounts_and_retains_image(self, temp_cache_d mounted = {paths.squashfs_mount_dir} process = AsyncMock() + process.pid = 999_999_999 process.communicate.return_value = (b"", b"") process.returncode = 0 @@ -910,6 +912,7 @@ async def mock_umount(*args, **kwargs): "create_subprocess_exec", side_effect=mock_umount, ), + patch("tracecat.sandbox.utils.os.killpg") as kill_group, ): async with cache.lease([artifact_uri]) as registry_paths: assert registry_paths == [paths.squashfs_mount_dir] @@ -918,6 +921,7 @@ async def mock_umount(*args, **kwargs): assert paths.squashfs_mount_dir not in mounted assert paths.squashfs_image_path.read_bytes() == b"squashfs" + kill_group.assert_called_once_with(process.pid, signal.SIGKILL) assert paths.squashfs_mount_dir.is_dir() @pytest.mark.anyio From dc4ee0c1c6660977949ce1b7f57ff38d07546185 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 19:51:27 -0400 Subject: [PATCH 090/161] fix(executor): discard incomplete startup cache entries --- .../test_registry_artifact_startup.py | 29 +++++++++++++++ .../executor/registry_artifact_storage.py | 36 ++++++++++++++++++- 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_startup.py b/tests/unit/executor/test_registry_artifact_startup.py index cd326b8e71..f1e1bbbb58 100644 --- a/tests/unit/executor/test_registry_artifact_startup.py +++ b/tests/unit/executor/test_registry_artifact_startup.py @@ -51,6 +51,35 @@ async def test_sweep_uses_entry_root_mtime_for_restart_safe_lru( assert old.is_dir() assert not new.exists() + @pytest.mark.anyio + async def test_sweep_removes_incomplete_shell_before_lru_trimming( + self, + temp_cache_dir: Path, + ) -> None: + """A crashed cold admission cannot displace a reusable cache entry.""" + cache = RegistryArtifactCache(temp_cache_dir) + reusable_image = write_image_entry( + temp_cache_dir, + "reusable", + size=16, + mtime=100.0, + ) + incomplete = cache._paths_for("incomplete") + incomplete.entry_dir.mkdir(parents=True) + incomplete.squashfs_mount_dir.mkdir() + os.utime(incomplete.entry_dir, (200.0, 200.0)) + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + await cache.ensure_swept() + + assert reusable_image.is_file() + assert not incomplete.entry_dir.exists() + assert not any(cache.trash_dir.iterdir()) + assert cache._budget_dirty is False + @pytest.mark.anyio async def test_sweep_tolerates_missing_cache_dir(self, temp_cache_dir): """A cache directory that does not exist yet is a no-op.""" diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 4e1e945d1c..4383e1c1fa 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -798,10 +798,44 @@ def _trim_startup_cache(self) -> bool: """ max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES + entries = self._scan_cache_entries() + for entry in tuple(entries.values()): + paths = self._paths_for(entry.cache_key) + if ( + paths.squashfs_mount_dir.is_mount() + or paths.squashfs_image_path.exists() + or paths.squashfs_extract_dir.exists() + or paths.tarball_target_dir.exists() + ): + continue + + try: + trash_path = _move_entry_to_trash( + paths.entry_dir, + self.trash_dir, + entry.cache_key, + ) + except OSError as e: + logger.warning( + "Failed to retire incomplete registry artifact during startup sweep", + cache_key=entry.cache_key, + entry_dir=str(paths.entry_dir), + error=str(e), + ) + return False + + del entries[entry.cache_key] + if not _delete_cache_path(trash_path): + return False + logger.info( + "Removed incomplete registry artifact during startup sweep", + cache_key=entry.cache_key, + size_bytes=entry.size_bytes, + ) + if max_entries <= 0 and max_bytes <= 0: return True - entries = self._scan_cache_entries() total_bytes = sum(entry.size_bytes for entry in entries.values()) # Mounted entries belong to a live process sharing this cache directory. candidates = sorted( From a00392773824a56dce60497c86b442b658743a7d Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 20:02:19 -0400 Subject: [PATCH 091/161] fix(executor): surface test cache sweep failures --- .../test_test_backend_no_registry_action.py | 63 ++++++++++++++++++- tracecat/executor/backends/test.py | 1 + 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_test_backend_no_registry_action.py b/tests/unit/executor/test_test_backend_no_registry_action.py index f954eb6ae9..01e83a6e23 100644 --- a/tests/unit/executor/test_test_backend_no_registry_action.py +++ b/tests/unit/executor/test_test_backend_no_registry_action.py @@ -13,7 +13,7 @@ import threading import uuid from collections.abc import AsyncIterator -from contextlib import asynccontextmanager +from contextlib import AsyncExitStack, asynccontextmanager from datetime import UTC, datetime from pathlib import Path @@ -271,6 +271,9 @@ class FakeRegistryArtifacts: def __init__(self) -> None: self.active = 0 + async def ensure_swept(self) -> None: + pass + @asynccontextmanager async def lease( self, artifact_uris: list[str] | None = None @@ -334,6 +337,61 @@ async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: finally: await backend.shutdown() + @pytest.mark.anyio + async def test_execute_surfaces_registry_cache_sweep_failure( + self, + test_role: Role, + test_run_action_input: RunActionInput, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Global cache inspection errors are not treated as one bad artifact.""" + artifact_uri = "s3://bucket/registry.tar.gz" + + class FakeRegistryArtifacts: + def __init__(self) -> None: + self.lease_attempted = False + + async def ensure_swept(self) -> None: + raise PermissionError("cannot inspect registry cache") + + @asynccontextmanager + async def lease( + self, artifact_uris: list[str] | None = None + ) -> AsyncIterator[list[Path]]: + del artifact_uris + self.lease_attempted = True + yield [] + + class FakeActionRunner: + def __init__(self) -> None: + self.registry_artifacts = FakeRegistryArtifacts() + + fake_runner = FakeActionRunner() + + async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: + return [artifact_uri] + + backend = TestBackend() + monkeypatch.setattr( + "tracecat.executor.backends.test.config.TRACECAT__LOCAL_REPOSITORY_ENABLED", + False, + ) + monkeypatch.setattr( + "tracecat.executor.backends.test.get_action_runner", + lambda: fake_runner, + ) + monkeypatch.setattr(backend, "_get_artifact_uris", _get_artifact_uris) + + with pytest.raises(PermissionError, match="cannot inspect registry cache"): + async with AsyncExitStack() as leases: + await backend._lease_registry_artifacts( + leases, + test_run_action_input, + test_role, + ) + + assert fake_runner.registry_artifacts.lease_attempted is False + @pytest.mark.anyio async def test_timed_out_sync_udf_keeps_artifact_lease_until_thread_finishes( self, @@ -354,6 +412,9 @@ class FakeRegistryArtifacts: def __init__(self) -> None: self.active = 0 + async def ensure_swept(self) -> None: + pass + @asynccontextmanager async def lease( self, artifact_uris: list[str] | None = None diff --git a/tracecat/executor/backends/test.py b/tracecat/executor/backends/test.py index 696141180c..8b70855721 100644 --- a/tracecat/executor/backends/test.py +++ b/tracecat/executor/backends/test.py @@ -309,6 +309,7 @@ async def _lease_registry_artifacts( return [] registry_artifacts = get_action_runner().registry_artifacts + await registry_artifacts.ensure_swept() extracted_paths: list[str] = [] for artifact_uri in artifact_uris: From e38b6fb84a25ef63d9efd7bc77905e7947c7e09e Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 20:04:22 -0400 Subject: [PATCH 092/161] perf(executor): reuse verified cache capacity headroom --- .../executor/test_registry_artifact_budget.py | 36 ++++++++++++++++++- .../executor/registry_artifact_storage.py | 25 +++++++++---- 2 files changed, 54 insertions(+), 7 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index e526e45988..022fa2d304 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -9,7 +9,7 @@ from collections.abc import Awaitable, Callable from pathlib import Path from typing import Literal -from unittest.mock import ANY, AsyncMock, MagicMock, patch +from unittest.mock import ANY, AsyncMock, MagicMock, call, patch import pytest @@ -116,6 +116,40 @@ async def test_admission_rounds_download_reservation_to_allocation_unit( max_bytes=8192, ) + @pytest.mark.anyio + async def test_admission_reuses_verified_capacity_headroom( + self, + temp_cache_dir: Path, + ) -> None: + """Chunked downloads rescan only after consuming known free bytes.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch(MAX_BYTES_CONFIG, 100), + patch( + "tracecat.executor.registry_artifact_storage." + "_filesystem_allocation_unit", + return_value=1, + ), + patch.object( + cache, + "_ensure_cache_capacity", + new_callable=AsyncMock, + side_effect=[10, 6], + ) as ensure_capacity, + ): + admission = cache._admission_for("new") + assert admission is not None + await admission.ensure_capacity(4) + await admission.ensure_capacity(6) + await admission.ensure_capacity(5) + await admission.ensure_capacity(2) + + assert ensure_capacity.await_args_list == [ + call(additional_bytes=4, protected_key="new", max_bytes=100), + call(additional_bytes=5, protected_key="new", max_bytes=100), + ] + def test_delete_cache_path_reports_directory_failure(self, temp_cache_dir): """Physical deletion failures are observable instead of suppressed.""" entry_dir = temp_cache_dir / "entry" diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 4383e1c1fa..fc77bf9e98 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -242,13 +242,20 @@ def _admission_for(self, cache_key: str) -> RegistryArtifactAdmission | None: return None allocation_unit = _filesystem_allocation_unit(self.cache_dir) + verified_headroom_bytes = 0 async def ensure_capacity(additional_bytes: int) -> None: - await self._ensure_cache_capacity( - additional_bytes=_allocated_size_bound( - additional_bytes, - allocation_unit=allocation_unit, - ), + nonlocal verified_headroom_bytes + allocated_bytes = _allocated_size_bound( + additional_bytes, + allocation_unit=allocation_unit, + ) + if allocated_bytes <= verified_headroom_bytes: + verified_headroom_bytes -= allocated_bytes + return + + verified_headroom_bytes = await self._ensure_cache_capacity( + additional_bytes=allocated_bytes, protected_key=cache_key, max_bytes=max_bytes, ) @@ -422,12 +429,16 @@ async def _ensure_cache_capacity( additional_bytes: int, protected_key: str, max_bytes: int, - ) -> None: + ) -> int: """Reserve peak bytes for a cold writer without exceeding the cap. The caller holds the admission lock and its key lock. Every normal budget pass takes the admission lock first, so acquiring the budget lock here cannot deadlock with eviction of the protected key. + + Returns the additional byte headroom proven by the same scan after the + requested allocation. The admission callback consumes that headroom + before rescanning, which keeps unknown-length chunked downloads cheap. """ if additional_bytes < 0: raise ValueError("additional_bytes must be non-negative") @@ -484,6 +495,8 @@ async def _ensure_cache_capacity( else: skipped.add(candidate.cache_key) + return max_bytes - total_bytes - additional_bytes + def _least_recently_used( self, entries: Iterable[RegistryArtifactCacheEntry], From 3fee0eea45a2a8eb7bced6de56293864a6c8d867 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 20:06:53 -0400 Subject: [PATCH 093/161] fix(executor): defer SquashFS staging cleanup --- .../test_registry_artifact_materialization.py | 59 +++++++++++++++++++ .../registry_artifact_materialization.py | 5 +- 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 5e02af5786..aa51e56135 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -672,6 +672,65 @@ def fail_download_unlink( assert cache._deferred_staging_cleanup == set() assert not deferred_path.exists() + @pytest.mark.anyio + async def test_failed_squashfs_unlink_is_deferred_after_concurrent_publish( + self, + temp_cache_dir: Path, + ) -> None: + """A losing SquashFS staging file remains retryable without masking success.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/concurrent.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) + ctx = cache._context_for(cache_key) + image_path = ctx.paths.squashfs_image_path + staging_paths: list[Path] = [] + real_rename = Path.rename + real_unlink = Path.unlink + + async def download( + artifact_uri: str, + output_path: Path, + *, + admission: object, + defer_cleanup: object, + ) -> None: + del artifact_uri, admission, defer_cleanup + output_path.write_bytes(b"loser") + staging_paths.append(output_path) + + def publish_concurrently(path: Path, target: Path) -> Path: + if path in staging_paths: + target.write_bytes(b"winner") + raise FileExistsError("published by another process") + return real_rename(path, target) + + def fail_staging_unlink(path: Path, missing_ok: bool = False) -> None: + if path in staging_paths: + raise PermissionError("cleanup denied") + real_unlink(path, missing_ok=missing_ok) + + with ( + patch( + "tracecat.executor.registry_artifact_materialization." + "_download_s3_artifact", + side_effect=download, + ), + patch.object(Path, "rename", publish_concurrently), + patch.object(Path, "unlink", fail_staging_unlink), + ): + await artifact.download(ctx, image_path) + + assert image_path.read_bytes() == b"winner" + assert len(staging_paths) == 1 + deferred_path = staging_paths[0] + assert deferred_path.exists() + assert cache._deferred_staging_cleanup == {deferred_path} + + assert cache._retry_deferred_staging_cleanup() is True + assert cache._deferred_staging_cleanup == set() + assert not deferred_path.exists() + @pytest.mark.parametrize("artifact_format", ["squashfs", "tarball"]) @pytest.mark.anyio async def test_repeatedly_cancelled_partial_cleanup_rejoins_thread( diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index bf09fb8775..a018d10fc8 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -339,7 +339,10 @@ async def download( raise return (time.monotonic() - download_start) * 1000 finally: - temp_image.unlink(missing_ok=True) + _remove_file_or_defer( + temp_image, + defer_cleanup=ctx.defer_cleanup, + ) async def mount( self, From 1447e740089e84fbffefe3597317a2a36bd25ec3 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 20:18:03 -0400 Subject: [PATCH 094/161] fix(executor): scope test caches to backend loops --- .../test_test_backend_no_registry_action.py | 40 ++++++++++++++++--- tracecat/executor/backends/base.py | 7 +++- tracecat/executor/backends/test.py | 20 +++++++++- 3 files changed, 58 insertions(+), 9 deletions(-) diff --git a/tests/unit/executor/test_test_backend_no_registry_action.py b/tests/unit/executor/test_test_backend_no_registry_action.py index 01e83a6e23..c3b6b66c7f 100644 --- a/tests/unit/executor/test_test_backend_no_registry_action.py +++ b/tests/unit/executor/test_test_backend_no_registry_action.py @@ -106,6 +106,31 @@ def test_run_action_input() -> RunActionInput: class TestTestBackendNoRegistryAction: """Test that TestBackend does not query RegistryActionsService.""" + def test_backend_instances_do_not_share_loop_affine_cache( + self, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + """Function-scoped backends can run in successive pytest event loops.""" + monkeypatch.setattr( + "tracecat.executor.backends.test.config" + ".TRACECAT__EXECUTOR_REGISTRY_CACHE_DIR", + str(tmp_path), + ) + first_backend = TestBackend() + second_backend = TestBackend() + + async def sweep(backend: TestBackend) -> None: + await backend._registry_artifact_cache().ensure_swept() + + asyncio.run(sweep(first_backend)) + asyncio.run(sweep(second_backend)) + + assert ( + first_backend._registry_artifact_cache() + is not second_backend._registry_artifact_cache() + ) + @pytest.mark.anyio async def test_execute_udf_without_db_lookup( self, @@ -306,8 +331,9 @@ async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: False, ) monkeypatch.setattr( - "tracecat.executor.backends.test.get_action_runner", - lambda: fake_runner, + backend, + "_registry_artifact_cache", + lambda: fake_runner.registry_artifacts, ) monkeypatch.setattr(backend, "_get_artifact_uris", _get_artifact_uris) monkeypatch.setattr( @@ -377,8 +403,9 @@ async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: False, ) monkeypatch.setattr( - "tracecat.executor.backends.test.get_action_runner", - lambda: fake_runner, + backend, + "_registry_artifact_cache", + lambda: fake_runner.registry_artifacts, ) monkeypatch.setattr(backend, "_get_artifact_uris", _get_artifact_uris) @@ -450,8 +477,9 @@ def blocking_udf(**_kwargs: object) -> str: False, ) monkeypatch.setattr( - "tracecat.executor.backends.test.get_action_runner", - lambda: fake_runner, + backend, + "_registry_artifact_cache", + lambda: fake_runner.registry_artifacts, ) monkeypatch.setattr(backend, "_get_artifact_uris", _get_artifact_uris) monkeypatch.setattr( diff --git a/tracecat/executor/backends/base.py b/tracecat/executor/backends/base.py index 082d4dfe7e..67f5169938 100644 --- a/tracecat/executor/backends/base.py +++ b/tracecat/executor/backends/base.py @@ -38,6 +38,7 @@ if TYPE_CHECKING: from tracecat.auth.types import Role from tracecat.dsl.schemas import RunActionInput + from tracecat.executor.registry_artifacts import RegistryArtifactCache from tracecat.executor.schemas import ResolvedContext @@ -149,7 +150,7 @@ async def _execute_run_python( # The lease is held for the whole sandbox run so cache eviction cannot # delete a directory the script is still importing from. - registry_artifacts = get_action_runner().registry_artifacts + registry_artifacts = self._registry_artifact_cache() async with registry_artifacts.lease(artifact_uris) as registry_paths: return await self._run_python_in_sandbox( script=script, @@ -159,6 +160,10 @@ async def _execute_run_python( resolved_context=resolved_context, ) + def _registry_artifact_cache(self) -> RegistryArtifactCache: + """Return the cache owned by this backend's executor process.""" + return get_action_runner().registry_artifacts + async def _run_python_in_sandbox( self, *, diff --git a/tracecat/executor/backends/test.py b/tracecat/executor/backends/test.py index 8b70855721..a9b6e42e13 100644 --- a/tracecat/executor/backends/test.py +++ b/tracecat/executor/backends/test.py @@ -22,6 +22,7 @@ import sys import threading from contextlib import AsyncExitStack, contextmanager +from pathlib import Path from typing import TYPE_CHECKING, Any from tracecat_registry import secrets as registry_secrets @@ -37,9 +38,9 @@ ctx_session_id, ) from tracecat.executor.action_gateway.config import action_gateway_socket_path -from tracecat.executor.action_runner import get_action_runner from tracecat.executor.backends.base import ExecutorBackend from tracecat.executor.backends.registry_helpers import get_registry_artifact_uris +from tracecat.executor.registry_artifacts import RegistryArtifactCache from tracecat.executor.schemas import ( ActionImplementation, ExecutorActionErrorInfo, @@ -107,6 +108,21 @@ def _temporary_sys_path(paths: list[str]) -> Iterator[None]: class TestBackend(ExecutorBackend): """In-process execution backend for tests only.""" + __test__ = False + + def __init__(self) -> None: + # Pytest creates a backend inside each function-scoped event loop. Keep + # its loop-affine cache on the same lifecycle instead of reusing the + # process-global ActionRunner cache across closed test loops. + self._owned_registry_artifacts: RegistryArtifactCache | None = None + + def _registry_artifact_cache(self) -> RegistryArtifactCache: + if self._owned_registry_artifacts is None: + self._owned_registry_artifacts = RegistryArtifactCache( + Path(config.TRACECAT__EXECUTOR_REGISTRY_CACHE_DIR) + ) + return self._owned_registry_artifacts + async def _execute( self, input: RunActionInput, @@ -308,7 +324,7 @@ async def _lease_registry_artifacts( logger.debug("No artifact URIs found, using empty paths") return [] - registry_artifacts = get_action_runner().registry_artifacts + registry_artifacts = self._registry_artifact_cache() await registry_artifacts.ensure_swept() extracted_paths: list[str] = [] From 995e5b65e44c8954945f5f277afe83af93c9c42d Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 20:22:06 -0400 Subject: [PATCH 095/161] refactor(executor): share cache eviction mechanics --- .../executor/registry_artifact_storage.py | 151 +++++++++++------- 1 file changed, 90 insertions(+), 61 deletions(-) diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index fc77bf9e98..d325d4764c 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -53,6 +53,15 @@ class RegistryArtifactEviction: reclaimed: bool +@dataclass(frozen=True, slots=True) +class _RegistryArtifactEvictionPass: + """Result of applying one cache limit policy to evictable entries.""" + + total_bytes: int + fits: bool + exhausted_candidates: bool + + @dataclass(frozen=True, slots=True) class RegistryArtifactCacheEntry: """Measured on-disk footprint and recency for one registry artifact key.""" @@ -375,13 +384,7 @@ async def _enforce_cache_budget_locked( protected_key: str | None, ) -> bool: """Enforce entry and byte limits while both cache-wide locks are held.""" - cleanup = asyncio.gather( - asyncio.to_thread(self._clear_work_dir, self.trash_dir), - asyncio.to_thread(self._retry_deferred_staging_cleanup), - ) - trash_clean, startup_clean = await _rejoin_future_on_cancel(cleanup) - cleanup_complete = trash_clean and startup_clean - if not cleanup_complete: + if not await self._cleanup_cache_work_dirs(): return False max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES @@ -392,37 +395,36 @@ async def _enforce_cache_budget_locked( entries = await asyncio.to_thread(self._scan_cache_entries) total_bytes = sum(entry.size_bytes for entry in entries.values()) protected = set() if protected_key is None else {protected_key} - skipped: set[str] = set() - - while (max_entries > 0 and len(entries) > max_entries) or ( - max_bytes > 0 and total_bytes > max_bytes - ): - candidate = self._least_recently_used( - entries.values(), - excluded=skipped | protected, - ) - if candidate is None: + eviction_pass = await self._evict_until_fits( + entries, + total_bytes=total_bytes, + excluded=protected, + max_entries=max_entries, + max_bytes=max_bytes, + ) + if not eviction_pass.fits: + if eviction_pass.exhausted_candidates: logger.warning( "Registry artifact cache is over budget but every entry is in use", cache_dir=str(self.cache_dir), entries=len(entries), max_entries=max_entries, - total_bytes=total_bytes, + total_bytes=eviction_pass.total_bytes, max_bytes=max_bytes, ) - return False - - eviction = await self._evict_entry(candidate.cache_key) - if eviction.retired: - del entries[candidate.cache_key] - if not eviction.reclaimed: - return False - total_bytes -= candidate.size_bytes - else: - skipped.add(candidate.cache_key) + return False return True + async def _cleanup_cache_work_dirs(self) -> bool: + """Rejoin cache-wide cleanup workers and report complete reclamation.""" + cleanup = asyncio.gather( + asyncio.to_thread(self._clear_work_dir, self.trash_dir), + asyncio.to_thread(self._retry_deferred_staging_cleanup), + ) + trash_clean, staging_clean = await _rejoin_future_on_cancel(cleanup) + return trash_clean and staging_clean + async def _ensure_cache_capacity( self, *, @@ -444,11 +446,7 @@ async def _ensure_cache_capacity( raise ValueError("additional_bytes must be non-negative") async with self._budget_lock: - cleanup = asyncio.gather( - asyncio.to_thread(self._clear_work_dir, self.trash_dir), - asyncio.to_thread(self._retry_deferred_staging_cleanup), - ) - trash_clean, startup_clean = await _rejoin_future_on_cancel(cleanup) + cleanup_complete = await self._cleanup_cache_work_dirs() entries = await asyncio.to_thread(self._scan_cache_entries) staging_bytes, trash_bytes = await asyncio.gather( asyncio.to_thread(_directory_footprint, self.staging_dir), @@ -459,43 +457,74 @@ async def _ensure_cache_capacity( + staging_bytes + trash_bytes ) - if ( - not (trash_clean and startup_clean) - and total_bytes + additional_bytes > max_bytes - ): + if not cleanup_complete and total_bytes + additional_bytes > max_bytes: raise RegistryArtifactCacheCapacityError( current_bytes=total_bytes, additional_bytes=additional_bytes, max_bytes=max_bytes, ) - skipped = {protected_key} + eviction_pass = await self._evict_until_fits( + entries, + total_bytes=total_bytes, + excluded={protected_key}, + max_entries=0, + max_bytes=max_bytes, + additional_bytes=additional_bytes, + ) + if not eviction_pass.fits: + raise RegistryArtifactCacheCapacityError( + current_bytes=eviction_pass.total_bytes, + additional_bytes=additional_bytes, + max_bytes=max_bytes, + ) + + return max_bytes - eviction_pass.total_bytes - additional_bytes + + async def _evict_until_fits( + self, + entries: dict[str, RegistryArtifactCacheEntry], + *, + total_bytes: int, + excluded: set[str], + max_entries: int, + max_bytes: int, + additional_bytes: int = 0, + ) -> _RegistryArtifactEvictionPass: + """Apply shared LRU retirement mechanics until the given limits fit.""" + skipped = set(excluded) - while total_bytes + additional_bytes > max_bytes: - candidate = self._least_recently_used( - entries.values(), - excluded=skipped, + while (max_entries > 0 and len(entries) > max_entries) or ( + max_bytes > 0 and total_bytes + additional_bytes > max_bytes + ): + candidate = self._least_recently_used( + entries.values(), + excluded=skipped, + ) + if candidate is None: + return _RegistryArtifactEvictionPass( + total_bytes=total_bytes, + fits=False, + exhausted_candidates=True, ) - if candidate is None: - raise RegistryArtifactCacheCapacityError( - current_bytes=total_bytes, - additional_bytes=additional_bytes, - max_bytes=max_bytes, + + eviction = await self._evict_entry(candidate.cache_key) + if eviction.retired: + del entries[candidate.cache_key] + if not eviction.reclaimed: + return _RegistryArtifactEvictionPass( + total_bytes=total_bytes, + fits=False, + exhausted_candidates=False, ) + total_bytes -= candidate.size_bytes + else: + skipped.add(candidate.cache_key) - eviction = await self._evict_entry(candidate.cache_key) - if eviction.retired: - del entries[candidate.cache_key] - if not eviction.reclaimed: - raise RegistryArtifactCacheCapacityError( - current_bytes=total_bytes, - additional_bytes=additional_bytes, - max_bytes=max_bytes, - ) - total_bytes -= candidate.size_bytes - else: - skipped.add(candidate.cache_key) - - return max_bytes - total_bytes - additional_bytes + return _RegistryArtifactEvictionPass( + total_bytes=total_bytes, + fits=True, + exhausted_candidates=False, + ) def _least_recently_used( self, From 7551b042f77e4953fb4760c45f8587dd69380a02 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 20:38:48 -0400 Subject: [PATCH 096/161] fix(executor): fail closed on mount inspection --- .../executor/test_registry_artifact_budget.py | 10 ++++- .../test_registry_artifact_eviction.py | 38 +++++++++++++++++-- .../executor/test_registry_artifact_leases.py | 30 ++++++++++++--- .../test_registry_artifact_materialization.py | 5 ++- .../test_registry_artifact_startup.py | 5 ++- .../registry_artifact_materialization.py | 9 +++-- tracecat/executor/registry_artifact_mounts.py | 30 +++++++++++++++ .../executor/registry_artifact_storage.py | 19 ++++++---- tracecat/executor/registry_artifacts.py | 3 +- 9 files changed, 124 insertions(+), 25 deletions(-) create mode 100644 tracecat/executor/registry_artifact_mounts.py diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 022fa2d304..a760f37405 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -936,7 +936,10 @@ async def mock_umount(*args, **kwargs): return process with ( - patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in mounted, + ), patch( "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/umount", @@ -985,7 +988,10 @@ async def flaky_unmount(mount_dir: Path) -> bool: return True with ( - patch.object(Path, "is_mount", lambda path: path in mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in mounted, + ), patch.object(cache, "_unmount", side_effect=flaky_unmount), ): async with cache.lease([artifact_uri]): diff --git a/tests/unit/executor/test_registry_artifact_eviction.py b/tests/unit/executor/test_registry_artifact_eviction.py index c51ebc2d01..7946ae8671 100644 --- a/tests/unit/executor/test_registry_artifact_eviction.py +++ b/tests/unit/executor/test_registry_artifact_eviction.py @@ -32,6 +32,29 @@ class TestRegistryArtifactCacheEviction: """Retire idle cache entries without disrupting live leases.""" + @pytest.mark.anyio + async def test_eviction_surfaces_unknown_mount_state(self, temp_cache_dir): + """Inspection failures cannot be mistaken for an unmounted entry.""" + cache = RegistryArtifactCache(temp_cache_dir) + paths = cache._paths_for("unknown-mount-state") + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + real_lstat = Path.lstat + + def fail_mount_lstat(path: Path): + if path == paths.squashfs_mount_dir: + raise PermissionError("mount inspection denied") + return real_lstat(path) + + with patch.object(Path, "lstat", fail_mount_lstat): + with pytest.raises(PermissionError, match="mount inspection denied"): + await cache._evict_entry("unknown-mount-state") + + assert paths.entry_dir.is_dir() + assert paths.squashfs_image_path.is_file() + assert paths.squashfs_mount_dir.is_dir() + @pytest.mark.anyio async def test_eviction_unmounts_before_deleting_the_image(self, temp_cache_dir): """Unlinking a mounted image would strand an open-file zombie.""" @@ -55,7 +78,10 @@ async def mock_umount(*args, **kwargs): return process with ( - patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in mounted, + ), patch( "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/umount", @@ -109,7 +135,10 @@ async def mock_umount(*args, **kwargs): return released_process with ( - patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in mounted, + ), patch( "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/umount", @@ -261,7 +290,10 @@ async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): process.returncode = 32 with ( - patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in mounted, + ), patch( "tracecat.executor.registry_artifact_materialization.shutil.which", return_value="/sbin/umount", diff --git a/tests/unit/executor/test_registry_artifact_leases.py b/tests/unit/executor/test_registry_artifact_leases.py index f84c7efe50..117e9a231a 100644 --- a/tests/unit/executor/test_registry_artifact_leases.py +++ b/tests/unit/executor/test_registry_artifact_leases.py @@ -338,7 +338,10 @@ async def hold_lease(index: int) -> None: await releases[index].wait() with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in harness.mounted, + ), patch(SQUASHFS_ENABLED_CONFIG, True), patch( "tracecat.executor.registry_artifact_materialization.shutil.which", @@ -400,7 +403,10 @@ async def hold_lease(index: int) -> None: await releases[index].wait() with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in harness.mounted, + ), patch(SQUASHFS_ENABLED_CONFIG, True), patch( "tracecat.executor.registry_artifact_materialization.shutil.which", @@ -569,7 +575,10 @@ async def new_holder() -> None: await release_newcomer.wait() with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in harness.mounted, + ), patch.object(cache, "_unmount", harness.unmount), patch.object( cache, @@ -632,7 +641,10 @@ async def fail_download( converge_cache_budget = AsyncMock() with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in harness.mounted, + ), patch(SQUASHFS_ENABLED_CONFIG, False), patch.object(TarballArtifact, "download", fail_download), patch.object(cache, "_unmount", harness.unmount), @@ -762,7 +774,10 @@ async def test_duplicate_uri_balances_each_acquisition_and_release( harness = SquashfsMountHarness(cache) with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in harness.mounted, + ), patch(SQUASHFS_ENABLED_CONFIG, True), patch( "tracecat.executor.registry_artifact_materialization.shutil.which", @@ -836,7 +851,10 @@ async def take_lease() -> None: leased_path_exists.append(registry_paths[0].is_dir()) with ( - patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in mounted, + ), patch(SQUASHFS_ENABLED_CONFIG, True), patch( "tracecat.executor.registry_artifact_materialization.shutil.which", diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index aa51e56135..5f6f299be2 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -1056,7 +1056,10 @@ async def test_loop_device_exhaustion_isolated_sticky_extraction_fallback( ) with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in harness.mounted, + ), patch(SQUASHFS_ENABLED_CONFIG, True), patch( "tracecat.executor.registry_artifact_materialization.shutil.which", diff --git a/tests/unit/executor/test_registry_artifact_startup.py b/tests/unit/executor/test_registry_artifact_startup.py index f1e1bbbb58..4c48e1ea4c 100644 --- a/tests/unit/executor/test_registry_artifact_startup.py +++ b/tests/unit/executor/test_registry_artifact_startup.py @@ -126,7 +126,10 @@ async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): with ( patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0), - patch.object(Path, "is_mount", lambda self: self == mount_dir), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path == mount_dir, + ), ): await cache.ensure_swept() diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index a018d10fc8..029a1edb9f 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -20,6 +20,7 @@ import tracecat_registry from tracecat import config +from tracecat.executor import registry_artifact_mounts from tracecat.logger import logger from tracecat.registry.artifact_keys import parse_s3_uri from tracecat.registry.constants import DEFAULT_REGISTRY_ORIGIN @@ -247,7 +248,7 @@ def format(self) -> RegistryArtifactFormat: def cached_path( self, ctx: RegistryArtifactMaterializationContext ) -> list[Path] | None: - if ctx.paths.squashfs_mount_dir.is_mount(): + if registry_artifact_mounts.is_mount(ctx.paths.squashfs_mount_dir): logger.debug( "Using cached SquashFS registry mount", cache_key=ctx.cache_key, @@ -363,7 +364,7 @@ async def mount( Exception: The image could not be downloaded or prepared. """ target_dir = ctx.paths.squashfs_mount_dir - if target_dir.is_mount(): + if registry_artifact_mounts.is_mount(target_dir): return target_dir ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) @@ -472,7 +473,7 @@ async def _mount_image(self, image_path: Path, target_dir: Path) -> None: Raises: SquashfsMountCommandError: The ``mount`` command failed. """ - if target_dir.is_mount(): + if registry_artifact_mounts.is_mount(target_dir): return proc = await asyncio.create_subprocess_exec( @@ -489,7 +490,7 @@ async def _mount_image(self, image_path: Path, target_dir: Path) -> None: ) stdout, stderr = await communicate_process_group(proc) - if proc.returncode == 0 or target_dir.is_mount(): + if proc.returncode == 0 or registry_artifact_mounts.is_mount(target_dir): return output = (stderr or stdout).decode(errors="replace").strip() diff --git a/tracecat/executor/registry_artifact_mounts.py b/tracecat/executor/registry_artifact_mounts.py new file mode 100644 index 0000000000..98db4250bc --- /dev/null +++ b/tracecat/executor/registry_artifact_mounts.py @@ -0,0 +1,30 @@ +"""Fail-closed mount-state inspection for registry artifact caches.""" + +from __future__ import annotations + +import os +import stat +from pathlib import Path + + +def is_mount(path: Path) -> bool: + """Return whether a path is mounted without hiding inspection failures. + + ``Path.is_mount()`` delegates to ``os.path.ismount()``, which converts every + ``OSError`` from ``lstat`` into ``False``. Cache cleanup must distinguish a + missing mount directory from an unreadable one so it never deletes the + backing image while mount state is unknown. + """ + try: + path_stat = path.lstat() + except FileNotFoundError: + return False + + if stat.S_ISLNK(path_stat.st_mode): + return False + + parent = Path(os.path.realpath(path / "..", strict=True)) + parent_stat = parent.lstat() + return ( + path_stat.st_dev != parent_stat.st_dev or path_stat.st_ino == parent_stat.st_ino + ) diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index d325d4764c..55a4b21ef1 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -12,6 +12,7 @@ from pathlib import Path from tracecat import config +from tracecat.executor import registry_artifact_mounts from tracecat.executor.registry_artifact_cache_state import ( _RegistryArtifactCacheState, ) @@ -294,7 +295,9 @@ async def _unmount_idle_entry(self, cache_key: str) -> None: mount_dir = self._paths_for(cache_key).squashfs_mount_dir try: - retry = self._refcount(cache_key) == 0 and mount_dir.is_mount() + retry = self._refcount( + cache_key + ) == 0 and registry_artifact_mounts.is_mount(mount_dir) except OSError as e: retry = True logger.warning( @@ -577,7 +580,7 @@ async def _unmount_entry(self, cache_key: str) -> bool: return False mount_dir = self._paths_for(cache_key).squashfs_mount_dir - if not mount_dir.is_mount(): + if not registry_artifact_mounts.is_mount(mount_dir): return False if not await self._unmount(mount_dir): logger.warning( @@ -626,9 +629,9 @@ async def _evict_entry(self, cache_key: str) -> RegistryArtifactEviction: paths = self._paths_for(cache_key) if not paths.entry_dir.exists(): return RegistryArtifactEviction(retired=True, reclaimed=True) - if paths.squashfs_mount_dir.is_mount() and not await self._unmount( + if registry_artifact_mounts.is_mount( paths.squashfs_mount_dir - ): + ) and not await self._unmount(paths.squashfs_mount_dir): logger.warning( "Failed to unmount registry artifact, skipping eviction", cache_key=cache_key, @@ -690,7 +693,7 @@ async def _unmount(self, mount_dir: Path) -> bool: start_new_session=True, ) stdout, stderr = await communicate_process_group(proc) - if proc.returncode == 0 or not mount_dir.is_mount(): + if proc.returncode == 0 or not registry_artifact_mounts.is_mount(mount_dir): return True logger.warning( @@ -844,7 +847,7 @@ def _trim_startup_cache(self) -> bool: for entry in tuple(entries.values()): paths = self._paths_for(entry.cache_key) if ( - paths.squashfs_mount_dir.is_mount() + registry_artifact_mounts.is_mount(paths.squashfs_mount_dir) or paths.squashfs_image_path.exists() or paths.squashfs_extract_dir.exists() or paths.tarball_target_dir.exists() @@ -884,7 +887,9 @@ def _trim_startup_cache(self) -> bool: ( entry for entry in entries.values() - if not self._paths_for(entry.cache_key).squashfs_mount_dir.is_mount() + if not registry_artifact_mounts.is_mount( + self._paths_for(entry.cache_key).squashfs_mount_dir + ) ), key=lambda entry: entry.last_used, ) diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index c20c1fe669..53d3252fe4 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -9,6 +9,7 @@ from pathlib import Path from tracecat import config +from tracecat.executor import registry_artifact_mounts from tracecat.executor.registry_artifact_cache_state import ( BASE_PYTHONPATH_DIR_NAME, CACHE_ENTRIES_DIR_NAME, @@ -353,7 +354,7 @@ def _remove_unpublished_entry( """ paths = ctx.paths try: - if paths.squashfs_mount_dir.is_mount(): + if registry_artifact_mounts.is_mount(paths.squashfs_mount_dir): return except OSError: return From a43aaeb7d4c17312855495bd121a4257b526ed1e Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 20:40:34 -0400 Subject: [PATCH 097/161] fix(executor): defer unusable SquashFS cleanup --- .../test_registry_artifact_materialization.py | 29 +++++++++++++++++++ .../registry_artifact_materialization.py | 17 ++++------- 2 files changed, 35 insertions(+), 11 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 5f6f299be2..58d715532d 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -731,6 +731,35 @@ def fail_staging_unlink(path: Path, missing_ok: bool = False) -> None: assert cache._deferred_staging_cleanup == set() assert not deferred_path.exists() + def test_failed_unusable_squashfs_unlink_is_deferred( + self, + temp_cache_dir: Path, + ) -> None: + """A failed canonical-image cleanup remains retryable after fallback.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/unusable.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) + ctx = cache._context_for(cache_key) + image_path = ctx.paths.squashfs_image_path + image_path.parent.mkdir(parents=True) + image_path.write_bytes(b"unusable") + real_unlink = Path.unlink + + def fail_image_unlink(path: Path, missing_ok: bool = False) -> None: + if path == image_path: + raise PermissionError("cleanup denied") + real_unlink(path, missing_ok=missing_ok) + + with patch.object(Path, "unlink", fail_image_unlink): + artifact.discard_failed_materialization(ctx) + + assert cache._deferred_staging_cleanup == {image_path} + assert image_path.is_file() + assert cache._retry_deferred_staging_cleanup() is True + assert cache._deferred_staging_cleanup == set() + assert not image_path.exists() + @pytest.mark.parametrize("artifact_format", ["squashfs", "tarball"]) @pytest.mark.anyio async def test_repeatedly_cancelled_partial_cleanup_rejoins_thread( diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 029a1edb9f..66e7c0ace1 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -197,13 +197,13 @@ def _remove_file_or_defer( *, defer_cleanup: Callable[[Path], None], ) -> None: - """Remove one staging file without masking the materialization outcome.""" + """Remove one artifact file without masking the materialization outcome.""" try: path.unlink(missing_ok=True) except OSError as e: defer_cleanup(path) logger.warning( - "Deferred failed registry artifact staging cleanup", + "Deferred failed registry artifact file cleanup", path=str(path), error=str(e), ) @@ -278,15 +278,10 @@ def discard_failed_materialization( ) return - try: - ctx.paths.squashfs_image_path.unlink(missing_ok=True) - except OSError as e: - logger.warning( - "Failed to discard unusable SquashFS candidate", - cache_key=ctx.cache_key, - artifact_uri=self.uri, - error=str(e), - ) + _remove_file_or_defer( + ctx.paths.squashfs_image_path, + defer_cleanup=ctx.defer_cleanup, + ) for directory in ( ctx.paths.squashfs_extract_dir, From dfe4baa1b3e526094a4fb21930eeb224d782c5df Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 20:43:15 -0400 Subject: [PATCH 098/161] refactor(executor): share cancellation rejoin helper --- tracecat/executor/backends/test.py | 20 ++++---------------- 1 file changed, 4 insertions(+), 16 deletions(-) diff --git a/tracecat/executor/backends/test.py b/tracecat/executor/backends/test.py index a9b6e42e13..6108f25f46 100644 --- a/tracecat/executor/backends/test.py +++ b/tracecat/executor/backends/test.py @@ -18,7 +18,6 @@ from __future__ import annotations import asyncio -import contextlib import sys import threading from contextlib import AsyncExitStack, contextmanager @@ -40,6 +39,9 @@ from tracecat.executor.action_gateway.config import action_gateway_socket_path from tracecat.executor.backends.base import ExecutorBackend from tracecat.executor.backends.registry_helpers import get_registry_artifact_uris +from tracecat.executor.registry_artifact_materialization import ( + _run_blocking_rejoin_on_cancel, +) from tracecat.executor.registry_artifacts import RegistryArtifactCache from tracecat.executor.schemas import ( ActionImplementation, @@ -264,21 +266,7 @@ async def _run_sync_udf( leases, temporary ``sys.path`` entries, and secret contexts alive until the function actually stops. """ - worker = asyncio.ensure_future(asyncio.to_thread(fn, **args)) - try: - return await asyncio.shield(worker) - except asyncio.CancelledError: - while not worker.done(): - try: - await asyncio.shield(worker) - except asyncio.CancelledError: - continue - except Exception: - break - if not worker.cancelled(): - with contextlib.suppress(Exception): - worker.result() - raise + return await _run_blocking_rejoin_on_cancel(lambda: fn(**args)) def _load_udf_callable(self, action_impl: ActionImplementation): """Load the UDF callable from action_impl metadata.""" From ca57de03f3000f49c10c7ced75425aa28fdf95ea Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 20:49:11 -0400 Subject: [PATCH 099/161] test(executor): mock fail-closed mount inspection --- tests/unit/test_action_runner.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index 536ceadb56..1b41cb152f 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -777,7 +777,10 @@ async def release_mount(mount_dir: Path) -> bool: return True with ( - patch.object(Path, "is_mount", lambda path: path in mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in mounted, + ), patch.object( action_runner, "_direct_subprocess_command", From 1951de677d0dc2dfc8088f5daf5f132258252546 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 20:50:19 -0400 Subject: [PATCH 100/161] test(sandbox): tolerate procfs stat exit race --- tests/unit/test_unsafe_pid_executor.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_unsafe_pid_executor.py b/tests/unit/test_unsafe_pid_executor.py index 1ecef7b382..25166ae7e2 100644 --- a/tests/unit/test_unsafe_pid_executor.py +++ b/tests/unit/test_unsafe_pid_executor.py @@ -23,8 +23,12 @@ def _process_is_running(pid: int) -> bool: return False stat_path = Path(f"/proc/{pid}/stat") - if not stat_path.exists(): - return True + try: + if not stat_path.exists(): + return True + except (FileNotFoundError, ProcessLookupError): + # Procfs can report ESRCH while resolving a process that just exited. + return False try: stat_fields = stat_path.read_text().split() except (FileNotFoundError, ProcessLookupError): @@ -92,6 +96,18 @@ def process_disappeared( assert not _process_is_running(123) + def test_process_probe_handles_procfs_stat_exit_race( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def process_disappeared(path: Path) -> bool: + del path + raise ProcessLookupError + + monkeypatch.setattr(os, "kill", lambda *_: None) + monkeypatch.setattr(Path, "exists", process_disappeared) + + assert not _process_is_running(123) + @pytest.mark.anyio async def test_build_execution_cmd_with_pid_namespace( self, executor: UnsafePidExecutor, monkeypatch: pytest.MonkeyPatch From e2bf986344d1ab83804701c8da990929215dc905 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 21:01:35 -0400 Subject: [PATCH 101/161] fix(executor): skip cache sweep for bundled registry --- .../test_test_backend_no_registry_action.py | 58 +++++++++++++++++++ tracecat/executor/backends/test.py | 4 +- 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_test_backend_no_registry_action.py b/tests/unit/executor/test_test_backend_no_registry_action.py index c3b6b66c7f..17de1926cc 100644 --- a/tests/unit/executor/test_test_backend_no_registry_action.py +++ b/tests/unit/executor/test_test_backend_no_registry_action.py @@ -29,6 +29,7 @@ RunContext, ) from tracecat.executor.backends.test import TestBackend +from tracecat.executor.registry_artifacts import bundled_builtin_registry_uri from tracecat.executor.schemas import ( ActionImplementation, ExecutorResult, @@ -419,6 +420,63 @@ async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: assert fake_runner.registry_artifacts.lease_attempted is False + @pytest.mark.anyio + async def test_builtin_only_execution_skips_registry_cache_sweep( + self, + test_role: Role, + test_run_action_input: RunActionInput, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Cache-free builtin execution is independent of cache inspection.""" + artifact_uri = bundled_builtin_registry_uri("1.2.3") + + class FakeRegistryArtifacts: + def __init__(self) -> None: + self.sweep_attempted = False + self.lease_attempted = False + + async def ensure_swept(self) -> None: + self.sweep_attempted = True + raise PermissionError("cannot inspect registry cache") + + @asynccontextmanager + async def lease( + self, artifact_uris: list[str] | None = None + ) -> AsyncIterator[list[Path]]: + assert artifact_uris == [artifact_uri] + self.lease_attempted = True + yield [] + + registry_artifacts = FakeRegistryArtifacts() + + async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: + return [artifact_uri] + + backend = TestBackend() + monkeypatch.setattr( + "tracecat.executor.backends.test.config.TRACECAT__LOCAL_REPOSITORY_ENABLED", + False, + ) + monkeypatch.setattr( + backend, + "_registry_artifact_cache", + lambda: registry_artifacts, + ) + monkeypatch.setattr(backend, "_get_artifact_uris", _get_artifact_uris) + + async with AsyncExitStack() as leases: + assert ( + await backend._lease_registry_artifacts( + leases, + test_run_action_input, + test_role, + ) + == [] + ) + + assert registry_artifacts.sweep_attempted is False + assert registry_artifacts.lease_attempted is True + @pytest.mark.anyio async def test_timed_out_sync_udf_keeps_artifact_lease_until_thread_finishes( self, diff --git a/tracecat/executor/backends/test.py b/tracecat/executor/backends/test.py index 6108f25f46..b8cfd23176 100644 --- a/tracecat/executor/backends/test.py +++ b/tracecat/executor/backends/test.py @@ -40,6 +40,7 @@ from tracecat.executor.backends.base import ExecutorBackend from tracecat.executor.backends.registry_helpers import get_registry_artifact_uris from tracecat.executor.registry_artifact_materialization import ( + _is_cache_entry_uri, _run_blocking_rejoin_on_cancel, ) from tracecat.executor.registry_artifacts import RegistryArtifactCache @@ -313,7 +314,8 @@ async def _lease_registry_artifacts( return [] registry_artifacts = self._registry_artifact_cache() - await registry_artifacts.ensure_swept() + if any(_is_cache_entry_uri(uri) for uri in artifact_uris): + await registry_artifacts.ensure_swept() extracted_paths: list[str] = [] for artifact_uri in artifact_uris: From 10b058f9fde9a02aea6d52af34f5dacd7c9f6421 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 21:20:45 -0400 Subject: [PATCH 102/161] fix(storage): rejoin cancelled file opens --- tests/unit/test_storage_blob.py | 72 +++++++++++++++++++++++++++++++++ tracecat/storage/blob.py | 56 ++++++++++++++++++++++++- 2 files changed, 126 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_storage_blob.py b/tests/unit/test_storage_blob.py index 53abd9a8a1..5cceec33e2 100644 --- a/tests/unit/test_storage_blob.py +++ b/tests/unit/test_storage_blob.py @@ -1077,6 +1077,78 @@ async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 assert not out.exists() assert not temp_path.exists() + @pytest.mark.anyio + async def test_download_file_to_path_cancellation_rejoins_active_open( + self, tmp_path: Path, monkeypatch + ): + """Cancellation cannot race cleanup against the aiofiles open worker.""" + open_started = threading.Event() + open_release = threading.Event() + open_finished = threading.Event() + close_finished = threading.Event() + temp_path = tmp_path / "out.bin.part" + + class DummyStream: + async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 + yield b"unused" + + class BlockingFile: + async def __aenter__(self): + def blocking_open(): + open_started.set() + open_release.wait() + temp_path.touch() + open_finished.set() + return self + + return await asyncio.to_thread(blocking_open) + + async def __aexit__(self, exc_type, exc, traceback): + del exc_type, exc, traceback + close_finished.set() + + @asynccontextmanager + async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 + yield DummyStream(), len(b"unused") + + monkeypatch.setattr( + "tracecat.storage.blob.open_download_stream", + _fake_open_download_stream, + ) + monkeypatch.setattr( + "tracecat.storage.blob.aiofiles.open", + lambda *args, **kwargs: BlockingFile(), + ) + + out = tmp_path / "out.bin" + download = asyncio.create_task( + download_file_to_path( + key="k", + bucket="b", + output_path=out, + ) + ) + try: + assert await asyncio.to_thread(open_started.wait, 1.0) + + download.cancel() + await asyncio.sleep(0) + assert not download.done() + + download.cancel() + await asyncio.sleep(0) + assert not download.done() + finally: + open_release.set() + + with pytest.raises(asyncio.CancelledError): + await download + + assert open_finished.is_set() + assert close_finished.is_set() + assert not out.exists() + assert not temp_path.exists() + @pytest.mark.anyio @patch("tracecat.storage.blob.get_storage_client") async def test_ensure_bucket_exists_create_error_propagates(self, mock_get_client): diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index 8aa43e32b1..674ad48623 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -11,7 +11,7 @@ from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass, field from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Protocol import aioboto3 import aiofiles @@ -36,6 +36,11 @@ DEFAULT_UPLOAD_MAX_CONCURRENCY = 4 DEFAULT_UPLOAD_MAX_IO_QUEUE_SIZE = 2 + +class _AsyncWritableFile(Protocol): + async def write(self, data: bytes, /) -> int: ... + + # Shared S3/MinIO client config: explicit standard-mode retries so transient # failures (throttling, 5xx, connection resets) are retried with backoff instead # of surfacing on the first error. @@ -789,6 +794,53 @@ async def download_file_to_path( hasher = hashlib.sha256() if expected_sha256 is not None else None bytes_written = 0 + @asynccontextmanager + async def open_file_rejoin_on_cancel() -> AsyncIterator[_AsyncWritableFile]: + """Keep the aiofiles open/close workers joined through cancellation.""" + opened: asyncio.Future[_AsyncWritableFile] = ( + asyncio.get_running_loop().create_future() + ) + close_file = asyncio.Event() + + async def file_lifecycle() -> None: + try: + async with aiofiles.open(temp_path, "wb", buffering=0) as file: + opened.set_result(file) + await close_file.wait() + except BaseException as e: + if not opened.done(): + opened.set_exception(e) + return + raise + + lifecycle = asyncio.create_task(file_lifecycle()) + operation_error: BaseException | None = None + try: + file = await asyncio.shield(opened) + yield file + except BaseException as e: + operation_error = e + raise + finally: + close_file.set() + pending_cancellation: asyncio.CancelledError | None = None + while not lifecycle.done(): + try: + await asyncio.shield(lifecycle) + except asyncio.CancelledError as e: + if lifecycle.cancelled(): + raise + pending_cancellation = e + + try: + lifecycle.result() + except BaseException as cleanup_error: + if operation_error is not None: + raise operation_error from cleanup_error + raise + if pending_cancellation is not None: + raise pending_cancellation + async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: """Write one chunk without abandoning the aiofiles worker thread.""" writer = asyncio.ensure_future(file.write(chunk)) @@ -848,7 +900,7 @@ async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: # Unbuffered writes keep the partial file's allocated size visible # to incremental capacity scans between unknown-length chunks. - async with aiofiles.open(temp_path, "wb", buffering=0) as f: + async with open_file_rejoin_on_cancel() as f: async for chunk in stream.iter_chunks(chunk_size=chunk_size): if not chunk: continue From 8203758d8cb58fa4a96e8ad83cf900a1e3d3d873 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 21:22:33 -0400 Subject: [PATCH 103/161] fix(executor): preserve exact artifact cache identity --- tests/unit/executor/test_registry_artifact_resolution.py | 9 +++++++++ tracecat/executor/registry_artifact_materialization.py | 5 ++--- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py index 049e562ba2..f4861d3a27 100644 --- a/tests/unit/executor/test_registry_artifact_resolution.py +++ b/tests/unit/executor/test_registry_artifact_resolution.py @@ -72,6 +72,15 @@ def test_compute_registry_artifact_cache_key_case_sensitive(self): assert key1 != key2 + def test_compute_registry_artifact_cache_key_whitespace_sensitive(self): + """Cache identity preserves whitespace that can belong to an S3 key.""" + key = compute_registry_artifact_cache_key("s3://bucket/path/file.tar.gz") + whitespace_key = compute_registry_artifact_cache_key( + "s3://bucket/path/file.tar.gz " + ) + + assert key != whitespace_key + def test_compute_registry_artifact_cache_key_empty(self): """Test that empty URI returns the base cache key.""" assert compute_registry_artifact_cache_key("") == "base" diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 66e7c0ace1..a99522bc89 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -716,9 +716,8 @@ def compute_registry_artifact_cache_key(artifact_uri: str) -> str: """Compute the local cache key for a registry artifact URI.""" if not artifact_uri: return "base" - # S3 keys are case-sensitive, so preserve URI case when hashing. - content = artifact_uri.strip() - return hashlib.sha256(content.encode()).hexdigest()[:16] + # S3 keys are byte-sensitive, so hash the exact URI used for retrieval. + return hashlib.sha256(artifact_uri.encode()).hexdigest()[:16] def bundled_builtin_registry_uri(version: str) -> str: From 0f7a32ecfbe9ed342ea9098e48bda211e6f1e2a4 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 21:23:33 -0400 Subject: [PATCH 104/161] fix(executor): tolerate non-UTF-8 SquashFS names --- .../unit/executor/test_registry_artifact_materialization.py | 5 +++++ tracecat/executor/registry_artifact_materialization.py | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 58d715532d..8491958aed 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -152,6 +152,11 @@ def test_squashfs_listing_size_bounds_each_inode_allocation(self) -> None: assert _squashfs_listing_size(listing, allocation_unit=4096) == 12_288 + def test_squashfs_listing_size_accepts_non_utf8_filenames(self) -> None: + listing = b"-rw-r--r-- 0/0 123 2026-01-01 00:00 squashfs-root/module-\xff.py" + + assert _squashfs_listing_size(listing, allocation_unit=4096) == 4096 + def test_tarball_size_bounds_each_member_allocation( self, temp_cache_dir: Path, diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index a99522bc89..02fbfde82d 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -850,7 +850,7 @@ def _tarball_extracted_size( def _squashfs_listing_size(output: bytes, *, allocation_unit: int = 1) -> int: """Bound allocated bytes from ``unsquashfs -lln`` output, failing closed.""" total_bytes = 0 - for raw_line in output.decode(errors="strict").splitlines(): + for raw_line in output.decode(errors="replace").splitlines(): line = raw_line.strip() if not line: continue From 9bd856fec9e1bdab3be88fd1ce2a57fb459b30ae Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 21:36:12 -0400 Subject: [PATCH 105/161] fix(executor): contain detached action descendants --- .../unit/executor/test_process_supervisor.py | 135 ++++++++++++++ tests/unit/test_action_runner.py | 13 ++ tracecat/executor/action_runner.py | 11 +- tracecat/executor/process_supervisor.py | 173 ++++++++++++++++++ tracecat/sandbox/utils.py | 24 ++- 5 files changed, 351 insertions(+), 5 deletions(-) create mode 100644 tests/unit/executor/test_process_supervisor.py create mode 100644 tracecat/executor/process_supervisor.py diff --git a/tests/unit/executor/test_process_supervisor.py b/tests/unit/executor/test_process_supervisor.py new file mode 100644 index 0000000000..87fd411165 --- /dev/null +++ b/tests/unit/executor/test_process_supervisor.py @@ -0,0 +1,135 @@ +"""Linux process-tree containment for direct registry actions.""" + +from __future__ import annotations + +import asyncio +import contextlib +import os +import signal +import sys +from pathlib import Path + +import pytest + +from tracecat.executor import process_supervisor +from tracecat.sandbox.utils import ( + communicate_process_group, + terminate_supervised_process, +) + +pytestmark = pytest.mark.skipif( + sys.platform != "linux", + reason="The direct action supervisor uses Linux prctl and procfs", +) + + +def _process_is_running(pid: int) -> bool: + """Return whether a process is alive, treating zombies as terminated.""" + try: + os.kill(pid, 0) + stat_path = Path(f"/proc/{pid}/stat") + return not stat_path.exists() or stat_path.read_text().split()[2] != "Z" + except (FileNotFoundError, ProcessLookupError): + return False + + +async def _wait_for_file(path: Path) -> None: + for _ in range(500): + if path.exists(): + return + await asyncio.sleep(0.01) + raise AssertionError(f"Timed out waiting for {path}") + + +def _write_action_script(path: Path) -> None: + path.write_text( + """ +import os +import subprocess +import sys +import time +from pathlib import Path + +pid_file = Path(sys.argv[1]) +mode = sys.argv[2] +child = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(30)"], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, +) +pid_file.write_text(f"{os.getpid()} {child.pid}") +if mode == "block": + time.sleep(30) +""".lstrip() + ) + + +async def _spawn_supervised_action( + tmp_path: Path, + *, + mode: str, +) -> tuple[asyncio.subprocess.Process, Path]: + action_script = tmp_path / f"action-{mode}.py" + pid_file = tmp_path / f"action-{mode}.pids" + _write_action_script(action_script) + supervisor_path = Path(process_supervisor.__file__) + process = await asyncio.create_subprocess_exec( + sys.executable, + str(supervisor_path), + sys.executable, + str(action_script), + str(pid_file), + mode, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + start_new_session=True, + ) + return process, pid_file + + +@pytest.mark.anyio +async def test_supervisor_reaps_detached_descendant_after_success( + tmp_path: Path, +) -> None: + process, pid_file = await _spawn_supervised_action(tmp_path, mode="success") + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=5) + + assert process.returncode == 0, stderr.decode() + assert stdout == b"" + _, detached_pid = (int(pid) for pid in pid_file.read_text().split()) + assert not _process_is_running(detached_pid) + + +@pytest.mark.anyio +async def test_supervisor_reaps_detached_descendant_before_cancellation_returns( + tmp_path: Path, +) -> None: + process, pid_file = await _spawn_supervised_action(tmp_path, mode="block") + await _wait_for_file(pid_file) + action_pid, detached_pid = (int(pid) for pid in pid_file.read_text().split()) + communication = asyncio.create_task( + communicate_process_group( + process, + timeout=30, + terminate=terminate_supervised_process, + ) + ) + + try: + await asyncio.sleep(0) + communication.cancel() + await asyncio.sleep(0) + communication.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(communication, timeout=5) + assert not _process_is_running(action_pid) + assert not _process_is_running(detached_pid) + finally: + for pid in (action_pid, detached_pid): + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index 1b41cb152f..1047ea8c1f 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -9,6 +9,7 @@ import contextlib import tempfile import uuid +from collections.abc import Awaitable, Callable from datetime import UTC, datetime from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch @@ -102,12 +103,15 @@ async def communicate( *, input: bytes | None = None, # noqa: A002 timeout: float | None = None, + terminate: Callable[[asyncio.subprocess.Process], Awaitable[None]] + | None = None, ) -> tuple[bytes, bytes]: if isinstance(process, asyncio.subprocess.Process): return await real_communication( process, input=input, timeout=timeout, + terminate=terminate, ) stdout, stderr = await asyncio.wait_for( process.communicate(input=input), @@ -529,6 +533,7 @@ async def test_execute_action_disables_new_privileges_for_direct_subprocess( temp_cache_dir, mock_run_action_input, mock_role, + mock_process_group_communication: AsyncMock, monkeypatch: pytest.MonkeyPatch, ): """Test direct subprocess execution disables new Linux privileges.""" @@ -577,6 +582,14 @@ async def create_subprocess_exec_side_effect(*args, **kwargs): # noqa: ARG001 ] assert captured_args[-2] == action_runner.sys.executable assert captured_args[-1].endswith("minimal_runner.py") + assert captured_args[-4] == action_runner.sys.executable + assert captured_args[-3].endswith("process_supervisor.py") + communication_call = mock_process_group_communication.await_args + assert communication_call is not None + assert ( + communication_call.kwargs["terminate"] + is action_runner.terminate_supervised_process + ) @pytest.mark.anyio async def test_execute_action_invalid_json_response( diff --git a/tracecat/executor/action_runner.py b/tracecat/executor/action_runner.py index 63b9c3606d..a285090588 100644 --- a/tracecat/executor/action_runner.py +++ b/tracecat/executor/action_runner.py @@ -45,7 +45,10 @@ from tracecat.logger import logger from tracecat.sandbox.executor import ActionSandboxConfig, NsjailExecutor from tracecat.sandbox.types import ResourceLimits -from tracecat.sandbox.utils import communicate_process_group +from tracecat.sandbox.utils import ( + communicate_process_group, + terminate_supervised_process, +) from tracecat.secrets.common import apply_masks, apply_masks_object if TYPE_CHECKING: @@ -118,11 +121,14 @@ def _direct_subprocess_command(minimal_runner_path: Path) -> list[str]: if setpriv is None: raise RuntimeError("setpriv is required for direct action subprocess isolation") + supervisor_path = Path(__file__).with_name("process_supervisor.py") return [ setpriv, "--no-new-privs", "--inh-caps=-all", "--ambient-caps=-all", + sys.executable, + str(supervisor_path), *runner_command, ] @@ -443,6 +449,9 @@ async def _execute_direct( proc, input=input_json, timeout=timeout, + terminate=( + terminate_supervised_process if sys.platform == "linux" else None + ), ) elapsed_ms = (time.monotonic() - start_time) * 1000 logger.info( diff --git a/tracecat/executor/process_supervisor.py b/tracecat/executor/process_supervisor.py new file mode 100644 index 0000000000..a53abbcfd1 --- /dev/null +++ b/tracecat/executor/process_supervisor.py @@ -0,0 +1,173 @@ +"""Contain descendants of one Linux direct-action subprocess. + +The outer process remains the child observed by ``ActionRunner``. A detached +subreaper monitor owns the actual action process and all descendants orphaned by +it. Closing the control pipe asks that monitor to kill and reap its complete +child tree before the outer process exits. + +This module intentionally uses only the Python standard library so invoking it +does not import the Tracecat application in an untrusted action process. +""" + +from __future__ import annotations + +import ctypes +import os +import select +import signal +import sys +from collections.abc import Sequence +from contextlib import suppress +from pathlib import Path +from types import FrameType + +_PR_SET_CHILD_SUBREAPER = 36 +_PARENT_POLL_INTERVAL_MS = 10 + + +def _exit_code(wait_status: int) -> int: + """Convert a wait status into a shell-compatible non-negative exit code.""" + code = os.waitstatus_to_exitcode(wait_status) + return code if code >= 0 else 128 - code + + +def _set_child_subreaper() -> None: + """Make this process the reparenting boundary for orphaned descendants.""" + libc = ctypes.CDLL(None, use_errno=True) + prctl = libc.prctl + prctl.argtypes = [ + ctypes.c_int, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ] + prctl.restype = ctypes.c_int + if prctl(_PR_SET_CHILD_SUBREAPER, 1, 0, 0, 0) != 0: + error_number = ctypes.get_errno() + raise OSError(error_number, os.strerror(error_number)) + + +def _direct_child_pids() -> list[int]: + """Return direct child PIDs from procfs for this single-threaded monitor.""" + children_path = Path(f"/proc/self/task/{os.getpid()}/children") + contents = children_path.read_text().strip() + return [int(pid) for pid in contents.split()] if contents else [] + + +def _waitpid(pid: int, options: int = 0) -> tuple[int, int]: + """Wait for a child while tolerating signal interruptions.""" + while True: + try: + return os.waitpid(pid, options) + except InterruptedError: + continue + + +def _kill_and_reap_children() -> None: + """Kill every adopted child, including descendants orphaned while reaping.""" + while child_pids := _direct_child_pids(): + for child_pid in child_pids: + with suppress(ProcessLookupError): + os.kill(child_pid, signal.SIGKILL) + for child_pid in child_pids: + with suppress(ChildProcessError): + _waitpid(child_pid) + + +def _exec(command: Sequence[str], control_fd: int) -> None: + """Replace this child with the actual action command.""" + os.close(control_fd) + try: + os.execvpe(command[0], list(command), os.environ) + except OSError as e: + message = f"Failed to execute supervised action: {e}\n".encode() + with suppress(OSError): + os.write(2, message) + os._exit(127) + + +def _run_monitor(control_fd: int, command: Sequence[str]) -> int: + """Run the action below a detached subreaper and contain its descendants.""" + try: + os.setsid() + _set_child_subreaper() + _direct_child_pids() # Fail before execution when procfs tracking is absent. + + action_pid = os.fork() + if action_pid == 0: + _exec(command, control_fd) + + poller = select.poll() + poller.register(control_fd, select.POLLIN | select.POLLHUP | select.POLLERR) + action_status: int | None = None + parent_closed = False + + while action_status is None and not parent_closed: + waited_pid, wait_status = _waitpid(action_pid, os.WNOHANG) + if waited_pid == action_pid: + action_status = wait_status + break + if poller.poll(_PARENT_POLL_INTERVAL_MS): + parent_closed = os.read(control_fd, 1) == b"" + + _kill_and_reap_children() + if parent_closed: + return 128 + signal.SIGTERM + if action_status is None: + raise RuntimeError("Supervised action exited without a wait status") + return _exit_code(action_status) + except BaseException as e: + with suppress(BaseException): + _kill_and_reap_children() + message = f"Direct action supervisor failed: {type(e).__name__}: {e}\n".encode() + with suppress(OSError): + os.write(2, message) + return 1 + finally: + with suppress(OSError): + os.close(control_fd) + + +def supervise(command: Sequence[str]) -> int: + """Run one command and return only after all of its descendants are reaped.""" + if sys.platform != "linux": + raise RuntimeError("The direct action process supervisor requires Linux") + if not command: + raise ValueError("A supervised command is required") + + control_read_fd, control_write_fd = os.pipe() + writer_open = True + + def request_cleanup(_signal: int, _frame: FrameType | None) -> None: + nonlocal writer_open + if writer_open: + with suppress(OSError): + os.close(control_write_fd) + writer_open = False + + signal.signal(signal.SIGTERM, request_cleanup) + signal.signal(signal.SIGINT, request_cleanup) + + monitor_pid = os.fork() + if monitor_pid == 0: + signal.signal(signal.SIGTERM, signal.SIG_DFL) + signal.signal(signal.SIGINT, signal.SIG_DFL) + os.close(control_write_fd) + os._exit(_run_monitor(control_read_fd, command)) + + os.close(control_read_fd) + try: + _, monitor_status = _waitpid(monitor_pid) + return _exit_code(monitor_status) + finally: + request_cleanup(signal.SIGTERM, None) + + +def main() -> int: + """CLI entry point.""" + return supervise(sys.argv[1:]) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tracecat/sandbox/utils.py b/tracecat/sandbox/utils.py index 680592fae5..a9b8dc0032 100644 --- a/tracecat/sandbox/utils.py +++ b/tracecat/sandbox/utils.py @@ -11,6 +11,7 @@ import shutil import signal import subprocess +from collections.abc import Awaitable, Callable from contextlib import suppress from pathlib import Path @@ -41,14 +42,26 @@ async def terminate_process_group(process: asyncio.subprocess.Process) -> None: await process.wait() +async def terminate_supervised_process(process: asyncio.subprocess.Process) -> None: + """Request descendant cleanup from a direct-action process supervisor.""" + if process.returncode is None: + with suppress(ProcessLookupError): + os.kill(process.pid, signal.SIGTERM) + # The supervisor exits only after its detached subreaper has killed and + # reaped the action's complete descendant tree. Do not impose a shorter + # wait that could release registry leases while cleanup is still active. + await process.wait() + + async def _finish_process_group_cleanup( process: asyncio.subprocess.Process, communicate_task: asyncio.Task[tuple[bytes | None, bytes | None]], - termination_task: asyncio.Task[None] | None, + termination_task: asyncio.Future[None] | None, + terminate: Callable[[asyncio.subprocess.Process], Awaitable[None]], ) -> None: """Finish process termination and consume the communication task.""" if termination_task is None: - termination_task = asyncio.create_task(terminate_process_group(process)) + termination_task = asyncio.ensure_future(terminate(process)) try: await termination_task finally: @@ -86,6 +99,7 @@ async def communicate_process_group( *, input: bytes | None = None, # noqa: A002 timeout: float | None = None, + terminate: Callable[[asyncio.subprocess.Process], Awaitable[None]] | None = None, ) -> tuple[bytes, bytes]: """Communicate with a process while containing its process group. @@ -97,14 +111,15 @@ async def communicate_process_group( task and is rejoined through repeated cancellation so callers cannot release resources while the process group is still alive. """ + terminator = terminate or terminate_process_group communicate_task = asyncio.create_task(process.communicate(input=input)) - termination_task: asyncio.Task[None] | None = None + termination_task: asyncio.Future[None] | None = None operation_error: BaseException | None = None try: async with asyncio.timeout(timeout): while process.returncode is None: await asyncio.sleep(_PROCESS_EXIT_POLL_INTERVAL_SECONDS) - termination_task = asyncio.create_task(terminate_process_group(process)) + termination_task = asyncio.ensure_future(terminator(process)) await asyncio.shield(termination_task) stdout, stderr = await communicate_task except BaseException as e: @@ -116,6 +131,7 @@ async def communicate_process_group( process, communicate_task, termination_task, + terminator, ) ) try: From b228c8b32b236aca04f3bedf66db79681cef5a86 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 21:45:45 -0400 Subject: [PATCH 106/161] fix(executor): sanitize supervisor failures --- tracecat/executor/process_supervisor.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/tracecat/executor/process_supervisor.py b/tracecat/executor/process_supervisor.py index a53abbcfd1..7680e0f30c 100644 --- a/tracecat/executor/process_supervisor.py +++ b/tracecat/executor/process_supervisor.py @@ -80,10 +80,9 @@ def _exec(command: Sequence[str], control_fd: int) -> None: os.close(control_fd) try: os.execvpe(command[0], list(command), os.environ) - except OSError as e: - message = f"Failed to execute supervised action: {e}\n".encode() + except OSError: with suppress(OSError): - os.write(2, message) + os.write(2, b"Failed to execute supervised action\n") os._exit(127) @@ -117,12 +116,11 @@ def _run_monitor(control_fd: int, command: Sequence[str]) -> int: if action_status is None: raise RuntimeError("Supervised action exited without a wait status") return _exit_code(action_status) - except BaseException as e: + except BaseException: with suppress(BaseException): _kill_and_reap_children() - message = f"Direct action supervisor failed: {type(e).__name__}: {e}\n".encode() with suppress(OSError): - os.write(2, message) + os.write(2, b"Direct action supervisor failed\n") return 1 finally: with suppress(OSError): From c0ba21813ca2fce0f7806990feb73d9cbc488489 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 21:46:55 -0400 Subject: [PATCH 107/161] test(executor): exercise supervised cancellation --- tests/unit/test_action_runner.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index 1047ea8c1f..d0767c4c57 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -33,7 +33,6 @@ from tracecat.executor.secret_preprocessors import SecretEnvProjection from tracecat.identifiers.workflow import WorkflowUUID from tracecat.registry.lock.types import RegistryLock -from tracecat.sandbox import utils as sandbox_utils from tracecat.sandbox.types import SandboxResult @@ -762,7 +761,7 @@ async def test_cancelled_action_reaps_child_before_releasing_mounted_artifact( ) real_create_subprocess_exec = asyncio.create_subprocess_exec - real_terminate_process_group = sandbox_utils.terminate_process_group + real_terminate_supervised_process = action_runner.terminate_supervised_process process_started = asyncio.Event() termination_started = asyncio.Event() finish_termination = asyncio.Event() @@ -780,7 +779,7 @@ async def controlled_termination( ) -> None: termination_started.set() await finish_termination.wait() - await real_terminate_process_group(requested_process) + await real_terminate_supervised_process(requested_process) async def release_mount(mount_dir: Path) -> bool: reaped_before_unmount.append( @@ -790,6 +789,7 @@ async def release_mount(mount_dir: Path) -> bool: return True with ( + patch.object(action_runner.sys, "platform", "linux"), patch( "tracecat.executor.registry_artifact_mounts.is_mount", lambda path: path in mounted, @@ -804,8 +804,8 @@ async def release_mount(mount_dir: Path) -> bool: side_effect=capture_subprocess, ), patch.object( - sandbox_utils, - "terminate_process_group", + action_runner, + "terminate_supervised_process", side_effect=controlled_termination, ), patch.object(cache, "_unmount", side_effect=release_mount), From f1c4151f52d30479aec4365671d1f1e1d062f2f3 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 21:49:23 -0400 Subject: [PATCH 108/161] fix(executor): redact registry artifact log URIs --- .../test_registry_artifact_resolution.py | 11 +++++++ tracecat/executor/backends/test.py | 3 +- .../registry_artifact_materialization.py | 31 +++++++++++++------ tracecat/executor/registry_artifacts.py | 14 +++++---- 4 files changed, 42 insertions(+), 17 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py index f4861d3a27..21a75e9a37 100644 --- a/tests/unit/executor/test_registry_artifact_resolution.py +++ b/tests/unit/executor/test_registry_artifact_resolution.py @@ -15,6 +15,7 @@ RegistryArtifactFormat, SquashfsArtifact, TarballArtifact, + _artifact_uri_for_logging, bundled_builtin_registry_uri, compute_registry_artifact_cache_key, ) @@ -85,6 +86,16 @@ def test_compute_registry_artifact_cache_key_empty(self): """Test that empty URI returns the base cache key.""" assert compute_registry_artifact_cache_key("") == "base" + def test_artifact_uri_for_logging_removes_credentials_and_signature(self): + uri = ( + "s3://access:secret@bucket/path/site-packages.squashfs" + "?X-Amz-Signature=signed-secret#fragment" + ) + + assert ( + _artifact_uri_for_logging(uri) == "s3://bucket/path/site-packages.squashfs" + ) + @pytest.mark.anyio async def test_download_artifact_uses_blob_download_file_to_path( self, temp_cache_dir diff --git a/tracecat/executor/backends/test.py b/tracecat/executor/backends/test.py index b8cfd23176..6fe5a0390a 100644 --- a/tracecat/executor/backends/test.py +++ b/tracecat/executor/backends/test.py @@ -40,6 +40,7 @@ from tracecat.executor.backends.base import ExecutorBackend from tracecat.executor.backends.registry_helpers import get_registry_artifact_uris from tracecat.executor.registry_artifact_materialization import ( + _artifact_uri_for_logging, _is_cache_entry_uri, _run_blocking_rejoin_on_cancel, ) @@ -326,7 +327,7 @@ async def _lease_registry_artifacts( except Exception as e: logger.warning( "Failed to materialize artifact for test execution", - artifact_uri=artifact_uri, + artifact_uri=_artifact_uri_for_logging(artifact_uri), error=str(e), ) continue diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 02fbfde82d..6ce1ae419b 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -15,6 +15,7 @@ from dataclasses import dataclass from enum import StrEnum from pathlib import Path, PurePosixPath +from urllib.parse import urlsplit, urlunsplit import httpx import tracecat_registry @@ -28,12 +29,22 @@ from tracecat.storage import blob __all__ = [ + "_artifact_uri_for_logging", "_is_cache_entry_uri", "_squashfs_sidecar_uri", "_tarball_uri_for_squashfs", ] +def _artifact_uri_for_logging(artifact_uri: str) -> str: + """Remove credentials, query parameters, and fragments from an artifact URI.""" + parsed = urlsplit(artifact_uri) + if not parsed.scheme or not parsed.hostname: + return "" + hostname = f"[{parsed.hostname}]" if ":" in parsed.hostname else parsed.hostname + return urlunsplit((parsed.scheme, hostname, parsed.path, "", "")) + + class RegistryArtifactFormat(StrEnum): """Executor-supported registry artifact encodings.""" @@ -273,7 +284,7 @@ def discard_failed_materialization( logger.warning( "Cannot determine whether failed SquashFS candidate is reusable", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), error=str(e), ) return @@ -302,7 +313,7 @@ async def materialize( logger.warning( "Failed to mount SquashFS registry artifact, trying extraction", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, error=str(e), ) @@ -368,7 +379,7 @@ async def mount( logger.info( "Materializing SquashFS registry artifact", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, ) start_time = time.monotonic() @@ -382,7 +393,7 @@ async def mount( logger.info( "SquashFS registry artifact mounted", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, download_ms=f"{download_elapsed:.1f}", mount_ms=f"{mount_elapsed:.1f}", @@ -404,7 +415,7 @@ async def extract( logger.info( "Extracting SquashFS registry artifact", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, ) start_time = time.monotonic() @@ -429,7 +440,7 @@ async def extract( logger.info( "SquashFS registry artifact extracted", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, download_ms=f"{download_elapsed:.1f}", extract_ms=f"{extract_elapsed:.1f}", @@ -440,7 +451,7 @@ async def extract( logger.debug( "SquashFS already extracted by another process", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, ) else: @@ -573,7 +584,7 @@ async def materialize( logger.info( "Materializing tarball registry artifact", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, ) start_time = time.monotonic() @@ -609,7 +620,7 @@ async def materialize( logger.info( "Tarball extracted and cached", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, download_ms=f"{download_elapsed:.1f}", extract_ms=f"{extract_elapsed:.1f}", @@ -620,7 +631,7 @@ async def materialize( logger.debug( "Tarball already extracted by another process", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, ) else: diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 53d3252fe4..682e3eb725 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -31,6 +31,7 @@ SquashfsMountCommandError, TarballArtifact, _artifact_format, + _artifact_uri_for_logging, _bundled_builtin_registry_import_paths, _bundled_builtin_registry_version, _download_s3_artifact, @@ -81,6 +82,7 @@ "SquashfsMountCommandError", "TarballArtifact", "_artifact_format", + "_artifact_uri_for_logging", "_allocated_stat_size", "_bundled_builtin_registry_import_paths", "_bundled_builtin_registry_version", @@ -273,7 +275,7 @@ async def _materialize_candidates( logger.info( "Trying registry artifact candidate", cache_key=cache_key, - artifact_uri=artifact.uri, + artifact_uri=_artifact_uri_for_logging(artifact.uri), artifact_format=artifact.format.value, candidate=index + 1, candidates=len(candidates), @@ -297,7 +299,7 @@ async def _materialize_candidates( logger.warning( "Failed to materialize registry artifact candidate, trying fallback", cache_key=cache_key, - artifact_uri=artifact.uri, + artifact_uri=_artifact_uri_for_logging(artifact.uri), artifact_format=artifact.format.value, error=str(e), ) @@ -439,16 +441,16 @@ async def _sidecar_exists( if await blob.file_exists(key=key, bucket=bucket): logger.debug( "Using registry artifact sidecar", - artifact_uri=base_uri, - sidecar_uri=sidecar_uri, + artifact_uri=_artifact_uri_for_logging(base_uri), + sidecar_uri=_artifact_uri_for_logging(sidecar_uri), artifact_format=artifact_format.value, ) return True except Exception as e: logger.warning( "Failed to check for registry artifact sidecar, falling back", - artifact_uri=base_uri, - sidecar_uri=sidecar_uri, + artifact_uri=_artifact_uri_for_logging(base_uri), + sidecar_uri=_artifact_uri_for_logging(sidecar_uri), artifact_format=artifact_format.value, error=str(e), ) From edef3cd1056ce5afcc32868b40e197659dc14194 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 21:50:36 -0400 Subject: [PATCH 109/161] test(executor): clean failed supervisor probes --- .../unit/executor/test_process_supervisor.py | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/tests/unit/executor/test_process_supervisor.py b/tests/unit/executor/test_process_supervisor.py index 87fd411165..47677c8540 100644 --- a/tests/unit/executor/test_process_supervisor.py +++ b/tests/unit/executor/test_process_supervisor.py @@ -94,12 +94,24 @@ async def test_supervisor_reaps_detached_descendant_after_success( tmp_path: Path, ) -> None: process, pid_file = await _spawn_supervised_action(tmp_path, mode="success") - stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=5) + tracked_pids: tuple[int, ...] = () + try: + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=5) - assert process.returncode == 0, stderr.decode() - assert stdout == b"" - _, detached_pid = (int(pid) for pid in pid_file.read_text().split()) - assert not _process_is_running(detached_pid) + assert process.returncode == 0, stderr.decode() + assert stdout == b"" + tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) + _, detached_pid = tracked_pids + assert not _process_is_running(detached_pid) + finally: + if not tracked_pids and pid_file.exists(): + tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) + for pid in tracked_pids: + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() @pytest.mark.anyio From d6af9138c70c0fa11ec08db7ff0a278e6f374c81 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 21:54:00 -0400 Subject: [PATCH 110/161] fix(executor): isolate direct action supervisor --- tests/unit/test_action_runner.py | 7 ++++--- tracecat/executor/action_runner.py | 3 +++ 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index d0767c4c57..c8d7a3a3b6 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -535,7 +535,7 @@ async def test_execute_action_disables_new_privileges_for_direct_subprocess( mock_process_group_communication: AsyncMock, monkeypatch: pytest.MonkeyPatch, ): - """Test direct subprocess execution disables new Linux privileges.""" + """Test Linux direct execution isolates and drops supervisor privileges.""" runner = ActionRunner(cache_dir=temp_cache_dir) base_dir = temp_cache_dir / "base" base_dir.mkdir() @@ -579,10 +579,11 @@ async def create_subprocess_exec_side_effect(*args, **kwargs): # noqa: ARG001 "--inh-caps=-all", "--ambient-caps=-all", ] + assert captured_args[-5] == action_runner.sys.executable + assert captured_args[-4] == "-I" + assert captured_args[-3].endswith("process_supervisor.py") assert captured_args[-2] == action_runner.sys.executable assert captured_args[-1].endswith("minimal_runner.py") - assert captured_args[-4] == action_runner.sys.executable - assert captured_args[-3].endswith("process_supervisor.py") communication_call = mock_process_group_communication.await_args assert communication_call is not None assert ( diff --git a/tracecat/executor/action_runner.py b/tracecat/executor/action_runner.py index a285090588..04ce830496 100644 --- a/tracecat/executor/action_runner.py +++ b/tracecat/executor/action_runner.py @@ -128,6 +128,9 @@ def _direct_subprocess_command(minimal_runner_path: Path) -> list[str]: "--inh-caps=-all", "--ambient-caps=-all", sys.executable, + # Keep registry-controlled PYTHONPATH out of the supervisor interpreter. + # Isolated mode leaves the environment intact for the nested action. + "-I", str(supervisor_path), *runner_command, ] From a4a78a77a42cca731fa4d12fb8228f94a50a936d Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 22:03:54 -0400 Subject: [PATCH 111/161] fix(executor): redact missing artifact errors --- .../unit/executor/test_registry_artifact_resolution.py | 10 +++++++++- tracecat/executor/registry_artifact_materialization.py | 2 +- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py index 21a75e9a37..0d5e9095e8 100644 --- a/tests/unit/executor/test_registry_artifact_resolution.py +++ b/tests/unit/executor/test_registry_artifact_resolution.py @@ -201,8 +201,12 @@ async def test_download_artifact_normalizes_missing_objects_to_http_404( ): """Preserve the missing-artifact error contract from presigned downloads.""" cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = ( + "s3://access:secret@bucket/path/site-packages.tar.gz" + "?X-Amz-Signature=signed-secret#fragment" + ) artifact = TarballArtifact( - uri="s3://bucket/path/site-packages.tar.gz", + uri=artifact_uri, cache_key="missing-test", ) ctx = cache._context_for(artifact.cache_key) @@ -218,6 +222,10 @@ async def test_download_artifact_normalizes_missing_objects_to_http_404( assert exc_info.value.response.status_code == 404 assert isinstance(exc_info.value.__cause__, FileNotFoundError) + assert str(exc_info.value) == ( + "Registry artifact not found: s3://bucket/path/site-packages.tar.gz" + ) + assert "secret" not in str(exc_info.value) @pytest.mark.anyio async def test_artifact_candidates_prefer_squashfs_sidecar(self, temp_cache_dir): diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 6ce1ae419b..fc6731bba5 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -717,7 +717,7 @@ async def _download_s3_artifact( request = httpx.Request("GET", artifact_uri) response = httpx.Response(status_code=404, request=request) raise httpx.HTTPStatusError( - f"Registry artifact not found: {artifact_uri}", + f"Registry artifact not found: {_artifact_uri_for_logging(artifact_uri)}", request=request, response=response, ) from e From 6afe59a3cd84b5c9a9ab4a575ad04068b6b610ca Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 22:04:48 -0400 Subject: [PATCH 112/161] fix(executor): redact cache cleanup failures --- tests/unit/executor/test_registry_artifact_leases.py | 12 ++++++------ tracecat/executor/registry_artifacts.py | 4 ++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_leases.py b/tests/unit/executor/test_registry_artifact_leases.py index 117e9a231a..7e8c1c400b 100644 --- a/tests/unit/executor/test_registry_artifact_leases.py +++ b/tests/unit/executor/test_registry_artifact_leases.py @@ -513,23 +513,23 @@ async def fail_cleanup( converge: bool, ) -> None: del idle_keys, converge - raise RuntimeError("cleanup failed") + raise RuntimeError( + "cleanup failed for s3://access:secret@bucket/path?signature=secret" + ) with ( patch.object(cache, "_finish_lease_cleanup", fail_cleanup), - patch( - "tracecat.executor.registry_artifacts.logger.exception" - ) as log_exception, + patch("tracecat.executor.registry_artifacts.logger.error") as log_error, ): async with cache.lease([artifact_uri]) as registry_paths: result = registry_paths assert result == [target_dir] assert cache._refcount(cache_key) == 0 - log_exception.assert_called_once_with( + log_error.assert_called_once_with( "Registry artifact lease cleanup failed; preserving caller outcome", cache_dir=str(temp_cache_dir), - error="cleanup failed", + error_type="RuntimeError", ) @pytest.mark.anyio diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 682e3eb725..984414f5a4 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -173,10 +173,10 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat raise pending_cancellation = e except Exception as e: - logger.exception( + logger.error( "Registry artifact lease cleanup failed; preserving caller outcome", cache_dir=str(self.cache_dir), - error=str(e), + error_type=type(e).__name__, ) break From 0f7a373225de0eb476f44f794f5d352188bddec6 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 22:05:48 -0400 Subject: [PATCH 113/161] test(executor): cover supervised action failures --- tests/unit/executor/test_process_supervisor.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/unit/executor/test_process_supervisor.py b/tests/unit/executor/test_process_supervisor.py index 47677c8540..069dc8461c 100644 --- a/tests/unit/executor/test_process_supervisor.py +++ b/tests/unit/executor/test_process_supervisor.py @@ -60,6 +60,8 @@ def _write_action_script(path: Path) -> None: start_new_session=True, ) pid_file.write_text(f"{os.getpid()} {child.pid}") +if mode == "failure": + raise SystemExit(23) if mode == "block": time.sleep(30) """.lstrip() @@ -90,15 +92,21 @@ async def _spawn_supervised_action( @pytest.mark.anyio -async def test_supervisor_reaps_detached_descendant_after_success( +@pytest.mark.parametrize( + ("mode", "expected_returncode"), + [("success", 0), ("failure", 23)], +) +async def test_supervisor_propagates_exit_and_reaps_detached_descendant( tmp_path: Path, + mode: str, + expected_returncode: int, ) -> None: - process, pid_file = await _spawn_supervised_action(tmp_path, mode="success") + process, pid_file = await _spawn_supervised_action(tmp_path, mode=mode) tracked_pids: tuple[int, ...] = () try: stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=5) - assert process.returncode == 0, stderr.decode() + assert process.returncode == expected_returncode, stderr.decode() assert stdout == b"" tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) _, detached_pid = tracked_pids From af91e3329d25cfdaea3791bd721d45953bbe3538 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 22:14:38 -0400 Subject: [PATCH 114/161] fix(executor): make artifact URI redaction total --- tests/unit/executor/test_registry_artifact_resolution.py | 5 +++++ tracecat/executor/registry_artifact_materialization.py | 5 ++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py index 0d5e9095e8..9ec61e281f 100644 --- a/tests/unit/executor/test_registry_artifact_resolution.py +++ b/tests/unit/executor/test_registry_artifact_resolution.py @@ -96,6 +96,11 @@ def test_artifact_uri_for_logging_removes_credentials_and_signature(self): _artifact_uri_for_logging(uri) == "s3://bucket/path/site-packages.squashfs" ) + def test_artifact_uri_for_logging_redacts_malformed_uri(self): + assert _artifact_uri_for_logging("s3://[malformed") == ( + "" + ) + @pytest.mark.anyio async def test_download_artifact_uses_blob_download_file_to_path( self, temp_cache_dir diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index fc6731bba5..9ac9c3c86e 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -38,7 +38,10 @@ def _artifact_uri_for_logging(artifact_uri: str) -> str: """Remove credentials, query parameters, and fragments from an artifact URI.""" - parsed = urlsplit(artifact_uri) + try: + parsed = urlsplit(artifact_uri) + except ValueError: + return "" if not parsed.scheme or not parsed.hostname: return "" hostname = f"[{parsed.hostname}]" if ":" in parsed.hostname else parsed.hostname From 3f7af1be27fecdb4e78a9e04013cf8886d3c6e0f Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 22:16:08 -0400 Subject: [PATCH 115/161] fix(executor): avoid staging retry collisions --- .../test_registry_artifact_materialization.py | 21 +++++++++++++++++++ .../registry_artifact_materialization.py | 10 +++++++-- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 8491958aed..44f45f4c2b 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -38,6 +38,27 @@ class TestRegistryArtifactMaterialization: """Materialize and reuse executor-local artifact formats.""" + def test_temp_path_avoids_deferred_staging_collision( + self, temp_cache_dir: Path + ) -> None: + cache = RegistryArtifactCache(temp_cache_dir) + artifact = TarballArtifact( + uri="s3://bucket/path/site-packages.tar.gz", + cache_key="staging-collision", + ) + ctx = cache._context_for(artifact.cache_key) + + with patch( + "tracecat.executor.registry_artifact_materialization.secrets.token_hex", + side_effect=["deferred", "deferred", "retry"], + ): + deferred_path = artifact._temp_path(ctx, ".tmp") + deferred_path.mkdir() + retry_path = artifact._temp_path(ctx, ".tmp") + + assert retry_path != deferred_path + assert retry_path.name.endswith(".retry.tmp") + @pytest.mark.anyio async def test_same_key_cold_fan_in_materializes_and_enforces_once( self, temp_cache_dir diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 9ac9c3c86e..b8489e1f3d 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -6,6 +6,7 @@ import contextlib import hashlib import os +import secrets import shutil import sysconfig import tarfile @@ -148,9 +149,14 @@ def _temp_path( ctx: RegistryArtifactMaterializationContext, suffix: str, ) -> Path: - unique_id = id(asyncio.current_task()) ctx.staging_dir.mkdir(parents=True, exist_ok=True) - return ctx.staging_dir / f"{self.cache_key}.{os.getpid()}.{unique_id}{suffix}" + while True: + attempt_id = secrets.token_hex(8) + candidate = ( + ctx.staging_dir / f"{self.cache_key}.{os.getpid()}.{attempt_id}{suffix}" + ) + if not os.path.lexists(candidate): + return candidate async def _rejoin_future_on_cancel[T](future: asyncio.Future[T]) -> T: From 3200fef135bc369c5dbda70afe0f783b25a6475f Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 22:17:51 -0400 Subject: [PATCH 116/161] fix(executor): omit artifact identifiers from logs --- .../test_registry_artifact_resolution.py | 17 ++++++++++------- .../registry_artifact_materialization.py | 5 ++--- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py index 9ec61e281f..27d3551817 100644 --- a/tests/unit/executor/test_registry_artifact_resolution.py +++ b/tests/unit/executor/test_registry_artifact_resolution.py @@ -86,14 +86,19 @@ def test_compute_registry_artifact_cache_key_empty(self): """Test that empty URI returns the base cache key.""" assert compute_registry_artifact_cache_key("") == "base" - def test_artifact_uri_for_logging_removes_credentials_and_signature(self): + def test_artifact_uri_for_logging_removes_identifiers_and_credentials(self): uri = ( - "s3://access:secret@bucket/path/site-packages.squashfs" + "s3://access:secret@bucket/org-id/repository-origin/1.2.3/" + "site-packages.squashfs" "?X-Amz-Signature=signed-secret#fragment" ) - assert ( - _artifact_uri_for_logging(uri) == "s3://bucket/path/site-packages.squashfs" + logged_uri = _artifact_uri_for_logging(uri) + + assert logged_uri == "s3://" + assert all( + value not in logged_uri + for value in ("secret", "bucket", "org-id", "repository-origin", "1.2.3") ) def test_artifact_uri_for_logging_redacts_malformed_uri(self): @@ -227,9 +232,7 @@ async def test_download_artifact_normalizes_missing_objects_to_http_404( assert exc_info.value.response.status_code == 404 assert isinstance(exc_info.value.__cause__, FileNotFoundError) - assert str(exc_info.value) == ( - "Registry artifact not found: s3://bucket/path/site-packages.tar.gz" - ) + assert str(exc_info.value) == "Registry artifact not found: s3://" assert "secret" not in str(exc_info.value) @pytest.mark.anyio diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index b8489e1f3d..f2acfa6a1a 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -38,15 +38,14 @@ def _artifact_uri_for_logging(artifact_uri: str) -> str: - """Remove credentials, query parameters, and fragments from an artifact URI.""" + """Retain only the non-sensitive scheme of an artifact URI.""" try: parsed = urlsplit(artifact_uri) except ValueError: return "" if not parsed.scheme or not parsed.hostname: return "" - hostname = f"[{parsed.hostname}]" if ":" in parsed.hostname else parsed.hostname - return urlunsplit((parsed.scheme, hostname, parsed.path, "", "")) + return urlunsplit((parsed.scheme, "", "", "", "")) class RegistryArtifactFormat(StrEnum): From b37e61b088e454a3e5a575c596033fc2d68d86c2 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 22:23:53 -0400 Subject: [PATCH 117/161] fix(executor): reject malformed cache hits --- .../test_registry_artifact_materialization.py | 21 +++++++++++++++++++ .../registry_artifact_materialization.py | 4 ++-- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 44f45f4c2b..da91f141f6 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -59,6 +59,27 @@ def test_temp_path_avoids_deferred_staging_collision( assert retry_path != deferred_path assert retry_path.name.endswith(".retry.tmp") + def test_cached_path_rejects_non_directory_extractions( + self, temp_cache_dir: Path + ) -> None: + cache = RegistryArtifactCache(temp_cache_dir) + ctx = cache._context_for("malformed-extractions") + ctx.paths.entry_dir.mkdir(parents=True) + ctx.paths.squashfs_extract_dir.write_bytes(b"not a directory") + ctx.paths.tarball_target_dir.write_bytes(b"not a directory") + + squashfs = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key=ctx.cache_key, + ) + tarball = TarballArtifact( + uri="s3://bucket/path/site-packages.tar.gz", + cache_key=ctx.cache_key, + ) + + assert squashfs.cached_path(ctx) is None + assert tarball.cached_path(ctx) is None + @pytest.mark.anyio async def test_same_key_cold_fan_in_materializes_and_enforces_once( self, temp_cache_dir diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index f2acfa6a1a..8d6311b681 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -273,7 +273,7 @@ def cached_path( cache_key=ctx.cache_key, ) return [ctx.paths.squashfs_mount_dir] - if ctx.paths.squashfs_extract_dir.exists(): + if ctx.paths.squashfs_extract_dir.is_dir(): logger.debug( "Using cached SquashFS registry extraction", cache_key=ctx.cache_key, @@ -577,7 +577,7 @@ def format(self) -> RegistryArtifactFormat: def cached_path( self, ctx: RegistryArtifactMaterializationContext ) -> list[Path] | None: - if ctx.paths.tarball_target_dir.exists(): + if ctx.paths.tarball_target_dir.is_dir(): logger.debug( "Using cached tarball extraction", cache_key=ctx.cache_key, From 6ee4dfb2bcd0b0ee900c1037d2df267d7fe85853 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 22:29:09 -0400 Subject: [PATCH 118/161] fix(executor): reclaim malformed cache targets --- .../test_registry_artifact_materialization.py | 2 + .../registry_artifact_materialization.py | 38 ++++++++++++++++--- 2 files changed, 35 insertions(+), 5 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index da91f141f6..3262c8dd2d 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -79,6 +79,8 @@ def test_cached_path_rejects_non_directory_extractions( assert squashfs.cached_path(ctx) is None assert tarball.cached_path(ctx) is None + assert not ctx.paths.squashfs_extract_dir.exists() + assert not ctx.paths.tarball_target_dir.exists() @pytest.mark.anyio async def test_same_key_cold_fan_in_materializes_and_enforces_once( diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 8d6311b681..33cdb2d282 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -228,6 +228,19 @@ def _remove_file_or_defer( ) +def _is_reusable_extraction_dir( + path: Path, + *, + defer_cleanup: Callable[[Path], None], +) -> bool: + """Accept canonical directories and reclaim malformed file or symlink targets.""" + if path.is_dir() and not path.is_symlink(): + return True + if os.path.lexists(path): + _remove_file_or_defer(path, defer_cleanup=defer_cleanup) + return False + + @dataclass(frozen=True, slots=True) class BuiltinArtifact(RegistryArtifact): """Current builtin registry package already installed in the executor image.""" @@ -273,7 +286,10 @@ def cached_path( cache_key=ctx.cache_key, ) return [ctx.paths.squashfs_mount_dir] - if ctx.paths.squashfs_extract_dir.is_dir(): + if _is_reusable_extraction_dir( + ctx.paths.squashfs_extract_dir, + defer_cleanup=ctx.defer_cleanup, + ): logger.debug( "Using cached SquashFS registry extraction", cache_key=ctx.cache_key, @@ -415,7 +431,10 @@ async def extract( image_path: Path, ) -> Path: target_dir = ctx.paths.squashfs_extract_dir - if target_dir.exists(): + if _is_reusable_extraction_dir( + target_dir, + defer_cleanup=ctx.defer_cleanup, + ): return target_dir ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) @@ -455,7 +474,10 @@ async def extract( total_ms=f"{total_elapsed:.1f}", ) except OSError: - if target_dir.exists(): + if _is_reusable_extraction_dir( + target_dir, + defer_cleanup=ctx.defer_cleanup, + ): logger.debug( "SquashFS already extracted by another process", cache_key=ctx.cache_key, @@ -577,7 +599,10 @@ def format(self) -> RegistryArtifactFormat: def cached_path( self, ctx: RegistryArtifactMaterializationContext ) -> list[Path] | None: - if ctx.paths.tarball_target_dir.is_dir(): + if _is_reusable_extraction_dir( + ctx.paths.tarball_target_dir, + defer_cleanup=ctx.defer_cleanup, + ): logger.debug( "Using cached tarball extraction", cache_key=ctx.cache_key, @@ -635,7 +660,10 @@ async def materialize( total_ms=f"{total_elapsed:.1f}", ) except OSError: - if target_dir.exists(): + if _is_reusable_extraction_dir( + target_dir, + defer_cleanup=ctx.defer_cleanup, + ): logger.debug( "Tarball already extracted by another process", cache_key=ctx.cache_key, From 6b06b17ffe8caa918a50a56a918edb41f6305b62 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 22:38:56 -0400 Subject: [PATCH 119/161] fix(executor): recover malformed squashfs images --- .../test_registry_artifact_materialization.py | 34 ++++++++++++++ .../test_registry_artifact_resolution.py | 27 +++++++++++ .../test_registry_artifact_startup.py | 14 ++++++ .../registry_artifact_materialization.py | 45 +++++++++++++++++-- .../executor/registry_artifact_storage.py | 8 ++-- tracecat/executor/registry_artifacts.py | 3 +- 6 files changed, 123 insertions(+), 8 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 3262c8dd2d..3b7410f968 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -82,6 +82,40 @@ def test_cached_path_rejects_non_directory_extractions( assert not ctx.paths.squashfs_extract_dir.exists() assert not ctx.paths.tarball_target_dir.exists() + @pytest.mark.anyio + async def test_download_reclaims_malformed_squashfs_image( + self, temp_cache_dir: Path + ) -> None: + cache = RegistryArtifactCache(temp_cache_dir) + artifact = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key="malformed-image", + ) + ctx = cache._context_for(artifact.cache_key) + image_path = ctx.paths.squashfs_image_path + image_path.mkdir(parents=True) + (image_path / "stale").write_bytes(b"stale") + + async def download_image( + artifact_uri: str, + output_path: Path, + *, + admission: object, + defer_cleanup: object, + ) -> None: + del artifact_uri, admission, defer_cleanup + output_path.write_bytes(b"fresh image") + + with patch( + "tracecat.executor.registry_artifact_materialization._download_s3_artifact", + download_image, + ): + await artifact.download(ctx, image_path) + + assert image_path.is_file() + assert not image_path.is_symlink() + assert image_path.read_bytes() == b"fresh image" + @pytest.mark.anyio async def test_same_key_cold_fan_in_materializes_and_enforces_once( self, temp_cache_dir diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py index 27d3551817..50dd9c9867 100644 --- a/tests/unit/executor/test_registry_artifact_resolution.py +++ b/tests/unit/executor/test_registry_artifact_resolution.py @@ -267,6 +267,33 @@ async def test_artifact_candidates_prefer_squashfs_sidecar(self, temp_cache_dir) bucket="bucket", ) + @pytest.mark.anyio + async def test_artifact_candidates_ignore_malformed_local_sidecar( + self, temp_cache_dir: Path + ) -> None: + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/site-packages.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + ctx = cache._context_for(cache_key) + ctx.paths.squashfs_image_path.mkdir(parents=True) + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + return_value=False, + ) as file_exists, + patch.object(cache, "_can_try_squashfs", return_value=True), + ): + candidates = await cache._artifact_candidates(ctx, artifact_uri) + + assert len(candidates) == 1 + assert isinstance(candidates[0], TarballArtifact) + file_exists.assert_awaited_once_with( + key="path/site-packages.squashfs", + bucket="bucket", + ) + @pytest.mark.anyio async def test_artifact_candidates_direct_squashfs_include_gzip_fallback( self, temp_cache_dir diff --git a/tests/unit/executor/test_registry_artifact_startup.py b/tests/unit/executor/test_registry_artifact_startup.py index 4c48e1ea4c..f90063347a 100644 --- a/tests/unit/executor/test_registry_artifact_startup.py +++ b/tests/unit/executor/test_registry_artifact_startup.py @@ -80,6 +80,20 @@ async def test_sweep_removes_incomplete_shell_before_lru_trimming( assert not any(cache.trash_dir.iterdir()) assert cache._budget_dirty is False + @pytest.mark.anyio + async def test_sweep_retires_malformed_squashfs_image( + self, temp_cache_dir: Path + ) -> None: + cache = RegistryArtifactCache(temp_cache_dir) + malformed = cache._paths_for("malformed-image") + malformed.squashfs_image_path.mkdir(parents=True) + (malformed.squashfs_image_path / "stale").write_bytes(b"stale") + + await cache.ensure_swept() + + assert not malformed.entry_dir.exists() + assert not any(cache.trash_dir.iterdir()) + @pytest.mark.anyio async def test_sweep_tolerates_missing_cache_dir(self, temp_cache_dir): """A cache directory that does not exist yet is a no-op.""" diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 33cdb2d282..d0a0df0fd2 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -228,13 +228,44 @@ def _remove_file_or_defer( ) +def _is_reusable_cache_file(path: Path) -> bool: + """Return whether a canonical cache file is regular and not a symlink.""" + return path.is_file() and not path.is_symlink() + + +def _is_reusable_cache_directory(path: Path) -> bool: + """Return whether a canonical cache directory is real and not a symlink.""" + return path.is_dir() and not path.is_symlink() + + +async def _reuse_or_reclaim_squashfs_image( + path: Path, + *, + defer_cleanup: Callable[[Path], None], +) -> bool: + """Reuse a safe image or reclaim a malformed target before downloading.""" + if _is_reusable_cache_file(path): + return True + if not os.path.lexists(path): + return False + + if _is_reusable_cache_directory(path): + await _remove_tree_rejoin_on_cancel(path, defer_cleanup=defer_cleanup) + else: + _remove_file_or_defer(path, defer_cleanup=defer_cleanup) + + if os.path.lexists(path): + raise OSError("Failed to reclaim malformed SquashFS image target") + return False + + def _is_reusable_extraction_dir( path: Path, *, defer_cleanup: Callable[[Path], None], ) -> bool: """Accept canonical directories and reclaim malformed file or symlink targets.""" - if path.is_dir() and not path.is_symlink(): + if _is_reusable_cache_directory(path): return True if os.path.lexists(path): _remove_file_or_defer(path, defer_cleanup=defer_cleanup) @@ -350,7 +381,10 @@ async def download( image_path: Path, ) -> float: """Ensure the SquashFS image exists locally and return download time.""" - if image_path.exists(): + if await _reuse_or_reclaim_squashfs_image( + image_path, + defer_cleanup=ctx.defer_cleanup, + ): return 0.0 image_path.parent.mkdir(parents=True, exist_ok=True) @@ -366,8 +400,11 @@ async def download( try: temp_image.rename(image_path) except OSError: - if not image_path.exists(): - raise + if not await _reuse_or_reclaim_squashfs_image( + image_path, + defer_cleanup=ctx.defer_cleanup, + ): + temp_image.rename(image_path) return (time.monotonic() - download_start) * 1000 finally: _remove_file_or_defer( diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 55a4b21ef1..23bbe59a12 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -19,6 +19,8 @@ from tracecat.executor.registry_artifact_materialization import ( RegistryArtifactAdmission, _allocated_size_bound, + _is_reusable_cache_directory, + _is_reusable_cache_file, _rejoin_future_on_cancel, _run_blocking_rejoin_on_cancel, ) @@ -848,9 +850,9 @@ def _trim_startup_cache(self) -> bool: paths = self._paths_for(entry.cache_key) if ( registry_artifact_mounts.is_mount(paths.squashfs_mount_dir) - or paths.squashfs_image_path.exists() - or paths.squashfs_extract_dir.exists() - or paths.tarball_target_dir.exists() + or _is_reusable_cache_file(paths.squashfs_image_path) + or _is_reusable_cache_directory(paths.squashfs_extract_dir) + or _is_reusable_cache_directory(paths.tarball_target_dir) ): continue diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 984414f5a4..5aa7e84a06 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -36,6 +36,7 @@ _bundled_builtin_registry_version, _download_s3_artifact, _is_cache_entry_uri, + _is_reusable_cache_file, _squashfs_listing_size, _squashfs_sidecar_uri, _tarball_extracted_size, @@ -401,7 +402,7 @@ async def _artifact_candidates( if self._can_try_squashfs(): squashfs_uri = _squashfs_sidecar_uri(artifact_uri) if squashfs_uri: - if ctx.paths.squashfs_image_path.exists(): + if _is_reusable_cache_file(ctx.paths.squashfs_image_path): candidates.append( SquashfsArtifact( uri=squashfs_uri, From c113e44d6a3b09c2406f994c46b0c7042d7c6371 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 22:45:16 -0400 Subject: [PATCH 120/161] fix(executor): redact registry blob downloads --- .../executor/test_registry_artifact_budget.py | 4 + .../test_registry_artifact_resolution.py | 2 + tests/unit/test_storage_blob.py | 118 ++++++++++++++++++ .../registry_artifact_materialization.py | 2 + tracecat/storage/blob.py | 90 ++++++++++--- 5 files changed, 202 insertions(+), 14 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index a760f37405..11f86010d5 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -226,8 +226,10 @@ async def download_file_to_path( max_bytes: int, ensure_capacity: Callable[[int], Awaitable[None]], defer_cleanup: Callable[[Path], None], + redact_log_identifiers: bool, ) -> int: del key, bucket, defer_cleanup + assert redact_log_identifiers is True nonlocal capacity_checked assert max_bytes == len(payload) + 33 await ensure_capacity(len(payload)) @@ -273,8 +275,10 @@ async def download_file_to_path( max_bytes: int, ensure_capacity: Callable[[int], Awaitable[None]], defer_cleanup: Callable[[Path], None], + redact_log_identifiers: bool, ) -> int: del key, bucket, defer_cleanup + assert redact_log_identifiers is True assert max_bytes == len(payload) + 256 await ensure_capacity(len(payload)) output_path.write_bytes(payload) diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py index 50dd9c9867..8f7b65cb62 100644 --- a/tests/unit/executor/test_registry_artifact_resolution.py +++ b/tests/unit/executor/test_registry_artifact_resolution.py @@ -125,8 +125,10 @@ async def mock_download_file_to_path( bucket: str, output_path: Path, defer_cleanup: Callable[[Path], None], + redact_log_identifiers: bool, ) -> None: assert defer_cleanup == ctx.defer_cleanup + assert redact_log_identifiers is True output_path.write_bytes(b"squashfs") with patch( diff --git a/tests/unit/test_storage_blob.py b/tests/unit/test_storage_blob.py index 5cceec33e2..b5bbdd8ec1 100644 --- a/tests/unit/test_storage_blob.py +++ b/tests/unit/test_storage_blob.py @@ -705,6 +705,45 @@ async def test_open_download_stream_yields_stream_and_length(self, mock_get_clie assert stream is mock_body assert length == 123 + @pytest.mark.anyio + @patch("tracecat.storage.blob.get_storage_client") + async def test_open_download_stream_redacts_sensitive_client_error( + self, mock_get_client + ): + mock_client = AsyncMock() + mock_get_client.return_value.__aenter__.return_value = mock_client + mock_client.get_object.side_effect = ClientError( + error_response={ + "Error": { + "Code": "AccessDenied", + "Message": "denied tenant-key/org/repository", + } + }, + operation_name="get_object", + ) + + with ( + patch("tracecat.storage.blob.logger.error") as log_error, + pytest.raises(blob_module.StorageDownloadError) as exc_info, + ): + async with open_download_stream( + key="tenant-key/org/repository", + bucket="tenant-bucket", + redact_log_identifiers=True, + ): + pass + + assert str(exc_info.value) == "Storage download failed" + assert exc_info.value.error_code == "AccessDenied" + assert exc_info.value.__cause__ is None + log_error.assert_called_once_with( + "Failed to open download stream", + key="", + bucket="", + error_code="AccessDenied", + error_type="ClientError", + ) + @pytest.mark.anyio async def test_download_file_to_path_writes_bytes( self, tmp_path: Path, monkeypatch @@ -741,6 +780,85 @@ async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 assert bytes_written == 11 assert out.read_bytes() == b"hello world" + @pytest.mark.anyio + async def test_download_file_to_path_redacts_sensitive_success_log( + self, tmp_path: Path, monkeypatch + ): + class DummyStream: + async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 + yield b"payload" + + @asynccontextmanager + async def _fake_open_download_stream( + *, + key: str, + bucket: str, + redact_log_identifiers: bool, + ): + assert key == "tenant-key/org/repository" + assert bucket == "tenant-bucket" + assert redact_log_identifiers is True + yield DummyStream(), 7 + + monkeypatch.setattr( + "tracecat.storage.blob.open_download_stream", + _fake_open_download_stream, + ) + out = tmp_path / "out.bin" + + with patch("tracecat.storage.blob.logger.debug") as log_debug: + await download_file_to_path( + key="tenant-key/org/repository", + bucket="tenant-bucket", + output_path=out, + redact_log_identifiers=True, + ) + + log_debug.assert_called_once_with( + "File streamed to disk successfully", + key="", + bucket="", + output_path=str(out), + size=7, + ) + + @pytest.mark.anyio + async def test_download_file_to_path_redacts_sensitive_capacity_error( + self, tmp_path: Path, monkeypatch + ): + class DummyStream: + async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 + yield b"payload" + + @asynccontextmanager + async def _fake_open_download_stream( + *, + key: str, # noqa: ARG001 + bucket: str, # noqa: ARG001 + redact_log_identifiers: bool, + ): + assert redact_log_identifiers is True + yield DummyStream(), 7 + + monkeypatch.setattr( + "tracecat.storage.blob.open_download_stream", + _fake_open_download_stream, + ) + + with pytest.raises(ValueError) as exc_info: + await download_file_to_path( + key="tenant-key/org/repository", + bucket="tenant-bucket", + output_path=tmp_path / "out.bin", + max_bytes=5, + redact_log_identifiers=True, + ) + + message = str(exc_info.value) + assert "tenant-key" not in message + assert "tenant-bucket" not in message + assert "/" in message + @pytest.mark.anyio async def test_download_file_to_path_max_bytes_refuses( self, tmp_path: Path, monkeypatch diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index d0a0df0fd2..89d2982086 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -776,6 +776,7 @@ async def _download_s3_artifact( bucket=bucket, output_path=output_path, defer_cleanup=defer_cleanup, + redact_log_identifiers=True, ) else: await blob.download_file_to_path( @@ -785,6 +786,7 @@ async def _download_s3_artifact( max_bytes=admission.max_bytes, ensure_capacity=admission.ensure_capacity, defer_cleanup=defer_cleanup, + redact_log_identifiers=True, ) except FileNotFoundError as e: request = httpx.Request("GET", artifact_uri) diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index 674ad48623..77c55d28ac 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -35,12 +35,33 @@ DEFAULT_UPLOAD_CHUNK_SIZE_BYTES = 8 * 1024 * 1024 # 8MB DEFAULT_UPLOAD_MAX_CONCURRENCY = 4 DEFAULT_UPLOAD_MAX_IO_QUEUE_SIZE = 2 +_REDACTED_STORAGE_IDENTIFIER = "" class _AsyncWritableFile(Protocol): async def write(self, data: bytes, /) -> int: ... +class StorageDownloadError(RuntimeError): + """A storage download failed without exposing sensitive object identifiers.""" + + def __init__(self, *, error_code: str | None) -> None: + super().__init__("Storage download failed") + self.error_code = error_code + + +def _download_log_identifiers( + key: str, + bucket: str, + *, + redact: bool, +) -> tuple[str, str]: + """Return storage identifiers that are safe for logs and error messages.""" + if redact: + return _REDACTED_STORAGE_IDENTIFIER, _REDACTED_STORAGE_IDENTIFIER + return key, bucket + + # Shared S3/MinIO client config: explicit standard-mode retries so transient # failures (throttling, 5xx, connection resets) are retried with backoff instead # of surfacing on the first error. @@ -708,6 +729,8 @@ async def download_file_range( async def open_download_stream( key: str, bucket: str, + *, + redact_log_identifiers: bool = False, ) -> AsyncIterator[tuple[StreamingBody, int | None]]: """Open a streaming download for an S3/MinIO object. @@ -723,14 +746,22 @@ async def open_download_stream( Args: key: The S3 object key. bucket: Bucket name (required). + redact_log_identifiers: Replace the key and bucket in logs and suppress + raw client-error text for sensitive internal objects. Yields: Tuple of (streaming body, content_length). Raises: ClientError: If the download fails. + StorageDownloadError: If a redacted download fails. FileNotFoundError: If the file doesn't exist. """ + log_key, log_bucket = _download_log_identifiers( + key, + bucket, + redact=redact_log_identifiers, + ) try: async with get_storage_client() as s3_client: response = await s3_client.get_object(Bucket=bucket, Key=key) @@ -739,17 +770,29 @@ async def open_download_stream( async with body: yield body, content_length except ClientError as e: - if e.response.get("Error", {}).get("Code") == "NoSuchKey": + error_code = e.response.get("Error", {}).get("Code") + if error_code == "NoSuchKey": logger.warning( "File not found in storage", - key=key, - bucket=bucket, + key=log_key, + bucket=log_bucket, ) + if redact_log_identifiers: + raise FileNotFoundError from None raise FileNotFoundError from e + if redact_log_identifiers: + logger.error( + "Failed to open download stream", + key=log_key, + bucket=log_bucket, + error_code=error_code, + error_type=type(e).__name__, + ) + raise StorageDownloadError(error_code=error_code) from None logger.error( "Failed to open download stream", - key=key, - bucket=bucket, + key=log_key, + bucket=log_bucket, error=str(e), ) raise @@ -765,6 +808,7 @@ async def download_file_to_path( expected_sha256: str | None = None, ensure_capacity: Callable[[int], Awaitable[None]] | None = None, defer_cleanup: Callable[[Path], None] | None = None, + redact_log_identifiers: bool = False, ) -> int: """Stream an S3/MinIO object to a local file. @@ -784,6 +828,8 @@ async def download_file_to_path( incrementally before each chunk is written. defer_cleanup: Optional callback that retains a partial-file path for a later cleanup retry when immediate deletion fails. + redact_log_identifiers: Replace the key and bucket in logs and generated + error messages for sensitive internal objects. Returns: Total bytes written. @@ -793,6 +839,11 @@ async def download_file_to_path( hasher = hashlib.sha256() if expected_sha256 is not None else None bytes_written = 0 + log_key, log_bucket = _download_log_identifiers( + key, + bucket, + redact=redact_log_identifiers, + ) @asynccontextmanager async def open_file_rejoin_on_cancel() -> AsyncIterator[_AsyncWritableFile]: @@ -865,9 +916,20 @@ async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: raise try: - async with open_download_stream(key=key, bucket=bucket) as ( - stream, - content_length, + download_stream = ( + open_download_stream( + key=key, + bucket=bucket, + redact_log_identifiers=True, + ) + if redact_log_identifiers + else open_download_stream(key=key, bucket=bucket) + ) + async with ( + download_stream as ( + stream, + content_length, + ) ): if ( max_bytes is not None @@ -875,7 +937,7 @@ async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: and content_length > max_bytes ): raise ValueError( - f"Refusing to download {bucket}/{key} to disk: " + f"Refusing to download {log_bucket}/{log_key} to disk: " f"ContentLength={content_length} exceeds max_bytes={max_bytes}" ) @@ -887,7 +949,7 @@ async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: if max_bytes is None: raise ValueError( "Cannot reserve disk capacity for a download without " - f"ContentLength or max_bytes: {bucket}/{key}" + f"ContentLength or max_bytes: {log_bucket}/{log_key}" ) grow_reservation_by_chunk = True else: @@ -907,7 +969,7 @@ async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: bytes_written += len(chunk) if download_limit is not None and bytes_written > download_limit: raise ValueError( - f"Refusing to download {bucket}/{key} to disk: " + f"Refusing to download {log_bucket}/{log_key} to disk: " f"bytes_written={bytes_written} exceeds " f"max_bytes={download_limit}" ) @@ -921,7 +983,7 @@ async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: actual_sha256 = hasher.hexdigest() if actual_sha256 != expected_sha256: raise ValueError( - f"Integrity check failed for {bucket}/{key}: " + f"Integrity check failed for {log_bucket}/{log_key}: " f"expected {expected_sha256}, got {actual_sha256}" ) @@ -940,8 +1002,8 @@ async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: logger.debug( "File streamed to disk successfully", - key=key, - bucket=bucket, + key=log_key, + bucket=log_bucket, output_path=str(output_path), size=bytes_written, ) From b8d2da0aa11acf014e4520a4b68ec1f3354d754c Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 22:56:03 -0400 Subject: [PATCH 121/161] fix(executor): reject unsafe squashfs mount paths --- .../test_registry_artifact_materialization.py | 45 +++++++++++++++++++ .../registry_artifact_materialization.py | 28 ++++++++++++ 2 files changed, 73 insertions(+) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 3b7410f968..7443a23f50 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -382,6 +382,51 @@ async def test_mount_squashfs_uses_hardened_read_only_options( ) terminate_process_group.assert_awaited_once_with(process) + @pytest.mark.parametrize("symlink_component", ["entry", "mount"]) + @pytest.mark.anyio + async def test_mount_squashfs_rejects_symlinked_cache_paths( + self, + temp_cache_dir: Path, + symlink_component: str, + ) -> None: + """A privileged mount cannot follow a cache path outside its entry.""" + cache = RegistryArtifactCache(temp_cache_dir) + ctx = cache._context_for("symlinked-mount") + artifact = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key=ctx.cache_key, + ) + outside_dir = temp_cache_dir / "outside" + outside_dir.mkdir(parents=True) + + if symlink_component == "entry": + ctx.paths.entry_dir.parent.mkdir(parents=True) + ctx.paths.entry_dir.symlink_to(outside_dir, target_is_directory=True) + else: + ctx.paths.entry_dir.mkdir(parents=True) + ctx.paths.squashfs_mount_dir.symlink_to( + outside_dir, + target_is_directory=True, + ) + + with ( + patch.object( + SquashfsArtifact, + "download", + new_callable=AsyncMock, + ) as download, + patch.object( + SquashfsArtifact, + "_mount_image", + new_callable=AsyncMock, + ) as mount_image, + pytest.raises(OSError, match="Unsafe .*SquashFS mount path"), + ): + await artifact.mount(ctx, ctx.paths.squashfs_image_path) + + download.assert_not_awaited() + mount_image.assert_not_awaited() + @pytest.mark.anyio async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): """Cancellation cannot leave an orphan mount process after lock release.""" diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 89d2982086..9328e5a2b3 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -238,6 +238,30 @@ def _is_reusable_cache_directory(path: Path) -> bool: return path.is_dir() and not path.is_symlink() +def _validate_squashfs_mount_paths(paths: RegistryArtifactPaths) -> None: + """Reject mount paths that can escape their canonical cache entry.""" + entry_dir = paths.entry_dir + mount_dir = paths.squashfs_mount_dir + if mount_dir.parent != entry_dir: + raise OSError("Unsafe SquashFS mount path outside its cache entry") + + for path in (entry_dir, mount_dir): + if os.path.lexists(path) and not _is_reusable_cache_directory(path): + raise OSError("Unsafe symlink or non-directory SquashFS mount path") + + if not entry_dir.exists() or not mount_dir.exists(): + return + + resolved_entries_dir = entry_dir.parent.resolve(strict=True) + resolved_entry_dir = entry_dir.resolve(strict=True) + resolved_mount_dir = mount_dir.resolve(strict=True) + if ( + resolved_entry_dir.parent != resolved_entries_dir + or resolved_mount_dir.parent != resolved_entry_dir + ): + raise OSError("Unsafe SquashFS mount path outside its cache entry") + + async def _reuse_or_reclaim_squashfs_image( path: Path, *, @@ -311,6 +335,7 @@ def format(self) -> RegistryArtifactFormat: def cached_path( self, ctx: RegistryArtifactMaterializationContext ) -> list[Path] | None: + _validate_squashfs_mount_paths(ctx.paths) if registry_artifact_mounts.is_mount(ctx.paths.squashfs_mount_dir): logger.debug( "Using cached SquashFS registry mount", @@ -431,11 +456,13 @@ async def mount( Exception: The image could not be downloaded or prepared. """ target_dir = ctx.paths.squashfs_mount_dir + _validate_squashfs_mount_paths(ctx.paths) if registry_artifact_mounts.is_mount(target_dir): return target_dir ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) target_dir.mkdir(parents=True, exist_ok=True) + _validate_squashfs_mount_paths(ctx.paths) logger.info( "Materializing SquashFS registry artifact", @@ -447,6 +474,7 @@ async def mount( download_elapsed = await self.download(ctx, image_path) mount_start = time.monotonic() + _validate_squashfs_mount_paths(ctx.paths) await self._mount_image(image_path, target_dir) mount_elapsed = (time.monotonic() - mount_start) * 1000 total_elapsed = (time.monotonic() - start_time) * 1000 From fa006797b9224aa53d6d4df546b93a4248def088 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 22:56:59 -0400 Subject: [PATCH 122/161] fix(executor): fail closed on squashfs listings --- .../test_registry_artifact_materialization.py | 17 +++++++++++++++++ .../registry_artifact_materialization.py | 4 ++++ 2 files changed, 21 insertions(+) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 7443a23f50..9920de4066 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -290,6 +290,23 @@ def test_squashfs_listing_size_rejects_unparseable_files(self) -> None: with pytest.raises(ValueError, match="Could not parse SquashFS listing"): _squashfs_listing_size(b"-rw-r--r-- malformed") + @pytest.mark.parametrize( + "listing", + [ + b"", + b"Parallel unsquashfs: Using 4 processors", + ], + ) + def test_squashfs_listing_size_rejects_listing_without_entries( + self, + listing: bytes, + ) -> None: + with pytest.raises( + ValueError, + match="Could not parse any SquashFS listing entries", + ): + _squashfs_listing_size(listing) + @pytest.mark.anyio async def test_materialize_mounts_squashfs_sidecar(self, temp_cache_dir): """Test that a SquashFS sidecar is mounted instead of extracting tarballs.""" diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 9328e5a2b3..35887ae298 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -964,6 +964,7 @@ def _tarball_extracted_size( def _squashfs_listing_size(output: bytes, *, allocation_unit: int = 1) -> int: """Bound allocated bytes from ``unsquashfs -lln`` output, failing closed.""" total_bytes = 0 + parsed_entries = 0 for raw_line in output.decode(errors="replace").splitlines(): line = raw_line.strip() if not line: @@ -974,8 +975,11 @@ def _squashfs_listing_size(output: bytes, *, allocation_unit: int = 1) -> int: continue if len(fields) < 5 or "/" not in fields[1] or not fields[2].isdigit(): raise ValueError(f"Could not parse SquashFS listing line: {line}") + parsed_entries += 1 total_bytes += _allocated_size_bound( int(fields[2]), allocation_unit=allocation_unit, ) + if parsed_entries == 0: + raise ValueError("Could not parse any SquashFS listing entries") return total_bytes From 2da290adfbac445d740dec82a8dbbc7730c806be Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:07:04 -0400 Subject: [PATCH 123/161] fix(executor): rescan writable artifact leases --- .../executor/test_registry_artifact_budget.py | 32 +++++++++++++++++++ .../test_test_backend_no_registry_action.py | 25 ++++++++++++--- tests/unit/test_action_runner.py | 16 ++++++++-- tracecat/executor/action_runner.py | 16 ++++++---- tracecat/executor/backends/test.py | 11 ++++++- tracecat/executor/registry_artifacts.py | 12 ++++++- 6 files changed, 96 insertions(+), 16 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 11f86010d5..ae82bfcd56 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -660,6 +660,38 @@ async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( assert cache._budget_dirty is False + @pytest.mark.anyio + async def test_releasing_a_mutable_cache_hit_rescans_post_admission_growth( + self, + temp_cache_dir: Path, + ) -> None: + """Direct consumers cannot grow a warm entry outside budget accounting.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/mutable-cached.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + target_dir = write_tarball_entry(temp_cache_dir, cache_key) + cache._budget_dirty = False + + with ( + patch(MAX_ENTRIES_CONFIG, 10), + patch(MAX_BYTES_CONFIG, 1_000_000), + patch.object( + cache, + "_scan_cache_entries", + wraps=cache._scan_cache_entries, + ) as scan_cache_entries, + ): + async with cache.lease( + [artifact_uri], + paths_may_be_modified=True, + ) as registry_paths: + assert registry_paths == [target_dir] + (target_dir / "action-output.bin").write_bytes(b"x" * 4096) + + assert scan_cache_entries.call_count == 1 + assert cache._budget_dirty is False + @pytest.mark.anyio async def test_successful_cold_admission_skips_release_rescan(self, temp_cache_dir): """A successful protected pass consumes the materialization dirty signal.""" diff --git a/tests/unit/executor/test_test_backend_no_registry_action.py b/tests/unit/executor/test_test_backend_no_registry_action.py index 17de1926cc..7a55cdac9a 100644 --- a/tests/unit/executor/test_test_backend_no_registry_action.py +++ b/tests/unit/executor/test_test_backend_no_registry_action.py @@ -302,8 +302,12 @@ async def ensure_swept(self) -> None: @asynccontextmanager async def lease( - self, artifact_uris: list[str] | None = None + self, + artifact_uris: list[str] | None = None, + *, + paths_may_be_modified: bool = False, ) -> AsyncIterator[list[Path]]: + assert paths_may_be_modified is True if broken_uri in (artifact_uris or []): raise RuntimeError("artifact unavailable") self.active += 1 @@ -383,9 +387,12 @@ async def ensure_swept(self) -> None: @asynccontextmanager async def lease( - self, artifact_uris: list[str] | None = None + self, + artifact_uris: list[str] | None = None, + *, + paths_may_be_modified: bool = False, ) -> AsyncIterator[list[Path]]: - del artifact_uris + del artifact_uris, paths_may_be_modified self.lease_attempted = True yield [] @@ -441,9 +448,13 @@ async def ensure_swept(self) -> None: @asynccontextmanager async def lease( - self, artifact_uris: list[str] | None = None + self, + artifact_uris: list[str] | None = None, + *, + paths_may_be_modified: bool = False, ) -> AsyncIterator[list[Path]]: assert artifact_uris == [artifact_uri] + assert paths_may_be_modified is False self.lease_attempted = True yield [] @@ -502,9 +513,13 @@ async def ensure_swept(self) -> None: @asynccontextmanager async def lease( - self, artifact_uris: list[str] | None = None + self, + artifact_uris: list[str] | None = None, + *, + paths_may_be_modified: bool = False, ) -> AsyncIterator[list[Path]]: assert artifact_uris == [artifact_uri] + assert paths_may_be_modified is True self.active += 1 try: yield [artifact_path] diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index c8d7a3a3b6..3ee69dd017 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -664,6 +664,7 @@ async def test_execute_action_holds_registry_lease_for_whole_subprocess( cache_key = compute_registry_artifact_cache_key(artifact_uri) entry_dir = runner.registry_artifacts._paths_for(cache_key).tarball_target_dir entry_dir.mkdir(parents=True) + await runner.registry_artifacts.ensure_swept() monkeypatch.setattr( action_runner.config, "TRACECAT__EXECUTOR_SANDBOX_ENABLED", False @@ -693,15 +694,23 @@ async def create_subprocess_exec_side_effect(*args, **kwargs): # noqa: ARG001 env = kwargs.get("env") assert isinstance(env, dict) registry_paths.append(env["PYTHONPATH"]) + (entry_dir / "action-output.bin").write_bytes(b"x" * 4096) mock_proc = AsyncMock() mock_proc.returncode = 0 mock_proc.communicate = AsyncMock(return_value=(success_response, b"")) return mock_proc - with patch( - "asyncio.create_subprocess_exec", - side_effect=create_subprocess_exec_side_effect, + with ( + patch( + "asyncio.create_subprocess_exec", + side_effect=create_subprocess_exec_side_effect, + ), + patch.object( + runner.registry_artifacts, + "_scan_cache_entries", + wraps=runner.registry_artifacts._scan_cache_entries, + ) as scan_cache_entries, ): result = await runner.execute_action( input=mock_run_action_input, @@ -715,6 +724,7 @@ async def create_subprocess_exec_side_effect(*args, **kwargs): # noqa: ARG001 assert refcounts == [1] assert registry_paths[0].startswith(str(entry_dir)) assert runner.registry_artifacts._refcount(cache_key) == 0 + assert scan_cache_entries.call_count == 1 @pytest.mark.anyio async def test_cancelled_action_reaps_child_before_releasing_mounted_artifact( diff --git a/tracecat/executor/action_runner.py b/tracecat/executor/action_runner.py index 04ce830496..63876bf19a 100644 --- a/tracecat/executor/action_runner.py +++ b/tracecat/executor/action_runner.py @@ -181,15 +181,19 @@ async def execute_action( """ timeout = timeout or config.TRACECAT__EXECUTOR_CLIENT_TIMEOUT + # Direct subprocesses receive host paths and can modify extracted + # artifacts. NsJail exposes the same paths through read-only bind mounts. + use_sandbox = force_sandbox or ( + config.TRACECAT__EXECUTOR_SANDBOX_ENABLED and _is_sandbox_available() + ) + # Materialize each registry artifact, collect paths in deterministic order. # The lease is held for the whole subprocess execution so cache eviction # cannot delete a directory the subprocess is still importing from. - async with self.registry_artifacts.lease(artifact_uris) as registry_paths: - # Check if sandbox execution is enabled and available - # force_sandbox=True overrides config (used by ephemeral backend) - use_sandbox = force_sandbox or ( - config.TRACECAT__EXECUTOR_SANDBOX_ENABLED and _is_sandbox_available() - ) + async with self.registry_artifacts.lease( + artifact_uris, + paths_may_be_modified=not use_sandbox, + ) as registry_paths: logger.debug( "Using sandbox execution", use_sandbox=use_sandbox, diff --git a/tracecat/executor/backends/test.py b/tracecat/executor/backends/test.py index 6fe5a0390a..2086a7b261 100644 --- a/tracecat/executor/backends/test.py +++ b/tracecat/executor/backends/test.py @@ -318,11 +318,18 @@ async def _lease_registry_artifacts( if any(_is_cache_entry_uri(uri) for uri in artifact_uris): await registry_artifacts.ensure_swept() extracted_paths: list[str] = [] + mutable_rescan_registered = False for artifact_uri in artifact_uris: + rescan_on_release = not mutable_rescan_registered and _is_cache_entry_uri( + artifact_uri + ) try: artifact_paths = await leases.enter_async_context( - registry_artifacts.lease([artifact_uri]) + registry_artifacts.lease( + [artifact_uri], + paths_may_be_modified=rescan_on_release, + ) ) except Exception as e: logger.warning( @@ -331,6 +338,8 @@ async def _lease_registry_artifacts( error=str(e), ) continue + if rescan_on_release: + mutable_rescan_registered = True extracted_paths.extend(str(path) for path in artifact_paths) logger.debug( diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 5aa7e84a06..2e93bc9775 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -107,7 +107,12 @@ class RegistryArtifactCache(_RegistryArtifactCacheStorage): """Materializes and leases executor-local registry artifact paths.""" @asynccontextmanager - async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Path]]: + async def lease( + self, + artifact_uris: list[str] | None, + *, + paths_may_be_modified: bool = False, + ) -> AsyncIterator[list[Path]]: """Materialize registry artifacts and pin them for the life of the context. Leased cache entries are never evicted, so callers may keep importing @@ -116,6 +121,9 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat Args: artifact_uris: Registry artifact URIs in deterministic PYTHONPATH order, or None to use the base PYTHONPATH directory. + paths_may_be_modified: Whether the consumer can write to returned + cache paths. Mutable leases re-arm byte-budget convergence when + execution ends so post-admission growth is measured. Yields: Importable Python paths for the requested artifacts. @@ -155,6 +163,8 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat lease_setup_complete = True yield registry_paths finally: + if paths_may_be_modified and lease_setup_complete and leased_keys: + self._budget_dirty = True idle_keys = [ cache_key for cache_key in leased_keys if self._release_lease(cache_key) ] From 2a76be8258bc748c59cc2d535129a177e6e5e764 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:10:12 -0400 Subject: [PATCH 124/161] fix(executor): contain registry cache entries --- .../test_registry_artifact_materialization.py | 40 ++++++++++++++++++- .../registry_artifact_materialization.py | 38 ++++++++++++++---- 2 files changed, 69 insertions(+), 9 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 9920de4066..4259de46da 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -116,6 +116,44 @@ async def download_image( assert not image_path.is_symlink() assert image_path.read_bytes() == b"fresh image" + @pytest.mark.anyio + async def test_tarball_rejects_symlinked_cache_entry_for_reuse_and_publish( + self, + temp_cache_dir: Path, + ) -> None: + """Tarball paths cannot escape through a symlinked cache entry root.""" + cache = RegistryArtifactCache(temp_cache_dir) + ctx = cache._context_for("symlinked-tarball-entry") + artifact = TarballArtifact( + uri="s3://bucket/path/site-packages.tar.gz", + cache_key=ctx.cache_key, + ) + outside_dir = temp_cache_dir / "outside-tarball-entry" + (outside_dir / "tarball").mkdir(parents=True) + ctx.paths.entry_dir.parent.mkdir(parents=True) + ctx.paths.entry_dir.symlink_to(outside_dir, target_is_directory=True) + + with pytest.raises(OSError, match="Unsafe .*registry cache entry path"): + artifact.cached_path(ctx) + + with ( + patch.object( + TarballArtifact, + "download", + new_callable=AsyncMock, + ) as download, + patch.object( + TarballArtifact, + "extract", + new_callable=AsyncMock, + ) as extract, + pytest.raises(OSError, match="Unsafe .*registry cache entry path"), + ): + await artifact.materialize(ctx) + + download.assert_not_awaited() + extract.assert_not_awaited() + @pytest.mark.anyio async def test_same_key_cold_fan_in_materializes_and_enforces_once( self, temp_cache_dir @@ -437,7 +475,7 @@ async def test_mount_squashfs_rejects_symlinked_cache_paths( "_mount_image", new_callable=AsyncMock, ) as mount_image, - pytest.raises(OSError, match="Unsafe .*SquashFS mount path"), + pytest.raises(OSError, match="Unsafe .* path"), ): await artifact.mount(ctx, ctx.paths.squashfs_image_path) diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 35887ae298..4023679235 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -238,27 +238,37 @@ def _is_reusable_cache_directory(path: Path) -> bool: return path.is_dir() and not path.is_symlink() +def _validate_cache_entry_path(paths: RegistryArtifactPaths) -> None: + """Reject an entry root redirected outside the configured entries directory.""" + entry_dir = paths.entry_dir + if os.path.lexists(entry_dir) and not _is_reusable_cache_directory(entry_dir): + raise OSError("Unsafe symlink or non-directory registry cache entry path") + if not entry_dir.exists(): + return + + resolved_entries_dir = entry_dir.parent.resolve(strict=True) + resolved_entry_dir = entry_dir.resolve(strict=True) + if resolved_entry_dir.parent != resolved_entries_dir: + raise OSError("Unsafe registry cache entry path outside the entries directory") + + def _validate_squashfs_mount_paths(paths: RegistryArtifactPaths) -> None: """Reject mount paths that can escape their canonical cache entry.""" entry_dir = paths.entry_dir mount_dir = paths.squashfs_mount_dir + _validate_cache_entry_path(paths) if mount_dir.parent != entry_dir: raise OSError("Unsafe SquashFS mount path outside its cache entry") - for path in (entry_dir, mount_dir): - if os.path.lexists(path) and not _is_reusable_cache_directory(path): - raise OSError("Unsafe symlink or non-directory SquashFS mount path") + if os.path.lexists(mount_dir) and not _is_reusable_cache_directory(mount_dir): + raise OSError("Unsafe symlink or non-directory SquashFS mount path") if not entry_dir.exists() or not mount_dir.exists(): return - resolved_entries_dir = entry_dir.parent.resolve(strict=True) resolved_entry_dir = entry_dir.resolve(strict=True) resolved_mount_dir = mount_dir.resolve(strict=True) - if ( - resolved_entry_dir.parent != resolved_entries_dir - or resolved_mount_dir.parent != resolved_entry_dir - ): + if resolved_mount_dir.parent != resolved_entry_dir: raise OSError("Unsafe SquashFS mount path outside its cache entry") @@ -406,6 +416,7 @@ async def download( image_path: Path, ) -> float: """Ensure the SquashFS image exists locally and return download time.""" + _validate_cache_entry_path(ctx.paths) if await _reuse_or_reclaim_squashfs_image( image_path, defer_cleanup=ctx.defer_cleanup, @@ -422,9 +433,11 @@ async def download( admission=ctx.admission, defer_cleanup=ctx.defer_cleanup, ) + _validate_cache_entry_path(ctx.paths) try: temp_image.rename(image_path) except OSError: + _validate_cache_entry_path(ctx.paths) if not await _reuse_or_reclaim_squashfs_image( image_path, defer_cleanup=ctx.defer_cleanup, @@ -496,6 +509,7 @@ async def extract( image_path: Path, ) -> Path: target_dir = ctx.paths.squashfs_extract_dir + _validate_cache_entry_path(ctx.paths) if _is_reusable_extraction_dir( target_dir, defer_cleanup=ctx.defer_cleanup, @@ -503,6 +517,7 @@ async def extract( return target_dir ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + _validate_cache_entry_path(ctx.paths) logger.info( "Extracting SquashFS registry artifact", @@ -526,6 +541,7 @@ async def extract( await self._extract_image(image_path, temp_dir) extract_elapsed = (time.monotonic() - extract_start) * 1000 + _validate_cache_entry_path(ctx.paths) try: temp_dir.rename(target_dir) total_elapsed = (time.monotonic() - start_time) * 1000 @@ -539,6 +555,7 @@ async def extract( total_ms=f"{total_elapsed:.1f}", ) except OSError: + _validate_cache_entry_path(ctx.paths) if _is_reusable_extraction_dir( target_dir, defer_cleanup=ctx.defer_cleanup, @@ -664,6 +681,7 @@ def format(self) -> RegistryArtifactFormat: def cached_path( self, ctx: RegistryArtifactMaterializationContext ) -> list[Path] | None: + _validate_cache_entry_path(ctx.paths) if _is_reusable_extraction_dir( ctx.paths.tarball_target_dir, defer_cleanup=ctx.defer_cleanup, @@ -679,6 +697,7 @@ async def materialize( self, ctx: RegistryArtifactMaterializationContext ) -> list[Path]: target_dir = ctx.paths.tarball_target_dir + _validate_cache_entry_path(ctx.paths) logger.info( "Materializing tarball registry artifact", cache_key=ctx.cache_key, @@ -692,6 +711,7 @@ async def materialize( try: ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + _validate_cache_entry_path(ctx.paths) download_start = time.monotonic() await self.download(ctx, temp_tarball) @@ -712,6 +732,7 @@ async def materialize( await self.extract(temp_tarball, temp_dir) extract_elapsed = (time.monotonic() - extract_start) * 1000 + _validate_cache_entry_path(ctx.paths) try: temp_dir.rename(target_dir) total_elapsed = (time.monotonic() - start_time) * 1000 @@ -725,6 +746,7 @@ async def materialize( total_ms=f"{total_elapsed:.1f}", ) except OSError: + _validate_cache_entry_path(ctx.paths) if _is_reusable_extraction_dir( target_dir, defer_cleanup=ctx.defer_cleanup, From d9b6915d1c213eb0e92890e07a6ec374fc899340 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:15:21 -0400 Subject: [PATCH 125/161] fix(sandbox): supervise unsafe pid fallback --- tests/unit/test_unsafe_pid_executor.py | 84 ++++++++++++++++++++++++- tracecat/executor/process_supervisor.py | 10 +-- tracecat/sandbox/unsafe_pid_executor.py | 43 +++++++++++-- tracecat/sandbox/utils.py | 2 +- 4 files changed, 126 insertions(+), 13 deletions(-) diff --git a/tests/unit/test_unsafe_pid_executor.py b/tests/unit/test_unsafe_pid_executor.py index 25166ae7e2..96ab11578f 100644 --- a/tests/unit/test_unsafe_pid_executor.py +++ b/tests/unit/test_unsafe_pid_executor.py @@ -74,6 +74,25 @@ def main(pid_file, wait): """ +def _detached_background_process_script() -> str: + return """ +import subprocess +import sys +from pathlib import Path + +def main(pid_file): + child = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(30)"], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + Path(pid_file).write_text(str(child.pid)) + return child.pid +""" + + class TestUnsafePidExecutor: @pytest.fixture def executor(self, tmp_path) -> UnsafePidExecutor: @@ -119,10 +138,35 @@ async def pid_namespace_available() -> bool: "tracecat.sandbox.unsafe_pid_executor.pid_namespace_available", pid_namespace_available, ) - cmd = await executor._build_execution_cmd( + command = await executor._build_execution_cmd( "python3", executor.cache_dir / "wrapper.py" ) - assert cmd[:4] == ["unshare", "--pid", "--fork", "--kill-child"] + assert command.argv[:4] == ["unshare", "--pid", "--fork", "--kill-child"] + assert command.supervised is False + + @pytest.mark.anyio + async def test_build_execution_cmd_without_pid_namespace_uses_linux_supervisor( + self, + executor: UnsafePidExecutor, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + async def pid_namespace_unavailable() -> bool: + return False + + monkeypatch.setattr( + unsafe_pid_executor, + "pid_namespace_available", + pid_namespace_unavailable, + ) + monkeypatch.setattr(unsafe_pid_executor.sys, "platform", "linux") + + wrapper_path = executor.cache_dir / "wrapper.py" + command = await executor._build_execution_cmd("python3", wrapper_path) + + assert command.argv[:2] == [sys.executable, "-I"] + assert Path(command.argv[2]).name == "process_supervisor.py" + assert command.argv[-2:] == ["python3", str(wrapper_path)] + assert command.supervised is True @pytest.mark.anyio async def test_pid_isolation_warning_logged_once( @@ -360,6 +404,42 @@ async def test_execute_kills_background_descendants_holding_output_pipes( with contextlib.suppress(ProcessLookupError): os.kill(child_pid, signal.SIGKILL) + @pytest.mark.skipif( + sys.platform != "linux", + reason="Detached fallback containment uses the Linux subreaper supervisor", + ) + @pytest.mark.anyio + async def test_execute_without_pid_namespace_reaps_detached_descendant( + self, + executor: UnsafePidExecutor, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + async def pid_namespace_unavailable() -> bool: + return False + + monkeypatch.setattr( + unsafe_pid_executor, + "pid_namespace_available", + pid_namespace_unavailable, + ) + pid_file = tmp_path / "detached-child.pid" + result = await asyncio.wait_for( + executor.execute( + script=_detached_background_process_script(), + inputs={"pid_file": str(pid_file)}, + ), + timeout=5, + ) + + assert result.success + child_pid = int(pid_file.read_text()) + try: + await _wait_for_process_exit(child_pid) + finally: + with contextlib.suppress(ProcessLookupError): + os.kill(child_pid, signal.SIGKILL) + @pytest.mark.anyio async def test_execute_kills_descendants_on_timeout( self, executor: UnsafePidExecutor, tmp_path: Path diff --git a/tracecat/executor/process_supervisor.py b/tracecat/executor/process_supervisor.py index 7680e0f30c..b5c94e1be3 100644 --- a/tracecat/executor/process_supervisor.py +++ b/tracecat/executor/process_supervisor.py @@ -1,6 +1,6 @@ -"""Contain descendants of one Linux direct-action subprocess. +"""Contain descendants of one Linux subprocess. -The outer process remains the child observed by ``ActionRunner``. A detached +The outer process remains the child observed by its caller. A detached subreaper monitor owns the actual action process and all descendants orphaned by it. Closing the control pipe asks that monitor to kill and reap its complete child tree before the outer process exits. @@ -82,7 +82,7 @@ def _exec(command: Sequence[str], control_fd: int) -> None: os.execvpe(command[0], list(command), os.environ) except OSError: with suppress(OSError): - os.write(2, b"Failed to execute supervised action\n") + os.write(2, b"Failed to execute supervised process\n") os._exit(127) @@ -120,7 +120,7 @@ def _run_monitor(control_fd: int, command: Sequence[str]) -> int: with suppress(BaseException): _kill_and_reap_children() with suppress(OSError): - os.write(2, b"Direct action supervisor failed\n") + os.write(2, b"Process supervisor failed\n") return 1 finally: with suppress(OSError): @@ -130,7 +130,7 @@ def _run_monitor(control_fd: int, command: Sequence[str]) -> int: def supervise(command: Sequence[str]) -> int: """Run one command and return only after all of its descendants are reaped.""" if sys.platform != "linux": - raise RuntimeError("The direct action process supervisor requires Linux") + raise RuntimeError("The process supervisor requires Linux") if not command: raise ValueError("A supervised command is required") diff --git a/tracecat/sandbox/unsafe_pid_executor.py b/tracecat/sandbox/unsafe_pid_executor.py index 2be02133cd..c308db0cbf 100644 --- a/tracecat/sandbox/unsafe_pid_executor.py +++ b/tracecat/sandbox/unsafe_pid_executor.py @@ -11,8 +11,10 @@ import logging import os import shutil +import sys import tempfile import time +from dataclasses import dataclass from pathlib import Path from typing import Any @@ -33,6 +35,7 @@ communicate_process_group, pid_namespace_available, pid_namespace_probe_error, + terminate_supervised_process, ) module_logger = logging.getLogger(__name__) @@ -240,6 +243,14 @@ def main(): ''' +@dataclass(frozen=True, slots=True) +class _ExecutionCommand: + """Command metadata for one unsafe executor subprocess.""" + + argv: list[str] + supervised: bool + + class UnsafePidExecutor: """Executor for Python scripts without nsjail, using subprocess isolation.""" @@ -307,17 +318,31 @@ def _with_python_paths( async def _build_execution_cmd( self, python_path: str, wrapper_path: Path - ) -> list[str]: + ) -> _ExecutionCommand: base_cmd = [python_path, str(wrapper_path)] if await pid_namespace_available(): - return ["unshare", "--pid", "--fork", "--kill-child", *base_cmd] + return _ExecutionCommand( + argv=["unshare", "--pid", "--fork", "--kill-child", *base_cmd], + supervised=False, + ) if not self._pid_isolation_warning_emitted: message = "PID namespace isolation unavailable; running script without PID isolation" logger.warning(message, reason=pid_namespace_probe_error()) module_logger.warning(message) self._pid_isolation_warning_emitted = True - return base_cmd + + if sys.platform == "linux": + supervisor_path = ( + Path(__file__).resolve().parents[1] + / "executor" + / "process_supervisor.py" + ) + return _ExecutionCommand( + argv=[sys.executable, "-I", str(supervisor_path), *base_cmd], + supervised=True, + ) + return _ExecutionCommand(argv=base_cmd, supervised=False) async def _create_venv(self, venv_path: Path) -> None: create_cmd = ["uv", "venv", str(venv_path), "--python", "3.12"] @@ -464,9 +489,12 @@ async def execute( if execution_env_vars: exec_env.update(execution_env_vars) - cmd = await self._build_execution_cmd(python_path, wrapper_path) + execution_command = await self._build_execution_cmd( + python_path, + wrapper_path, + ) process = await asyncio.create_subprocess_exec( - *cmd, + *execution_command.argv, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=str(work_dir), @@ -477,6 +505,11 @@ async def execute( stdout_bytes, stderr_bytes = await communicate_process_group( process, timeout=timeout_seconds, + terminate=( + terminate_supervised_process + if execution_command.supervised + else None + ), ) except TimeoutError as e: raise SandboxTimeoutError( diff --git a/tracecat/sandbox/utils.py b/tracecat/sandbox/utils.py index a9b8dc0032..fc63c8a14c 100644 --- a/tracecat/sandbox/utils.py +++ b/tracecat/sandbox/utils.py @@ -43,7 +43,7 @@ async def terminate_process_group(process: asyncio.subprocess.Process) -> None: async def terminate_supervised_process(process: asyncio.subprocess.Process) -> None: - """Request descendant cleanup from a direct-action process supervisor.""" + """Request descendant cleanup from a Linux process supervisor.""" if process.returncode is None: with suppress(ProcessLookupError): os.kill(process.pid, signal.SIGTERM) From c678aedb8c2ccd65b5c4db5fbe435d6e42d4a2b4 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:18:39 -0400 Subject: [PATCH 126/161] test(executor): synchronize udf timeout probe --- .../test_test_backend_no_registry_action.py | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/tests/unit/executor/test_test_backend_no_registry_action.py b/tests/unit/executor/test_test_backend_no_registry_action.py index 7a55cdac9a..fdd08769bb 100644 --- a/tests/unit/executor/test_test_backend_no_registry_action.py +++ b/tests/unit/executor/test_test_backend_no_registry_action.py @@ -12,7 +12,7 @@ import sys import threading import uuid -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Awaitable from contextlib import AsyncExitStack, asynccontextmanager from datetime import UTC, datetime from pathlib import Path @@ -503,6 +503,7 @@ async def test_timed_out_sync_udf_keeps_artifact_lease_until_thread_finishes( artifact_uri = "s3://bucket/sync-timeout.tar.gz" worker_started = threading.Event() finish_worker = threading.Event() + timeout_triggered = asyncio.Event() class FakeRegistryArtifacts: def __init__(self) -> None: @@ -540,6 +541,21 @@ def blocking_udf(**_kwargs: object) -> str: assert finish_worker.wait(timeout=5) return "finished" + async def wait_for_after_worker_started[T]( + awaitable: Awaitable[T], + timeout: float | None, + ) -> T: + """Drive wait_for cancellation only after the UDF thread exists.""" + del timeout + task = asyncio.ensure_future(awaitable) + assert await asyncio.to_thread(worker_started.wait, 1) + task.cancel() + timeout_triggered.set() + try: + return await task + except asyncio.CancelledError as e: + raise TimeoutError from e + backend = TestBackend() await backend.start() execution: asyncio.Task[ExecutorResult] | None = None @@ -560,17 +576,17 @@ def blocking_udf(**_kwargs: object) -> str: "_load_udf_callable", lambda _action_impl: blocking_udf, ) + monkeypatch.setattr(asyncio, "wait_for", wait_for_after_worker_started) execution = asyncio.create_task( backend.execute( input=test_run_action_input, role=test_role, resolved_context=test_resolved_context, - timeout=0.01, + timeout=30.0, ) ) - assert await asyncio.to_thread(worker_started.wait, 1) - await asyncio.sleep(0.05) + await timeout_triggered.wait() assert not execution.done() assert fake_runner.registry_artifacts.active == 1 From 434c507e9106266158ddf63eb8345f1ac199842a Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:30:00 -0400 Subject: [PATCH 127/161] fix(executor): contain registry cache work paths --- .../test_registry_artifact_materialization.py | 20 +++++++++++++++ .../test_registry_artifact_startup.py | 25 +++++++++++++++++++ .../registry_artifact_materialization.py | 15 +++++++++++ .../executor/registry_artifact_storage.py | 5 ++++ 4 files changed, 65 insertions(+) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 4259de46da..a2eef24f37 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -38,6 +38,26 @@ class TestRegistryArtifactMaterialization: """Materialize and reuse executor-local artifact formats.""" + def test_temp_path_rejects_symlinked_staging_root( + self, temp_cache_dir: Path + ) -> None: + cache_dir = temp_cache_dir / "cache" + cache_dir.mkdir() + cache = RegistryArtifactCache(cache_dir) + artifact = TarballArtifact( + uri="s3://bucket/path/site-packages.tar.gz", + cache_key="symlinked-staging", + ) + ctx = cache._context_for(artifact.cache_key) + outside_dir = temp_cache_dir / "outside-staging" + outside_dir.mkdir() + ctx.staging_dir.symlink_to(outside_dir, target_is_directory=True) + + with pytest.raises(OSError, match="Unsafe .*registry cache work path"): + artifact._temp_path(ctx, ".tmp") + + assert not any(outside_dir.iterdir()) + def test_temp_path_avoids_deferred_staging_collision( self, temp_cache_dir: Path ) -> None: diff --git a/tests/unit/executor/test_registry_artifact_startup.py b/tests/unit/executor/test_registry_artifact_startup.py index f90063347a..02a8def9df 100644 --- a/tests/unit/executor/test_registry_artifact_startup.py +++ b/tests/unit/executor/test_registry_artifact_startup.py @@ -127,6 +127,31 @@ async def test_sweep_removes_orphaned_work(self, temp_cache_dir): assert unrelated_file.read_text() == "keep" assert entry_dir.is_dir() + @pytest.mark.anyio + @pytest.mark.parametrize("work_dir_name", ["staging", "trash"]) + async def test_sweep_rejects_symlinked_work_directory( + self, + temp_cache_dir: Path, + work_dir_name: str, + ) -> None: + """Startup cleanup cannot follow a cache work root outside the cache.""" + cache_dir = temp_cache_dir / "cache" + cache_dir.mkdir() + cache = RegistryArtifactCache(cache_dir) + outside_dir = temp_cache_dir / f"outside-{work_dir_name}" + outside_dir.mkdir() + outside_file = outside_dir / "keep.txt" + outside_file.write_text("keep") + getattr(cache, f"{work_dir_name}_dir").symlink_to( + outside_dir, + target_is_directory=True, + ) + + with pytest.raises(OSError, match="Unsafe .*registry cache work path"): + await cache.ensure_swept() + + assert outside_file.read_text() == "keep" + @pytest.mark.anyio async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): """A live mountpoint belongs to a running process and must survive.""" diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 4023679235..a97b945832 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -148,7 +148,9 @@ def _temp_path( ctx: RegistryArtifactMaterializationContext, suffix: str, ) -> Path: + _validate_cache_work_directory(ctx.staging_dir) ctx.staging_dir.mkdir(parents=True, exist_ok=True) + _validate_cache_work_directory(ctx.staging_dir) while True: attempt_id = secrets.token_hex(8) candidate = ( @@ -238,6 +240,19 @@ def _is_reusable_cache_directory(path: Path) -> bool: return path.is_dir() and not path.is_symlink() +def _validate_cache_work_directory(path: Path) -> None: + """Reject a staging or trash root redirected outside its cache root.""" + if os.path.lexists(path) and not _is_reusable_cache_directory(path): + raise OSError("Unsafe symlink or non-directory registry cache work path") + if not path.exists(): + return + + resolved_cache_dir = path.parent.resolve(strict=True) + resolved_work_dir = path.resolve(strict=True) + if resolved_work_dir.parent != resolved_cache_dir: + raise OSError("Unsafe registry cache work path outside the cache directory") + + def _validate_cache_entry_path(paths: RegistryArtifactPaths) -> None: """Reject an entry root redirected outside the configured entries directory.""" entry_dir = paths.entry_dir diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 23bbe59a12..9f5ff1c0a7 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -23,6 +23,7 @@ _is_reusable_cache_file, _rejoin_future_on_cancel, _run_blocking_rejoin_on_cancel, + _validate_cache_work_directory, ) from tracecat.logger import logger from tracecat.sandbox.utils import communicate_process_group @@ -190,7 +191,9 @@ async def _delete_cache_path_off_loop(path: Path) -> bool: def _unique_work_path(root: Path, cache_key: str) -> Path: """Return a unique path beneath a cache work directory.""" + _validate_cache_work_directory(root) root.mkdir(parents=True, exist_ok=True) + _validate_cache_work_directory(root) unique_id = time.time_ns() while True: path = root / f"{cache_key}.{os.getpid()}.{unique_id}" @@ -803,6 +806,7 @@ def _clear_work_dir( remember_failures: bool = False, ) -> bool: """Best-effort remove every child of a staging or trash directory.""" + _validate_cache_work_directory(work_dir) try: paths = list(work_dir.iterdir()) except FileNotFoundError: @@ -832,6 +836,7 @@ def _clear_work_dir( def _retry_deferred_staging_cleanup(self) -> bool: """Retry exact failed paths without sweeping live staging work.""" + _validate_cache_work_directory(self.staging_dir) for path in tuple(self._deferred_staging_cleanup): if _delete_cache_path(path): self._deferred_staging_cleanup.discard(path) From c88c48e9ae4c6f8637f11e6b421c03160d2af7b4 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:31:46 -0400 Subject: [PATCH 128/161] fix(executor): contain registry cache entries root --- .../test_registry_artifact_materialization.py | 40 +++++++++++++++++++ .../test_registry_artifact_startup.py | 21 ++++++++++ .../registry_artifact_materialization.py | 17 +++++++- .../executor/registry_artifact_storage.py | 5 +++ 4 files changed, 82 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index a2eef24f37..2f261a5913 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -174,6 +174,46 @@ async def test_tarball_rejects_symlinked_cache_entry_for_reuse_and_publish( download.assert_not_awaited() extract.assert_not_awaited() + @pytest.mark.anyio + async def test_tarball_rejects_symlinked_entries_root_before_creation( + self, + temp_cache_dir: Path, + ) -> None: + """An absent cache key cannot be created through a redirected entries root.""" + cache_dir = temp_cache_dir / "cache" + cache_dir.mkdir() + cache = RegistryArtifactCache(cache_dir) + ctx = cache._context_for("symlinked-entries-root") + artifact = TarballArtifact( + uri="s3://bucket/path/site-packages.tar.gz", + cache_key=ctx.cache_key, + ) + outside_dir = temp_cache_dir / "outside-entries" + outside_dir.mkdir() + cache.entries_dir.symlink_to(outside_dir, target_is_directory=True) + + with pytest.raises(OSError, match="Unsafe .*registry cache entries path"): + artifact.cached_path(ctx) + + with ( + patch.object( + TarballArtifact, + "download", + new_callable=AsyncMock, + ) as download, + patch.object( + TarballArtifact, + "extract", + new_callable=AsyncMock, + ) as extract, + pytest.raises(OSError, match="Unsafe .*registry cache entries path"), + ): + await artifact.materialize(ctx) + + download.assert_not_awaited() + extract.assert_not_awaited() + assert not any(outside_dir.iterdir()) + @pytest.mark.anyio async def test_same_key_cold_fan_in_materializes_and_enforces_once( self, temp_cache_dir diff --git a/tests/unit/executor/test_registry_artifact_startup.py b/tests/unit/executor/test_registry_artifact_startup.py index 02a8def9df..cdc33d3a62 100644 --- a/tests/unit/executor/test_registry_artifact_startup.py +++ b/tests/unit/executor/test_registry_artifact_startup.py @@ -152,6 +152,27 @@ async def test_sweep_rejects_symlinked_work_directory( assert outside_file.read_text() == "keep" + @pytest.mark.anyio + async def test_sweep_rejects_symlinked_entries_directory( + self, + temp_cache_dir: Path, + ) -> None: + """Startup inspection cannot retire entries outside the configured cache.""" + cache_dir = temp_cache_dir / "cache" + cache_dir.mkdir() + cache = RegistryArtifactCache(cache_dir) + outside_dir = temp_cache_dir / "outside-entries" + outside_entry = write_tarball_entry(outside_dir, "keep") + cache.entries_dir.symlink_to( + outside_dir / "entries", + target_is_directory=True, + ) + + with pytest.raises(OSError, match="Unsafe .*registry cache entries path"): + await cache.ensure_swept() + + assert outside_entry.is_dir() + @pytest.mark.anyio async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): """A live mountpoint belongs to a running process and must survive.""" diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index a97b945832..4f6cd766bb 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -253,15 +253,30 @@ def _validate_cache_work_directory(path: Path) -> None: raise OSError("Unsafe registry cache work path outside the cache directory") +def _validate_cache_entries_directory(entries_dir: Path) -> None: + """Reject an entries root redirected outside its configured cache root.""" + if os.path.lexists(entries_dir) and not _is_reusable_cache_directory(entries_dir): + raise OSError("Unsafe symlink or non-directory registry cache entries path") + if not entries_dir.exists(): + return + + resolved_cache_dir = entries_dir.parent.resolve(strict=True) + resolved_entries_dir = entries_dir.resolve(strict=True) + if resolved_entries_dir.parent != resolved_cache_dir: + raise OSError("Unsafe registry cache entries path outside the cache directory") + + def _validate_cache_entry_path(paths: RegistryArtifactPaths) -> None: """Reject an entry root redirected outside the configured entries directory.""" entry_dir = paths.entry_dir + entries_dir = entry_dir.parent + _validate_cache_entries_directory(entries_dir) if os.path.lexists(entry_dir) and not _is_reusable_cache_directory(entry_dir): raise OSError("Unsafe symlink or non-directory registry cache entry path") if not entry_dir.exists(): return - resolved_entries_dir = entry_dir.parent.resolve(strict=True) + resolved_entries_dir = entries_dir.resolve(strict=True) resolved_entry_dir = entry_dir.resolve(strict=True) if resolved_entry_dir.parent != resolved_entries_dir: raise OSError("Unsafe registry cache entry path outside the entries directory") diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 9f5ff1c0a7..db369426dc 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -23,6 +23,8 @@ _is_reusable_cache_file, _rejoin_future_on_cancel, _run_blocking_rejoin_on_cancel, + _validate_cache_entries_directory, + _validate_cache_entry_path, _validate_cache_work_directory, ) from tracecat.logger import logger @@ -632,6 +634,7 @@ async def _evict_entry(self, cache_key: str) -> RegistryArtifactEviction: return RegistryArtifactEviction(retired=False, reclaimed=False) paths = self._paths_for(cache_key) + _validate_cache_entry_path(paths) if not paths.entry_dir.exists(): return RegistryArtifactEviction(retired=True, reclaimed=True) if registry_artifact_mounts.is_mount( @@ -644,6 +647,7 @@ async def _evict_entry(self, cache_key: str) -> RegistryArtifactEviction: ) return RegistryArtifactEviction(retired=False, reclaimed=False) + _validate_cache_entry_path(paths) try: trash_path = _move_entry_to_trash( paths.entry_dir, @@ -717,6 +721,7 @@ def _scan_cache_entries(self) -> dict[str, RegistryArtifactCacheEntry]: def _discover_cache_keys(self) -> set[str]: """Return cache keys represented by atomic entry directories.""" + _validate_cache_entries_directory(self.entries_dir) try: entries = list(os.scandir(self.entries_dir)) except FileNotFoundError: From 15be3fde924c6602f3457741cf7517ccc23931df Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:34:13 -0400 Subject: [PATCH 129/161] fix(executor): survive process monitor termination --- .../unit/executor/test_process_supervisor.py | 34 ++++++++++++++++++- tracecat/executor/process_supervisor.py | 7 ++++ 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_process_supervisor.py b/tests/unit/executor/test_process_supervisor.py index 069dc8461c..6314eb4f8c 100644 --- a/tests/unit/executor/test_process_supervisor.py +++ b/tests/unit/executor/test_process_supervisor.py @@ -45,6 +45,7 @@ def _write_action_script(path: Path) -> None: path.write_text( """ import os +import signal import subprocess import sys import time @@ -62,7 +63,9 @@ def _write_action_script(path: Path) -> None: pid_file.write_text(f"{os.getpid()} {child.pid}") if mode == "failure": raise SystemExit(23) -if mode == "block": +if mode == "kill-monitor": + os.kill(os.getppid(), signal.SIGKILL) +if mode in {"block", "kill-monitor"}: time.sleep(30) """.lstrip() ) @@ -153,3 +156,32 @@ async def test_supervisor_reaps_detached_descendant_before_cancellation_returns( with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) await process.wait() + + +@pytest.mark.anyio +async def test_supervisor_reaps_action_that_kills_its_monitor( + tmp_path: Path, +) -> None: + process, pid_file = await _spawn_supervised_action( + tmp_path, + mode="kill-monitor", + ) + tracked_pids: tuple[int, ...] = () + try: + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=5) + + assert process.returncode == 128 + signal.SIGKILL, stderr.decode() + assert stdout == b"" + tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) + action_pid, detached_pid = tracked_pids + assert not _process_is_running(action_pid) + assert not _process_is_running(detached_pid) + finally: + if not tracked_pids and pid_file.exists(): + tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) + for pid in tracked_pids: + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() diff --git a/tracecat/executor/process_supervisor.py b/tracecat/executor/process_supervisor.py index b5c94e1be3..4e8865b37a 100644 --- a/tracecat/executor/process_supervisor.py +++ b/tracecat/executor/process_supervisor.py @@ -134,6 +134,12 @@ def supervise(command: Sequence[str]) -> int: if not command: raise ValueError("A supervised command is required") + # The monitor is intentionally visible as the action's parent. Make this + # outer process a second subreaper so killing that monitor only reparents + # the action tree here, where it is still killed and reaped before return. + _set_child_subreaper() + _direct_child_pids() # Fail before execution when procfs tracking is absent. + control_read_fd, control_write_fd = os.pipe() writer_open = True @@ -160,6 +166,7 @@ def request_cleanup(_signal: int, _frame: FrameType | None) -> None: return _exit_code(monitor_status) finally: request_cleanup(signal.SIGTERM, None) + _kill_and_reap_children() def main() -> int: From 2b72e2c8b0674b3f79ae726efdc920138c161c1b Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:41:30 -0400 Subject: [PATCH 130/161] fix(executor): rescan run-python artifact leases --- tests/unit/executor/test_run_python_sdk_context.py | 8 +++++++- tracecat/executor/backends/base.py | 9 +++++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/tests/unit/executor/test_run_python_sdk_context.py b/tests/unit/executor/test_run_python_sdk_context.py index 9adb24b45d..bdef85587c 100644 --- a/tests/unit/executor/test_run_python_sdk_context.py +++ b/tests/unit/executor/test_run_python_sdk_context.py @@ -473,13 +473,18 @@ class _FakeRunPythonRegistryArtifacts: def __init__(self, paths: list[Path]) -> None: self.paths = paths self.artifact_uris: list[str] | None = None + self.paths_may_be_modified: bool | None = None self.leased = False @asynccontextmanager async def lease( - self, artifact_uris: list[str] | None = None + self, + artifact_uris: list[str] | None = None, + *, + paths_may_be_modified: bool = False, ) -> AsyncIterator[list[Path]]: self.artifact_uris = artifact_uris + self.paths_may_be_modified = paths_may_be_modified self.leased = True try: yield self.paths @@ -1232,6 +1237,7 @@ async def run_python(self, **kwargs: Any) -> dict[str, bool]: assert result.type == "success" assert leased_during_run == [True] assert fake_runner.registry_artifacts.leased is False + assert fake_runner.registry_artifacts.paths_may_be_modified is True @pytest.mark.anyio diff --git a/tracecat/executor/backends/base.py b/tracecat/executor/backends/base.py index 67f5169938..0507882b9d 100644 --- a/tracecat/executor/backends/base.py +++ b/tracecat/executor/backends/base.py @@ -149,9 +149,14 @@ async def _execute_run_python( ) # The lease is held for the whole sandbox run so cache eviction cannot - # delete a directory the script is still importing from. + # delete a directory the script is still importing from. SandboxService + # may select UnsafePidExecutor, which exposes these host paths writable, + # so conservatively rescan their footprint after every successful run. registry_artifacts = self._registry_artifact_cache() - async with registry_artifacts.lease(artifact_uris) as registry_paths: + async with registry_artifacts.lease( + artifact_uris, + paths_may_be_modified=True, + ) as registry_paths: return await self._run_python_in_sandbox( script=script, args=args, From c32ecdc6f21a022647ebcfb0d8e90b85cd64bf49 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:42:48 -0400 Subject: [PATCH 131/161] fix(executor): redact sidecar lookup failures --- .../test_registry_artifact_resolution.py | 37 +++++++++++++++++++ tracecat/executor/registry_artifacts.py | 2 +- 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py index 8f7b65cb62..06dd954b90 100644 --- a/tests/unit/executor/test_registry_artifact_resolution.py +++ b/tests/unit/executor/test_registry_artifact_resolution.py @@ -352,6 +352,43 @@ async def test_artifact_candidates_fall_back_to_gzip(self, temp_cache_dir): assert artifact.uri == "s3://bucket/path/site-packages.tar.gz" assert artifact.format == RegistryArtifactFormat.TAR_GZ + @pytest.mark.anyio + async def test_sidecar_lookup_failure_redacts_exception_text( + self, + temp_cache_dir: Path, + ) -> None: + """SDK exception strings cannot leak registry identifiers into logs.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://sensitive-bucket/org/repo/site-packages.tar.gz" + sidecar_uri = "s3://sensitive-bucket/org/repo/site-packages.squashfs" + ctx = cache._context_for(compute_registry_artifact_cache_key(artifact_uri)) + lookup_error = ConnectionError( + f"Could not connect to endpoint URL: {sidecar_uri}" + ) + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + side_effect=lookup_error, + ), + patch.object(cache, "_can_try_squashfs", return_value=True), + patch("tracecat.executor.registry_artifacts.logger.warning") as warning, + ): + candidates = await cache._artifact_candidates(ctx, artifact_uri) + + assert len(candidates) == 1 + assert isinstance(candidates[0], TarballArtifact) + warning.assert_called_once_with( + "Failed to check for registry artifact sidecar, falling back", + artifact_uri="s3://", + sidecar_uri="s3://", + artifact_format=RegistryArtifactFormat.SQUASHFS.value, + error_type="ConnectionError", + ) + assert "sensitive-bucket" not in repr(warning.call_args) + assert "org/repo" not in repr(warning.call_args) + def test_can_try_squashfs_does_not_require_mount_binary(self, temp_cache_dir): """Prefer SquashFS whenever enabled; extraction may work without mounts.""" cache = RegistryArtifactCache(temp_cache_dir) diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 2e93bc9775..c560d842e7 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -463,7 +463,7 @@ async def _sidecar_exists( artifact_uri=_artifact_uri_for_logging(base_uri), sidecar_uri=_artifact_uri_for_logging(sidecar_uri), artifact_format=artifact_format.value, - error=str(e), + error_type=type(e).__name__, ) return False From 1ea56a3d998257d3ecde026952eb95a8c544da5b Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:54:45 -0400 Subject: [PATCH 132/161] fix(storage): redact transport download failures --- tests/unit/test_storage_blob.py | 38 ++++++++++++++++++++++++++++++++- tracecat/storage/blob.py | 13 ++++++++++- 2 files changed, 49 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_storage_blob.py b/tests/unit/test_storage_blob.py index b5bbdd8ec1..2885151cb4 100644 --- a/tests/unit/test_storage_blob.py +++ b/tests/unit/test_storage_blob.py @@ -9,7 +9,7 @@ from urllib.parse import urlparse import pytest -from botocore.exceptions import ClientError +from botocore.exceptions import ClientError, EndpointConnectionError from tracecat.storage import blob as blob_module from tracecat.storage.blob import ( @@ -744,6 +744,42 @@ async def test_open_download_stream_redacts_sensitive_client_error( error_type="ClientError", ) + @pytest.mark.anyio + @patch("tracecat.storage.blob.get_storage_client") + async def test_open_download_stream_redacts_sensitive_transport_error( + self, + mock_get_client, + ) -> None: + mock_client = AsyncMock() + mock_get_client.return_value.__aenter__.return_value = mock_client + sensitive_endpoint = "https://tenant-bucket.invalid/tenant-key/org/repository" + mock_client.get_object.side_effect = EndpointConnectionError( + endpoint_url=sensitive_endpoint + ) + + with ( + patch("tracecat.storage.blob.logger.error") as log_error, + pytest.raises(blob_module.StorageDownloadError) as exc_info, + ): + async with open_download_stream( + key="tenant-key/org/repository", + bucket="tenant-bucket", + redact_log_identifiers=True, + ): + pass + + assert str(exc_info.value) == "Storage download failed" + assert exc_info.value.error_code is None + assert exc_info.value.__cause__ is None + log_error.assert_called_once_with( + "Failed to open download stream", + key="", + bucket="", + error_code=None, + error_type="EndpointConnectionError", + ) + assert sensitive_endpoint not in repr(log_error.call_args) + @pytest.mark.anyio async def test_download_file_to_path_writes_bytes( self, tmp_path: Path, monkeypatch diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index 77c55d28ac..7c2cd09fab 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -17,7 +17,7 @@ import aiofiles from aiobotocore.config import AioConfig from boto3.s3.transfer import TransferConfig -from botocore.exceptions import ClientError +from botocore.exceptions import BotoCoreError, ClientError from tracecat import config from tracecat.logger import logger @@ -796,6 +796,17 @@ async def open_download_stream( error=str(e), ) raise + except BotoCoreError as e: + if not redact_log_identifiers: + raise + logger.error( + "Failed to open download stream", + key=log_key, + bucket=log_bucket, + error_code=None, + error_type=type(e).__name__, + ) + raise StorageDownloadError(error_code=None) from None async def download_file_to_path( From 11e88768848849ded45c8062fe6b432f5fd07564 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sun, 2 Aug 2026 00:08:27 -0400 Subject: [PATCH 133/161] fix(executor): support portable procfs child discovery --- .../unit/executor/test_process_supervisor.py | 22 ++++++++++ tracecat/executor/process_supervisor.py | 43 ++++++++++++++++++- 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/tests/unit/executor/test_process_supervisor.py b/tests/unit/executor/test_process_supervisor.py index 6314eb4f8c..2cc5615830 100644 --- a/tests/unit/executor/test_process_supervisor.py +++ b/tests/unit/executor/test_process_supervisor.py @@ -71,6 +71,28 @@ def _write_action_script(path: Path) -> None: ) +def test_direct_child_discovery_falls_back_to_proc_stat(monkeypatch) -> None: + child_pid = os.fork() + if child_pid == 0: + signal.pause() + os._exit(0) + + def missing_children_file() -> list[int]: + raise FileNotFoundError + + monkeypatch.setattr( + process_supervisor, + "_direct_child_pids_from_children_file", + missing_children_file, + ) + try: + assert child_pid in process_supervisor._direct_child_pids() + finally: + with contextlib.suppress(ProcessLookupError): + os.kill(child_pid, signal.SIGKILL) + os.waitpid(child_pid, 0) + + async def _spawn_supervised_action( tmp_path: Path, *, diff --git a/tracecat/executor/process_supervisor.py b/tracecat/executor/process_supervisor.py index 4e8865b37a..ae59c47411 100644 --- a/tracecat/executor/process_supervisor.py +++ b/tracecat/executor/process_supervisor.py @@ -48,13 +48,52 @@ def _set_child_subreaper() -> None: raise OSError(error_number, os.strerror(error_number)) -def _direct_child_pids() -> list[int]: - """Return direct child PIDs from procfs for this single-threaded monitor.""" +def _direct_child_pids_from_children_file() -> list[int]: + """Return direct child PIDs using procfs's optional children file.""" children_path = Path(f"/proc/self/task/{os.getpid()}/children") contents = children_path.read_text().strip() return [int(pid) for pid in contents.split()] if contents else [] +def _proc_parent_pid(stat_path: Path) -> int: + """Read one process's parent PID from its procfs stat record.""" + contents = stat_path.read_text() + closing_paren = contents.rfind(")") + fields = contents[closing_paren + 1 :].split() + if closing_paren < 0 or len(fields) < 2: + raise RuntimeError("Malformed procfs stat record") + return int(fields[1]) + + +def _direct_child_pids_from_proc_stat() -> list[int]: + """Return direct child PIDs by scanning portable procfs stat records.""" + parent_pid = os.getpid() + _proc_parent_pid(Path("/proc/self/stat")) # Verify procfs is readable. + child_pids: list[int] = [] + + with os.scandir("/proc") as entries: + for entry in entries: + if not entry.name.isdecimal(): + continue + try: + candidate_pid = int(entry.name) + candidate_parent_pid = _proc_parent_pid(Path(entry.path) / "stat") + except (FileNotFoundError, PermissionError, ProcessLookupError): + continue + if candidate_parent_pid == parent_pid: + child_pids.append(candidate_pid) + + return child_pids + + +def _direct_child_pids() -> list[int]: + """Return direct child PIDs from either available procfs interface.""" + try: + return _direct_child_pids_from_children_file() + except OSError: + return _direct_child_pids_from_proc_stat() + + def _waitpid(pid: int, options: int = 0) -> tuple[int, int]: """Wait for a child while tolerating signal interruptions.""" while True: From a754b13d49ae233173ba6f7c5af7f5d5ad74c482 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sun, 2 Aug 2026 00:11:27 -0400 Subject: [PATCH 134/161] fix(executor): measure complete mutable cache entries --- .../executor/test_registry_artifact_budget.py | 33 ++++++++++-- .../executor/registry_artifact_storage.py | 50 ++++++++++--------- 2 files changed, 55 insertions(+), 28 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index ae82bfcd56..1b4deed6ff 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -86,6 +86,28 @@ def test_directory_footprint_counts_hard_linked_inode_once( assert allocated_size.call_count == 2 + def test_directory_footprint_prunes_directory_contents( + self, + temp_cache_dir: Path, + ) -> None: + mounted = temp_cache_dir / "mount" + mounted.mkdir() + (mounted / "module.py").write_text("x") + + with patch( + "tracecat.executor.registry_artifact_storage._allocated_stat_size", + return_value=4096, + ) as allocated_size: + assert ( + _directory_footprint( + temp_cache_dir, + pruned_directories=(mounted,), + ) + == 2 * 4096 + ) + + assert allocated_size.call_count == 2 + @pytest.mark.anyio async def test_admission_rounds_download_reservation_to_allocation_unit( self, @@ -661,21 +683,23 @@ async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( assert cache._budget_dirty is False @pytest.mark.anyio - async def test_releasing_a_mutable_cache_hit_rescans_post_admission_growth( + async def test_releasing_a_mutable_cache_hit_counts_unknown_entry_growth( self, temp_cache_dir: Path, ) -> None: - """Direct consumers cannot grow a warm entry outside budget accounting.""" + """Direct consumers cannot grow unknown entry paths outside the budget.""" cache = RegistryArtifactCache(temp_cache_dir) await cache.ensure_swept() artifact_uri = "s3://bucket/mutable-cached.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) target_dir = write_tarball_entry(temp_cache_dir, cache_key) + entry_dir = target_dir.parent + initial_size = cache._measure_entry(cache_key).size_bytes cache._budget_dirty = False with ( patch(MAX_ENTRIES_CONFIG, 10), - patch(MAX_BYTES_CONFIG, 1_000_000), + patch(MAX_BYTES_CONFIG, initial_size), patch.object( cache, "_scan_cache_entries", @@ -687,9 +711,10 @@ async def test_releasing_a_mutable_cache_hit_rescans_post_admission_growth( paths_may_be_modified=True, ) as registry_paths: assert registry_paths == [target_dir] - (target_dir / "action-output.bin").write_bytes(b"x" * 4096) + (entry_dir / "action-output.bin").write_bytes(b"x" * 4096) assert scan_cache_entries.call_count == 1 + assert not entry_dir.exists() assert cache._budget_dirty is False @pytest.mark.anyio diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index db369426dc..b6c8b34d52 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -113,11 +113,14 @@ def _directory_footprint( directory: Path, *, allocation_unit: int | None = None, + pruned_directories: Iterable[Path] = (), ) -> int: """Return the allocated footprint of a cache directory tree. Args: directory: Cache directory to measure. + pruned_directories: Directories whose own inodes are counted without + descending into their contents. Returns: Total allocated bytes of unique contained inodes, or zero when the @@ -132,6 +135,7 @@ def raise_walk_error(error: OSError) -> None: total_bytes = 0 seen_inodes: set[tuple[int, int]] = set() + pruned_paths = frozenset(pruned_directories) def allocated_inode_size(file_stat: os.stat_result) -> int: inode_key = (file_stat.st_dev, file_stat.st_ino) @@ -146,6 +150,7 @@ def allocated_inode_size(file_stat: os.stat_result) -> int: try: walker = os.walk(directory, onerror=raise_walk_error) for root, dirs, files in walker: + root_path = Path(root) try: total_bytes += allocated_inode_size(os.lstat(root)) except FileNotFoundError: @@ -157,13 +162,18 @@ def allocated_inode_size(file_stat: os.stat_result) -> int: ) except FileNotFoundError: continue + traversed_directories: list[str] = [] for directory_name in dirs: + child_path = root_path / directory_name try: - directory_stat = os.lstat(os.path.join(root, directory_name)) + directory_stat = child_path.lstat() except FileNotFoundError: continue - if stat.S_ISLNK(directory_stat.st_mode): + if stat.S_ISLNK(directory_stat.st_mode) or child_path in pruned_paths: total_bytes += allocated_inode_size(directory_stat) + continue + traversed_directories.append(directory_name) + dirs[:] = traversed_directories except FileNotFoundError: return 0 return total_bytes @@ -736,32 +746,24 @@ def _discover_cache_keys(self) -> set[str]: def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: """Measure the on-disk footprint and recency of one cache entry. - Mounted contents are excluded because the image already accounts for - their backing bytes. The entry root, mount-point inode, and image are - measured individually so a concurrent eviction cannot fail the scan. + Every entry-owned inode is included, including paths created by mutable + consumers. Active mounted contents are pruned because the image already + accounts for their backing bytes. """ paths = self._paths_for(cache_key) allocation_unit = _filesystem_allocation_unit(self.cache_dir) - size_bytes = 0 - - for path in ( - paths.entry_dir, - paths.squashfs_image_path, - paths.squashfs_mount_dir, - ): - try: - size_bytes += _allocated_stat_size( - path.lstat(), - allocation_unit=allocation_unit, - ) - except FileNotFoundError: - continue - - for directory in (paths.squashfs_extract_dir, paths.tarball_target_dir): - size_bytes += _directory_footprint( - directory, - allocation_unit=allocation_unit, + try: + mount_is_active = registry_artifact_mounts.is_mount( + paths.squashfs_mount_dir ) + except FileNotFoundError: + mount_is_active = False + pruned_directories = (paths.squashfs_mount_dir,) if mount_is_active else () + size_bytes = _directory_footprint( + paths.entry_dir, + allocation_unit=allocation_unit, + pruned_directories=pruned_directories, + ) try: last_used = paths.entry_dir.stat().st_mtime From 867029d7944a6d65c316f10a8882cd12704246d9 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sun, 2 Aug 2026 00:25:33 -0400 Subject: [PATCH 135/161] fix(executor): redact artifact fallback failures --- .../test_registry_artifact_resolution.py | 43 +++++++++++++++++++ tracecat/executor/registry_artifacts.py | 2 +- 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py index 06dd954b90..75257696d6 100644 --- a/tests/unit/executor/test_registry_artifact_resolution.py +++ b/tests/unit/executor/test_registry_artifact_resolution.py @@ -389,6 +389,49 @@ async def test_sidecar_lookup_failure_redacts_exception_text( assert "sensitive-bucket" not in repr(warning.call_args) assert "org/repo" not in repr(warning.call_args) + @pytest.mark.anyio + async def test_materialization_fallback_redacts_malformed_uri_error( + self, + temp_cache_dir: Path, + ) -> None: + """Malformed candidate URIs cannot leak identifiers through errors.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = ( + "https://tenant-bucket.invalid/org/repository/site-packages.squashfs" + ) + cache_key = compute_registry_artifact_cache_key(artifact_uri) + ctx = cache._context_for(cache_key) + candidates = await cache._artifact_candidates(ctx, artifact_uri) + fallback_path = temp_cache_dir / "fallback" + + with ( + patch( + "tracecat.executor.registry_artifact_materialization." + "config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", + False, + ), + patch.object( + TarballArtifact, + "materialize", + new_callable=AsyncMock, + return_value=[fallback_path], + ) as materialize_fallback, + patch("tracecat.executor.registry_artifacts.logger.warning") as warning, + ): + registry_paths = await cache._materialize_candidates(ctx, candidates) + + assert registry_paths == [fallback_path] + materialize_fallback.assert_awaited_once() + warning.assert_called_once_with( + "Failed to materialize registry artifact candidate, trying fallback", + cache_key=cache_key, + artifact_uri="https://", + artifact_format=RegistryArtifactFormat.SQUASHFS.value, + error_type="ValueError", + ) + assert "tenant-bucket" not in repr(warning.call_args) + assert "org/repository" not in repr(warning.call_args) + def test_can_try_squashfs_does_not_require_mount_binary(self, temp_cache_dir): """Prefer SquashFS whenever enabled; extraction may work without mounts.""" cache = RegistryArtifactCache(temp_cache_dir) diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index c560d842e7..f442b61564 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -312,7 +312,7 @@ async def _materialize_candidates( cache_key=cache_key, artifact_uri=_artifact_uri_for_logging(artifact.uri), artifact_format=artifact.format.value, - error=str(e), + error_type=type(e).__name__, ) raise RuntimeError(f"No registry artifact candidates for {ctx.cache_key}") From 46be136a95a867558c8db3fefca877427f7733a0 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sun, 2 Aug 2026 00:38:37 -0400 Subject: [PATCH 136/161] fix(executor): reject impossible cache reservations --- .../executor/test_registry_artifact_budget.py | 17 ++++++++++++++-- .../executor/registry_artifact_storage.py | 20 +++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 1b4deed6ff..5d57901efc 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -279,11 +279,17 @@ async def download_file_to_path( assert (registry_paths[0] / "module.py").read_bytes() == b"x" * 32 @pytest.mark.anyio - async def test_compression_heavy_tarball_is_rejected_before_extraction( + async def test_impossible_tarball_reservation_preserves_warm_entry( self, temp_cache_dir: Path ) -> None: - """Compressed bytes plus declared extraction cannot exceed the cache cap.""" + """Impossible extraction cannot evict warm entries before rejection.""" cache = RegistryArtifactCache(temp_cache_dir) + warm = write_image_entry( + temp_cache_dir, + "warm", + size=64, + mtime=100.0, + ) artifact_uri = "s3://bucket/compression-heavy.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) payload = tarball_payload(size=4096) @@ -319,6 +325,11 @@ async def download_file_to_path( "extract", new_callable=AsyncMock, ) as extract, + patch.object( + cache, + "_evict_entry", + wraps=cache._evict_entry, + ) as evict_entry, ): with pytest.raises(RegistryArtifactCacheCapacityError) as raised: async with cache.lease([artifact_uri]): @@ -327,6 +338,8 @@ async def download_file_to_path( assert raised.value.additional_bytes == 4097 assert raised.value.max_bytes == max_bytes extract.assert_not_awaited() + evict_entry.assert_not_awaited() + assert warm.is_file() assert not cache._paths_for(cache_key).entry_dir.exists() assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index b6c8b34d52..9dc8d561c0 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -477,6 +477,26 @@ async def _ensure_cache_capacity( + staging_bytes + trash_bytes ) + non_evictable_bytes = ( + staging_bytes + + trash_bytes + + sum( + entry.size_bytes + for entry in entries.values() + if entry.cache_key == protected_key + or self._refcount(entry.cache_key) > 0 + or ( + (runtime := self._runtime.get(entry.cache_key)) is not None + and runtime.lock.locked() + ) + ) + ) + if non_evictable_bytes + additional_bytes > max_bytes: + raise RegistryArtifactCacheCapacityError( + current_bytes=non_evictable_bytes, + additional_bytes=additional_bytes, + max_bytes=max_bytes, + ) if not cleanup_complete and total_bytes + additional_bytes > max_bytes: raise RegistryArtifactCacheCapacityError( current_bytes=total_bytes, From ee182e93d3ff0a2a273cdb7e443e1a64567730f7 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sun, 2 Aug 2026 00:49:00 -0400 Subject: [PATCH 137/161] fix(executor): sanitize invalid artifact URIs --- .../test_registry_artifact_resolution.py | 20 ++++++++++++++++++- .../registry_artifact_materialization.py | 9 ++++++++- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py index 75257696d6..e41d1442dc 100644 --- a/tests/unit/executor/test_registry_artifact_resolution.py +++ b/tests/unit/executor/test_registry_artifact_resolution.py @@ -427,11 +427,29 @@ async def test_materialization_fallback_redacts_malformed_uri_error( cache_key=cache_key, artifact_uri="https://", artifact_format=RegistryArtifactFormat.SQUASHFS.value, - error_type="ValueError", + error_type="RegistryArtifactUriError", ) assert "tenant-bucket" not in repr(warning.call_args) assert "org/repository" not in repr(warning.call_args) + @pytest.mark.anyio + async def test_final_candidate_malformed_uri_error_is_sanitized( + self, + temp_cache_dir: Path, + ) -> None: + """A final candidate cannot expose its malformed URI to callers.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "https://tenant-bucket.invalid/org/repository/artifact.tar.gz" + + with pytest.raises(ValueError) as exc_info: + async with cache.lease([artifact_uri]): + pass + + assert str(exc_info.value) == "Invalid registry artifact URI" + assert exc_info.value.__cause__ is None + assert "tenant-bucket" not in repr(exc_info.value) + assert "org/repository" not in repr(exc_info.value) + def test_can_try_squashfs_does_not_require_mount_binary(self, temp_cache_dir): """Prefer SquashFS whenever enabled; extraction may work without mounts.""" cache = RegistryArtifactCache(temp_cache_dir) diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index 4f6cd766bb..cdecabf563 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -88,6 +88,10 @@ class SquashfsMountCommandError(RuntimeError): """ +class RegistryArtifactUriError(ValueError): + """A registry artifact URI is malformed, with identifiers suppressed.""" + + @dataclass(frozen=True, slots=True) class RegistryArtifactAdmission: """Byte-bound admission hook shared by one cold materialization.""" @@ -848,7 +852,10 @@ async def _download_s3_artifact( defer_cleanup: Callable[[Path], None], ) -> None: """Download an S3 registry artifact to a local path.""" - bucket, key = parse_s3_uri(artifact_uri) + try: + bucket, key = parse_s3_uri(artifact_uri) + except ValueError: + raise RegistryArtifactUriError("Invalid registry artifact URI") from None try: if admission is None: await blob.download_file_to_path( From 4ae52601d1384739ca25dcb7607cdb471a7e441c Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sun, 2 Aug 2026 00:53:54 -0400 Subject: [PATCH 138/161] fix(executor): reserve tar directory metadata --- .../executor/test_registry_artifact_budget.py | 10 +++-- .../test_registry_artifact_materialization.py | 30 +++++++++++++-- .../registry_artifact_materialization.py | 38 +++++++++++++++++++ 3 files changed, 71 insertions(+), 7 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index 5d57901efc..c0de89b22b 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -237,7 +237,8 @@ async def test_cold_download_reserves_space_before_writing( artifact_uri = "s3://bucket/new.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) payload = tarball_payload(size=32) - max_bytes = len(payload) + 33 + extracted_size = 74 # File, root, and directory entry at unit size 1. + max_bytes = len(payload) + extracted_size capacity_checked = False async def download_file_to_path( @@ -253,7 +254,7 @@ async def download_file_to_path( del key, bucket, defer_cleanup assert redact_log_identifiers is True nonlocal capacity_checked - assert max_bytes == len(payload) + 33 + assert max_bytes == len(payload) + extracted_size await ensure_capacity(len(payload)) capacity_checked = True assert not idle.exists() @@ -335,7 +336,7 @@ async def download_file_to_path( async with cache.lease([artifact_uri]): pass - assert raised.value.additional_bytes == 4097 + assert raised.value.additional_bytes == 4138 assert raised.value.max_bytes == max_bytes extract.assert_not_awaited() evict_entry.assert_not_awaited() @@ -352,7 +353,8 @@ async def test_failed_squashfs_bytes_do_not_block_tarball_fallback( artifact_uri = "s3://bucket/path/site-packages.squashfs" cache_key = compute_registry_artifact_cache_key(artifact_uri) payload = tarball_payload(size=32) - max_bytes = len(payload) + 33 + extracted_size = 74 # File, root, and directory entry at unit size 1. + max_bytes = len(payload) + extracted_size async def fail_after_squashfs_download( self: SquashfsArtifact, diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 2f261a5913..8cc8944060 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -344,7 +344,7 @@ def test_tarball_size_bounds_each_member_allocation( member.size = 1 tar.addfile(member, io.BytesIO(b"x")) - assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 16_384 + assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 28_672 def test_tarball_size_includes_extraction_root( self, @@ -357,7 +357,7 @@ def test_tarball_size_includes_extraction_root( member.size = 0 tar.addfile(member, io.BytesIO()) - assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 8192 + assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 12_288 def test_tarball_size_does_not_duplicate_explicit_root( self, @@ -382,7 +382,31 @@ def test_tarball_size_includes_implicit_parent_directories( member.size = 0 tar.addfile(member, io.BytesIO()) - assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 20_480 + assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 36_864 + + def test_tarball_size_reserves_growing_directory_metadata( + self, + temp_cache_dir: Path, + ) -> None: + """Many child entries reserve more than one directory block.""" + tarball_path = temp_cache_dir / "wide-directory.tar.gz" + child_count = 1000 + with tarfile.open(tarball_path, "w:gz") as tar: + for index in range(child_count): + member = tarfile.TarInfo(f"module-{index:04d}.py") + member.size = 0 + tar.addfile(member, io.BytesIO()) + + file_bytes = child_count * 4096 + root_bytes = 4096 + directory_entry_bytes = child_count * 4096 + assert ( + _tarball_extracted_size( + tarball_path, + allocation_unit=4096, + ) + == file_bytes + root_bytes + directory_entry_bytes + ) def test_squashfs_listing_size_rejects_unparseable_files(self) -> None: with pytest.raises(ValueError, match="Could not parse SquashFS listing"): diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index cdecabf563..ef12e0ca21 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -983,6 +983,27 @@ def _allocated_size_bound(size_bytes: int, *, allocation_unit: int) -> int: ) +_DIRECTORY_ENTRY_OVERHEAD_BYTES = 32 +"""Conservative per-child filesystem directory-record overhead.""" + + +def _directory_entry_size_bound( + child_name: str, + *, + allocation_unit: int, +) -> int: + """Reserve directory storage for one unique child name. + + Each child receives at least one allocation unit. This intentionally + overbounds filesystem-specific records and index blocks without assuming a + particular executor filesystem layout. + """ + return _allocated_size_bound( + _DIRECTORY_ENTRY_OVERHEAD_BYTES + len(os.fsencode(child_name)), + allocation_unit=allocation_unit, + ) + + def _tarball_extracted_size( tarball_path: Path, *, @@ -994,6 +1015,13 @@ def _tarball_extracted_size( has_explicit_root_directory = False required_parent_dirs: set[PurePosixPath] = set() explicit_dirs: set[PurePosixPath] = set() + directory_children: dict[PurePosixPath, set[str]] = {} + + def record_directory_child(path: PurePosixPath) -> None: + if path == root_path: + return + directory_children.setdefault(path.parent, set()).add(path.name) + with tarfile.open(tarball_path, "r:gz") as tar: for member in tar: if member.size < 0: @@ -1005,6 +1033,7 @@ def _tarball_extracted_size( allocation_unit=allocation_unit, ) member_path = PurePosixPath(member.name) + record_directory_child(member_path) if member.isdir(): explicit_dirs.add(member_path) has_explicit_root_directory |= member_path == root_path @@ -1012,11 +1041,20 @@ def _tarball_extracted_size( if parent == root_path: break required_parent_dirs.add(parent) + record_directory_child(parent) # Extraction creates a target root even when the tar manifest omits it. if not has_explicit_root_directory: total_bytes += _allocated_size_bound(0, allocation_unit=allocation_unit) implicit_parent_dirs = required_parent_dirs - explicit_dirs total_bytes += len(implicit_parent_dirs) * allocation_unit + total_bytes += sum( + _directory_entry_size_bound( + child_name, + allocation_unit=allocation_unit, + ) + for child_names in directory_children.values() + for child_name in child_names + ) return total_bytes From 5e402c6bad485ab4380a304b13a1c28dae9be2ac Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sun, 2 Aug 2026 01:03:20 -0400 Subject: [PATCH 139/161] fix(executor): contain sidecar URI parse failures --- .../test_registry_artifact_resolution.py | 35 +++++++++++++++++++ tracecat/executor/registry_artifacts.py | 2 +- 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py index e41d1442dc..bab757c885 100644 --- a/tests/unit/executor/test_registry_artifact_resolution.py +++ b/tests/unit/executor/test_registry_artifact_resolution.py @@ -389,6 +389,41 @@ async def test_sidecar_lookup_failure_redacts_exception_text( assert "sensitive-bucket" not in repr(warning.call_args) assert "org/repo" not in repr(warning.call_args) + @pytest.mark.anyio + async def test_sidecar_parse_failure_uses_redacted_fallback( + self, + temp_cache_dir: Path, + ) -> None: + """Malformed sidecar URIs cannot escape the redacted fallback path.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = ( + "https://tenant-bucket.invalid/org/repository/site-packages.tar.gz" + ) + ctx = cache._context_for(compute_registry_artifact_cache_key(artifact_uri)) + + with ( + patch.object(cache, "_can_try_squashfs", return_value=True), + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + ) as file_exists, + patch("tracecat.executor.registry_artifacts.logger.warning") as warning, + ): + candidates = await cache._artifact_candidates(ctx, artifact_uri) + + assert len(candidates) == 1 + assert isinstance(candidates[0], TarballArtifact) + file_exists.assert_not_awaited() + warning.assert_called_once_with( + "Failed to check for registry artifact sidecar, falling back", + artifact_uri="https://", + sidecar_uri="https://", + artifact_format=RegistryArtifactFormat.SQUASHFS.value, + error_type="ValueError", + ) + assert "tenant-bucket" not in repr(warning.call_args) + assert "org/repository" not in repr(warning.call_args) + @pytest.mark.anyio async def test_materialization_fallback_redacts_malformed_uri_error( self, diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index f442b61564..b4412f6845 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -447,8 +447,8 @@ async def _sidecar_exists( artifact_format: RegistryArtifactFormat, ) -> bool: """Return whether a registry sidecar exists, logging lookup failures.""" - bucket, key = parse_s3_uri(sidecar_uri) try: + bucket, key = parse_s3_uri(sidecar_uri) if await blob.file_exists(key=key, bucket=bucket): logger.debug( "Using registry artifact sidecar", From 7d6382548be96554ce2f3ecd952d6b3b5138617d Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sun, 2 Aug 2026 01:07:12 -0400 Subject: [PATCH 140/161] fix(executor): reserve SquashFS directory metadata --- .../test_registry_artifact_materialization.py | 30 ++++++++++- .../registry_artifact_materialization.py | 51 +++++++++++++++++-- 2 files changed, 76 insertions(+), 5 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 8cc8944060..6db5c8ff92 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -326,12 +326,38 @@ def test_squashfs_listing_size_bounds_each_inode_allocation(self) -> None: ] ) - assert _squashfs_listing_size(listing, allocation_unit=4096) == 12_288 + assert _squashfs_listing_size(listing, allocation_unit=4096) == 20_480 def test_squashfs_listing_size_accepts_non_utf8_filenames(self) -> None: listing = b"-rw-r--r-- 0/0 123 2026-01-01 00:00 squashfs-root/module-\xff.py" - assert _squashfs_listing_size(listing, allocation_unit=4096) == 4096 + assert _squashfs_listing_size(listing, allocation_unit=4096) == 12_288 + + def test_squashfs_listing_size_reserves_growing_directory_metadata( + self, + ) -> None: + """Wide extracted directories reserve metadata for every child.""" + child_count = 1000 + listing = b"\n".join( + [ + b"drwxr-xr-x 0/0 64 2026-01-01 00:00 squashfs-root", + *( + f"-rw-r--r-- 0/0 0 2026-01-01 00:00 " + f"squashfs-root/module-{index:04d}.py".encode() + for index in range(child_count) + ), + ] + ) + + inode_bytes = (child_count + 1) * 4096 + directory_entry_bytes = child_count * 4096 + assert ( + _squashfs_listing_size( + listing, + allocation_unit=4096, + ) + == inode_bytes + directory_entry_bytes + ) def test_tarball_size_bounds_each_member_allocation( self, diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index ef12e0ca21..de84688b3d 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -1062,21 +1062,66 @@ def _squashfs_listing_size(output: bytes, *, allocation_unit: int = 1) -> int: """Bound allocated bytes from ``unsquashfs -lln`` output, failing closed.""" total_bytes = 0 parsed_entries = 0 + root_path = PurePosixPath(".") + listing_root: str | None = None + required_dirs: set[PurePosixPath] = {root_path} + listed_dirs: set[PurePosixPath] = set() + directory_children: dict[PurePosixPath, set[str]] = {} + + def record_directory_child(path: PurePosixPath) -> None: + if path == root_path: + return + directory_children.setdefault(path.parent, set()).add(path.name) + for raw_line in output.decode(errors="replace").splitlines(): line = raw_line.strip() if not line: continue - fields = line.split(maxsplit=4) + fields = line.split(maxsplit=5) mode = fields[0] if len(mode) != 10 or mode[0] not in "bcdlps-": continue - if len(fields) < 5 or "/" not in fields[1] or not fields[2].isdigit(): - raise ValueError(f"Could not parse SquashFS listing line: {line}") + if len(fields) < 6 or "/" not in fields[1] or not fields[2].isdigit(): + raise ValueError("Could not parse SquashFS listing line") + + listed_path_text = fields[5] + if mode[0] == "l": + listed_path_text = listed_path_text.split(" -> ", maxsplit=1)[0] + listed_path = PurePosixPath(listed_path_text) + if listed_path.is_absolute() or not listed_path.parts: + raise ValueError("Could not parse SquashFS listing path") + if listing_root is None: + listing_root = listed_path.parts[0] + elif listed_path.parts[0] != listing_root: + raise ValueError("Inconsistent SquashFS listing root") + + relative_parts = listed_path.parts[1:] + entry_path = PurePosixPath(*relative_parts) if relative_parts else root_path + if entry_path == root_path and mode[0] != "d": + raise ValueError("Could not parse SquashFS listing root") + parsed_entries += 1 total_bytes += _allocated_size_bound( int(fields[2]), allocation_unit=allocation_unit, ) + record_directory_child(entry_path) + if mode[0] == "d": + listed_dirs.add(entry_path) + for parent in entry_path.parents: + required_dirs.add(parent) + if parent == root_path: + break + record_directory_child(parent) if parsed_entries == 0: raise ValueError("Could not parse any SquashFS listing entries") + total_bytes += len(required_dirs - listed_dirs) * allocation_unit + total_bytes += sum( + _directory_entry_size_bound( + child_name, + allocation_unit=allocation_unit, + ) + for child_names in directory_children.values() + for child_name in child_names + ) return total_bytes From de633a366a09cd071e337fd6fc6ad2f12d6a469e Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sun, 2 Aug 2026 01:21:05 -0400 Subject: [PATCH 141/161] fix(executor): recover stopped process supervisors --- .../unit/executor/test_process_supervisor.py | 48 ++++++++++++++++++- tracecat/executor/process_supervisor.py | 14 +++++- tracecat/sandbox/utils.py | 5 ++ 3 files changed, 64 insertions(+), 3 deletions(-) diff --git a/tests/unit/executor/test_process_supervisor.py b/tests/unit/executor/test_process_supervisor.py index 2cc5615830..cbafd6b644 100644 --- a/tests/unit/executor/test_process_supervisor.py +++ b/tests/unit/executor/test_process_supervisor.py @@ -65,7 +65,19 @@ def _write_action_script(path: Path) -> None: raise SystemExit(23) if mode == "kill-monitor": os.kill(os.getppid(), signal.SIGKILL) -if mode in {"block", "kill-monitor"}: +if mode in {"stop-monitor", "stop-supervisors"}: + monitor_pid = os.getppid() + outer_pid = int( + next( + line.removeprefix("PPid:").strip() + for line in Path(f"/proc/{monitor_pid}/status").read_text().splitlines() + if line.startswith("PPid:") + ) + ) + os.kill(monitor_pid, signal.SIGSTOP) + if mode == "stop-supervisors": + os.kill(outer_pid, signal.SIGSTOP) +if mode in {"block", "kill-monitor", "stop-monitor", "stop-supervisors"}: time.sleep(30) """.lstrip() ) @@ -207,3 +219,37 @@ async def test_supervisor_reaps_action_that_kills_its_monitor( with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) await process.wait() + + +@pytest.mark.anyio +@pytest.mark.parametrize("mode", ["stop-monitor", "stop-supervisors"]) +async def test_termination_recovers_stopped_supervisor_tree( + tmp_path: Path, + mode: str, +) -> None: + process, pid_file = await _spawn_supervised_action(tmp_path, mode=mode) + await _wait_for_file(pid_file) + tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) + action_pid, detached_pid = tracked_pids + + try: + with pytest.raises(TimeoutError): + await asyncio.wait_for( + communicate_process_group( + process, + timeout=0.05, + terminate=terminate_supervised_process, + ), + timeout=5, + ) + + assert process.returncode is not None + assert not _process_is_running(action_pid) + assert not _process_is_running(detached_pid) + finally: + for pid in tracked_pids: + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() diff --git a/tracecat/executor/process_supervisor.py b/tracecat/executor/process_supervisor.py index ae59c47411..712d44e5d4 100644 --- a/tracecat/executor/process_supervisor.py +++ b/tracecat/executor/process_supervisor.py @@ -181,14 +181,24 @@ def supervise(command: Sequence[str]) -> int: control_read_fd, control_write_fd = os.pipe() writer_open = True + monitor_pid: int | None = None - def request_cleanup(_signal: int, _frame: FrameType | None) -> None: + def close_control_pipe() -> None: nonlocal writer_open if writer_open: with suppress(OSError): os.close(control_write_fd) writer_open = False + def request_cleanup(_signal: int, _frame: FrameType | None) -> None: + close_control_pipe() + # The action shares the monitor's UID and can stop it. SIGKILL the + # monitor session so the outer subreaper adopts and reaps every member, + # including descendants that detached into their own sessions. + if monitor_pid is not None: + with suppress(ProcessLookupError): + os.killpg(monitor_pid, signal.SIGKILL) + signal.signal(signal.SIGTERM, request_cleanup) signal.signal(signal.SIGINT, request_cleanup) @@ -204,7 +214,7 @@ def request_cleanup(_signal: int, _frame: FrameType | None) -> None: _, monitor_status = _waitpid(monitor_pid) return _exit_code(monitor_status) finally: - request_cleanup(signal.SIGTERM, None) + close_control_pipe() _kill_and_reap_children() diff --git a/tracecat/sandbox/utils.py b/tracecat/sandbox/utils.py index fc63c8a14c..f75fd0773a 100644 --- a/tracecat/sandbox/utils.py +++ b/tracecat/sandbox/utils.py @@ -47,6 +47,11 @@ async def terminate_supervised_process(process: asyncio.subprocess.Process) -> N if process.returncode is None: with suppress(ProcessLookupError): os.kill(process.pid, signal.SIGTERM) + # An untrusted action runs under the supervisor's UID and can stop the + # outer supervisor. SIGTERM remains pending for a stopped process, so + # resume it before waiting for its cleanup handler to run. + with suppress(ProcessLookupError): + os.kill(process.pid, signal.SIGCONT) # The supervisor exits only after its detached subreaper has killed and # reaped the action's complete descendant tree. Do not impose a shorter # wait that could release registry leases while cleanup is still active. From 33962afc87deb643df16e3a36066d74233f65cc7 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sun, 2 Aug 2026 01:37:10 -0400 Subject: [PATCH 142/161] fix(executor): count cache structural metadata --- .../executor/test_registry_artifact_budget.py | 69 ++++++++++++++++ .../executor/registry_artifact_storage.py | 82 ++++++++++++++++--- 2 files changed, 138 insertions(+), 13 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py index c0de89b22b..160bb79756 100644 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -108,6 +108,75 @@ def test_directory_footprint_prunes_directory_contents( assert allocated_size.call_count == 2 + def test_directory_footprint_can_exclude_root_inode( + self, + temp_cache_dir: Path, + ) -> None: + nested = temp_cache_dir / "nested" + nested.mkdir() + (nested / "module.py").write_text("x") + + with patch( + "tracecat.executor.registry_artifact_storage._allocated_stat_size", + return_value=4096, + ) as allocated_size: + assert _directory_footprint(temp_cache_dir, include_root=False) == 2 * 4096 + + assert allocated_size.call_count == 2 + + def test_cache_structure_counts_roots_without_subtree_contents( + self, + temp_cache_dir: Path, + ) -> None: + cache = RegistryArtifactCache(temp_cache_dir) + for root in (cache.entries_dir, cache.staging_dir, cache.trash_dir): + child = root / "child" + child.mkdir(parents=True) + (child / "payload").write_text("x") + base_dir = temp_cache_dir / "base" + base_dir.mkdir() + (base_dir / "payload").write_text("x") + + with patch( + "tracecat.executor.registry_artifact_storage._allocated_stat_size", + return_value=4096, + ) as allocated_size: + assert cache._cache_structural_footprint() == 6 * 4096 + + # Cache, entries, staging, and trash roots plus the unpruned base tree. + assert allocated_size.call_count == 6 + + @pytest.mark.anyio + async def test_admission_counts_non_evictable_cache_structure( + self, + temp_cache_dir: Path, + ) -> None: + cache = RegistryArtifactCache(temp_cache_dir) + + with patch.object(cache, "_cache_structural_footprint", return_value=5): + with pytest.raises(RegistryArtifactCacheCapacityError) as raised: + await cache._ensure_cache_capacity( + additional_bytes=0, + protected_key="new", + max_bytes=4, + ) + + assert raised.value.current_bytes == 5 + assert raised.value.additional_bytes == 0 + + @pytest.mark.anyio + async def test_enforcement_counts_cache_structure( + self, + temp_cache_dir: Path, + ) -> None: + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch(MAX_BYTES_CONFIG, 4), + patch.object(cache, "_cache_structural_footprint", return_value=5), + ): + assert await cache._enforce_cache_budget() is False + @pytest.mark.anyio async def test_admission_rounds_download_reservation_to_allocation_unit( self, diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 9dc8d561c0..6e4360689b 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -114,6 +114,7 @@ def _directory_footprint( *, allocation_unit: int | None = None, pruned_directories: Iterable[Path] = (), + include_root: bool = True, ) -> int: """Return the allocated footprint of a cache directory tree. @@ -121,6 +122,8 @@ def _directory_footprint( directory: Cache directory to measure. pruned_directories: Directories whose own inodes are counted without descending into their contents. + include_root: Whether to count the root directory's inode. Disable this + when a separate structural scan already owns that inode. Returns: Total allocated bytes of unique contained inodes, or zero when the @@ -151,10 +154,11 @@ def allocated_inode_size(file_stat: os.stat_result) -> int: walker = os.walk(directory, onerror=raise_walk_error) for root, dirs, files in walker: root_path = Path(root) - try: - total_bytes += allocated_inode_size(os.lstat(root)) - except FileNotFoundError: - continue + if include_root or root_path != directory: + try: + total_bytes += allocated_inode_size(os.lstat(root)) + except FileNotFoundError: + continue for file_name in files: try: total_bytes += allocated_inode_size( @@ -224,6 +228,22 @@ def _move_entry_to_trash(entry_dir: Path, trash_dir: Path, cache_key: str) -> Pa class _RegistryArtifactCacheStorage(_RegistryArtifactCacheState): """Adds startup recovery, eviction, and byte-budget enforcement.""" + def _cache_structural_footprint(self) -> int: + """Measure cache roots and non-entry data exactly once. + + Entry, staging, and trash contents are measured separately. Pruning + those trees here retains their root-directory blocks, including growth + caused by child names, while avoiding double-counting their contents. + """ + return _directory_footprint( + self.cache_dir, + pruned_directories=( + self.entries_dir, + self.staging_dir, + self.trash_dir, + ), + ) + async def ensure_swept(self) -> None: """Run the startup sweep exactly once successfully, off the event loop. @@ -412,8 +432,26 @@ async def _enforce_cache_budget_locked( if max_entries <= 0 and max_bytes <= 0: return True - entries = await asyncio.to_thread(self._scan_cache_entries) - total_bytes = sum(entry.size_bytes for entry in entries.values()) + entries, structural_bytes, staging_bytes, trash_bytes = await asyncio.gather( + asyncio.to_thread(self._scan_cache_entries), + asyncio.to_thread(self._cache_structural_footprint), + asyncio.to_thread( + _directory_footprint, + self.staging_dir, + include_root=False, + ), + asyncio.to_thread( + _directory_footprint, + self.trash_dir, + include_root=False, + ), + ) + total_bytes = ( + structural_bytes + + staging_bytes + + trash_bytes + + sum(entry.size_bytes for entry in entries.values()) + ) protected = set() if protected_key is None else {protected_key} eviction_pass = await self._evict_until_fits( entries, @@ -467,18 +505,34 @@ async def _ensure_cache_capacity( async with self._budget_lock: cleanup_complete = await self._cleanup_cache_work_dirs() - entries = await asyncio.to_thread(self._scan_cache_entries) - staging_bytes, trash_bytes = await asyncio.gather( - asyncio.to_thread(_directory_footprint, self.staging_dir), - asyncio.to_thread(_directory_footprint, self.trash_dir), + ( + entries, + structural_bytes, + staging_bytes, + trash_bytes, + ) = await asyncio.gather( + asyncio.to_thread(self._scan_cache_entries), + asyncio.to_thread(self._cache_structural_footprint), + asyncio.to_thread( + _directory_footprint, + self.staging_dir, + include_root=False, + ), + asyncio.to_thread( + _directory_footprint, + self.trash_dir, + include_root=False, + ), ) total_bytes = ( - sum(entry.size_bytes for entry in entries.values()) + structural_bytes + staging_bytes + trash_bytes + + sum(entry.size_bytes for entry in entries.values()) ) non_evictable_bytes = ( - staging_bytes + structural_bytes + + staging_bytes + trash_bytes + sum( entry.size_bytes @@ -915,7 +969,9 @@ def _trim_startup_cache(self) -> bool: if max_entries <= 0 and max_bytes <= 0: return True - total_bytes = sum(entry.size_bytes for entry in entries.values()) + total_bytes = self._cache_structural_footprint() + sum( + entry.size_bytes for entry in entries.values() + ) # Mounted entries belong to a live process sharing this cache directory. candidates = sorted( ( From 56bf7b1056b812da903759d40c52993bf54242f2 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sun, 2 Aug 2026 01:48:22 -0400 Subject: [PATCH 143/161] fix(executor): sanitize archive extraction failures --- .../test_registry_artifact_materialization.py | 26 +++++++++++++++++++ .../registry_artifact_materialization.py | 25 +++++++++++++----- tracecat/executor/registry_artifacts.py | 2 ++ 3 files changed, 47 insertions(+), 6 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 6db5c8ff92..4de4f76904 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -16,6 +16,7 @@ from tracecat.executor.registry_artifacts import ( SQUASHFS_MOUNT_OPTIONS, RegistryArtifactCache, + RegistryArtifactExtractionError, RegistryArtifactMaterializationContext, SquashfsArtifact, SquashfsMountCommandError, @@ -811,6 +812,31 @@ def blocking_extractall(*args: object, **kwargs: object) -> None: assert first_cancellation_propagated_early is False assert second_cancellation_propagated_early is False + @pytest.mark.anyio + async def test_tarball_extract_sanitizes_member_failures( + self, + temp_cache_dir: Path, + ) -> None: + sensitive_member = "../../synthetic-customer/repository/module.py" + artifact = TarballArtifact( + uri="s3://bucket/path/site-packages.tar.gz", + cache_key="malformed-tarball", + ) + tarball_path = temp_cache_dir / "artifact.tar.gz" + target_dir = temp_cache_dir / "target" + target_dir.mkdir() + with tarfile.open(tarball_path, "w:gz") as tar: + member = tarfile.TarInfo(sensitive_member) + member.size = 1 + tar.addfile(member, io.BytesIO(b"x")) + + with pytest.raises(RegistryArtifactExtractionError) as raised: + await artifact.extract(tarball_path, target_dir) + + assert str(raised.value) == "Registry artifact extraction failed" + assert sensitive_member not in str(raised.value) + assert raised.value.__cause__ is None + @pytest.mark.anyio async def test_repeatedly_cancelled_tarball_size_scan_rejoins_thread( self, temp_cache_dir diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index de84688b3d..a55a407043 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -92,6 +92,13 @@ class RegistryArtifactUriError(ValueError): """A registry artifact URI is malformed, with identifiers suppressed.""" +class RegistryArtifactExtractionError(RuntimeError): + """A registry archive could not be inspected or extracted safely.""" + + def __init__(self) -> None: + super().__init__("Registry artifact extraction failed") + + @dataclass(frozen=True, slots=True) class RegistryArtifactAdmission: """Byte-bound admission hook shared by one cold materialization.""" @@ -753,12 +760,15 @@ async def materialize( admission = ctx.admission if admission is not None: - extracted_size = await _run_blocking_rejoin_on_cancel( - lambda: _tarball_extracted_size( - temp_tarball, - allocation_unit=admission.allocation_unit, + try: + extracted_size = await _run_blocking_rejoin_on_cancel( + lambda: _tarball_extracted_size( + temp_tarball, + allocation_unit=admission.allocation_unit, + ) ) - ) + except Exception: + raise RegistryArtifactExtractionError() from None await admission.ensure_capacity(extracted_size) extract_start = time.monotonic() @@ -835,7 +845,10 @@ def _do_extract() -> None: raise ValueError(f"Unsupported tarball format: {tarball_path}") - await _run_blocking_rejoin_on_cancel(_do_extract) + try: + await _run_blocking_rejoin_on_cancel(_do_extract) + except Exception: + raise RegistryArtifactExtractionError() from None logger.debug( "Tarball extracted", diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index b4412f6845..4156cd1dc5 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -24,6 +24,7 @@ BuiltinArtifact, RegistryArtifact, RegistryArtifactAdmission, + RegistryArtifactExtractionError, RegistryArtifactFormat, RegistryArtifactMaterializationContext, RegistryArtifactPaths, @@ -75,6 +76,7 @@ "RegistryArtifactCacheEntry", "RegistryArtifactCacheLoopError", "RegistryArtifactEviction", + "RegistryArtifactExtractionError", "RegistryArtifactFormat", "RegistryArtifactMaterializationContext", "RegistryArtifactPaths", From 51c3bd63204541e45213ac9389f5e51f2138ca26 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sun, 2 Aug 2026 01:59:16 -0400 Subject: [PATCH 144/161] fix(executor): sanitize SquashFS extraction failures --- .../test_registry_artifact_materialization.py | 47 +++++++++++++++++++ .../registry_artifact_materialization.py | 15 +++--- 2 files changed, 55 insertions(+), 7 deletions(-) diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py index 4de4f76904..33792d1fbb 100644 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ b/tests/unit/executor/test_registry_artifact_materialization.py @@ -702,6 +702,53 @@ async def create_sleep_subprocess( assert production_reaped is True assert captured.returncode is not None + @pytest.mark.parametrize("operation", ["extract", "size-command", "size-parse"]) + @pytest.mark.anyio + async def test_squashfs_failures_sanitize_subprocess_output( + self, + temp_cache_dir: Path, + operation: str, + ) -> None: + sensitive_output = b"synthetic-customer/repository/member.py" + artifact = SquashfsArtifact( + uri="s3://bucket/path/custom.squashfs", + cache_key="malformed-squashfs", + ) + image_path = temp_cache_dir / "image.squashfs" + image_path.write_bytes(b"squashfs") + target_dir = temp_cache_dir / "target" + target_dir.mkdir() + process = AsyncMock() + process.returncode = 0 if operation == "size-parse" else 1 + stdout = sensitive_output if operation == "size-parse" else b"" + stderr = b"" if operation == "size-parse" else sensitive_output + + with ( + patch( + "tracecat.executor.registry_artifact_materialization.shutil.which", + return_value="/usr/bin/unsquashfs", + ), + patch( + "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ), + patch( + "tracecat.executor.registry_artifact_materialization.communicate_process_group", + new_callable=AsyncMock, + return_value=(stdout, stderr), + ), + pytest.raises(RegistryArtifactExtractionError) as raised, + ): + if operation == "extract": + await artifact._extract_image(image_path, target_dir) + else: + await artifact._squashfs_extracted_size(image_path) + + assert str(raised.value) == "Registry artifact extraction failed" + assert sensitive_output.decode() not in str(raised.value) + assert raised.value.__cause__ is None + @pytest.mark.parametrize("operation", ["mount", "extract", "size"]) @pytest.mark.anyio async def test_repeated_cancellation_reaps_squashfs_subprocess( diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py index a55a407043..111064205c 100644 --- a/tracecat/executor/registry_artifact_materialization.py +++ b/tracecat/executor/registry_artifact_materialization.py @@ -676,13 +676,12 @@ async def _extract_image(self, image_path: Path, target_dir: Path) -> None: stderr=asyncio.subprocess.PIPE, start_new_session=True, ) - stdout, stderr = await communicate_process_group(proc) + await communicate_process_group(proc) if proc.returncode == 0: return - output = (stderr or stdout).decode(errors="replace").strip() - raise RuntimeError(output or "unsquashfs command failed") + raise RegistryArtifactExtractionError() async def _squashfs_extracted_size( self, @@ -703,12 +702,14 @@ async def _squashfs_extracted_size( stderr=asyncio.subprocess.PIPE, start_new_session=True, ) - stdout, stderr = await communicate_process_group(proc) + stdout, _ = await communicate_process_group(proc) if proc.returncode != 0: - output = (stderr or stdout).decode(errors="replace").strip() - raise RuntimeError(output or "unsquashfs listing failed") - return _squashfs_listing_size(stdout, allocation_unit=allocation_unit) + raise RegistryArtifactExtractionError() + try: + return _squashfs_listing_size(stdout, allocation_unit=allocation_unit) + except Exception: + raise RegistryArtifactExtractionError() from None @dataclass(frozen=True, slots=True) From 23b8a40995358dd960473b576d24089814d1727f Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Sun, 2 Aug 2026 02:13:03 -0400 Subject: [PATCH 145/161] fix(executor): recover monitor after supervisor death --- .../unit/executor/test_process_supervisor.py | 72 ++++++++++++++++--- tracecat/executor/process_supervisor.py | 57 +++++++++++++-- 2 files changed, 114 insertions(+), 15 deletions(-) diff --git a/tests/unit/executor/test_process_supervisor.py b/tests/unit/executor/test_process_supervisor.py index cbafd6b644..a63a953dd0 100644 --- a/tests/unit/executor/test_process_supervisor.py +++ b/tests/unit/executor/test_process_supervisor.py @@ -60,24 +60,36 @@ def _write_action_script(path: Path) -> None: stderr=subprocess.DEVNULL, start_new_session=True, ) -pid_file.write_text(f"{os.getpid()} {child.pid}") +monitor_pid = os.getppid() +outer_pid = int( + next( + line.removeprefix("PPid:").strip() + for line in Path(f"/proc/{monitor_pid}/status").read_text().splitlines() + if line.startswith("PPid:") + ) +) +if mode == "stop-monitor-kill-supervisor": + pid_file.write_text(f"{os.getpid()} {child.pid} {monitor_pid} {outer_pid}") +else: + pid_file.write_text(f"{os.getpid()} {child.pid}") if mode == "failure": raise SystemExit(23) if mode == "kill-monitor": - os.kill(os.getppid(), signal.SIGKILL) + os.kill(monitor_pid, signal.SIGKILL) if mode in {"stop-monitor", "stop-supervisors"}: - monitor_pid = os.getppid() - outer_pid = int( - next( - line.removeprefix("PPid:").strip() - for line in Path(f"/proc/{monitor_pid}/status").read_text().splitlines() - if line.startswith("PPid:") - ) - ) os.kill(monitor_pid, signal.SIGSTOP) if mode == "stop-supervisors": os.kill(outer_pid, signal.SIGSTOP) -if mode in {"block", "kill-monitor", "stop-monitor", "stop-supervisors"}: +if mode == "stop-monitor-kill-supervisor": + os.kill(monitor_pid, signal.SIGSTOP) + os.kill(outer_pid, signal.SIGKILL) +if mode in { + "block", + "kill-monitor", + "stop-monitor", + "stop-supervisors", + "stop-monitor-kill-supervisor", +}: time.sleep(30) """.lstrip() ) @@ -253,3 +265,41 @@ async def test_termination_recovers_stopped_supervisor_tree( with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) await process.wait() + + +@pytest.mark.anyio +async def test_parent_death_resumes_monitor_and_reaps_action_tree( + tmp_path: Path, +) -> None: + process, pid_file = await _spawn_supervised_action( + tmp_path, + mode="stop-monitor-kill-supervisor", + ) + await _wait_for_file(pid_file) + action_pid, detached_pid, monitor_pid, outer_pid = ( + int(pid) for pid in pid_file.read_text().split() + ) + assert outer_pid == process.pid + + try: + stdout, stderr = await asyncio.wait_for( + communicate_process_group( + process, + timeout=2, + terminate=terminate_supervised_process, + ), + timeout=5, + ) + + assert process.returncode == -signal.SIGKILL, stderr.decode() + assert stdout == b"" + assert not _process_is_running(monitor_pid) + assert not _process_is_running(action_pid) + assert not _process_is_running(detached_pid) + finally: + for pid in (monitor_pid, action_pid, detached_pid): + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() diff --git a/tracecat/executor/process_supervisor.py b/tracecat/executor/process_supervisor.py index 712d44e5d4..2820815eaa 100644 --- a/tracecat/executor/process_supervisor.py +++ b/tracecat/executor/process_supervisor.py @@ -22,6 +22,7 @@ from types import FrameType _PR_SET_CHILD_SUBREAPER = 36 +_PR_SET_PDEATHSIG = 1 _PARENT_POLL_INTERVAL_MS = 10 @@ -48,6 +49,23 @@ def _set_child_subreaper() -> None: raise OSError(error_number, os.strerror(error_number)) +def _set_parent_death_signal(signal_number: int) -> None: + """Ask Linux to signal this process when its current parent exits.""" + libc = ctypes.CDLL(None, use_errno=True) + prctl = libc.prctl + prctl.argtypes = [ + ctypes.c_int, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ] + prctl.restype = ctypes.c_int + if prctl(_PR_SET_PDEATHSIG, signal_number, 0, 0, 0) != 0: + error_number = ctypes.get_errno() + raise OSError(error_number, os.strerror(error_number)) + + def _direct_child_pids_from_children_file() -> list[int]: """Return direct child PIDs using procfs's optional children file.""" children_path = Path(f"/proc/self/task/{os.getpid()}/children") @@ -125,9 +143,31 @@ def _exec(command: Sequence[str], control_fd: int) -> None: os._exit(127) -def _run_monitor(control_fd: int, command: Sequence[str]) -> int: +def _run_monitor( + control_fd: int, + command: Sequence[str], + *, + supervisor_pid: int, +) -> int: """Run the action below a detached subreaper and contain its descendants.""" try: + parent_cleanup_requested = False + + def request_parent_cleanup( + _signal: int, + _frame: FrameType | None, + ) -> None: + nonlocal parent_cleanup_requested + parent_cleanup_requested = True + + # The action can stop this same-UID monitor and kill its outer parent. + # SIGCONT both resumes the monitor and records that it must clean up. + signal.signal(signal.SIGCONT, request_parent_cleanup) + _set_parent_death_signal(signal.SIGCONT) + # Close the fork-to-prctl race before any untrusted command is started. + if os.getppid() != supervisor_pid or parent_cleanup_requested: + return 128 + signal.SIGTERM + os.setsid() _set_child_subreaper() _direct_child_pids() # Fail before execution when procfs tracking is absent. @@ -141,7 +181,9 @@ def _run_monitor(control_fd: int, command: Sequence[str]) -> int: action_status: int | None = None parent_closed = False - while action_status is None and not parent_closed: + while ( + action_status is None and not parent_closed and not parent_cleanup_requested + ): waited_pid, wait_status = _waitpid(action_pid, os.WNOHANG) if waited_pid == action_pid: action_status = wait_status @@ -150,7 +192,7 @@ def _run_monitor(control_fd: int, command: Sequence[str]) -> int: parent_closed = os.read(control_fd, 1) == b"" _kill_and_reap_children() - if parent_closed: + if parent_closed or parent_cleanup_requested: return 128 + signal.SIGTERM if action_status is None: raise RuntimeError("Supervised action exited without a wait status") @@ -202,12 +244,19 @@ def request_cleanup(_signal: int, _frame: FrameType | None) -> None: signal.signal(signal.SIGTERM, request_cleanup) signal.signal(signal.SIGINT, request_cleanup) + supervisor_pid = os.getpid() monitor_pid = os.fork() if monitor_pid == 0: signal.signal(signal.SIGTERM, signal.SIG_DFL) signal.signal(signal.SIGINT, signal.SIG_DFL) os.close(control_write_fd) - os._exit(_run_monitor(control_read_fd, command)) + os._exit( + _run_monitor( + control_read_fd, + command, + supervisor_pid=supervisor_pid, + ) + ) os.close(control_read_fd) try: From 074149341c990d2881b39b33a39f8af8b008fa4e Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:32:59 -0400 Subject: [PATCH 146/161] refactor(executor): trim registry cache review scope --- ...registry_artifact_cache_mount_lifecycle.py | 12 +- ...registry_artifact_cache_temporal_worker.py | 12 +- tests/unit/executor/conftest.py | 46 - .../registry_artifact_test_helpers.py | 199 - .../unit/executor/test_process_supervisor.py | 305 -- .../executor/test_registry_artifact_budget.py | 1284 ------- .../test_registry_artifact_eviction.py | 344 -- .../executor/test_registry_artifact_leases.py | 927 ----- .../test_registry_artifact_materialization.py | 1546 -------- .../test_registry_artifact_resolution.py | 546 --- .../test_registry_artifact_startup.py | 533 --- .../executor/test_run_python_sdk_context.py | 8 +- .../test_test_backend_no_registry_action.py | 202 +- tests/unit/test_action_runner.py | 69 +- tests/unit/test_executor_sandbox_nsjail.py | 18 +- ...arball.py => test_multitenant_registry.py} | 59 +- tests/unit/test_registry_artifacts.py | 3418 +++++++++++++++++ tests/unit/test_sandbox_utils.py | 103 - tests/unit/test_storage_blob.py | 387 +- tests/unit/test_unsafe_pid_executor.py | 189 +- tracecat/executor/action_runner.py | 30 +- tracecat/executor/backends/base.py | 16 +- tracecat/executor/backends/test.py | 57 +- tracecat/executor/process_supervisor.py | 276 -- .../executor/registry_artifact_cache_state.py | 218 -- .../registry_artifact_materialization.py | 1141 ------ tracecat/executor/registry_artifact_mounts.py | 30 - .../executor/registry_artifact_storage.py | 1022 ----- tracecat/executor/registry_artifacts.py | 1781 ++++++++- tracecat/sandbox/unsafe_pid_executor.py | 66 +- tracecat/sandbox/utils.py | 94 +- tracecat/storage/blob.py | 212 +- 32 files changed, 5207 insertions(+), 9943 deletions(-) delete mode 100644 tests/unit/executor/conftest.py delete mode 100644 tests/unit/executor/registry_artifact_test_helpers.py delete mode 100644 tests/unit/executor/test_process_supervisor.py delete mode 100644 tests/unit/executor/test_registry_artifact_budget.py delete mode 100644 tests/unit/executor/test_registry_artifact_eviction.py delete mode 100644 tests/unit/executor/test_registry_artifact_leases.py delete mode 100644 tests/unit/executor/test_registry_artifact_materialization.py delete mode 100644 tests/unit/executor/test_registry_artifact_resolution.py delete mode 100644 tests/unit/executor/test_registry_artifact_startup.py rename tests/unit/{executor/test_registry_artifact_tarball.py => test_multitenant_registry.py} (83%) create mode 100644 tests/unit/test_registry_artifacts.py delete mode 100644 tests/unit/test_sandbox_utils.py delete mode 100644 tracecat/executor/process_supervisor.py delete mode 100644 tracecat/executor/registry_artifact_cache_state.py delete mode 100644 tracecat/executor/registry_artifact_materialization.py delete mode 100644 tracecat/executor/registry_artifact_mounts.py delete mode 100644 tracecat/executor/registry_artifact_storage.py diff --git a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py index 639539d632..6982282480 100644 --- a/tests/integration/test_registry_artifact_cache_mount_lifecycle.py +++ b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py @@ -184,8 +184,9 @@ def test_registry_artifact_cache_mount_lifecycle() -> None: assert payload["converged_loop_device_released"] is True assert payload["converged_entries_remaining"] == 2 - # The startup sweep trims retained images to budget. + # The startup sweep trims to budget and removes stale mount directories. assert payload["startup_sweep_trimmed"] is True + assert payload["startup_sweep_removed_stale_mount_dir"] is True def _squashfs_mounts(cache_dir: Path) -> dict[str, str]: @@ -400,7 +401,7 @@ async def hold_concurrent_lease(index: int) -> None: ) payload["converged_entries_remaining"] = len(cache._discover_cache_keys()) - # (f) The startup sweep trims retained images to budget. + # (f) The startup sweep trims to budget and drops stale mount directories. sweep_dir = root / "sweep-cache" sweep_dir.mkdir() sweep_cache = RegistryArtifactCache(sweep_dir) @@ -411,6 +412,12 @@ async def hold_concurrent_lease(index: int) -> None: image_path = sweep_paths.squashfs_image_path image_path.write_bytes(b"x" * 4096) os.utime(sweep_paths.entry_dir, (100.0 + index, 100.0 + index)) + stale_mount_dir = sweep_cache._paths_for(sweep_keys[0]).squashfs_mount_dir + stale_mount_dir.mkdir() + os.utime( + sweep_cache._paths_for(sweep_keys[0]).entry_dir, + (100.0, 100.0), + ) config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES = 1 await sweep_cache.ensure_swept() @@ -420,6 +427,7 @@ async def hold_concurrent_lease(index: int) -> None: not evicted_paths.squashfs_image_path.exists() and retained_paths.squashfs_image_path.exists() ) + payload["startup_sweep_removed_stale_mount_dir"] = not stale_mount_dir.exists() # Leave no mounts behind for the container teardown. for cache_key in sorted(cache._discover_cache_keys()): diff --git a/tests/integration/test_registry_artifact_cache_temporal_worker.py b/tests/integration/test_registry_artifact_cache_temporal_worker.py index aa122cdc71..47fae9f691 100644 --- a/tests/integration/test_registry_artifact_cache_temporal_worker.py +++ b/tests/integration/test_registry_artifact_cache_temporal_worker.py @@ -35,8 +35,9 @@ async def temporal_env() -> AsyncGenerator[WorkflowEnvironment, None]: @dataclass(frozen=True, slots=True) class _ProbeResult: - """Runtime observations from one Temporal activity.""" + """Runtime identity observed by one Temporal activity.""" + cache_instance_id: int event_loop_id: int thread_id: int registry_path: str @@ -72,6 +73,7 @@ async def run(self, index: int) -> _ProbeResult: try: await self.release.wait() return _ProbeResult( + cache_instance_id=id(self.cache), event_loop_id=id(asyncio.get_running_loop()), thread_id=threading.get_ident(), registry_path=str(registry_paths[0]), @@ -110,9 +112,10 @@ async def test_one_temporal_worker_uses_one_cache_loop_and_thread( """Protect the production async-activity ownership contract. A real Temporal worker must schedule overlapping cache users on one event - loop and thread and balance their leases. The explicit production-activity - assertion also prevents action execution from quietly moving into Temporal's - synchronous thread pool, the historical failure mode for async storage state. + loop and thread while sharing one process-wide cache instance. The explicit + production-activity assertion also prevents action execution from quietly + moving into Temporal's synchronous thread pool, the historical failure mode + for process-wide async storage state. """ assert inspect.iscoroutinefunction(ExecutorActivities.execute_action_activity) @@ -163,6 +166,7 @@ async def test_one_temporal_worker_uses_one_cache_loop_and_thread( assert probe.peak_holders == activity_count assert probe.active_holders == 0 + assert {result.cache_instance_id for result in results} == {id(cache)} assert len({result.event_loop_id for result in results}) == 1 assert len({result.thread_id for result in results}) == 1 assert {result.registry_path for result in results} == {str(registry_path)} diff --git a/tests/unit/executor/conftest.py b/tests/unit/executor/conftest.py deleted file mode 100644 index 0bfce17b4f..0000000000 --- a/tests/unit/executor/conftest.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Fixtures for executor registry artifact tests.""" - -from __future__ import annotations - -import os -import stat -import tempfile -from collections.abc import Iterator -from pathlib import Path - -import pytest - -from tracecat.executor import registry_artifact_storage - - -@pytest.fixture(autouse=True) -def logical_cache_sizes(monkeypatch: pytest.MonkeyPatch) -> None: - """Keep synthetic byte budgets independent from host filesystem block sizes.""" - - def logical_stat_size( - file_stat: os.stat_result, - *, - allocation_unit: int, - ) -> int: - del allocation_unit - if stat.S_ISDIR(file_stat.st_mode): - return 0 - return file_stat.st_size - - monkeypatch.setattr( - registry_artifact_storage, - "_allocated_stat_size", - logical_stat_size, - ) - monkeypatch.setattr( - registry_artifact_storage, - "_filesystem_allocation_unit", - lambda _path: 1, - ) - - -@pytest.fixture -def temp_cache_dir() -> Iterator[Path]: - """Create an isolated registry artifact cache directory.""" - with tempfile.TemporaryDirectory() as tmpdir: - yield Path(tmpdir) diff --git a/tests/unit/executor/registry_artifact_test_helpers.py b/tests/unit/executor/registry_artifact_test_helpers.py deleted file mode 100644 index b41c94b117..0000000000 --- a/tests/unit/executor/registry_artifact_test_helpers.py +++ /dev/null @@ -1,199 +0,0 @@ -"""Shared fixtures and fakes for registry artifact cache tests.""" - -from __future__ import annotations - -import asyncio -import io -import os -import tarfile -from dataclasses import dataclass, field -from pathlib import Path - -from tracecat.executor.registry_artifacts import ( - RegistryArtifactCache, - RegistryArtifactMaterializationContext, - SquashfsMountCommandError, -) - -MAX_ENTRIES_CONFIG = ( - "tracecat.executor.registry_artifacts.config" - ".TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES" -) -MAX_BYTES_CONFIG = ( - "tracecat.executor.registry_artifacts.config" - ".TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES" -) -SQUASHFS_ENABLED_CONFIG = ( - "tracecat.executor.registry_artifacts.config" - ".TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED" -) - - -def write_tarball_entry(cache_dir: Path, cache_key: str) -> Path: - """Create a materialized tarball cache entry on disk.""" - target_dir = cache_dir / "entries" / cache_key / "tarball" - target_dir.mkdir(parents=True) - (target_dir / "module.py").write_text("VALUE = 1") - return target_dir - - -def write_image_entry( - cache_dir: Path, cache_key: str, *, size: int, mtime: float -) -> Path: - """Create a downloaded SquashFS image cache entry with a fixed mtime.""" - entry_dir = cache_dir / "entries" / cache_key - entry_dir.mkdir(parents=True, exist_ok=True) - image_path = entry_dir / "image.squashfs" - image_path.write_bytes(b"x" * size) - os.utime(image_path, (mtime, mtime)) - os.utime(entry_dir, (mtime, mtime)) - return image_path - - -def tarball_payload(*, size: int) -> bytes: - """Return a gzip tarball containing one synthetic regular file.""" - payload = b"x" * size - output = io.BytesIO() - with tarfile.open(fileobj=output, mode="w:gz") as tar: - member = tarfile.TarInfo("module.py") - member.size = len(payload) - tar.addfile(member, io.BytesIO(payload)) - return output.getvalue() - - -@dataclass(slots=True) -class SquashfsMountHarness: - """Model mount ownership without consuming host loop devices.""" - - cache: RegistryArtifactCache - failed_mount_keys: set[str] = field(default_factory=set) - mounted: set[Path] = field(default_factory=set) - mount_attempts: list[str] = field(default_factory=list) - extraction_attempts: list[str] = field(default_factory=list) - unmounts: list[Path] = field(default_factory=list) - - def seed_mount(self, cache_key: str) -> Path: - """Create one reusable image and mark its mount directory active.""" - paths = self.cache._paths_for(cache_key) - paths.entry_dir.mkdir(parents=True, exist_ok=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir(exist_ok=True) - (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") - self.mounted.add(paths.squashfs_mount_dir) - return paths.squashfs_mount_dir - - async def mount( - self, - ctx: RegistryArtifactMaterializationContext, - image_path: Path, - ) -> Path: - """Record a mount attempt, failing selected keys like loop exhaustion.""" - self.mount_attempts.append(ctx.cache_key) - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - image_path.write_bytes(b"squashfs") - if ctx.cache_key in self.failed_mount_keys: - raise SquashfsMountCommandError("no free loop device") - ctx.paths.squashfs_mount_dir.mkdir(exist_ok=True) - (ctx.paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") - self.mounted.add(ctx.paths.squashfs_mount_dir) - return ctx.paths.squashfs_mount_dir - - async def extract( - self, - ctx: RegistryArtifactMaterializationContext, - image_path: Path, - ) -> Path: - """Publish an extracted fallback from the image retained by mount.""" - assert image_path.is_file() - self.extraction_attempts.append(ctx.cache_key) - ctx.paths.squashfs_extract_dir.mkdir(parents=True, exist_ok=True) - (ctx.paths.squashfs_extract_dir / "module.py").write_text("VALUE = 1") - return ctx.paths.squashfs_extract_dir - - async def unmount(self, mount_dir: Path) -> bool: - """Record final-release unmounts and release the modeled loop device.""" - self.unmounts.append(mount_dir) - self.mounted.discard(mount_dir) - return True - - -async def lease_paths( - cache: RegistryArtifactCache, - artifact_uri: str, -) -> list[Path]: - """Return paths from the cache's public lease API.""" - async with cache.lease([artifact_uri]) as paths: - return paths - - -class BlockingSubprocess: - """Fake subprocess that blocks in communicate until it is cancelled.""" - - def __init__(self, *, block_wait: bool = False) -> None: - self.pid = 999_999_999 - self.communicate_started = asyncio.Event() - self.wait_started = asyncio.Event() - self.release_wait = asyncio.Event() - self.cleanup_calls: list[str] = [] - self.returncode: int | None = None - self._block_wait = block_wait - - async def communicate( - self, - input: bytes | None = None, # noqa: A002 - ) -> tuple[bytes, bytes]: - """Block until the task awaiting subprocess completion is cancelled.""" - del input - self.communicate_started.set() - await asyncio.Event().wait() - return b"", b"" - - def kill(self) -> None: - """Record that the subprocess was killed.""" - self.cleanup_calls.append("kill") - self.returncode = -9 - - async def wait(self) -> int: - """Record that the killed subprocess was reaped.""" - self.cleanup_calls.append("wait") - self.wait_started.set() - if self._block_wait: - await self.release_wait.wait() - return -9 - - -class CapturedSubprocess: - """Capture cleanup of a real subprocess used by cancellation tests.""" - - def __init__(self, process: asyncio.subprocess.Process) -> None: - self.process = process - self.killed = False - self.reaped = False - - @property - def returncode(self) -> int | None: - """Return the wrapped subprocess exit status.""" - return self.process.returncode - - @property - def pid(self) -> int: - """Return the wrapped subprocess process-group identifier.""" - return self.process.pid - - async def communicate( - self, - input: bytes | None = None, # noqa: A002 - ) -> tuple[bytes, bytes]: - """Wait for the wrapped subprocess and collect its output.""" - return await self.process.communicate(input=input) - - def kill(self) -> None: - """Kill the wrapped subprocess and record the signal.""" - self.killed = True - self.process.kill() - - async def wait(self) -> int: - """Reap the wrapped subprocess and record completion.""" - returncode = await self.process.wait() - self.reaped = True - return returncode diff --git a/tests/unit/executor/test_process_supervisor.py b/tests/unit/executor/test_process_supervisor.py deleted file mode 100644 index a63a953dd0..0000000000 --- a/tests/unit/executor/test_process_supervisor.py +++ /dev/null @@ -1,305 +0,0 @@ -"""Linux process-tree containment for direct registry actions.""" - -from __future__ import annotations - -import asyncio -import contextlib -import os -import signal -import sys -from pathlib import Path - -import pytest - -from tracecat.executor import process_supervisor -from tracecat.sandbox.utils import ( - communicate_process_group, - terminate_supervised_process, -) - -pytestmark = pytest.mark.skipif( - sys.platform != "linux", - reason="The direct action supervisor uses Linux prctl and procfs", -) - - -def _process_is_running(pid: int) -> bool: - """Return whether a process is alive, treating zombies as terminated.""" - try: - os.kill(pid, 0) - stat_path = Path(f"/proc/{pid}/stat") - return not stat_path.exists() or stat_path.read_text().split()[2] != "Z" - except (FileNotFoundError, ProcessLookupError): - return False - - -async def _wait_for_file(path: Path) -> None: - for _ in range(500): - if path.exists(): - return - await asyncio.sleep(0.01) - raise AssertionError(f"Timed out waiting for {path}") - - -def _write_action_script(path: Path) -> None: - path.write_text( - """ -import os -import signal -import subprocess -import sys -import time -from pathlib import Path - -pid_file = Path(sys.argv[1]) -mode = sys.argv[2] -child = subprocess.Popen( - [sys.executable, "-c", "import time; time.sleep(30)"], - stdin=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - start_new_session=True, -) -monitor_pid = os.getppid() -outer_pid = int( - next( - line.removeprefix("PPid:").strip() - for line in Path(f"/proc/{monitor_pid}/status").read_text().splitlines() - if line.startswith("PPid:") - ) -) -if mode == "stop-monitor-kill-supervisor": - pid_file.write_text(f"{os.getpid()} {child.pid} {monitor_pid} {outer_pid}") -else: - pid_file.write_text(f"{os.getpid()} {child.pid}") -if mode == "failure": - raise SystemExit(23) -if mode == "kill-monitor": - os.kill(monitor_pid, signal.SIGKILL) -if mode in {"stop-monitor", "stop-supervisors"}: - os.kill(monitor_pid, signal.SIGSTOP) - if mode == "stop-supervisors": - os.kill(outer_pid, signal.SIGSTOP) -if mode == "stop-monitor-kill-supervisor": - os.kill(monitor_pid, signal.SIGSTOP) - os.kill(outer_pid, signal.SIGKILL) -if mode in { - "block", - "kill-monitor", - "stop-monitor", - "stop-supervisors", - "stop-monitor-kill-supervisor", -}: - time.sleep(30) -""".lstrip() - ) - - -def test_direct_child_discovery_falls_back_to_proc_stat(monkeypatch) -> None: - child_pid = os.fork() - if child_pid == 0: - signal.pause() - os._exit(0) - - def missing_children_file() -> list[int]: - raise FileNotFoundError - - monkeypatch.setattr( - process_supervisor, - "_direct_child_pids_from_children_file", - missing_children_file, - ) - try: - assert child_pid in process_supervisor._direct_child_pids() - finally: - with contextlib.suppress(ProcessLookupError): - os.kill(child_pid, signal.SIGKILL) - os.waitpid(child_pid, 0) - - -async def _spawn_supervised_action( - tmp_path: Path, - *, - mode: str, -) -> tuple[asyncio.subprocess.Process, Path]: - action_script = tmp_path / f"action-{mode}.py" - pid_file = tmp_path / f"action-{mode}.pids" - _write_action_script(action_script) - supervisor_path = Path(process_supervisor.__file__) - process = await asyncio.create_subprocess_exec( - sys.executable, - str(supervisor_path), - sys.executable, - str(action_script), - str(pid_file), - mode, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - start_new_session=True, - ) - return process, pid_file - - -@pytest.mark.anyio -@pytest.mark.parametrize( - ("mode", "expected_returncode"), - [("success", 0), ("failure", 23)], -) -async def test_supervisor_propagates_exit_and_reaps_detached_descendant( - tmp_path: Path, - mode: str, - expected_returncode: int, -) -> None: - process, pid_file = await _spawn_supervised_action(tmp_path, mode=mode) - tracked_pids: tuple[int, ...] = () - try: - stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=5) - - assert process.returncode == expected_returncode, stderr.decode() - assert stdout == b"" - tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) - _, detached_pid = tracked_pids - assert not _process_is_running(detached_pid) - finally: - if not tracked_pids and pid_file.exists(): - tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) - for pid in tracked_pids: - with contextlib.suppress(ProcessLookupError): - os.kill(pid, signal.SIGKILL) - with contextlib.suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) - await process.wait() - - -@pytest.mark.anyio -async def test_supervisor_reaps_detached_descendant_before_cancellation_returns( - tmp_path: Path, -) -> None: - process, pid_file = await _spawn_supervised_action(tmp_path, mode="block") - await _wait_for_file(pid_file) - action_pid, detached_pid = (int(pid) for pid in pid_file.read_text().split()) - communication = asyncio.create_task( - communicate_process_group( - process, - timeout=30, - terminate=terminate_supervised_process, - ) - ) - - try: - await asyncio.sleep(0) - communication.cancel() - await asyncio.sleep(0) - communication.cancel() - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(communication, timeout=5) - assert not _process_is_running(action_pid) - assert not _process_is_running(detached_pid) - finally: - for pid in (action_pid, detached_pid): - with contextlib.suppress(ProcessLookupError): - os.kill(pid, signal.SIGKILL) - with contextlib.suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) - await process.wait() - - -@pytest.mark.anyio -async def test_supervisor_reaps_action_that_kills_its_monitor( - tmp_path: Path, -) -> None: - process, pid_file = await _spawn_supervised_action( - tmp_path, - mode="kill-monitor", - ) - tracked_pids: tuple[int, ...] = () - try: - stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=5) - - assert process.returncode == 128 + signal.SIGKILL, stderr.decode() - assert stdout == b"" - tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) - action_pid, detached_pid = tracked_pids - assert not _process_is_running(action_pid) - assert not _process_is_running(detached_pid) - finally: - if not tracked_pids and pid_file.exists(): - tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) - for pid in tracked_pids: - with contextlib.suppress(ProcessLookupError): - os.kill(pid, signal.SIGKILL) - with contextlib.suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) - await process.wait() - - -@pytest.mark.anyio -@pytest.mark.parametrize("mode", ["stop-monitor", "stop-supervisors"]) -async def test_termination_recovers_stopped_supervisor_tree( - tmp_path: Path, - mode: str, -) -> None: - process, pid_file = await _spawn_supervised_action(tmp_path, mode=mode) - await _wait_for_file(pid_file) - tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) - action_pid, detached_pid = tracked_pids - - try: - with pytest.raises(TimeoutError): - await asyncio.wait_for( - communicate_process_group( - process, - timeout=0.05, - terminate=terminate_supervised_process, - ), - timeout=5, - ) - - assert process.returncode is not None - assert not _process_is_running(action_pid) - assert not _process_is_running(detached_pid) - finally: - for pid in tracked_pids: - with contextlib.suppress(ProcessLookupError): - os.kill(pid, signal.SIGKILL) - with contextlib.suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) - await process.wait() - - -@pytest.mark.anyio -async def test_parent_death_resumes_monitor_and_reaps_action_tree( - tmp_path: Path, -) -> None: - process, pid_file = await _spawn_supervised_action( - tmp_path, - mode="stop-monitor-kill-supervisor", - ) - await _wait_for_file(pid_file) - action_pid, detached_pid, monitor_pid, outer_pid = ( - int(pid) for pid in pid_file.read_text().split() - ) - assert outer_pid == process.pid - - try: - stdout, stderr = await asyncio.wait_for( - communicate_process_group( - process, - timeout=2, - terminate=terminate_supervised_process, - ), - timeout=5, - ) - - assert process.returncode == -signal.SIGKILL, stderr.decode() - assert stdout == b"" - assert not _process_is_running(monitor_pid) - assert not _process_is_running(action_pid) - assert not _process_is_running(detached_pid) - finally: - for pid in (monitor_pid, action_pid, detached_pid): - with contextlib.suppress(ProcessLookupError): - os.kill(pid, signal.SIGKILL) - with contextlib.suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) - await process.wait() diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py deleted file mode 100644 index 160bb79756..0000000000 --- a/tests/unit/executor/test_registry_artifact_budget.py +++ /dev/null @@ -1,1284 +0,0 @@ -"""Registry artifact byte admission and budget convergence tests.""" - -from __future__ import annotations - -import asyncio -import os -import signal -import threading -from collections.abc import Awaitable, Callable -from pathlib import Path -from typing import Literal -from unittest.mock import ANY, AsyncMock, MagicMock, call, patch - -import pytest - -from tracecat.executor.registry_artifacts import ( - RegistryArtifactCache, - RegistryArtifactCacheCapacityError, - RegistryArtifactEviction, - RegistryArtifactMaterializationContext, - SquashfsArtifact, - TarballArtifact, - _allocated_stat_size, - _delete_cache_path, - _directory_footprint, - compute_registry_artifact_cache_key, -) - -from .registry_artifact_test_helpers import ( - MAX_BYTES_CONFIG, - MAX_ENTRIES_CONFIG, - SQUASHFS_ENABLED_CONFIG, - lease_paths, - tarball_payload, - write_image_entry, - write_tarball_entry, -) - - -class TestRegistryArtifactCacheBudget: - """Enforce peak and steady-state cache capacity.""" - - def test_allocated_stat_size_uses_filesystem_blocks(self) -> None: - file_stat = MagicMock(spec=os.stat_result) - file_stat.st_blocks = 7 - file_stat.st_size = 1 - - assert _allocated_stat_size(file_stat, allocation_unit=512) == 7 * 512 - - def test_allocated_stat_size_charges_zero_block_inode(self) -> None: - file_stat = MagicMock(spec=os.stat_result) - file_stat.st_blocks = 0 - file_stat.st_size = 0 - - assert _allocated_stat_size(file_stat, allocation_unit=4096) == 4096 - - def test_directory_footprint_includes_directory_inodes( - self, - temp_cache_dir: Path, - ) -> None: - nested = temp_cache_dir / "nested" - nested.mkdir() - (nested / "module.py").write_text("x") - - with patch( - "tracecat.executor.registry_artifact_storage._allocated_stat_size", - return_value=4096, - ) as allocated_size: - assert _directory_footprint(temp_cache_dir) == 3 * 4096 - - assert allocated_size.call_count == 3 - - def test_directory_footprint_counts_hard_linked_inode_once( - self, - temp_cache_dir: Path, - ) -> None: - payload = temp_cache_dir / "payload" - payload.write_text("x") - os.link(payload, temp_cache_dir / "payload-link") - - with patch( - "tracecat.executor.registry_artifact_storage._allocated_stat_size", - return_value=4096, - ) as allocated_size: - assert _directory_footprint(temp_cache_dir) == 2 * 4096 - - assert allocated_size.call_count == 2 - - def test_directory_footprint_prunes_directory_contents( - self, - temp_cache_dir: Path, - ) -> None: - mounted = temp_cache_dir / "mount" - mounted.mkdir() - (mounted / "module.py").write_text("x") - - with patch( - "tracecat.executor.registry_artifact_storage._allocated_stat_size", - return_value=4096, - ) as allocated_size: - assert ( - _directory_footprint( - temp_cache_dir, - pruned_directories=(mounted,), - ) - == 2 * 4096 - ) - - assert allocated_size.call_count == 2 - - def test_directory_footprint_can_exclude_root_inode( - self, - temp_cache_dir: Path, - ) -> None: - nested = temp_cache_dir / "nested" - nested.mkdir() - (nested / "module.py").write_text("x") - - with patch( - "tracecat.executor.registry_artifact_storage._allocated_stat_size", - return_value=4096, - ) as allocated_size: - assert _directory_footprint(temp_cache_dir, include_root=False) == 2 * 4096 - - assert allocated_size.call_count == 2 - - def test_cache_structure_counts_roots_without_subtree_contents( - self, - temp_cache_dir: Path, - ) -> None: - cache = RegistryArtifactCache(temp_cache_dir) - for root in (cache.entries_dir, cache.staging_dir, cache.trash_dir): - child = root / "child" - child.mkdir(parents=True) - (child / "payload").write_text("x") - base_dir = temp_cache_dir / "base" - base_dir.mkdir() - (base_dir / "payload").write_text("x") - - with patch( - "tracecat.executor.registry_artifact_storage._allocated_stat_size", - return_value=4096, - ) as allocated_size: - assert cache._cache_structural_footprint() == 6 * 4096 - - # Cache, entries, staging, and trash roots plus the unpruned base tree. - assert allocated_size.call_count == 6 - - @pytest.mark.anyio - async def test_admission_counts_non_evictable_cache_structure( - self, - temp_cache_dir: Path, - ) -> None: - cache = RegistryArtifactCache(temp_cache_dir) - - with patch.object(cache, "_cache_structural_footprint", return_value=5): - with pytest.raises(RegistryArtifactCacheCapacityError) as raised: - await cache._ensure_cache_capacity( - additional_bytes=0, - protected_key="new", - max_bytes=4, - ) - - assert raised.value.current_bytes == 5 - assert raised.value.additional_bytes == 0 - - @pytest.mark.anyio - async def test_enforcement_counts_cache_structure( - self, - temp_cache_dir: Path, - ) -> None: - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch(MAX_BYTES_CONFIG, 4), - patch.object(cache, "_cache_structural_footprint", return_value=5), - ): - assert await cache._enforce_cache_budget() is False - - @pytest.mark.anyio - async def test_admission_rounds_download_reservation_to_allocation_unit( - self, - temp_cache_dir: Path, - ) -> None: - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch(MAX_BYTES_CONFIG, 8192), - patch( - "tracecat.executor.registry_artifact_storage." - "_filesystem_allocation_unit", - return_value=4096, - ), - patch.object( - cache, - "_ensure_cache_capacity", - new_callable=AsyncMock, - ) as ensure_capacity, - ): - admission = cache._admission_for("new") - assert admission is not None - await admission.ensure_capacity(1) - - ensure_capacity.assert_awaited_once_with( - additional_bytes=4096, - protected_key="new", - max_bytes=8192, - ) - - @pytest.mark.anyio - async def test_admission_reuses_verified_capacity_headroom( - self, - temp_cache_dir: Path, - ) -> None: - """Chunked downloads rescan only after consuming known free bytes.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch(MAX_BYTES_CONFIG, 100), - patch( - "tracecat.executor.registry_artifact_storage." - "_filesystem_allocation_unit", - return_value=1, - ), - patch.object( - cache, - "_ensure_cache_capacity", - new_callable=AsyncMock, - side_effect=[10, 6], - ) as ensure_capacity, - ): - admission = cache._admission_for("new") - assert admission is not None - await admission.ensure_capacity(4) - await admission.ensure_capacity(6) - await admission.ensure_capacity(5) - await admission.ensure_capacity(2) - - assert ensure_capacity.await_args_list == [ - call(additional_bytes=4, protected_key="new", max_bytes=100), - call(additional_bytes=5, protected_key="new", max_bytes=100), - ] - - def test_delete_cache_path_reports_directory_failure(self, temp_cache_dir): - """Physical deletion failures are observable instead of suppressed.""" - entry_dir = temp_cache_dir / "entry" - entry_dir.mkdir() - - with ( - patch( - "tracecat.executor.registry_artifact_storage.shutil.rmtree", - side_effect=OSError("permission denied"), - ), - patch( - "tracecat.executor.registry_artifact_storage.logger.warning" - ) as warning, - ): - deleted = _delete_cache_path(entry_dir) - - assert deleted is False - assert entry_dir.is_dir() - warning.assert_called_once_with( - "Failed to delete registry artifact cache path", - path=str(entry_dir), - error="permission denied", - ) - - @pytest.mark.anyio - async def test_leased_entry_survives_eviction_of_idle_entry(self, temp_cache_dir): - """Eviction must never remove an entry a live action is importing from.""" - cache = RegistryArtifactCache(temp_cache_dir) - leased_uri = "s3://bucket/leased.tar.gz" - idle_uri = "s3://bucket/idle.tar.gz" - new_uri = "s3://bucket/new.tar.gz" - leased_dir = write_tarball_entry( - temp_cache_dir, compute_registry_artifact_cache_key(leased_uri) - ) - idle_dir = write_tarball_entry( - temp_cache_dir, compute_registry_artifact_cache_key(idle_uri) - ) - - async def mock_download(self, ctx, path): - path.write_bytes(tarball_payload(size=1)) - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_text("VALUE = 2") - - with ( - patch(MAX_ENTRIES_CONFIG, 2), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - async with cache.lease([leased_uri]): - await lease_paths(cache, new_uri) - - assert leased_dir.is_dir() - assert not idle_dir.exists() - - @pytest.mark.anyio - async def test_cold_download_reserves_space_before_writing( - self, temp_cache_dir: Path - ) -> None: - """Admission evicts idle bytes before a new download enters staging.""" - cache = RegistryArtifactCache(temp_cache_dir) - idle = write_image_entry(temp_cache_dir, "idle", size=80, mtime=100.0) - artifact_uri = "s3://bucket/new.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - payload = tarball_payload(size=32) - extracted_size = 74 # File, root, and directory entry at unit size 1. - max_bytes = len(payload) + extracted_size - capacity_checked = False - - async def download_file_to_path( - *, - key: str, - bucket: str, - output_path: Path, - max_bytes: int, - ensure_capacity: Callable[[int], Awaitable[None]], - defer_cleanup: Callable[[Path], None], - redact_log_identifiers: bool, - ) -> int: - del key, bucket, defer_cleanup - assert redact_log_identifiers is True - nonlocal capacity_checked - assert max_bytes == len(payload) + extracted_size - await ensure_capacity(len(payload)) - capacity_checked = True - assert not idle.exists() - output_path.write_bytes(payload) - return len(payload) - - with ( - patch(SQUASHFS_ENABLED_CONFIG, False), - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, max_bytes), - patch( - "tracecat.executor.registry_artifacts.blob.download_file_to_path", - side_effect=download_file_to_path, - ), - ): - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [ - cache._paths_for(cache_key).tarball_target_dir - ] - - assert capacity_checked is True - assert not idle.exists() - assert (registry_paths[0] / "module.py").read_bytes() == b"x" * 32 - - @pytest.mark.anyio - async def test_impossible_tarball_reservation_preserves_warm_entry( - self, temp_cache_dir: Path - ) -> None: - """Impossible extraction cannot evict warm entries before rejection.""" - cache = RegistryArtifactCache(temp_cache_dir) - warm = write_image_entry( - temp_cache_dir, - "warm", - size=64, - mtime=100.0, - ) - artifact_uri = "s3://bucket/compression-heavy.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - payload = tarball_payload(size=4096) - max_bytes = len(payload) + 256 - - async def download_file_to_path( - *, - key: str, - bucket: str, - output_path: Path, - max_bytes: int, - ensure_capacity: Callable[[int], Awaitable[None]], - defer_cleanup: Callable[[Path], None], - redact_log_identifiers: bool, - ) -> int: - del key, bucket, defer_cleanup - assert redact_log_identifiers is True - assert max_bytes == len(payload) + 256 - await ensure_capacity(len(payload)) - output_path.write_bytes(payload) - return len(payload) - - with ( - patch(SQUASHFS_ENABLED_CONFIG, False), - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, max_bytes), - patch( - "tracecat.executor.registry_artifacts.blob.download_file_to_path", - side_effect=download_file_to_path, - ), - patch.object( - TarballArtifact, - "extract", - new_callable=AsyncMock, - ) as extract, - patch.object( - cache, - "_evict_entry", - wraps=cache._evict_entry, - ) as evict_entry, - ): - with pytest.raises(RegistryArtifactCacheCapacityError) as raised: - async with cache.lease([artifact_uri]): - pass - - assert raised.value.additional_bytes == 4138 - assert raised.value.max_bytes == max_bytes - extract.assert_not_awaited() - evict_entry.assert_not_awaited() - assert warm.is_file() - assert not cache._paths_for(cache_key).entry_dir.exists() - assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) - - @pytest.mark.anyio - async def test_failed_squashfs_bytes_do_not_block_tarball_fallback( - self, temp_cache_dir: Path - ) -> None: - """An unusable image cannot consume the tarball fallback's budget.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/path/site-packages.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - payload = tarball_payload(size=32) - extracted_size = 74 # File, root, and directory entry at unit size 1. - max_bytes = len(payload) + extracted_size - - async def fail_after_squashfs_download( - self: SquashfsArtifact, - ctx: RegistryArtifactMaterializationContext, - ) -> list[Path]: - del self - assert ctx.admission is not None - await ctx.admission.ensure_capacity(max_bytes) - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - ctx.paths.squashfs_image_path.write_bytes(b"x" * max_bytes) - raise RuntimeError("unusable SquashFS image") - - async def download_tarball( - self: TarballArtifact, - ctx: RegistryArtifactMaterializationContext, - path: Path, - ) -> None: - del self - assert ctx.admission is not None - await ctx.admission.ensure_capacity(len(payload)) - path.write_bytes(payload) - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, max_bytes), - patch.object( - SquashfsArtifact, - "materialize", - fail_after_squashfs_download, - ), - patch.object(TarballArtifact, "download", download_tarball), - ): - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [ - cache._paths_for(cache_key).tarball_target_dir - ] - - paths = cache._paths_for(cache_key) - assert not paths.squashfs_image_path.exists() - assert (paths.tarball_target_dir / "module.py").read_bytes() == b"x" * 32 - - @pytest.mark.anyio - async def test_squashfs_expansion_is_rejected_before_extraction( - self, temp_cache_dir: Path - ) -> None: - """SquashFS metadata is accounted before unsquashfs writes scratch.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/oversized.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - - async def download( - self: SquashfsArtifact, - ctx: RegistryArtifactMaterializationContext, - image_path: Path, - ) -> float: - del self, ctx - image_path.parent.mkdir(parents=True, exist_ok=True) - image_path.write_bytes(b"image") - return 0.0 - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 100), - patch.object( - RegistryArtifactMaterializationContext, - "can_mount_squashfs", - return_value=False, - ), - patch.object(SquashfsArtifact, "download", download), - patch.object( - SquashfsArtifact, - "_squashfs_extracted_size", - new_callable=AsyncMock, - return_value=101, - ), - patch.object( - SquashfsArtifact, - "_extract_image", - new_callable=AsyncMock, - ) as extract_image, - ): - with pytest.raises(RegistryArtifactCacheCapacityError) as raised: - async with cache.lease([artifact_uri]): - pass - - assert raised.value.additional_bytes == 101 - extract_image.assert_not_awaited() - paths = cache._paths_for(cache_key) - assert paths.squashfs_image_path.read_bytes() == b"image" - assert not paths.squashfs_extract_dir.exists() - - @pytest.mark.anyio - async def test_successful_admission_enforces_actual_size_before_yield( - self, temp_cache_dir - ): - """Post-publication enforcement sees the new entry's actual size.""" - cache = RegistryArtifactCache(temp_cache_dir) - idle = write_image_entry(temp_cache_dir, "idle", size=4096, mtime=100.0) - new_uri = "s3://bucket/new.tar.gz" - new_key = compute_registry_artifact_cache_key(new_uri) - - async def mock_download(self, ctx, path): - path.write_bytes(tarball_payload(size=1)) - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_bytes(b"x" * 4096) - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 6000), - patch( - "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", - return_value=4096, - ), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - async with cache.lease([new_uri]) as registry_paths: - assert registry_paths == [cache._paths_for(new_key).tarball_target_dir] - assert not idle.exists() - - assert not idle.exists() - assert cache._paths_for(new_key).tarball_target_dir.is_dir() - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_deletion_failure_does_not_block_materialization( - self, temp_cache_dir - ): - """Failed cleanup is reported while cache admission remains fail-open.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - write_image_entry(temp_cache_dir, "idle", size=4096, mtime=100.0) - artifact_uri = "s3://bucket/new.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - - async def mock_download(self, ctx, path): - path.write_bytes(tarball_payload(size=1)) - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_text("VALUE = 2") - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - return_value=False, - ), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - patch( - "tracecat.executor.registry_artifact_storage.logger.warning" - ) as warning, - ): - registry_paths = await lease_paths(cache, artifact_uri) - - assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] - assert registry_paths[0].is_dir() - assert cache._budget_dirty is True - warning.assert_any_call( - "Registry artifact eviction remains pending physical deletion", - cache_key="idle", - trash_path=ANY, - ) - - @pytest.mark.anyio - async def test_failed_physical_delete_retries_without_extra_eviction( - self, temp_cache_dir - ): - """Failed byte reclamation stops eviction until trash cleanup succeeds.""" - cache = RegistryArtifactCache(temp_cache_dir) - oldest = write_image_entry( - temp_cache_dir, - "oldest", - size=16, - mtime=100.0, - ) - older = write_image_entry( - temp_cache_dir, - "older", - size=16, - mtime=200.0, - ) - newest = write_image_entry( - temp_cache_dir, - "newest", - size=16, - mtime=300.0, - ) - real_delete = _delete_cache_path - failed_once = False - - def fail_once(path: Path) -> bool: - nonlocal failed_once - if not failed_once: - failed_once = True - return False - return real_delete(path) - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 16), - patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - side_effect=fail_once, - ), - ): - assert await cache._enforce_cache_budget() is False - assert not oldest.exists() - assert older.exists() - assert newest.exists() - assert cache._discover_cache_keys() == {"older", "newest"} - assert len(tuple(cache.trash_dir.iterdir())) == 1 - - assert await cache._enforce_cache_budget() is True - - assert not older.exists() - assert newest.exists() - assert not any(cache.trash_dir.iterdir()) - - @pytest.mark.anyio - async def test_undeletable_trash_does_not_evict_warm_entries_for_admission( - self, temp_cache_dir - ): - """Unreclaimed trash blocks an oversized write without extra eviction.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - warm = write_image_entry(temp_cache_dir, "warm", size=32, mtime=100.0) - stale_trash = cache.trash_dir / "stale" - stale_trash.mkdir(parents=True) - (stale_trash / "image.squashfs").write_bytes(b"x" * 32) - real_delete = _delete_cache_path - - def fail_stale_trash(path: Path) -> bool: - if path == stale_trash: - return False - return real_delete(path) - - with patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - side_effect=fail_stale_trash, - ): - with pytest.raises(RegistryArtifactCacheCapacityError) as raised: - await cache._ensure_cache_capacity( - additional_bytes=16, - protected_key="new", - max_bytes=64, - ) - - assert raised.value.current_bytes == 64 - assert warm.exists() - assert stale_trash.exists() - - @pytest.mark.anyio - async def test_rename_failure_does_not_block_materialization(self, temp_cache_dir): - """A failed atomic retirement keeps the old entry and admits new work.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - idle = write_image_entry(temp_cache_dir, "idle", size=16, mtime=100.0) - artifact_uri = "s3://bucket/new-after-rename-failure.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - - async def mock_download(self, ctx, path): - path.write_bytes(tarball_payload(size=1)) - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_text("VALUE = 2") - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - patch( - "tracecat.executor.registry_artifact_storage._move_entry_to_trash", - side_effect=OSError("rename failed"), - ), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - registry_paths = await lease_paths(cache, artifact_uri) - - assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] - assert registry_paths[0].is_dir() - assert idle.is_file() - assert cache._budget_dirty is True - - @pytest.mark.anyio - async def test_cache_scan_failure_does_not_block_materialization( - self, temp_cache_dir - ): - """Maintenance errors are observable while artifact admission stays open.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/new-after-scan-failure.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - - async def mock_download(self, ctx, path): - path.write_bytes(tarball_payload(size=1)) - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_text("VALUE = 2") - - with ( - patch.object( - cache, - "_enforce_cache_budget", - new_callable=AsyncMock, - side_effect=PermissionError("denied"), - ), - patch( - "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", - return_value=9, - ), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [ - cache._paths_for(cache_key).tarball_target_dir - ] - - @pytest.mark.anyio - async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( - self, temp_cache_dir - ): - """Steady-state cache hits must not pay for a full cache scan.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/cached.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - write_tarball_entry(temp_cache_dir, cache_key) - cache._budget_dirty = False - - with patch.object( - cache, - "_scan_cache_entries", - side_effect=AssertionError("cache hits must not scan the cache dir"), - ): - async with cache.lease([artifact_uri]): - pass - - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_releasing_a_mutable_cache_hit_counts_unknown_entry_growth( - self, - temp_cache_dir: Path, - ) -> None: - """Direct consumers cannot grow unknown entry paths outside the budget.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/mutable-cached.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - target_dir = write_tarball_entry(temp_cache_dir, cache_key) - entry_dir = target_dir.parent - initial_size = cache._measure_entry(cache_key).size_bytes - cache._budget_dirty = False - - with ( - patch(MAX_ENTRIES_CONFIG, 10), - patch(MAX_BYTES_CONFIG, initial_size), - patch.object( - cache, - "_scan_cache_entries", - wraps=cache._scan_cache_entries, - ) as scan_cache_entries, - ): - async with cache.lease( - [artifact_uri], - paths_may_be_modified=True, - ) as registry_paths: - assert registry_paths == [target_dir] - (entry_dir / "action-output.bin").write_bytes(b"x" * 4096) - - assert scan_cache_entries.call_count == 1 - assert not entry_dir.exists() - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_successful_cold_admission_skips_release_rescan(self, temp_cache_dir): - """A successful protected pass consumes the materialization dirty signal.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/new-with-one-budget-pass.tar.gz" - - async def mock_download(self, ctx, path): - del self, ctx - path.write_bytes(tarball_payload(size=1)) - - with ( - patch(MAX_ENTRIES_CONFIG, 10), - patch(MAX_BYTES_CONFIG, 0), - patch(SQUASHFS_ENABLED_CONFIG, False), - patch.object(TarballArtifact, "download", mock_download), - patch.object( - cache, - "_scan_cache_entries", - wraps=cache._scan_cache_entries, - ) as scan_cache_entries, - ): - async with cache.lease([artifact_uri]): - assert cache._budget_dirty is False - - assert scan_cache_entries.call_count == 1 - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_failed_cold_admission_keeps_warm_lru(self, temp_cache_dir): - """A missing cold artifact must not evict an existing warm entry.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - warm_uri = "s3://bucket/warm.tar.gz" - warm_key = compute_registry_artifact_cache_key(warm_uri) - warm_dir = write_tarball_entry(temp_cache_dir, warm_key) - missing_uri = "s3://bucket/missing.tar.gz" - missing_key = compute_registry_artifact_cache_key(missing_uri) - - async def mock_download(self, ctx, path): - raise FileNotFoundError("missing artifact") - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - patch(SQUASHFS_ENABLED_CONFIG, False), - patch.object(TarballArtifact, "download", mock_download), - ): - with pytest.raises(FileNotFoundError, match="missing artifact"): - async with cache.lease([missing_uri]): - pytest.fail("failed admission must not yield a lease") - - assert warm_dir.is_dir() - assert not cache._paths_for(missing_key).entry_dir.exists() - - @pytest.mark.anyio - async def test_non_final_release_defers_retry_while_cache_stays_over_budget( - self, temp_cache_dir - ): - """A non-final release defers rescanning until an entry becomes idle.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/pinned.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - write_image_entry(temp_cache_dir, cache_key, size=4096, mtime=100.0) - write_tarball_entry(temp_cache_dir, cache_key) - cache._budget_dirty = True - # A second holder keeps the entry pinned past the inner lease. - cache._acquire_lease(cache_key) - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 1), - patch.object( - cache, - "_scan_cache_entries", - side_effect=AssertionError("non-final releases must not rescan"), - ), - ): - async with cache.lease([artifact_uri]): - pass - - assert cache._budget_dirty is True - - @pytest.mark.anyio - async def test_convergence_rescans_after_concurrent_materialization( - self, temp_cache_dir - ): - """A materialization during a budget scan must schedule another scan.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/concurrent.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - scan_started = asyncio.Event() - materialized = asyncio.Event() - convergence_scans = 0 - - async def mock_enforce_cache_budget( - *, protected_key: str | None = None - ) -> bool: - nonlocal convergence_scans - if protected_key is not None: - return True - - convergence_scans += 1 - if convergence_scans == 1: - scan_started.set() - await materialized.wait() - return True - - async def mock_download(self, ctx, path): - path.write_bytes(tarball_payload(size=1)) - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_text("VALUE = 1") - - cache._budget_dirty = True - with ( - patch.object( - cache, - "_enforce_cache_budget", - side_effect=mock_enforce_cache_budget, - ), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - convergence = asyncio.create_task(cache._converge_cache_budget()) - await scan_started.wait() - registry_paths = await lease_paths(cache, artifact_uri) - materialized.set() - await convergence - - assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] - assert convergence_scans == 2 - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_cancelled_convergence_rearms_budget_dirty( - self, temp_cache_dir - ) -> None: - """Cancellation must preserve the need for a later budget pass.""" - cache = RegistryArtifactCache(temp_cache_dir) - enforcement_started = asyncio.Event() - finish_enforcement = asyncio.Event() - - async def mock_enforce_cache_budget( - *, protected_key: str | None = None - ) -> bool: - del protected_key - enforcement_started.set() - await finish_enforcement.wait() - return True - - cache._budget_dirty = True - with patch.object( - cache, - "_enforce_cache_budget", - side_effect=mock_enforce_cache_budget, - ): - convergence = asyncio.create_task(cache._converge_cache_budget()) - await enforcement_started.wait() - convergence.cancel() - - with pytest.raises(asyncio.CancelledError): - await convergence - - assert cache._budget_dirty is True - - @pytest.mark.parametrize("operation", ["budget", "admission"]) - @pytest.mark.anyio - async def test_cancelled_cleanup_rejoins_workers_before_releasing_locks( - self, - temp_cache_dir: Path, - operation: Literal["budget", "admission"], - ) -> None: - """Cache-wide locks outlive repeatedly cancelled cleanup workers.""" - cache = RegistryArtifactCache(temp_cache_dir) - cleanup_started = threading.Event() - cleanup_release = threading.Event() - cleanup_finished = threading.Event() - - def blocking_clear(work_dir: Path) -> bool: - assert work_dir == cache.trash_dir - cleanup_started.set() - cleanup_release.wait(timeout=5) - cleanup_finished.set() - return True - - async def run_operation() -> object: - if operation == "budget": - return await cache._enforce_cache_budget() - await cache._ensure_cache_capacity( - additional_bytes=0, - protected_key="pending", - max_bytes=1, - ) - return None - - with ( - patch.object(cache, "_clear_work_dir", side_effect=blocking_clear), - patch.object( - cache, - "_retry_deferred_staging_cleanup", - return_value=True, - ), - ): - running = asyncio.create_task(run_operation()) - try: - assert await asyncio.to_thread(cleanup_started.wait, 1) - running.cancel() - await asyncio.sleep(0) - running.cancel() - await asyncio.sleep(0) - - assert not running.done() - assert cache._budget_lock.locked() - assert cache._admission_lock.locked() is (operation == "budget") - finally: - cleanup_release.set() - - with pytest.raises(asyncio.CancelledError): - await running - - assert cleanup_finished.is_set() - assert not cache._budget_lock.locked() - assert not cache._admission_lock.locked() - - @pytest.mark.parametrize( - "oldest_has_tarball", - [False, True], - ids=["squashfs-only", "tarball-bearing"], - ) - @pytest.mark.anyio - async def test_enforce_budget_evicts_least_recently_used_until_under_max_bytes( - self, temp_cache_dir, oldest_has_tarball: bool - ): - """Size eviction stops once the cache is within budget.""" - cache = RegistryArtifactCache(temp_cache_dir) - oldest = write_image_entry(temp_cache_dir, "oldest", size=4096, mtime=100.0) - if oldest_has_tarball: - write_tarball_entry(temp_cache_dir, "oldest") - oldest_entry = cache._paths_for("oldest").entry_dir - os.utime(oldest_entry, (100.0, 100.0)) - older = write_image_entry(temp_cache_dir, "older", size=4096, mtime=200.0) - newest = write_image_entry(temp_cache_dir, "newest", size=4096, mtime=300.0) - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 9000), - ): - within_budget = await cache._enforce_cache_budget(protected_key="pending") - - assert within_budget is True - assert not oldest.exists() - assert not cache._paths_for("oldest").tarball_target_dir.exists() - assert older.exists() - assert newest.exists() - - @pytest.mark.anyio - async def test_final_lease_release_unmounts_and_retains_image(self, temp_cache_dir): - """An idle entry releases its loop device without deleting its image.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/path/site-packages.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - paths = cache._paths_for(cache_key) - paths.entry_dir.mkdir(parents=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir() - mounted = {paths.squashfs_mount_dir} - - process = AsyncMock() - process.pid = 999_999_999 - process.communicate.return_value = (b"", b"") - process.returncode = 0 - - async def mock_umount(*args, **kwargs): - assert kwargs["start_new_session"] is True - mounted.discard(paths.squashfs_mount_dir) - return process - - with ( - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path in mounted, - ), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/umount", - ), - patch.object( - asyncio, - "create_subprocess_exec", - side_effect=mock_umount, - ), - patch("tracecat.sandbox.utils.os.killpg") as kill_group, - ): - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [paths.squashfs_mount_dir] - assert paths.squashfs_mount_dir in mounted - - assert paths.squashfs_mount_dir not in mounted - - assert paths.squashfs_image_path.read_bytes() == b"squashfs" - kill_group.assert_called_once_with(process.pid, signal.SIGKILL) - assert paths.squashfs_mount_dir.is_dir() - - @pytest.mark.anyio - async def test_failed_final_release_unmount_retries_on_later_cleanup( - self, temp_cache_dir - ): - """A transient unmount failure is retried after another lease release.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/path/retry-unmount.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - paths = cache._paths_for(cache_key) - paths.entry_dir.mkdir(parents=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir() - mounted = {paths.squashfs_mount_dir} - retry_uri = "s3://bucket/path/retry-trigger.tar.gz" - retry_key = compute_registry_artifact_cache_key(retry_uri) - write_tarball_entry(temp_cache_dir, retry_key) - attempts: list[Path] = [] - - async def flaky_unmount(mount_dir: Path) -> bool: - attempts.append(mount_dir) - if len(attempts) == 1: - return False - mounted.discard(mount_dir) - return True - - with ( - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path in mounted, - ), - patch.object(cache, "_unmount", side_effect=flaky_unmount), - ): - async with cache.lease([artifact_uri]): - pass - - assert cache._failed_unmounts == {cache_key} - assert paths.squashfs_mount_dir in mounted - - async with cache.lease([retry_uri]): - pass - - assert attempts == [paths.squashfs_mount_dir, paths.squashfs_mount_dir] - assert cache._failed_unmounts == set() - assert paths.squashfs_mount_dir not in mounted - - @pytest.mark.anyio - async def test_concurrent_budget_passes_only_evict_once(self, temp_cache_dir): - """A waiting budget pass must re-scan after the active pass evicts.""" - cache = RegistryArtifactCache(temp_cache_dir) - - class ObservedLock(asyncio.Lock): - def __init__(self) -> None: - super().__init__() - self.second_acquire_started = asyncio.Event() - self._acquire_attempts = 0 - - async def acquire(self) -> Literal[True]: - self._acquire_attempts += 1 - if self._acquire_attempts == 2: - self.second_acquire_started.set() - return await super().acquire() - - admission_lock = ObservedLock() - cache._admission_lock = admission_lock - oldest = write_image_entry(temp_cache_dir, "oldest", size=16, mtime=100.0) - retained = write_image_entry(temp_cache_dir, "retained", size=16, mtime=200.0) - original_scan = cache._scan_cache_entries - first_scan_started = threading.Event() - second_scan_started = threading.Event() - release_first_scan = threading.Event() - release_second_scan = threading.Event() - scan_count_lock = threading.Lock() - scan_count = 0 - - def controlled_scan(): - nonlocal scan_count - entries = original_scan() - with scan_count_lock: - scan_index = scan_count - scan_count += 1 - if scan_index == 0: - first_scan_started.set() - release_first_scan.wait(timeout=5) - elif scan_index == 1: - second_scan_started.set() - release_second_scan.wait(timeout=5) - return entries - - eviction_started = asyncio.Event() - finish_eviction = asyncio.Event() - evicted_keys: list[str] = [] - - async def controlled_evict( - cache_key: str, - ) -> RegistryArtifactEviction: - if cache_key == "oldest": - if eviction_started.is_set(): - return RegistryArtifactEviction( - retired=False, - reclaimed=False, - ) - eviction_started.set() - await finish_eviction.wait() - _delete_cache_path(cache._paths_for("oldest").entry_dir) - else: - _delete_cache_path(cache._paths_for("retained").entry_dir) - evicted_keys.append(cache_key) - return RegistryArtifactEviction(retired=True, reclaimed=True) - - with ( - patch.object(cache, "_scan_cache_entries", side_effect=controlled_scan), - patch.object(cache, "_evict_entry", side_effect=controlled_evict), - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - ): - first_pass = asyncio.create_task(cache._enforce_cache_budget()) - assert await asyncio.to_thread(first_scan_started.wait, 1) - second_pass = asyncio.create_task(cache._enforce_cache_budget()) - await admission_lock.second_acquire_started.wait() - assert not second_scan_started.is_set() - - release_first_scan.set() - await asyncio.wait_for(eviction_started.wait(), timeout=1) - assert not second_scan_started.is_set() - finish_eviction.set() - assert await asyncio.to_thread(second_scan_started.wait, 1) - release_second_scan.set() - await asyncio.gather(first_pass, second_pass) - - assert scan_count == 2 - assert evicted_keys == ["oldest"] - assert not oldest.exists() - assert retained.exists() - - @pytest.mark.anyio - async def test_enforce_budget_ignores_missing_protected_key(self, temp_cache_dir): - """Only successfully published entries count against the entry budget.""" - cache = RegistryArtifactCache(temp_cache_dir) - existing = write_image_entry(temp_cache_dir, "existing", size=16, mtime=100.0) - - with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): - within_budget = await cache._enforce_cache_budget(protected_key="missing") - - assert within_budget is True - assert existing.exists() - - @pytest.mark.anyio - async def test_enforce_budget_never_evicts_the_protected_key(self, temp_cache_dir): - """The newly materialized key is exempt even when it is the LRU entry.""" - cache = RegistryArtifactCache(temp_cache_dir) - protected = write_image_entry( - temp_cache_dir, "protected", size=4096, mtime=100.0 - ) - other = write_image_entry(temp_cache_dir, "other", size=4096, mtime=300.0) - - with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): - await cache._enforce_cache_budget(protected_key="protected") - - assert protected.exists() - assert not other.exists() - - @pytest.mark.anyio - async def test_enforce_budget_proceeds_over_budget_when_everything_is_leased( - self, temp_cache_dir - ): - """An over-budget cache must degrade, not fail the action.""" - cache = RegistryArtifactCache(temp_cache_dir) - leased = write_image_entry(temp_cache_dir, "leased", size=4096, mtime=100.0) - cache._acquire_lease("leased") - - with patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 1): - within_budget = await cache._enforce_cache_budget(protected_key="missing") - - assert within_budget is False - assert leased.exists() diff --git a/tests/unit/executor/test_registry_artifact_eviction.py b/tests/unit/executor/test_registry_artifact_eviction.py deleted file mode 100644 index 7946ae8671..0000000000 --- a/tests/unit/executor/test_registry_artifact_eviction.py +++ /dev/null @@ -1,344 +0,0 @@ -"""Registry artifact retirement and unmount tests.""" - -from __future__ import annotations - -import asyncio -import os -import signal -import threading -from pathlib import Path -from unittest.mock import AsyncMock, patch - -import pytest - -from tracecat.executor.registry_artifacts import ( - RegistryArtifactCache, - RegistryArtifactEviction, - TarballArtifact, - _delete_cache_path, - compute_registry_artifact_cache_key, -) - -from .registry_artifact_test_helpers import ( - MAX_BYTES_CONFIG, - MAX_ENTRIES_CONFIG, - BlockingSubprocess, - tarball_payload, - write_image_entry, - write_tarball_entry, -) - - -class TestRegistryArtifactCacheEviction: - """Retire idle cache entries without disrupting live leases.""" - - @pytest.mark.anyio - async def test_eviction_surfaces_unknown_mount_state(self, temp_cache_dir): - """Inspection failures cannot be mistaken for an unmounted entry.""" - cache = RegistryArtifactCache(temp_cache_dir) - paths = cache._paths_for("unknown-mount-state") - paths.entry_dir.mkdir(parents=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir() - real_lstat = Path.lstat - - def fail_mount_lstat(path: Path): - if path == paths.squashfs_mount_dir: - raise PermissionError("mount inspection denied") - return real_lstat(path) - - with patch.object(Path, "lstat", fail_mount_lstat): - with pytest.raises(PermissionError, match="mount inspection denied"): - await cache._evict_entry("unknown-mount-state") - - assert paths.entry_dir.is_dir() - assert paths.squashfs_image_path.is_file() - assert paths.squashfs_mount_dir.is_dir() - - @pytest.mark.anyio - async def test_eviction_unmounts_before_deleting_the_image(self, temp_cache_dir): - """Unlinking a mounted image would strand an open-file zombie.""" - cache = RegistryArtifactCache(temp_cache_dir) - paths = cache._paths_for("mounted") - paths.entry_dir.mkdir(parents=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir() - mounted = {paths.squashfs_mount_dir} - image_present_at_umount: list[bool] = [] - - process = AsyncMock() - process.pid = 999_999_999 - process.communicate.return_value = (b"", b"") - process.returncode = 0 - - async def mock_umount(*args, **kwargs): - assert kwargs["start_new_session"] is True - image_present_at_umount.append(paths.squashfs_image_path.exists()) - mounted.discard(paths.squashfs_mount_dir) - return process - - with ( - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path in mounted, - ), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/umount", - ), - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - side_effect=mock_umount, - ) as create_subprocess_exec, - patch("tracecat.sandbox.utils.os.killpg") as kill_group, - ): - evicted = await cache._evict_entry("mounted") - - assert evicted == RegistryArtifactEviction(retired=True, reclaimed=True) - assert image_present_at_umount == [True] - assert not paths.squashfs_image_path.exists() - assert not paths.squashfs_mount_dir.exists() - create_subprocess_exec.assert_called_once() - assert create_subprocess_exec.call_args.args == ( - "/sbin/umount", - str(paths.squashfs_mount_dir), - ) - kill_group.assert_called_once_with(process.pid, signal.SIGKILL) - - @pytest.mark.anyio - async def test_repeatedly_cancelled_unmount_reaps_before_releasing_key_lock( - self, temp_cache_dir - ): - """Repeated cancellation leaves a consistent entry for next admission.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/path/cancelled-unmount.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - paths = cache._paths_for(cache_key) - paths.entry_dir.mkdir(parents=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir() - (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") - mounted = {paths.squashfs_mount_dir} - blocked_process = BlockingSubprocess(block_wait=True) - released_process = AsyncMock() - released_process.communicate.return_value = (b"", b"") - released_process.returncode = 0 - unmount_attempts = 0 - - async def mock_umount(*args, **kwargs): - nonlocal unmount_attempts - assert kwargs["start_new_session"] is True - unmount_attempts += 1 - if unmount_attempts == 1: - return blocked_process - mounted.discard(paths.squashfs_mount_dir) - return released_process - - with ( - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path in mounted, - ), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/umount", - ), - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - side_effect=mock_umount, - ), - patch("tracecat.sandbox.utils.os.killpg") as kill_group, - ): - eviction = asyncio.create_task(cache._evict_entry(cache_key)) - await blocked_process.communicate_started.wait() - eviction.cancel() - await blocked_process.wait_started.wait() - - eviction.cancel() - done, _ = await asyncio.wait({eviction}, timeout=0.05) - assert not done - blocked_process.release_wait.set() - with pytest.raises(asyncio.CancelledError): - await eviction - - assert blocked_process.cleanup_calls == ["kill", "wait"] - kill_group.assert_called_once_with( - blocked_process.pid, - signal.SIGKILL, - ) - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [paths.squashfs_mount_dir] - assert registry_paths[0].is_dir() - assert (registry_paths[0] / "module.py").read_text() == "VALUE = 1" - - assert paths.squashfs_image_path.is_file() - assert paths.squashfs_mount_dir.is_dir() - assert paths.squashfs_mount_dir not in mounted - - @pytest.mark.anyio - async def test_cancelled_background_deletion_leaves_a_clean_miss( - self, temp_cache_dir - ): - """Cancellation cannot expose live paths that deletion still owns.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/cancelled-eviction.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - original_target = write_tarball_entry(temp_cache_dir, cache_key) - delete_started = threading.Event() - finish_delete = threading.Event() - delete_finished = threading.Event() - doomed: list[Path] = [] - - def blocked_delete(path: Path) -> bool: - doomed.append(path) - delete_started.set() - finish_delete.wait(timeout=5) - deleted = _delete_cache_path(path) - delete_finished.set() - return deleted - - async def mock_download(self, ctx, path): - path.write_bytes(tarball_payload(size=1)) - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_text("VALUE = 2") - - with ( - patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - side_effect=blocked_delete, - ), - patch( - "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", - return_value=9, - ), - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - eviction = asyncio.create_task(cache._evict_entry(cache_key)) - assert await asyncio.to_thread(delete_started.wait, 1) - eviction.cancel() - try: - await asyncio.sleep(0) - eviction.cancel() - await asyncio.sleep(0) - assert not eviction.done() - assert not original_target.exists() - assert (doomed[0] / "tarball").is_dir() - assert cache._discover_cache_keys() == set() - finally: - finish_delete.set() - - with pytest.raises(asyncio.CancelledError): - await eviction - assert await asyncio.to_thread(delete_finished.wait, 1) - assert not doomed[0].exists() - - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [original_target] - assert (original_target / "module.py").read_text() == "VALUE = 2" - - assert original_target.is_dir() - - @pytest.mark.anyio - async def test_doomed_eviction_names_are_startup_scratch(self, temp_cache_dir): - """Every renamed entry root is invisible and reclaimed on startup.""" - cache = RegistryArtifactCache(temp_cache_dir) - cache_key = "squashfs-doomed" - paths = cache._paths_for(cache_key) - paths.entry_dir.mkdir(parents=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir() - paths.squashfs_extract_dir.mkdir() - paths.tarball_target_dir.mkdir() - - with patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - return_value=True, - ) as delete_cache_path: - assert await cache._evict_entry(cache_key) == RegistryArtifactEviction( - retired=True, reclaimed=True - ) - - delete_cache_path.assert_called_once() - trash_path = delete_cache_path.call_args.args[0] - assert trash_path.parent == cache.trash_dir - assert trash_path.exists() - assert cache._discover_cache_keys() == set() - - cache._sweep_startup_state() - - assert not trash_path.exists() - - @pytest.mark.anyio - async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): - """A failed unmount skips the key instead of forcing a lazy detach.""" - cache = RegistryArtifactCache(temp_cache_dir) - stuck = cache._paths_for("stuck") - stuck.entry_dir.mkdir(parents=True) - stuck.squashfs_image_path.write_bytes(b"squashfs") - stuck.squashfs_mount_dir.mkdir() - os.utime(stuck.squashfs_image_path, (100.0, 100.0)) - os.utime(stuck.entry_dir, (100.0, 100.0)) - idle = write_image_entry(temp_cache_dir, "idle", size=16, mtime=300.0) - mounted = {stuck.squashfs_mount_dir} - - process = AsyncMock() - process.pid = 999_999_999 - process.communicate.return_value = (b"", b"target is busy") - process.returncode = 32 - - with ( - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path in mounted, - ), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/umount", - ), - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - new_callable=AsyncMock, - return_value=process, - ) as create_subprocess_exec, - patch("tracecat.sandbox.utils.os.killpg") as kill_group, - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - ): - await cache._enforce_cache_budget(protected_key="pending") - - assert stuck.squashfs_image_path.exists() - assert stuck.squashfs_mount_dir.exists() - assert not idle.exists() - create_subprocess_exec.assert_awaited_once() - process.communicate.assert_awaited_once() - kill_group.assert_called_once_with(process.pid, signal.SIGKILL) - - @pytest.mark.anyio - async def test_eviction_discards_idle_runtime_state(self, temp_cache_dir): - """An evicted key releases runtime state after every lock user exits.""" - cache = RegistryArtifactCache(temp_cache_dir) - write_tarball_entry(temp_cache_dir, "bookkeeping") - cache._acquire_lease("bookkeeping") - cache._release_lease("bookkeeping") - - assert await cache._evict_entry("bookkeeping") == RegistryArtifactEviction( - retired=True, reclaimed=True - ) - assert "bookkeeping" not in cache._runtime - - @pytest.mark.anyio - async def test_eviction_skips_busy_key(self, temp_cache_dir): - """A key another task is materializing is never evicted underneath it.""" - cache = RegistryArtifactCache(temp_cache_dir) - target_dir = write_tarball_entry(temp_cache_dir, "busy") - lock = cache._runtime_for("busy").lock - - async with lock: - assert await cache._evict_entry("busy") == RegistryArtifactEviction( - retired=False, reclaimed=False - ) - - assert target_dir.is_dir() diff --git a/tests/unit/executor/test_registry_artifact_leases.py b/tests/unit/executor/test_registry_artifact_leases.py deleted file mode 100644 index 7e8c1c400b..0000000000 --- a/tests/unit/executor/test_registry_artifact_leases.py +++ /dev/null @@ -1,927 +0,0 @@ -"""Registry artifact lease lifetime and concurrency tests.""" - -from __future__ import annotations - -import asyncio -import os -import signal -from pathlib import Path -from unittest.mock import AsyncMock, patch - -import pytest -import tracecat_registry - -from tracecat.executor.registry_artifacts import ( - RegistryArtifactCache, - RegistryArtifactCacheLoopError, - RegistryArtifactEviction, - RegistryArtifactMaterializationContext, - SquashfsArtifact, - TarballArtifact, - bundled_builtin_registry_uri, - compute_registry_artifact_cache_key, -) - -from .registry_artifact_test_helpers import ( - MAX_BYTES_CONFIG, - MAX_ENTRIES_CONFIG, - SQUASHFS_ENABLED_CONFIG, - SquashfsMountHarness, - write_image_entry, - write_tarball_entry, -) - - -class TestRegistryArtifactCacheLease: - """Tests for lease-based pinning of registry artifact cache entries.""" - - @pytest.mark.anyio - async def test_cache_rejects_use_from_a_second_event_loop( - self, temp_cache_dir: Path - ) -> None: - """Protect the process-wide cache from thread-local Temporal loops. - - The cache owns asyncio locks, tasks, and refcounts as one unit. A typed - ownership failure prevents a future synchronous activity from silently - sharing only part of that state across thread-local event loops, which - previously caused intermittent cross-loop failures in storage caches. - """ - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - - def use_cache_from_another_thread() -> None: - asyncio.run(cache.ensure_swept()) - - with pytest.raises(RegistryArtifactCacheLoopError): - await asyncio.to_thread(use_cache_from_another_thread) - - # A rejected caller must not poison the owning loop's cache. - async with cache.lease(None) as registry_paths: - assert registry_paths == [temp_cache_dir / "base"] - - def test_touch_entry_refreshes_tarball_root_mtime(self, temp_cache_dir): - """Touching a tarball-only entry persists its restart-safe recency.""" - cache = RegistryArtifactCache(temp_cache_dir) - cache_key = "tarball-only" - write_tarball_entry(temp_cache_dir, cache_key) - entry_dir = cache._paths_for(cache_key).entry_dir - os.utime(entry_dir, (100.0, 100.0)) - - cache._touch_entry(cache_key) - - assert entry_dir.stat().st_mtime > 100.0 - - @pytest.mark.anyio - async def test_lease_refcounts_and_touches_image_mtime(self, temp_cache_dir): - """Acquire and final release persist the restart-safe LRU timestamp.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/leased.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - target_dir = write_tarball_entry(temp_cache_dir, cache_key) - image_path = write_image_entry(temp_cache_dir, cache_key, size=16, mtime=100.0) - entry_dir = cache._paths_for(cache_key).entry_dir - - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [target_dir] - assert cache._refcount(cache_key) == 1 - assert entry_dir.stat().st_mtime > 100.0 - os.utime(entry_dir, (100.0, 100.0)) - - assert cache._refcount(cache_key) == 0 - assert entry_dir.stat().st_mtime > 100.0 - assert image_path.is_file() - - @pytest.mark.anyio - async def test_lease_releases_refcount_when_materialization_fails( - self, temp_cache_dir - ): - """A failed materialization must not leak a pin or empty cache entry.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/broken.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - - async def mock_download(self, ctx, path): - raise RuntimeError("download failed") - - with patch.object(TarballArtifact, "download", mock_download): - with pytest.raises(RuntimeError, match="download failed"): - async with cache.lease([artifact_uri]): - pass - - assert cache._refcount(cache_key) == 0 - assert not cache._paths_for(cache_key).entry_dir.exists() - assert cache_key not in cache._discover_cache_keys() - assert cache_key not in cache._runtime - - @pytest.mark.anyio - async def test_distinct_failed_keys_do_not_accumulate_runtime_states( - self, temp_cache_dir: Path - ) -> None: - """Failed cold keys release lock state after their last waiter exits.""" - cache = RegistryArtifactCache(temp_cache_dir) - - async def fail_download(self, ctx, path): - del self, ctx, path - raise RuntimeError("download failed") - - with patch.object(TarballArtifact, "download", fail_download): - for index in range(100): - with pytest.raises(RuntimeError, match="download failed"): - async with cache.lease([f"s3://bucket/broken-{index}.tar.gz"]): - pass - - assert cache._runtime == {} - - @pytest.mark.anyio - async def test_failed_first_admission_converges_deposited_image( - self, temp_cache_dir: Path - ) -> None: - """A failed first artifact must not strand an over-budget image.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/failed-first.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) - - async def fail_after_deposit( - artifact: SquashfsArtifact, - ctx: RegistryArtifactMaterializationContext, - ) -> list[Path]: - del artifact - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - ctx.paths.squashfs_image_path.write_bytes(b"reusable image") - raise RuntimeError("mount failed") - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 1), - patch.object( - cache, - "_artifact_candidates", - new_callable=AsyncMock, - return_value=[artifact], - ), - patch.object(SquashfsArtifact, "materialize", fail_after_deposit), - ): - with pytest.raises(RuntimeError, match="mount failed"): - async with cache.lease([artifact_uri]): - pass - - assert cache._refcount(cache_key) == 0 - assert not cache._paths_for(cache_key).entry_dir.exists() - assert cache._discover_cache_keys() == set() - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_cancelled_lease_admission_releases_refcount(self, temp_cache_dir): - """Cancellation during candidate lookup must not leak a permanent pin.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/path/site-packages.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - lookup_started = asyncio.Event() - finish_lookup = asyncio.Event() - - async def blocked_sidecar_lookup(**kwargs): - lookup_started.set() - await finish_lookup.wait() - return False - - async def take_lease() -> None: - async with cache.lease([artifact_uri]): - pass - - with ( - patch(SQUASHFS_ENABLED_CONFIG, True), - patch.object(cache, "_sidecar_exists", blocked_sidecar_lookup), - ): - acquisition = asyncio.create_task(take_lease()) - await lookup_started.wait() - assert cache._refcount(cache_key) == 1 - acquisition.cancel() - with pytest.raises(asyncio.CancelledError): - await acquisition - - assert cache._refcount(cache_key) == 0 - assert not cache._paths_for(cache_key).entry_dir.exists() - - @pytest.mark.anyio - async def test_repeated_cancellation_finishes_acquisition_rollback( - self, - temp_cache_dir: Path, - ) -> None: - """A second cancellation cannot abandon the final rollback unmount.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/path/site-packages.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - lookup_started = asyncio.Event() - rollback_started = asyncio.Event() - finish_rollback = asyncio.Event() - - async def blocked_sidecar_lookup(**kwargs): - del kwargs - lookup_started.set() - await asyncio.Event().wait() - - async def blocked_unmount(requested_key: str) -> None: - assert requested_key == cache_key - rollback_started.set() - await finish_rollback.wait() - - with ( - patch(SQUASHFS_ENABLED_CONFIG, True), - patch.object(cache, "_sidecar_exists", blocked_sidecar_lookup), - patch.object(cache, "_unmount_idle_entry", blocked_unmount), - ): - acquisition = asyncio.create_task(cache._lease_artifact(artifact_uri)) - await lookup_started.wait() - acquisition.cancel() - await rollback_started.wait() - - acquisition.cancel() - await asyncio.sleep(0) - assert not acquisition.done() - - finish_rollback.set() - with pytest.raises(asyncio.CancelledError): - await acquisition - - assert cache._refcount(cache_key) == 0 - - @pytest.mark.anyio - async def test_cancelled_waiter_preserves_existing_same_key_lease( - self, temp_cache_dir: Path - ) -> None: - """A waiter cancelled before admission cannot release another holder.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/path/shared.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - lock = cache._runtime_for(cache_key).lock - - async def take_lease() -> None: - async with cache.lease([artifact_uri]): - pytest.fail("cancelled waiter must not enter the lease context") - - unmount_idle_entry = AsyncMock() - with patch.object(cache, "_unmount_idle_entry", unmount_idle_entry): - async with lock: - cache._acquire_lease(cache_key) - try: - waiter = asyncio.create_task(take_lease()) - await asyncio.sleep(0) - assert not waiter.done() - - waiter.cancel() - with pytest.raises(asyncio.CancelledError): - await waiter - - assert cache._refcount(cache_key) == 1 - unmount_idle_entry.assert_not_awaited() - finally: - cache._release_lease(cache_key) - - @pytest.mark.anyio - async def test_lease_without_uris_returns_base_pythonpath_dir(self, temp_cache_dir): - """No artifact URIs still yields the base PYTHONPATH directory.""" - cache = RegistryArtifactCache(temp_cache_dir) - - async with cache.lease(None) as registry_paths: - assert registry_paths == [temp_cache_dir / "base"] - assert registry_paths[0].is_dir() - - assert cache._runtime == {} - - @pytest.mark.anyio - async def test_lease_preserves_uri_order(self, temp_cache_dir): - """Multiple artifacts keep their deterministic PYTHONPATH order.""" - cache = RegistryArtifactCache(temp_cache_dir) - uris = ["s3://bucket/first.tar.gz", "s3://bucket/second.tar.gz"] - expected = [ - write_tarball_entry( - temp_cache_dir, compute_registry_artifact_cache_key(uri) - ) - for uri in uris - ] - - async with cache.lease(uris) as registry_paths: - assert registry_paths == expected - - @pytest.mark.anyio - async def test_overlapping_same_key_leases_share_one_mount_until_final_release( - self, temp_cache_dir: Path - ) -> None: - """Protect refcount and loop-device lifetime under heavy fan-in. - - Every holder must observe one published mount, intermediate releases - must leave that mount usable, and exactly the final release may reclaim - its loop device while retaining the downloaded image for a remount. - """ - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/shared.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - harness = SquashfsMountHarness(cache) - holder_count = 32 - entered = 0 - all_entered = asyncio.Event() - releases = [asyncio.Event() for _ in range(holder_count)] - - async def hold_lease(index: int) -> None: - nonlocal entered - async with cache.lease([artifact_uri]) as registry_paths: - assert registry_paths == [ - cache._paths_for(cache_key).squashfs_mount_dir - ] - entered += 1 - if entered == holder_count: - all_entered.set() - await releases[index].wait() - - with ( - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path in harness.mounted, - ), - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", harness.mount), - patch.object(SquashfsArtifact, "extract", harness.extract), - patch.object(cache, "_unmount", harness.unmount), - ): - holders = [ - asyncio.create_task(hold_lease(index)) for index in range(holder_count) - ] - await asyncio.wait_for(all_entered.wait(), timeout=5) - - mount_dir = cache._paths_for(cache_key).squashfs_mount_dir - assert cache._refcount(cache_key) == holder_count - assert harness.mount_attempts == [cache_key] - assert harness.extraction_attempts == [] - assert mount_dir in harness.mounted - - for release in releases[:-1]: - release.set() - await asyncio.gather(*holders[:-1]) - - assert cache._refcount(cache_key) == 1 - assert mount_dir in harness.mounted - assert harness.unmounts == [] - - releases[-1].set() - await holders[-1] - - assert cache._refcount(cache_key) == 0 - assert harness.unmounts == [mount_dir] - assert mount_dir not in harness.mounted - assert cache._paths_for(cache_key).squashfs_image_path.is_file() - assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) - assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) - - @pytest.mark.anyio - async def test_cancelling_one_holder_preserves_sibling_leases( - self, temp_cache_dir: Path - ) -> None: - """Protect sibling actions from another holder's cancellation. - - Cancellation must release exactly one pin without unmounting the shared - artifact beneath surviving actions; the last surviving holder remains - solely responsible for loop-device reclamation. - """ - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/cancel-one.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - harness = SquashfsMountHarness(cache) - entered = [asyncio.Event() for _ in range(3)] - releases = [asyncio.Event() for _ in range(3)] - - async def hold_lease(index: int) -> None: - async with cache.lease([artifact_uri]): - entered[index].set() - await releases[index].wait() - - with ( - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path in harness.mounted, - ), - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", harness.mount), - patch.object(SquashfsArtifact, "extract", harness.extract), - patch.object(cache, "_unmount", harness.unmount), - ): - holders = [asyncio.create_task(hold_lease(index)) for index in range(3)] - await asyncio.wait_for( - asyncio.gather(*(event.wait() for event in entered)), - timeout=5, - ) - - holders[0].cancel() - with pytest.raises(asyncio.CancelledError): - await holders[0] - - mount_dir = cache._paths_for(cache_key).squashfs_mount_dir - assert cache._refcount(cache_key) == 2 - assert mount_dir in harness.mounted - assert harness.unmounts == [] - - releases[1].set() - await holders[1] - assert cache._refcount(cache_key) == 1 - assert mount_dir in harness.mounted - assert harness.unmounts == [] - - releases[2].set() - await holders[2] - - assert cache._refcount(cache_key) == 0 - assert harness.mount_attempts == [cache_key] - assert harness.unmounts == [mount_dir] - - @pytest.mark.anyio - async def test_repeated_cancellation_finishes_all_lease_cleanup( - self, temp_cache_dir: Path - ) -> None: - """Every unmount and budget pass finishes before cancellation propagates.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uris = [ - "s3://bucket/first.tar.gz", - "s3://bucket/second.tar.gz", - ] - cache_keys = [compute_registry_artifact_cache_key(uri) for uri in artifact_uris] - for cache_key in cache_keys: - write_tarball_entry(temp_cache_dir, cache_key) - - lease_entered = asyncio.Event() - first_unmount_started = asyncio.Event() - finish_first_unmount = asyncio.Event() - cleanup_calls: list[str] = [] - - async def mock_unmount_idle_entry(cache_key: str) -> None: - cleanup_calls.append(cache_key) - if cache_key == cache_keys[0]: - first_unmount_started.set() - await finish_first_unmount.wait() - - async def mock_converge_cache_budget() -> None: - cleanup_calls.append("converge") - - async def hold_lease() -> None: - async with cache.lease(artifact_uris): - lease_entered.set() - await asyncio.Event().wait() - - with ( - patch.object(cache, "_unmount_idle_entry", mock_unmount_idle_entry), - patch.object(cache, "_converge_cache_budget", mock_converge_cache_budget), - ): - holder = asyncio.create_task(hold_lease()) - await lease_entered.wait() - holder.cancel() - await first_unmount_started.wait() - - holder.cancel() - await asyncio.sleep(0) - assert not holder.done() - - finish_first_unmount.set() - with pytest.raises(asyncio.CancelledError): - await holder - - assert cleanup_calls == [*cache_keys, "converge"] - assert all(cache._refcount(cache_key) == 0 for cache_key in cache_keys) - - @pytest.mark.anyio - async def test_cleanup_failure_preserves_successful_lease_outcome( - self, temp_cache_dir: Path - ) -> None: - """Post-lease maintenance cannot replace a successful caller result.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/cleanup-failure.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - target_dir = write_tarball_entry(temp_cache_dir, cache_key) - - async def fail_cleanup( - idle_keys: list[str], - *, - converge: bool, - ) -> None: - del idle_keys, converge - raise RuntimeError( - "cleanup failed for s3://access:secret@bucket/path?signature=secret" - ) - - with ( - patch.object(cache, "_finish_lease_cleanup", fail_cleanup), - patch("tracecat.executor.registry_artifacts.logger.error") as log_error, - ): - async with cache.lease([artifact_uri]) as registry_paths: - result = registry_paths - - assert result == [target_dir] - assert cache._refcount(cache_key) == 0 - log_error.assert_called_once_with( - "Registry artifact lease cleanup failed; preserving caller outcome", - cache_dir=str(temp_cache_dir), - error_type="RuntimeError", - ) - - @pytest.mark.anyio - async def test_new_lease_racing_final_release_prevents_stale_unmount( - self, temp_cache_dir: Path - ) -> None: - """Protect the zero-refcount-to-unmount handoff from a new acquisition. - - A lease can arrive after the old holder decrements to zero but before - unmount takes the key lock. The lock-time refcount recheck must preserve - the mount for that newcomer instead of returning a path being torn down. - """ - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/release-race.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - harness = SquashfsMountHarness(cache) - mount_dir = harness.seed_mount(cache_key) - first_entered = asyncio.Event() - release_first = asyncio.Event() - release_reached_unmount = asyncio.Event() - allow_unmount_recheck = asyncio.Event() - newcomer_entered = asyncio.Event() - release_newcomer = asyncio.Event() - original_unmount_idle_entry = cache._unmount_idle_entry - unmount_requests = 0 - - async def pause_first_unmount_request(requested_key: str) -> None: - nonlocal unmount_requests - unmount_requests += 1 - if unmount_requests == 1: - release_reached_unmount.set() - await allow_unmount_recheck.wait() - await original_unmount_idle_entry(requested_key) - - async def old_holder() -> None: - async with cache.lease([artifact_uri]): - first_entered.set() - await release_first.wait() - - async def new_holder() -> None: - async with cache.lease([artifact_uri]): - newcomer_entered.set() - await release_newcomer.wait() - - with ( - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path in harness.mounted, - ), - patch.object(cache, "_unmount", harness.unmount), - patch.object( - cache, - "_unmount_idle_entry", - pause_first_unmount_request, - ), - ): - old_task = asyncio.create_task(old_holder()) - await first_entered.wait() - release_first.set() - await release_reached_unmount.wait() - assert cache._refcount(cache_key) == 0 - - new_task = asyncio.create_task(new_holder()) - await newcomer_entered.wait() - assert cache._refcount(cache_key) == 1 - - allow_unmount_recheck.set() - await old_task - - assert mount_dir in harness.mounted - assert harness.unmounts == [] - - release_newcomer.set() - await new_task - - assert cache._refcount(cache_key) == 0 - assert harness.unmounts == [mount_dir] - - @pytest.mark.anyio - async def test_partial_multi_artifact_failure_rolls_back_prior_leases( - self, temp_cache_dir: Path - ) -> None: - """Protect sequential multi-artifact admission from partial pin leaks. - - If a middle artifact fails, every earlier acquisition must be released - and unmounted, the failed key must leave no shell, later artifacts must - never be acquired, and the release path must still request convergence. - """ - cache = RegistryArtifactCache(temp_cache_dir) - first_uri = "s3://bucket/first.squashfs" - failed_uri = "s3://bucket/failed.tar.gz" - untouched_uri = "s3://bucket/untouched.tar.gz" - first_key = compute_registry_artifact_cache_key(first_uri) - failed_key = compute_registry_artifact_cache_key(failed_uri) - untouched_key = compute_registry_artifact_cache_key(untouched_uri) - harness = SquashfsMountHarness(cache) - first_mount = harness.seed_mount(first_key) - untouched_path = write_tarball_entry(temp_cache_dir, untouched_key) - - async def fail_download( - artifact: TarballArtifact, - ctx: RegistryArtifactMaterializationContext, - output_path: Path, - ) -> None: - del artifact, ctx, output_path - raise RuntimeError("download failed") - - tracked_lease_artifact = AsyncMock(wraps=cache._lease_artifact) - converge_cache_budget = AsyncMock() - - with ( - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path in harness.mounted, - ), - patch(SQUASHFS_ENABLED_CONFIG, False), - patch.object(TarballArtifact, "download", fail_download), - patch.object(cache, "_unmount", harness.unmount), - patch.object(cache, "_lease_artifact", tracked_lease_artifact), - patch.object( - cache, - "_converge_cache_budget", - converge_cache_budget, - ), - ): - with pytest.raises(RuntimeError, match="download failed"): - async with cache.lease([first_uri, failed_uri, untouched_uri]): - pass - - requested_uris = [ - await_call.args[0] for await_call in tracked_lease_artifact.await_args_list - ] - assert requested_uris == [first_uri, failed_uri] - assert cache._refcount(first_key) == 0 - assert cache._refcount(failed_key) == 0 - assert cache._refcount(untouched_key) == 0 - assert harness.unmounts == [first_mount] - assert cache._paths_for(first_key).squashfs_image_path.is_file() - assert not cache._paths_for(failed_key).entry_dir.exists() - assert untouched_path.is_dir() - converge_cache_budget.assert_awaited_once_with() - assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) - assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) - - @pytest.mark.anyio - async def test_waiter_retries_after_first_same_key_materializer_fails( - self, temp_cache_dir: Path - ) -> None: - """Protect same-key waiters from a failed cold-cache publisher. - - The first materializer may fail after writing scratch. Its waiter must - retry against a clean staging area, publish the sole valid entry, and - leave neither a leaked pin nor an empty LRU-visible cache shell. - """ - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/retry-after-failure.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - first_download_started = asyncio.Event() - fail_first_download = asyncio.Event() - retry_download_started = asyncio.Event() - allow_retry_download = asyncio.Event() - download_attempts = 0 - retry_saw_clean_staging = False - - async def controlled_download( - artifact: TarballArtifact, - ctx: RegistryArtifactMaterializationContext, - output_path: Path, - ) -> None: - del artifact, ctx - nonlocal download_attempts, retry_saw_clean_staging - download_attempts += 1 - if download_attempts == 1: - output_path.write_bytes(b"partial") - first_download_started.set() - await fail_first_download.wait() - raise RuntimeError("first publisher failed") - - retry_saw_clean_staging = not any(cache.staging_dir.iterdir()) - retry_download_started.set() - await allow_retry_download.wait() - output_path.write_bytes(b"complete") - - async def mock_extract( - artifact: TarballArtifact, - tarball_path: Path, - target_dir: Path, - ) -> None: - del artifact - assert tarball_path.read_bytes() == b"complete" - (target_dir / "module.py").write_text("VALUE = 1") - - async def take_lease() -> list[Path]: - async with cache.lease([artifact_uri]) as registry_paths: - return registry_paths - - with ( - patch(SQUASHFS_ENABLED_CONFIG, False), - patch( - "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", - return_value=1, - ), - patch.object(TarballArtifact, "download", controlled_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - first = asyncio.create_task(take_lease()) - await first_download_started.wait() - waiter = asyncio.create_task(take_lease()) - await asyncio.sleep(0) - assert download_attempts == 1 - - fail_first_download.set() - await retry_download_started.wait() - allow_retry_download.set() - registry_paths = await waiter - with pytest.raises(RuntimeError, match="first publisher failed"): - await first - - target_dir = cache._paths_for(cache_key).tarball_target_dir - assert registry_paths == [target_dir] - assert (target_dir / "module.py").read_text() == "VALUE = 1" - assert download_attempts == 2 - assert retry_saw_clean_staging is True - assert cache._discover_cache_keys() == {cache_key} - assert cache._refcount(cache_key) == 0 - assert not any(cache.staging_dir.iterdir()) - assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) - - @pytest.mark.anyio - async def test_duplicate_uri_balances_each_acquisition_and_release( - self, temp_cache_dir: Path - ) -> None: - """Protect list bookkeeping when one lease requests the same URI twice. - - Duplicate PYTHONPATH entries intentionally acquire two pins. Both must - be released, while materialization and final unmount still occur once; - accidental deduplication on only one side would leak or underflow pins. - """ - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/duplicate.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - harness = SquashfsMountHarness(cache) - - with ( - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path in harness.mounted, - ), - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", harness.mount), - patch.object(SquashfsArtifact, "extract", harness.extract), - patch.object(cache, "_unmount", harness.unmount), - ): - async with cache.lease([artifact_uri, artifact_uri]) as registry_paths: - mount_dir = cache._paths_for(cache_key).squashfs_mount_dir - assert registry_paths == [mount_dir, mount_dir] - assert cache._refcount(cache_key) == 2 - assert harness.mount_attempts == [cache_key] - - assert cache._refcount(cache_key) == 0 - assert harness.unmounts == [mount_dir] - assert cache._paths_for(cache_key).squashfs_image_path.is_file() - - @pytest.mark.anyio - async def test_lease_is_never_admitted_across_an_in_flight_eviction( - self, temp_cache_dir - ): - """A lease must not return a mount an in-flight eviction is deleting.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/path/site-packages.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - paths = cache._paths_for(cache_key) - paths.entry_dir.mkdir(parents=True) - paths.squashfs_image_path.write_bytes(b"squashfs") - paths.squashfs_mount_dir.mkdir() - mounted = {paths.squashfs_mount_dir} - umount_started = asyncio.Event() - finish_umount = asyncio.Event() - lease_waiting_for_key = asyncio.Event() - remounts: list[str] = [] - original_runtime_for = cache._runtime_for - - def observed_runtime_for(requested_key: str): - runtime = original_runtime_for(requested_key) - if requested_key == cache_key and umount_started.is_set(): - lease_waiting_for_key.set() - return runtime - - umount_process = AsyncMock() - umount_process.pid = 999_999_999 - umount_process.communicate.return_value = (b"", b"") - umount_process.returncode = 0 - - async def mock_umount(*args, **kwargs): - umount_started.set() - await finish_umount.wait() - mounted.discard(paths.squashfs_mount_dir) - return umount_process - - async def mock_mount(self, ctx, image_path): - remounts.append(ctx.cache_key) - target_dir = ctx.paths.squashfs_mount_dir - target_dir.mkdir(parents=True, exist_ok=True) - (target_dir / "module.py").write_text("VALUE = 1") - mounted.add(target_dir) - return target_dir - - leased_paths: list[Path] = [] - leased_path_exists: list[bool] = [] - - async def take_lease() -> None: - async with cache.lease([artifact_uri]) as registry_paths: - leased_paths.extend(registry_paths) - leased_path_exists.append(registry_paths[0].is_dir()) - - with ( - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path in mounted, - ), - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/umount", - ), - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - side_effect=mock_umount, - ), - patch("tracecat.sandbox.utils.os.killpg") as kill_group, - patch.object(cache, "_runtime_for", side_effect=observed_runtime_for), - patch.object(SquashfsArtifact, "mount", mock_mount), - # This test targets the per-key eviction/lease handoff. Keep the - # lease's budget pass from concurrently sweeping the same trash - # path and turning physical reclamation into a two-deleter race. - patch.object( - cache, - "_enforce_cache_budget", - new_callable=AsyncMock, - return_value=True, - ), - ): - eviction = asyncio.create_task(cache._evict_entry(cache_key)) - await umount_started.wait() - lease = asyncio.create_task(take_lease()) - await lease_waiting_for_key.wait() - assert cache._runtime[cache_key].users == 2 - finish_umount.set() - evicted, _ = await asyncio.gather(eviction, lease) - - assert evicted == RegistryArtifactEviction(retired=True, reclaimed=True) - # The lease waited for the eviction and re-materialized the entry. - assert remounts == [cache_key] - assert leased_paths == [paths.squashfs_mount_dir] - assert leased_path_exists == [True] - assert (paths.squashfs_mount_dir / "module.py").read_text() == "VALUE = 1" - assert kill_group.call_count == 2 - kill_group.assert_called_with(umount_process.pid, signal.SIGKILL) - - @pytest.mark.anyio - async def test_builtin_artifact_is_exempt_from_cache_accounting( - self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch - ): - """The bundled builtin registry is never a cache entry.""" - version = "1.2.3" - site_packages = temp_cache_dir / "venv" / "site-packages" - package_dir = site_packages / "tracecat_registry" - package_dir.mkdir(parents=True) - package_file = package_dir / "__init__.py" - package_file.write_text("") - - monkeypatch.setattr(tracecat_registry, "__version__", version) - monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) - monkeypatch.setattr( - "tracecat.executor.registry_artifact_materialization.sysconfig.get_path", - lambda name: str(site_packages) if name == "purelib" else None, - ) - - cache = RegistryArtifactCache(temp_cache_dir) - - with patch.object( - cache, - "_enforce_cache_budget", - new_callable=AsyncMock, - ) as enforce_cache_budget: - async with cache.lease([bundled_builtin_registry_uri(version)]) as paths: - assert paths == [site_packages.resolve()] - assert cache._runtime == {} - - enforce_cache_budget.assert_not_awaited() diff --git a/tests/unit/executor/test_registry_artifact_materialization.py b/tests/unit/executor/test_registry_artifact_materialization.py deleted file mode 100644 index 33792d1fbb..0000000000 --- a/tests/unit/executor/test_registry_artifact_materialization.py +++ /dev/null @@ -1,1546 +0,0 @@ -"""Artifact selection, download, and materialization tests.""" - -from __future__ import annotations - -import asyncio -import io -import shutil -import signal -import tarfile -import threading -from pathlib import Path -from unittest.mock import AsyncMock, patch - -import pytest - -from tracecat.executor.registry_artifacts import ( - SQUASHFS_MOUNT_OPTIONS, - RegistryArtifactCache, - RegistryArtifactExtractionError, - RegistryArtifactMaterializationContext, - SquashfsArtifact, - SquashfsMountCommandError, - TarballArtifact, - _squashfs_listing_size, - _tarball_extracted_size, - compute_registry_artifact_cache_key, -) - -from .registry_artifact_test_helpers import ( - SQUASHFS_ENABLED_CONFIG, - BlockingSubprocess, - CapturedSubprocess, - SquashfsMountHarness, - lease_paths, - tarball_payload, -) - - -class TestRegistryArtifactMaterialization: - """Materialize and reuse executor-local artifact formats.""" - - def test_temp_path_rejects_symlinked_staging_root( - self, temp_cache_dir: Path - ) -> None: - cache_dir = temp_cache_dir / "cache" - cache_dir.mkdir() - cache = RegistryArtifactCache(cache_dir) - artifact = TarballArtifact( - uri="s3://bucket/path/site-packages.tar.gz", - cache_key="symlinked-staging", - ) - ctx = cache._context_for(artifact.cache_key) - outside_dir = temp_cache_dir / "outside-staging" - outside_dir.mkdir() - ctx.staging_dir.symlink_to(outside_dir, target_is_directory=True) - - with pytest.raises(OSError, match="Unsafe .*registry cache work path"): - artifact._temp_path(ctx, ".tmp") - - assert not any(outside_dir.iterdir()) - - def test_temp_path_avoids_deferred_staging_collision( - self, temp_cache_dir: Path - ) -> None: - cache = RegistryArtifactCache(temp_cache_dir) - artifact = TarballArtifact( - uri="s3://bucket/path/site-packages.tar.gz", - cache_key="staging-collision", - ) - ctx = cache._context_for(artifact.cache_key) - - with patch( - "tracecat.executor.registry_artifact_materialization.secrets.token_hex", - side_effect=["deferred", "deferred", "retry"], - ): - deferred_path = artifact._temp_path(ctx, ".tmp") - deferred_path.mkdir() - retry_path = artifact._temp_path(ctx, ".tmp") - - assert retry_path != deferred_path - assert retry_path.name.endswith(".retry.tmp") - - def test_cached_path_rejects_non_directory_extractions( - self, temp_cache_dir: Path - ) -> None: - cache = RegistryArtifactCache(temp_cache_dir) - ctx = cache._context_for("malformed-extractions") - ctx.paths.entry_dir.mkdir(parents=True) - ctx.paths.squashfs_extract_dir.write_bytes(b"not a directory") - ctx.paths.tarball_target_dir.write_bytes(b"not a directory") - - squashfs = SquashfsArtifact( - uri="s3://bucket/path/site-packages.squashfs", - cache_key=ctx.cache_key, - ) - tarball = TarballArtifact( - uri="s3://bucket/path/site-packages.tar.gz", - cache_key=ctx.cache_key, - ) - - assert squashfs.cached_path(ctx) is None - assert tarball.cached_path(ctx) is None - assert not ctx.paths.squashfs_extract_dir.exists() - assert not ctx.paths.tarball_target_dir.exists() - - @pytest.mark.anyio - async def test_download_reclaims_malformed_squashfs_image( - self, temp_cache_dir: Path - ) -> None: - cache = RegistryArtifactCache(temp_cache_dir) - artifact = SquashfsArtifact( - uri="s3://bucket/path/site-packages.squashfs", - cache_key="malformed-image", - ) - ctx = cache._context_for(artifact.cache_key) - image_path = ctx.paths.squashfs_image_path - image_path.mkdir(parents=True) - (image_path / "stale").write_bytes(b"stale") - - async def download_image( - artifact_uri: str, - output_path: Path, - *, - admission: object, - defer_cleanup: object, - ) -> None: - del artifact_uri, admission, defer_cleanup - output_path.write_bytes(b"fresh image") - - with patch( - "tracecat.executor.registry_artifact_materialization._download_s3_artifact", - download_image, - ): - await artifact.download(ctx, image_path) - - assert image_path.is_file() - assert not image_path.is_symlink() - assert image_path.read_bytes() == b"fresh image" - - @pytest.mark.anyio - async def test_tarball_rejects_symlinked_cache_entry_for_reuse_and_publish( - self, - temp_cache_dir: Path, - ) -> None: - """Tarball paths cannot escape through a symlinked cache entry root.""" - cache = RegistryArtifactCache(temp_cache_dir) - ctx = cache._context_for("symlinked-tarball-entry") - artifact = TarballArtifact( - uri="s3://bucket/path/site-packages.tar.gz", - cache_key=ctx.cache_key, - ) - outside_dir = temp_cache_dir / "outside-tarball-entry" - (outside_dir / "tarball").mkdir(parents=True) - ctx.paths.entry_dir.parent.mkdir(parents=True) - ctx.paths.entry_dir.symlink_to(outside_dir, target_is_directory=True) - - with pytest.raises(OSError, match="Unsafe .*registry cache entry path"): - artifact.cached_path(ctx) - - with ( - patch.object( - TarballArtifact, - "download", - new_callable=AsyncMock, - ) as download, - patch.object( - TarballArtifact, - "extract", - new_callable=AsyncMock, - ) as extract, - pytest.raises(OSError, match="Unsafe .*registry cache entry path"), - ): - await artifact.materialize(ctx) - - download.assert_not_awaited() - extract.assert_not_awaited() - - @pytest.mark.anyio - async def test_tarball_rejects_symlinked_entries_root_before_creation( - self, - temp_cache_dir: Path, - ) -> None: - """An absent cache key cannot be created through a redirected entries root.""" - cache_dir = temp_cache_dir / "cache" - cache_dir.mkdir() - cache = RegistryArtifactCache(cache_dir) - ctx = cache._context_for("symlinked-entries-root") - artifact = TarballArtifact( - uri="s3://bucket/path/site-packages.tar.gz", - cache_key=ctx.cache_key, - ) - outside_dir = temp_cache_dir / "outside-entries" - outside_dir.mkdir() - cache.entries_dir.symlink_to(outside_dir, target_is_directory=True) - - with pytest.raises(OSError, match="Unsafe .*registry cache entries path"): - artifact.cached_path(ctx) - - with ( - patch.object( - TarballArtifact, - "download", - new_callable=AsyncMock, - ) as download, - patch.object( - TarballArtifact, - "extract", - new_callable=AsyncMock, - ) as extract, - pytest.raises(OSError, match="Unsafe .*registry cache entries path"), - ): - await artifact.materialize(ctx) - - download.assert_not_awaited() - extract.assert_not_awaited() - assert not any(outside_dir.iterdir()) - - @pytest.mark.anyio - async def test_same_key_cold_fan_in_materializes_and_enforces_once( - self, temp_cache_dir - ): - """Same-key waiters share candidate lookup, materialization, and enforcement.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - artifact_uri = "s3://bucket/path/site-packages.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - materialization_started = asyncio.Event() - finish_materialization = asyncio.Event() - - async def mock_materialize( - self: TarballArtifact, - ctx: RegistryArtifactMaterializationContext, - ) -> list[Path]: - materialization_started.set() - await finish_materialization.wait() - ctx.paths.tarball_target_dir.mkdir(parents=True) - return [ctx.paths.tarball_target_dir] - - async def take_lease() -> list[Path]: - async with cache.lease([artifact_uri]) as paths: - return paths - - with ( - patch.object( - cache, - "_sidecar_exists", - new_callable=AsyncMock, - return_value=False, - ) as sidecar_exists, - patch.object( - cache, - "_enforce_cache_budget", - new_callable=AsyncMock, - return_value=True, - ) as enforce_cache_budget, - patch.object(cache, "_converge_cache_budget", new_callable=AsyncMock), - patch.object(TarballArtifact, "materialize", mock_materialize), - ): - leases = [asyncio.create_task(take_lease()) for _ in range(5)] - await materialization_started.wait() - await asyncio.sleep(0) - finish_materialization.set() - results = await asyncio.gather(*leases) - - expected_paths = [cache._paths_for(cache_key).tarball_target_dir] - assert results == [expected_paths] * 5 - sidecar_exists.assert_awaited_once() - enforce_cache_budget.assert_awaited_once_with(protected_key=cache_key) - - @pytest.mark.anyio - async def test_different_cold_keys_serialize_materialization( - self, temp_cache_dir: Path - ) -> None: - """Distinct cold keys cannot multiply staging and extraction peaks.""" - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - uris = ["s3://bucket/first.tar.gz", "s3://bucket/second.tar.gz"] - first_started = asyncio.Event() - release_first = asyncio.Event() - second_started = asyncio.Event() - started_keys: list[str] = [] - - async def controlled_materialize( - self: TarballArtifact, - ctx: RegistryArtifactMaterializationContext, - ) -> list[Path]: - del self - started_keys.append(ctx.cache_key) - if len(started_keys) == 1: - first_started.set() - await release_first.wait() - else: - second_started.set() - ctx.paths.tarball_target_dir.mkdir(parents=True) - return [ctx.paths.tarball_target_dir] - - async def take_lease(uri: str) -> None: - async with cache.lease([uri]): - pass - - with ( - patch(SQUASHFS_ENABLED_CONFIG, False), - patch.object(TarballArtifact, "materialize", controlled_materialize), - patch.object(cache, "_enforce_cache_budget", new_callable=AsyncMock), - patch.object(cache, "_converge_cache_budget", new_callable=AsyncMock), - ): - first = asyncio.create_task(take_lease(uris[0])) - await first_started.wait() - second = asyncio.create_task(take_lease(uris[1])) - await asyncio.sleep(0) - assert not second_started.is_set() - - release_first.set() - await second_started.wait() - await asyncio.gather(first, second) - - assert started_keys == [ - compute_registry_artifact_cache_key(uri) for uri in uris - ] - - def test_squashfs_listing_size_bounds_each_inode_allocation(self) -> None: - listing = b"\n".join( - [ - b"drwxr-xr-x 0/0 64 2026-01-01 00:00 squashfs-root", - b"-rw-r--r-- 0/0 123 2026-01-01 00:00 squashfs-root/module.py", - b"lrwxrwxrwx 0/0 9 2026-01-01 00:00 squashfs-root/current -> module.py", - ] - ) - - assert _squashfs_listing_size(listing, allocation_unit=4096) == 20_480 - - def test_squashfs_listing_size_accepts_non_utf8_filenames(self) -> None: - listing = b"-rw-r--r-- 0/0 123 2026-01-01 00:00 squashfs-root/module-\xff.py" - - assert _squashfs_listing_size(listing, allocation_unit=4096) == 12_288 - - def test_squashfs_listing_size_reserves_growing_directory_metadata( - self, - ) -> None: - """Wide extracted directories reserve metadata for every child.""" - child_count = 1000 - listing = b"\n".join( - [ - b"drwxr-xr-x 0/0 64 2026-01-01 00:00 squashfs-root", - *( - f"-rw-r--r-- 0/0 0 2026-01-01 00:00 " - f"squashfs-root/module-{index:04d}.py".encode() - for index in range(child_count) - ), - ] - ) - - inode_bytes = (child_count + 1) * 4096 - directory_entry_bytes = child_count * 4096 - assert ( - _squashfs_listing_size( - listing, - allocation_unit=4096, - ) - == inode_bytes + directory_entry_bytes - ) - - def test_tarball_size_bounds_each_member_allocation( - self, - temp_cache_dir: Path, - ) -> None: - tarball_path = temp_cache_dir / "many-small-files.tar.gz" - with tarfile.open(tarball_path, "w:gz") as tar: - for index in range(3): - member = tarfile.TarInfo(f"module-{index}.py") - member.size = 1 - tar.addfile(member, io.BytesIO(b"x")) - - assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 28_672 - - def test_tarball_size_includes_extraction_root( - self, - temp_cache_dir: Path, - ) -> None: - """Extraction reserves its root even when the manifest omits it.""" - tarball_path = temp_cache_dir / "implicit-root.tar.gz" - with tarfile.open(tarball_path, "w:gz") as tar: - member = tarfile.TarInfo("module.py") - member.size = 0 - tar.addfile(member, io.BytesIO()) - - assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 12_288 - - def test_tarball_size_does_not_duplicate_explicit_root( - self, - temp_cache_dir: Path, - ) -> None: - tarball_path = temp_cache_dir / "explicit-root.tar.gz" - with tarfile.open(tarball_path, "w:gz") as tar: - root = tarfile.TarInfo(".") - root.type = tarfile.DIRTYPE - tar.addfile(root) - - assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 4096 - - def test_tarball_size_includes_implicit_parent_directories( - self, - temp_cache_dir: Path, - ) -> None: - """Extraction reserves directories omitted from the tar manifest.""" - tarball_path = temp_cache_dir / "implicit-directories.tar.gz" - with tarfile.open(tarball_path, "w:gz") as tar: - member = tarfile.TarInfo("one/two/three/module.py") - member.size = 0 - tar.addfile(member, io.BytesIO()) - - assert _tarball_extracted_size(tarball_path, allocation_unit=4096) == 36_864 - - def test_tarball_size_reserves_growing_directory_metadata( - self, - temp_cache_dir: Path, - ) -> None: - """Many child entries reserve more than one directory block.""" - tarball_path = temp_cache_dir / "wide-directory.tar.gz" - child_count = 1000 - with tarfile.open(tarball_path, "w:gz") as tar: - for index in range(child_count): - member = tarfile.TarInfo(f"module-{index:04d}.py") - member.size = 0 - tar.addfile(member, io.BytesIO()) - - file_bytes = child_count * 4096 - root_bytes = 4096 - directory_entry_bytes = child_count * 4096 - assert ( - _tarball_extracted_size( - tarball_path, - allocation_unit=4096, - ) - == file_bytes + root_bytes + directory_entry_bytes - ) - - def test_squashfs_listing_size_rejects_unparseable_files(self) -> None: - with pytest.raises(ValueError, match="Could not parse SquashFS listing"): - _squashfs_listing_size(b"-rw-r--r-- malformed") - - @pytest.mark.parametrize( - "listing", - [ - b"", - b"Parallel unsquashfs: Using 4 processors", - ], - ) - def test_squashfs_listing_size_rejects_listing_without_entries( - self, - listing: bytes, - ) -> None: - with pytest.raises( - ValueError, - match="Could not parse any SquashFS listing entries", - ): - _squashfs_listing_size(listing) - - @pytest.mark.anyio - async def test_materialize_mounts_squashfs_sidecar(self, temp_cache_dir): - """Test that a SquashFS sidecar is mounted instead of extracting tarballs.""" - cache = RegistryArtifactCache(temp_cache_dir) - - async def mock_mount(self, ctx, image_path): - assert image_path.name.endswith(".squashfs") - target_dir = ctx.paths.squashfs_mount_dir - target_dir.mkdir(parents=True, exist_ok=True) - (target_dir / "module.py").write_text("VALUE = 1") - return target_dir - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - return_value=True, - ), - patch( - "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", - True, - ), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", mock_mount), - patch.object( - TarballArtifact, - "materialize", - new_callable=AsyncMock, - ) as tarball_materialize, - ): - result = await lease_paths( - cache, - "s3://bucket/path/site-packages.tar.gz", - ) - - assert len(result) == 1 - assert (result[0] / "module.py").read_text() == "VALUE = 1" - tarball_materialize.assert_not_awaited() - - @pytest.mark.anyio - async def test_mount_squashfs_uses_hardened_read_only_options( - self, - temp_cache_dir, - ): - """Test that SquashFS images are mounted read-only without device/setuid bits.""" - cache_key = "cache-key" - cache = RegistryArtifactCache(temp_cache_dir) - ctx = cache._context_for(cache_key) - artifact = SquashfsArtifact( - uri="s3://bucket/path/site-packages.squashfs", - cache_key=cache_key, - ) - image_path = ctx.paths.squashfs_image_path - target_dir = ctx.paths.squashfs_mount_dir - ctx.paths.entry_dir.mkdir(parents=True) - image_path.write_bytes(b"squashfs") - target_dir.mkdir() - process = AsyncMock() - process.pid = 1234 - process.communicate.return_value = (b"", b"") - process.returncode = 0 - - with ( - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - new_callable=AsyncMock, - return_value=process, - ) as create_subprocess_exec, - patch( - "tracecat.sandbox.utils.terminate_process_group", - new_callable=AsyncMock, - ) as terminate_process_group, - ): - await artifact.mount(ctx, image_path) - - create_subprocess_exec.assert_awaited_once_with( - "mount", - "-t", - "squashfs", - "-o", - SQUASHFS_MOUNT_OPTIONS, - str(image_path), - str(target_dir), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - start_new_session=True, - ) - terminate_process_group.assert_awaited_once_with(process) - - @pytest.mark.parametrize("symlink_component", ["entry", "mount"]) - @pytest.mark.anyio - async def test_mount_squashfs_rejects_symlinked_cache_paths( - self, - temp_cache_dir: Path, - symlink_component: str, - ) -> None: - """A privileged mount cannot follow a cache path outside its entry.""" - cache = RegistryArtifactCache(temp_cache_dir) - ctx = cache._context_for("symlinked-mount") - artifact = SquashfsArtifact( - uri="s3://bucket/path/site-packages.squashfs", - cache_key=ctx.cache_key, - ) - outside_dir = temp_cache_dir / "outside" - outside_dir.mkdir(parents=True) - - if symlink_component == "entry": - ctx.paths.entry_dir.parent.mkdir(parents=True) - ctx.paths.entry_dir.symlink_to(outside_dir, target_is_directory=True) - else: - ctx.paths.entry_dir.mkdir(parents=True) - ctx.paths.squashfs_mount_dir.symlink_to( - outside_dir, - target_is_directory=True, - ) - - with ( - patch.object( - SquashfsArtifact, - "download", - new_callable=AsyncMock, - ) as download, - patch.object( - SquashfsArtifact, - "_mount_image", - new_callable=AsyncMock, - ) as mount_image, - pytest.raises(OSError, match="Unsafe .* path"), - ): - await artifact.mount(ctx, ctx.paths.squashfs_image_path) - - download.assert_not_awaited() - mount_image.assert_not_awaited() - - @pytest.mark.anyio - async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): - """Cancellation cannot leave an orphan mount process after lock release.""" - cache_key = "cancelled-mount" - cache = RegistryArtifactCache(temp_cache_dir) - ctx = cache._context_for(cache_key) - artifact = SquashfsArtifact( - uri="s3://bucket/path/site-packages.squashfs", - cache_key=cache_key, - ) - image_path = ctx.paths.squashfs_image_path - target_dir = ctx.paths.squashfs_mount_dir - ctx.paths.entry_dir.mkdir(parents=True) - image_path.write_bytes(b"squashfs") - target_dir.mkdir() - process = BlockingSubprocess() - - with ( - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - new_callable=AsyncMock, - return_value=process, - ), - patch("tracecat.sandbox.utils.os.killpg") as kill_group, - ): - mounting = asyncio.create_task( - artifact._mount_image(image_path, target_dir) - ) - await process.communicate_started.wait() - mounting.cancel() - - with pytest.raises(asyncio.CancelledError): - await mounting - - assert process.cleanup_calls == ["kill", "wait"] - kill_group.assert_called_once_with(process.pid, signal.SIGKILL) - assert target_dir.is_dir() - assert not target_dir.is_mount() - - @pytest.mark.anyio - async def test_cancelled_squashfs_extract_kills_and_reaps_subprocess( - self, temp_cache_dir - ): - """Cancellation cannot leave unsquashfs writing into scratch.""" - cache_key = "cancelled-extract" - cache = RegistryArtifactCache(temp_cache_dir) - ctx = cache._context_for(cache_key) - artifact = SquashfsArtifact( - uri="s3://bucket/path/site-packages.squashfs", - cache_key=cache_key, - ) - image_path = ctx.paths.squashfs_image_path - target_dir = ctx.paths.squashfs_extract_dir - ctx.paths.entry_dir.mkdir(parents=True) - image_path.write_bytes(b"squashfs") - target_dir.mkdir() - real_create_subprocess_exec = asyncio.create_subprocess_exec - process_started = asyncio.Event() - captured_processes: list[CapturedSubprocess] = [] - - async def create_sleep_subprocess( - *args: object, **kwargs: object - ) -> CapturedSubprocess: - del args, kwargs - process = await real_create_subprocess_exec( - "/bin/sleep", - "30", - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - start_new_session=True, - ) - captured = CapturedSubprocess(process) - captured_processes.append(captured) - process_started.set() - return captured - - with ( - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/usr/bin/unsquashfs", - ), - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - side_effect=create_sleep_subprocess, - ), - ): - extracting = asyncio.create_task( - artifact._extract_image(image_path, target_dir) - ) - await process_started.wait() - captured = captured_processes[0] - extracting.cancel() - - production_killed = False - production_reaped = False - try: - with pytest.raises(asyncio.CancelledError): - await extracting - production_killed = captured.killed - production_reaped = captured.reaped - finally: - if captured.returncode is None: - captured.process.kill() - await captured.process.wait() - - assert production_killed is True - assert production_reaped is True - assert captured.returncode is not None - - @pytest.mark.parametrize("operation", ["extract", "size-command", "size-parse"]) - @pytest.mark.anyio - async def test_squashfs_failures_sanitize_subprocess_output( - self, - temp_cache_dir: Path, - operation: str, - ) -> None: - sensitive_output = b"synthetic-customer/repository/member.py" - artifact = SquashfsArtifact( - uri="s3://bucket/path/custom.squashfs", - cache_key="malformed-squashfs", - ) - image_path = temp_cache_dir / "image.squashfs" - image_path.write_bytes(b"squashfs") - target_dir = temp_cache_dir / "target" - target_dir.mkdir() - process = AsyncMock() - process.returncode = 0 if operation == "size-parse" else 1 - stdout = sensitive_output if operation == "size-parse" else b"" - stderr = b"" if operation == "size-parse" else sensitive_output - - with ( - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/usr/bin/unsquashfs", - ), - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - new_callable=AsyncMock, - return_value=process, - ), - patch( - "tracecat.executor.registry_artifact_materialization.communicate_process_group", - new_callable=AsyncMock, - return_value=(stdout, stderr), - ), - pytest.raises(RegistryArtifactExtractionError) as raised, - ): - if operation == "extract": - await artifact._extract_image(image_path, target_dir) - else: - await artifact._squashfs_extracted_size(image_path) - - assert str(raised.value) == "Registry artifact extraction failed" - assert sensitive_output.decode() not in str(raised.value) - assert raised.value.__cause__ is None - - @pytest.mark.parametrize("operation", ["mount", "extract", "size"]) - @pytest.mark.anyio - async def test_repeated_cancellation_reaps_squashfs_subprocess( - self, - temp_cache_dir: Path, - operation: str, - ) -> None: - """A second cancellation cannot abandon a killed SquashFS child.""" - artifact = SquashfsArtifact( - uri="s3://bucket/path/site-packages.squashfs", - cache_key="repeated-subprocess-cancellation", - ) - image_path = temp_cache_dir / "image.squashfs" - image_path.write_bytes(b"squashfs") - target_dir = temp_cache_dir / "target" - target_dir.mkdir() - process = BlockingSubprocess(block_wait=True) - - with ( - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/usr/bin/unsquashfs", - ), - patch( - "tracecat.executor.registry_artifact_materialization.asyncio.create_subprocess_exec", - new_callable=AsyncMock, - return_value=process, - ) as create_subprocess_exec, - patch("tracecat.sandbox.utils.os.killpg") as kill_group, - ): - if operation == "mount": - running = asyncio.create_task( - artifact._mount_image(image_path, target_dir) - ) - elif operation == "extract": - running = asyncio.create_task( - artifact._extract_image(image_path, target_dir) - ) - else: - running = asyncio.create_task( - artifact._squashfs_extracted_size(image_path) - ) - - await process.communicate_started.wait() - running.cancel() - await process.wait_started.wait() - - running.cancel() - done, _ = await asyncio.wait({running}, timeout=0.05) - second_cancellation_propagated_early = bool(done) - process.release_wait.set() - - with pytest.raises(asyncio.CancelledError): - await running - - assert second_cancellation_propagated_early is False - assert process.cleanup_calls == ["kill", "wait"] - await_args = create_subprocess_exec.await_args - assert await_args is not None - assert await_args.kwargs["start_new_session"] is True - kill_group.assert_called_once_with(process.pid, signal.SIGKILL) - - @pytest.mark.anyio - async def test_repeatedly_cancelled_tarball_extract_rejoins_thread( - self, temp_cache_dir - ): - """Repeated cancellation waits until the tar extractor stops writing.""" - artifact = TarballArtifact( - uri="s3://bucket/path/site-packages.tar.gz", - cache_key="cancelled-tarball-extract", - ) - tarball_path = temp_cache_dir / "artifact.tar.gz" - target_dir = temp_cache_dir / "target" - target_dir.mkdir() - with tarfile.open(tarball_path, "w:gz"): - pass - - extraction_started = threading.Event() - extraction_release = threading.Event() - extraction_finished = threading.Event() - - def blocking_extractall(*args: object, **kwargs: object) -> None: - del args, kwargs - extraction_started.set() - extraction_release.wait() - extraction_finished.set() - - with patch.object( - tarfile.TarFile, - "extractall", - new=blocking_extractall, - ): - extracting = asyncio.create_task(artifact.extract(tarball_path, target_dir)) - assert await asyncio.to_thread(extraction_started.wait, 1) - extracting.cancel() - done, _ = await asyncio.wait({extracting}, timeout=0.05) - first_cancellation_propagated_early = bool(done) - - extracting.cancel() - done, _ = await asyncio.wait({extracting}, timeout=0.05) - second_cancellation_propagated_early = bool(done) - extraction_release.set() - - with pytest.raises(asyncio.CancelledError): - await extracting - - assert extraction_finished.is_set() - assert first_cancellation_propagated_early is False - assert second_cancellation_propagated_early is False - - @pytest.mark.anyio - async def test_tarball_extract_sanitizes_member_failures( - self, - temp_cache_dir: Path, - ) -> None: - sensitive_member = "../../synthetic-customer/repository/module.py" - artifact = TarballArtifact( - uri="s3://bucket/path/site-packages.tar.gz", - cache_key="malformed-tarball", - ) - tarball_path = temp_cache_dir / "artifact.tar.gz" - target_dir = temp_cache_dir / "target" - target_dir.mkdir() - with tarfile.open(tarball_path, "w:gz") as tar: - member = tarfile.TarInfo(sensitive_member) - member.size = 1 - tar.addfile(member, io.BytesIO(b"x")) - - with pytest.raises(RegistryArtifactExtractionError) as raised: - await artifact.extract(tarball_path, target_dir) - - assert str(raised.value) == "Registry artifact extraction failed" - assert sensitive_member not in str(raised.value) - assert raised.value.__cause__ is None - - @pytest.mark.anyio - async def test_repeatedly_cancelled_tarball_size_scan_rejoins_thread( - self, temp_cache_dir - ): - """Cancellation cannot unlink a tarball while its size scan is running.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/path/slow-size-scan.tar.gz" - downloaded_paths: list[Path] = [] - scan_started = threading.Event() - scan_release = threading.Event() - scan_finished = threading.Event() - input_present_at_finish: list[bool] = [] - - async def mock_download(self, ctx, path): - del self, ctx - path.write_bytes(tarball_payload(size=1)) - downloaded_paths.append(path) - - def blocking_size_scan(path: Path, *, allocation_unit: int) -> int: - assert allocation_unit == 1 - scan_started.set() - scan_release.wait() - input_present_at_finish.append(path.exists()) - scan_finished.set() - return 1 - - with ( - patch(SQUASHFS_ENABLED_CONFIG, False), - patch.object(TarballArtifact, "download", mock_download), - patch( - "tracecat.executor.registry_artifact_materialization._tarball_extracted_size", - side_effect=blocking_size_scan, - ), - patch.object( - TarballArtifact, - "extract", - new_callable=AsyncMock, - ) as extract, - ): - materializing = asyncio.create_task(lease_paths(cache, artifact_uri)) - assert await asyncio.to_thread(scan_started.wait, 1) - materializing.cancel() - done, _ = await asyncio.wait({materializing}, timeout=0.05) - first_cancellation_propagated_early = bool(done) - - materializing.cancel() - done, _ = await asyncio.wait({materializing}, timeout=0.05) - second_cancellation_propagated_early = bool(done) - assert downloaded_paths[0].exists() - scan_release.set() - - with pytest.raises(asyncio.CancelledError): - await materializing - - assert scan_finished.is_set() - assert input_present_at_finish == [True] - assert not downloaded_paths[0].exists() - assert first_cancellation_propagated_early is False - assert second_cancellation_propagated_early is False - extract.assert_not_awaited() - - @pytest.mark.anyio - async def test_failed_partial_cleanup_is_deferred_for_capacity_retry( - self, temp_cache_dir - ): - """A failed staging-tree deletion remains discoverable and retryable.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/path/failed-cleanup.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - artifact = TarballArtifact(uri=artifact_uri, cache_key=cache_key) - ctx = cache._context_for(cache_key) - cleanup_attempts: list[Path] = [] - - async def download(self, ctx, path): - del self, ctx - path.write_bytes(b"archive") - - async def fail_extract(self, tarball_path, target_dir): - del self, tarball_path - (target_dir / "partial.py").write_text("partial") - raise RuntimeError("extraction failed") - - def fail_cleanup(path: Path) -> None: - cleanup_attempts.append(path) - raise PermissionError("cleanup denied") - - with ( - patch.object(TarballArtifact, "download", download), - patch.object(TarballArtifact, "extract", fail_extract), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.rmtree", - side_effect=fail_cleanup, - ), - ): - with pytest.raises(RuntimeError, match="extraction failed"): - await artifact.materialize(ctx) - - assert len(cleanup_attempts) == 1 - assert set(cleanup_attempts) == cache._deferred_staging_cleanup - deferred_path = cleanup_attempts[0] - assert deferred_path.is_dir() - - assert cache._retry_deferred_staging_cleanup() is True - assert cache._deferred_staging_cleanup == set() - assert not deferred_path.exists() - - @pytest.mark.anyio - async def test_failed_tarball_unlink_is_deferred_without_masking_success( - self, temp_cache_dir - ): - """A failed tarball unlink preserves success and remains retryable.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/path/failed-tarball-cleanup.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - artifact = TarballArtifact(uri=artifact_uri, cache_key=cache_key) - ctx = cache._context_for(cache_key) - downloaded_paths: list[Path] = [] - real_unlink = Path.unlink - - async def download(self, ctx, path): - del self, ctx - path.write_bytes(b"archive") - downloaded_paths.append(path) - - async def extract(self, tarball_path, target_dir): - del self, tarball_path - (target_dir / "module.py").write_text("VALUE = 1") - - def fail_download_unlink( - path: Path, - missing_ok: bool = False, - ) -> None: - if path in downloaded_paths: - raise PermissionError("cleanup denied") - real_unlink(path, missing_ok=missing_ok) - - with ( - patch.object(TarballArtifact, "download", download), - patch.object(TarballArtifact, "extract", extract), - patch.object(Path, "unlink", fail_download_unlink), - ): - result = await artifact.materialize(ctx) - - assert result == [ctx.paths.tarball_target_dir] - assert (result[0] / "module.py").read_text() == "VALUE = 1" - assert len(downloaded_paths) == 1 - deferred_path = downloaded_paths[0] - assert cache._deferred_staging_cleanup == {deferred_path} - assert deferred_path.exists() - - assert cache._retry_deferred_staging_cleanup() is True - assert cache._deferred_staging_cleanup == set() - assert not deferred_path.exists() - - @pytest.mark.anyio - async def test_failed_squashfs_unlink_is_deferred_after_concurrent_publish( - self, - temp_cache_dir: Path, - ) -> None: - """A losing SquashFS staging file remains retryable without masking success.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/path/concurrent.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) - ctx = cache._context_for(cache_key) - image_path = ctx.paths.squashfs_image_path - staging_paths: list[Path] = [] - real_rename = Path.rename - real_unlink = Path.unlink - - async def download( - artifact_uri: str, - output_path: Path, - *, - admission: object, - defer_cleanup: object, - ) -> None: - del artifact_uri, admission, defer_cleanup - output_path.write_bytes(b"loser") - staging_paths.append(output_path) - - def publish_concurrently(path: Path, target: Path) -> Path: - if path in staging_paths: - target.write_bytes(b"winner") - raise FileExistsError("published by another process") - return real_rename(path, target) - - def fail_staging_unlink(path: Path, missing_ok: bool = False) -> None: - if path in staging_paths: - raise PermissionError("cleanup denied") - real_unlink(path, missing_ok=missing_ok) - - with ( - patch( - "tracecat.executor.registry_artifact_materialization." - "_download_s3_artifact", - side_effect=download, - ), - patch.object(Path, "rename", publish_concurrently), - patch.object(Path, "unlink", fail_staging_unlink), - ): - await artifact.download(ctx, image_path) - - assert image_path.read_bytes() == b"winner" - assert len(staging_paths) == 1 - deferred_path = staging_paths[0] - assert deferred_path.exists() - assert cache._deferred_staging_cleanup == {deferred_path} - - assert cache._retry_deferred_staging_cleanup() is True - assert cache._deferred_staging_cleanup == set() - assert not deferred_path.exists() - - def test_failed_unusable_squashfs_unlink_is_deferred( - self, - temp_cache_dir: Path, - ) -> None: - """A failed canonical-image cleanup remains retryable after fallback.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/path/unusable.squashfs" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) - ctx = cache._context_for(cache_key) - image_path = ctx.paths.squashfs_image_path - image_path.parent.mkdir(parents=True) - image_path.write_bytes(b"unusable") - real_unlink = Path.unlink - - def fail_image_unlink(path: Path, missing_ok: bool = False) -> None: - if path == image_path: - raise PermissionError("cleanup denied") - real_unlink(path, missing_ok=missing_ok) - - with patch.object(Path, "unlink", fail_image_unlink): - artifact.discard_failed_materialization(ctx) - - assert cache._deferred_staging_cleanup == {image_path} - assert image_path.is_file() - assert cache._retry_deferred_staging_cleanup() is True - assert cache._deferred_staging_cleanup == set() - assert not image_path.exists() - - @pytest.mark.parametrize("artifact_format", ["squashfs", "tarball"]) - @pytest.mark.anyio - async def test_repeatedly_cancelled_partial_cleanup_rejoins_thread( - self, - temp_cache_dir, - artifact_format: str, - ): - """Large partial environments are deleted off-loop before cancellation.""" - cache = RegistryArtifactCache(temp_cache_dir) - cleanup_started = threading.Event() - cleanup_release = threading.Event() - cleanup_finished = threading.Event() - real_rmtree = shutil.rmtree - - def blocking_rmtree(path: Path, *, ignore_errors: bool = False) -> None: - cleanup_started.set() - cleanup_release.wait() - real_rmtree(path, ignore_errors=ignore_errors) - cleanup_finished.set() - - async def assert_cleanup(task: asyncio.Task[list[Path]]) -> None: - assert await asyncio.to_thread(cleanup_started.wait, 1) - task.cancel() - done, _ = await asyncio.wait({task}, timeout=0.05) - first_cancellation_propagated_early = bool(done) - - task.cancel() - done, _ = await asyncio.wait({task}, timeout=0.05) - second_cancellation_propagated_early = bool(done) - cleanup_release.set() - - with pytest.raises(asyncio.CancelledError): - await task - - assert cleanup_finished.is_set() - assert first_cancellation_propagated_early is False - assert second_cancellation_propagated_early is False - assert not cache.staging_dir.exists() or not any( - cache.staging_dir.iterdir() - ) - - async def fail_tarball_extract(self, tarball_path, target_dir): - del self, tarball_path - target_dir.mkdir(parents=True) - (target_dir / "partial.py").write_text("partial") - raise RuntimeError("tarball extraction failed") - - async def download_tarball(self, ctx, path): - del self, ctx - path.write_bytes(tarball_payload(size=1)) - - async def download_squashfs(self, ctx, image_path): - del self, ctx - image_path.parent.mkdir(parents=True, exist_ok=True) - image_path.write_bytes(b"image") - return 0.0 - - async def fail_squashfs_extract(self, image_path, target_dir): - del self, image_path - target_dir.mkdir(parents=True) - (target_dir / "partial.py").write_text("partial") - raise RuntimeError("SquashFS extraction failed") - - with patch( - "tracecat.executor.registry_artifact_materialization.shutil.rmtree", - side_effect=blocking_rmtree, - ): - if artifact_format == "tarball": - with ( - patch(SQUASHFS_ENABLED_CONFIG, False), - patch.object(TarballArtifact, "download", download_tarball), - patch.object(TarballArtifact, "extract", fail_tarball_extract), - ): - task = asyncio.create_task( - lease_paths(cache, "s3://bucket/partial.tar.gz") - ) - await assert_cleanup(task) - else: - with ( - patch.object( - RegistryArtifactMaterializationContext, - "can_mount_squashfs", - return_value=False, - ), - patch.object(SquashfsArtifact, "download", download_squashfs), - patch.object( - SquashfsArtifact, - "_squashfs_extracted_size", - new_callable=AsyncMock, - return_value=1, - ), - patch.object( - SquashfsArtifact, - "_extract_image", - fail_squashfs_extract, - ), - ): - task = asyncio.create_task( - lease_paths(cache, "s3://bucket/partial.squashfs") - ) - await assert_cleanup(task) - - @pytest.mark.anyio - async def test_materialize_extracts_squashfs_when_mount_fails(self, temp_cache_dir): - """Test that SquashFS mount failures fall back to unsquashfs extraction.""" - cache = RegistryArtifactCache(temp_cache_dir) - - async def mock_mount(self, ctx, image_path): - raise SquashfsMountCommandError("operation not permitted") - - async def mock_extract(self, ctx, image_path): - target_dir = ctx.paths.squashfs_extract_dir - target_dir.mkdir(parents=True, exist_ok=True) - (target_dir / "module.py").write_text("VALUE = 1") - return target_dir - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - return_value=True, - ), - patch( - "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", - True, - ), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", mock_mount), - patch.object(SquashfsArtifact, "extract", mock_extract), - patch.object( - TarballArtifact, - "materialize", - new_callable=AsyncMock, - ) as tarball_materialize, - ): - result = await lease_paths( - cache, - "s3://bucket/path/site-packages.tar.gz", - ) - - assert len(result) == 1 - assert (result[0] / "module.py").read_text() == "VALUE = 1" - assert result[0].name == "extracted" - tarball_materialize.assert_not_awaited() - - @pytest.mark.anyio - async def test_materialize_extracts_squashfs_without_mount_binary( - self, temp_cache_dir - ): - """Test that SquashFS is still preferred when only unsquashfs is available.""" - cache = RegistryArtifactCache(temp_cache_dir) - - async def mock_extract(self, ctx, image_path): - target_dir = ctx.paths.squashfs_extract_dir - target_dir.mkdir(parents=True, exist_ok=True) - (target_dir / "module.py").write_text("VALUE = 1") - return target_dir - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - return_value=True, - ), - patch( - "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", - True, - ), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value=None, - ), - patch.object(SquashfsArtifact, "extract", mock_extract), - ): - result = await lease_paths( - cache, - "s3://bucket/path/site-packages.tar.gz", - ) - - assert len(result) == 1 - assert (result[0] / "module.py").read_text() == "VALUE = 1" - assert result[0].name == "extracted" - - @pytest.mark.anyio - async def test_materialize_falls_back_to_gzip_when_squashfs_extract_fails( - self, temp_cache_dir - ): - """Test that legacy gzip remains the final compatibility fallback.""" - cache = RegistryArtifactCache(temp_cache_dir) - source = temp_cache_dir / "source" - source.mkdir() - (source / "module.py").write_text("VALUE = 1") - - async def mock_tarball_download(self, ctx, path): - with tarfile.open(path, "w:gz") as tar: - tar.add(source / "module.py", arcname="module.py") - - async def mock_mount(self, ctx, image_path): - raise SquashfsMountCommandError("operation not permitted") - - async def mock_extract(self, ctx, image_path): - raise RuntimeError("unsquashfs unavailable") - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - side_effect=[True, False], - ), - patch( - "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", - True, - ), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", mock_mount), - patch.object(SquashfsArtifact, "extract", mock_extract), - patch.object(TarballArtifact, "download", mock_tarball_download), - ): - result = await lease_paths( - cache, - "s3://bucket/path/site-packages.tar.gz", - ) - - assert len(result) == 1 - assert (result[0] / "module.py").read_text() == "VALUE = 1" - assert result[0].name == "tarball" - - @pytest.mark.anyio - async def test_materialize_treats_unknown_suffix_as_gzip(self, temp_cache_dir): - """Test that existing gzip artifacts can use arbitrary S3 key suffixes.""" - cache = RegistryArtifactCache(temp_cache_dir) - source = temp_cache_dir / "source" - source.mkdir() - (source / "module.py").write_text("VALUE = 1") - - async def mock_download(self, ctx, path): - assert path.name.endswith(".tar.gz") - with tarfile.open(path, "w:gz") as tar: - tar.add(source / "module.py", arcname="module.py") - - with patch.object(TarballArtifact, "download", mock_download): - result = await lease_paths( - cache, - "s3://bucket/path/custom-key", - ) - - assert len(result) == 1 - assert (result[0] / "module.py").read_text() == "VALUE = 1" - - @pytest.mark.anyio - async def test_materialize_caches_result(self, temp_cache_dir): - """Test that tarball extraction is cached.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/test.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - target_dir = cache._paths_for(cache_key).tarball_target_dir - target_dir.mkdir(parents=True) - - result = await lease_paths(cache, artifact_uri) - - assert result == [target_dir] - - @pytest.mark.anyio - async def test_materialize_concurrent_requests(self, temp_cache_dir): - """Test that concurrent requests for same artifact do not race.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/test.tar.gz" - download_count = 0 - - async def mock_download(self, ctx, path): - nonlocal download_count - download_count += 1 - await asyncio.sleep(0.1) - path.write_bytes(tarball_payload(size=1)) - - async def mock_extract(self, tarball_path, target_dir): - (target_dir / "extracted.txt").write_text("extracted") - - with ( - patch.object(TarballArtifact, "download", mock_download), - patch.object(TarballArtifact, "extract", mock_extract), - ): - results = await asyncio.gather( - lease_paths(cache, artifact_uri), - lease_paths(cache, artifact_uri), - lease_paths(cache, artifact_uri), - ) - - assert all(r == results[0] for r in results) - assert download_count == 1 - - -class TestSquashfsMountPolicy: - """Tests for per-artifact SquashFS fallback.""" - - @pytest.mark.anyio - async def test_loop_device_exhaustion_isolated_sticky_extraction_fallback( - self, temp_cache_dir: Path - ) -> None: - """Protect the cache's fail-open policy when loop devices are saturated. - - A mount-command failure must extract only the affected cold artifact, - preserve already-leased mounts, reuse that extraction on later leases, - and still let unrelated artifacts attempt mounting. This models loop - exhaustion deterministically without consuming host-global devices. - """ - cache = RegistryArtifactCache(temp_cache_dir) - held_uri = "s3://bucket/already-mounted.squashfs" - saturated_uri = "s3://bucket/no-loop-available.squashfs" - later_uri = "s3://bucket/later-artifact.squashfs" - held_key = compute_registry_artifact_cache_key(held_uri) - saturated_key = compute_registry_artifact_cache_key(saturated_uri) - later_key = compute_registry_artifact_cache_key(later_uri) - harness = SquashfsMountHarness( - cache, - failed_mount_keys={saturated_key}, - ) - - with ( - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path in harness.mounted, - ), - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", harness.mount), - patch.object(SquashfsArtifact, "extract", harness.extract), - patch.object(cache, "_unmount", harness.unmount), - ): - async with cache.lease([held_uri]) as held_paths: - held_mount = cache._paths_for(held_key).squashfs_mount_dir - assert held_paths == [held_mount] - - async with cache.lease([saturated_uri]) as saturated_paths: - extracted = cache._paths_for(saturated_key).squashfs_extract_dir - assert saturated_paths == [extracted] - assert held_mount in harness.mounted - assert cache._refcount(held_key) == 1 - assert harness.unmounts == [] - - # The extracted directory is the cache entry's sticky path; - # freeing a loop elsewhere does not trigger a remount attempt. - async with cache.lease([saturated_uri]) as cached_paths: - assert cached_paths == [extracted] - - async with cache.lease([later_uri]) as later_paths: - later_mount = cache._paths_for(later_key).squashfs_mount_dir - assert later_paths == [later_mount] - assert held_mount in harness.mounted - - assert held_mount in harness.mounted - - assert harness.mount_attempts == [held_key, saturated_key, later_key] - assert harness.extraction_attempts == [saturated_key] - assert harness.unmounts == [later_mount, held_mount] - assert cache._paths_for(saturated_key).squashfs_image_path.is_file() - assert cache._paths_for(saturated_key).squashfs_extract_dir.is_dir() - assert cache._refcount(held_key) == 0 - assert cache._refcount(saturated_key) == 0 - assert cache._refcount(later_key) == 0 - - @pytest.mark.anyio - async def test_mount_failure_does_not_disable_later_artifacts( - self, temp_cache_dir - ) -> None: - """One failed mount falls back without changing later mount attempts.""" - cache = RegistryArtifactCache(temp_cache_dir) - first_ctx = cache._context_for("first") - second_ctx = cache._context_for("second") - first = SquashfsArtifact( - uri="s3://bucket/first.squashfs", - cache_key="first", - ) - second = SquashfsArtifact( - uri="s3://bucket/second.squashfs", - cache_key="second", - ) - mount_attempts: list[str] = [] - - async def mock_mount(self, ctx, image_path): - del image_path - mount_attempts.append(ctx.cache_key) - if ctx.cache_key == "first": - raise SquashfsMountCommandError("operation not permitted") - ctx.paths.squashfs_mount_dir.mkdir(parents=True) - return ctx.paths.squashfs_mount_dir - - async def mock_extract(self, ctx, image_path): - del image_path - ctx.paths.squashfs_extract_dir.mkdir(parents=True) - return ctx.paths.squashfs_extract_dir - - with ( - patch(SQUASHFS_ENABLED_CONFIG, True), - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value="/sbin/mount", - ), - patch.object(SquashfsArtifact, "mount", mock_mount), - patch.object(SquashfsArtifact, "extract", mock_extract), - ): - assert await first.materialize(first_ctx) == [ - first_ctx.paths.squashfs_extract_dir - ] - assert await second.materialize(second_ctx) == [ - second_ctx.paths.squashfs_mount_dir - ] - - assert mount_attempts == ["first", "second"] diff --git a/tests/unit/executor/test_registry_artifact_resolution.py b/tests/unit/executor/test_registry_artifact_resolution.py deleted file mode 100644 index bab757c885..0000000000 --- a/tests/unit/executor/test_registry_artifact_resolution.py +++ /dev/null @@ -1,546 +0,0 @@ -"""Registry artifact URI, key, and candidate resolution tests.""" - -from __future__ import annotations - -from collections.abc import Callable -from pathlib import Path -from unittest.mock import AsyncMock, patch - -import httpx -import pytest -import tracecat_registry - -from tracecat.executor.registry_artifacts import ( - RegistryArtifactCache, - RegistryArtifactFormat, - SquashfsArtifact, - TarballArtifact, - _artifact_uri_for_logging, - bundled_builtin_registry_uri, - compute_registry_artifact_cache_key, -) -from tracecat.registry.artifact_keys import parse_s3_uri - - -class TestParseS3Uri: - """Tests for parse_s3_uri function.""" - - def test_valid_uri(self): - """Test parsing a valid S3 URI.""" - bucket, key = parse_s3_uri("s3://my-bucket/path/to/file.tar.gz") - assert bucket == "my-bucket" - assert key == "path/to/file.tar.gz" - - def test_uri_with_nested_path(self): - """Test parsing URI with deeply nested path.""" - bucket, key = parse_s3_uri("s3://bucket/a/b/c/d/e/file.tar.gz") - assert bucket == "bucket" - assert key == "a/b/c/d/e/file.tar.gz" - - def test_invalid_uri_no_prefix(self): - """Test that non-S3 URIs raise ValueError.""" - with pytest.raises(ValueError, match="Invalid S3 URI"): - parse_s3_uri("https://bucket/key") - - def test_invalid_uri_no_key(self): - """Test that URIs without keys raise ValueError.""" - with pytest.raises(ValueError, match="Invalid S3 URI"): - parse_s3_uri("s3://bucket") - - def test_invalid_uri_empty_bucket(self): - """Test that URIs with empty bucket raise ValueError.""" - with pytest.raises(ValueError, match="Invalid S3 URI"): - parse_s3_uri("s3:///key") - - -class TestRegistryArtifactResolution: - """Resolve artifact identities and preferred formats.""" - - def test_compute_registry_artifact_cache_key_deterministic(self): - """Test that cache key computation is deterministic.""" - uri = "s3://bucket/path/to/registry-v1.2.3.tar.gz" - - key1 = compute_registry_artifact_cache_key(uri) - key2 = compute_registry_artifact_cache_key(uri) - - assert key1 == key2 - assert len(key1) == 16 - - def test_compute_registry_artifact_cache_key_case_sensitive(self): - """Test that cache key is case-sensitive because S3 keys are case-sensitive.""" - key1 = compute_registry_artifact_cache_key("s3://BUCKET/PATH/FILE.tar.gz") - key2 = compute_registry_artifact_cache_key("s3://bucket/path/file.tar.gz") - - assert key1 != key2 - - def test_compute_registry_artifact_cache_key_whitespace_sensitive(self): - """Cache identity preserves whitespace that can belong to an S3 key.""" - key = compute_registry_artifact_cache_key("s3://bucket/path/file.tar.gz") - whitespace_key = compute_registry_artifact_cache_key( - "s3://bucket/path/file.tar.gz " - ) - - assert key != whitespace_key - - def test_compute_registry_artifact_cache_key_empty(self): - """Test that empty URI returns the base cache key.""" - assert compute_registry_artifact_cache_key("") == "base" - - def test_artifact_uri_for_logging_removes_identifiers_and_credentials(self): - uri = ( - "s3://access:secret@bucket/org-id/repository-origin/1.2.3/" - "site-packages.squashfs" - "?X-Amz-Signature=signed-secret#fragment" - ) - - logged_uri = _artifact_uri_for_logging(uri) - - assert logged_uri == "s3://" - assert all( - value not in logged_uri - for value in ("secret", "bucket", "org-id", "repository-origin", "1.2.3") - ) - - def test_artifact_uri_for_logging_redacts_malformed_uri(self): - assert _artifact_uri_for_logging("s3://[malformed") == ( - "" - ) - - @pytest.mark.anyio - async def test_download_artifact_uses_blob_download_file_to_path( - self, temp_cache_dir - ): - """Test that artifact downloads stay behind the blob storage helper.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact = SquashfsArtifact( - uri="s3://bucket/path/site-packages.squashfs", - cache_key="download-test", - ) - ctx = cache._context_for(artifact.cache_key) - output_path = temp_cache_dir / "artifact.squashfs" - - async def mock_download_file_to_path( - *, - key: str, - bucket: str, - output_path: Path, - defer_cleanup: Callable[[Path], None], - redact_log_identifiers: bool, - ) -> None: - assert defer_cleanup == ctx.defer_cleanup - assert redact_log_identifiers is True - output_path.write_bytes(b"squashfs") - - with patch( - "tracecat.executor.registry_artifacts.blob.download_file_to_path", - new_callable=AsyncMock, - side_effect=mock_download_file_to_path, - ) as download_file_to_path: - await artifact.download(ctx, output_path) - - download_file_to_path.assert_awaited_once() - await_args = download_file_to_path.await_args - assert await_args is not None - assert await_args.kwargs["key"] == "path/site-packages.squashfs" - assert await_args.kwargs["bucket"] == "bucket" - assert output_path.read_bytes() == b"squashfs" - - @pytest.mark.anyio - async def test_lease_uses_bundled_current_builtin( - self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch - ): - """In-tree builtin registry returns only the installed site-packages.""" - version = "1.2.3" - site_packages = temp_cache_dir / "venv" / "site-packages" - package_dir = site_packages / "tracecat_registry" - package_dir.mkdir(parents=True) - package_file = package_dir / "__init__.py" - package_file.write_text("") - - monkeypatch.setattr(tracecat_registry, "__version__", version) - monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) - monkeypatch.setattr( - "tracecat.executor.registry_artifact_materialization.sysconfig.get_path", - lambda name: str(site_packages) if name == "purelib" else None, - ) - - cache = RegistryArtifactCache(temp_cache_dir) - async with cache.lease([bundled_builtin_registry_uri(version)]) as result: - assert result == [site_packages.resolve()] - - @pytest.mark.anyio - async def test_lease_exposes_editable_builtin_parent( - self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch - ): - """Editable builtin registry exposes the package wrapper + site-packages.""" - version = "1.2.3" - site_packages = temp_cache_dir / "venv" / "site-packages" - dependency_dir = site_packages / "orjson" - dependency_dir.mkdir(parents=True) - (dependency_dir / "__init__.py").write_text("VALUE = 1\n") - source_root = temp_cache_dir / "src" / "tracecat-registry" - package_dir = source_root / "tracecat_registry" - package_dir.mkdir(parents=True) - package_file = package_dir / "__init__.py" - package_file.write_text("__version__ = '1.2.3'\n") - - monkeypatch.setattr(tracecat_registry, "__version__", version) - monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) - monkeypatch.setattr( - "tracecat.executor.registry_artifact_materialization.sysconfig.get_path", - lambda name: str(site_packages) if name == "purelib" else None, - ) - - cache = RegistryArtifactCache(temp_cache_dir) - async with cache.lease([bundled_builtin_registry_uri(version)]) as result: - assert result == [source_root.resolve(), site_packages.resolve()] - - @pytest.mark.anyio - async def test_lease_rejects_stale_bundled_builtin( - self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch - ): - """Bundled pseudo-URIs must match this executor's installed package.""" - monkeypatch.setattr(tracecat_registry, "__version__", "1.2.3") - - cache = RegistryArtifactCache(temp_cache_dir) - with pytest.raises(RuntimeError, match="does not match installed version"): - async with cache.lease([bundled_builtin_registry_uri("1.2.4")]): - pass - - @pytest.mark.anyio - async def test_download_artifact_normalizes_missing_objects_to_http_404( - self, temp_cache_dir - ): - """Preserve the missing-artifact error contract from presigned downloads.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = ( - "s3://access:secret@bucket/path/site-packages.tar.gz" - "?X-Amz-Signature=signed-secret#fragment" - ) - artifact = TarballArtifact( - uri=artifact_uri, - cache_key="missing-test", - ) - ctx = cache._context_for(artifact.cache_key) - output_path = temp_cache_dir / "artifact.tar.gz" - - with patch( - "tracecat.executor.registry_artifacts.blob.download_file_to_path", - new_callable=AsyncMock, - side_effect=FileNotFoundError, - ): - with pytest.raises(httpx.HTTPStatusError) as exc_info: - await artifact.download(ctx, output_path) - - assert exc_info.value.response.status_code == 404 - assert isinstance(exc_info.value.__cause__, FileNotFoundError) - assert str(exc_info.value) == "Registry artifact not found: s3://" - assert "secret" not in str(exc_info.value) - - @pytest.mark.anyio - async def test_artifact_candidates_prefer_squashfs_sidecar(self, temp_cache_dir): - """Test that gzip tarballs prefer a sibling SquashFS sidecar.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - return_value=True, - ) as file_exists, - patch.object(cache, "_can_try_squashfs", return_value=True), - ): - cache_key = compute_registry_artifact_cache_key( - "s3://bucket/path/site-packages.tar.gz" - ) - ctx = cache._context_for(cache_key) - candidates = await cache._artifact_candidates( - ctx, "s3://bucket/path/site-packages.tar.gz" - ) - - artifact = candidates[0] - assert len(candidates) == 2 - assert isinstance(artifact, SquashfsArtifact) - assert isinstance(candidates[1], TarballArtifact) - assert artifact.uri == "s3://bucket/path/site-packages.squashfs" - assert artifact.format == RegistryArtifactFormat.SQUASHFS - file_exists.assert_awaited_once_with( - key="path/site-packages.squashfs", - bucket="bucket", - ) - - @pytest.mark.anyio - async def test_artifact_candidates_ignore_malformed_local_sidecar( - self, temp_cache_dir: Path - ) -> None: - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://bucket/path/site-packages.tar.gz" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - ctx = cache._context_for(cache_key) - ctx.paths.squashfs_image_path.mkdir(parents=True) - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - return_value=False, - ) as file_exists, - patch.object(cache, "_can_try_squashfs", return_value=True), - ): - candidates = await cache._artifact_candidates(ctx, artifact_uri) - - assert len(candidates) == 1 - assert isinstance(candidates[0], TarballArtifact) - file_exists.assert_awaited_once_with( - key="path/site-packages.squashfs", - bucket="bucket", - ) - - @pytest.mark.anyio - async def test_artifact_candidates_direct_squashfs_include_gzip_fallback( - self, temp_cache_dir - ): - """Test direct SquashFS URIs fall back to sibling gzip tarballs.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with patch.object(cache, "_can_try_squashfs") as can_try_squashfs: - cache_key = compute_registry_artifact_cache_key( - "s3://bucket/path/site-packages.squashfs" - ) - ctx = cache._context_for(cache_key) - candidates = await cache._artifact_candidates( - ctx, - "s3://bucket/path/site-packages.squashfs", - ) - - assert isinstance(candidates[0], SquashfsArtifact) - assert isinstance(candidates[1], TarballArtifact) - assert [artifact.uri for artifact in candidates] == [ - "s3://bucket/path/site-packages.squashfs", - "s3://bucket/path/site-packages.tar.gz", - ] - assert [artifact.format for artifact in candidates] == [ - RegistryArtifactFormat.SQUASHFS, - RegistryArtifactFormat.TAR_GZ, - ] - can_try_squashfs.assert_not_called() - - @pytest.mark.anyio - async def test_artifact_candidates_fall_back_to_gzip(self, temp_cache_dir): - """Test that gzip tarballs are used when no sidecar exists.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - return_value=False, - ), - patch.object(cache, "_can_try_squashfs", return_value=True), - ): - cache_key = compute_registry_artifact_cache_key( - "s3://bucket/path/site-packages.tar.gz" - ) - ctx = cache._context_for(cache_key) - candidates = await cache._artifact_candidates( - ctx, "s3://bucket/path/site-packages.tar.gz" - ) - - artifact = candidates[0] - assert len(candidates) == 1 - assert isinstance(artifact, TarballArtifact) - assert artifact.uri == "s3://bucket/path/site-packages.tar.gz" - assert artifact.format == RegistryArtifactFormat.TAR_GZ - - @pytest.mark.anyio - async def test_sidecar_lookup_failure_redacts_exception_text( - self, - temp_cache_dir: Path, - ) -> None: - """SDK exception strings cannot leak registry identifiers into logs.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "s3://sensitive-bucket/org/repo/site-packages.tar.gz" - sidecar_uri = "s3://sensitive-bucket/org/repo/site-packages.squashfs" - ctx = cache._context_for(compute_registry_artifact_cache_key(artifact_uri)) - lookup_error = ConnectionError( - f"Could not connect to endpoint URL: {sidecar_uri}" - ) - - with ( - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - side_effect=lookup_error, - ), - patch.object(cache, "_can_try_squashfs", return_value=True), - patch("tracecat.executor.registry_artifacts.logger.warning") as warning, - ): - candidates = await cache._artifact_candidates(ctx, artifact_uri) - - assert len(candidates) == 1 - assert isinstance(candidates[0], TarballArtifact) - warning.assert_called_once_with( - "Failed to check for registry artifact sidecar, falling back", - artifact_uri="s3://", - sidecar_uri="s3://", - artifact_format=RegistryArtifactFormat.SQUASHFS.value, - error_type="ConnectionError", - ) - assert "sensitive-bucket" not in repr(warning.call_args) - assert "org/repo" not in repr(warning.call_args) - - @pytest.mark.anyio - async def test_sidecar_parse_failure_uses_redacted_fallback( - self, - temp_cache_dir: Path, - ) -> None: - """Malformed sidecar URIs cannot escape the redacted fallback path.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = ( - "https://tenant-bucket.invalid/org/repository/site-packages.tar.gz" - ) - ctx = cache._context_for(compute_registry_artifact_cache_key(artifact_uri)) - - with ( - patch.object(cache, "_can_try_squashfs", return_value=True), - patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - ) as file_exists, - patch("tracecat.executor.registry_artifacts.logger.warning") as warning, - ): - candidates = await cache._artifact_candidates(ctx, artifact_uri) - - assert len(candidates) == 1 - assert isinstance(candidates[0], TarballArtifact) - file_exists.assert_not_awaited() - warning.assert_called_once_with( - "Failed to check for registry artifact sidecar, falling back", - artifact_uri="https://", - sidecar_uri="https://", - artifact_format=RegistryArtifactFormat.SQUASHFS.value, - error_type="ValueError", - ) - assert "tenant-bucket" not in repr(warning.call_args) - assert "org/repository" not in repr(warning.call_args) - - @pytest.mark.anyio - async def test_materialization_fallback_redacts_malformed_uri_error( - self, - temp_cache_dir: Path, - ) -> None: - """Malformed candidate URIs cannot leak identifiers through errors.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = ( - "https://tenant-bucket.invalid/org/repository/site-packages.squashfs" - ) - cache_key = compute_registry_artifact_cache_key(artifact_uri) - ctx = cache._context_for(cache_key) - candidates = await cache._artifact_candidates(ctx, artifact_uri) - fallback_path = temp_cache_dir / "fallback" - - with ( - patch( - "tracecat.executor.registry_artifact_materialization." - "config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", - False, - ), - patch.object( - TarballArtifact, - "materialize", - new_callable=AsyncMock, - return_value=[fallback_path], - ) as materialize_fallback, - patch("tracecat.executor.registry_artifacts.logger.warning") as warning, - ): - registry_paths = await cache._materialize_candidates(ctx, candidates) - - assert registry_paths == [fallback_path] - materialize_fallback.assert_awaited_once() - warning.assert_called_once_with( - "Failed to materialize registry artifact candidate, trying fallback", - cache_key=cache_key, - artifact_uri="https://", - artifact_format=RegistryArtifactFormat.SQUASHFS.value, - error_type="RegistryArtifactUriError", - ) - assert "tenant-bucket" not in repr(warning.call_args) - assert "org/repository" not in repr(warning.call_args) - - @pytest.mark.anyio - async def test_final_candidate_malformed_uri_error_is_sanitized( - self, - temp_cache_dir: Path, - ) -> None: - """A final candidate cannot expose its malformed URI to callers.""" - cache = RegistryArtifactCache(temp_cache_dir) - artifact_uri = "https://tenant-bucket.invalid/org/repository/artifact.tar.gz" - - with pytest.raises(ValueError) as exc_info: - async with cache.lease([artifact_uri]): - pass - - assert str(exc_info.value) == "Invalid registry artifact URI" - assert exc_info.value.__cause__ is None - assert "tenant-bucket" not in repr(exc_info.value) - assert "org/repository" not in repr(exc_info.value) - - def test_can_try_squashfs_does_not_require_mount_binary(self, temp_cache_dir): - """Prefer SquashFS whenever enabled; extraction may work without mounts.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", - return_value=None, - ), - patch( - "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", - True, - ), - ): - ctx = cache._context_for("squashfs-test") - assert cache._can_try_squashfs() is True - assert ctx.can_mount_squashfs() is False - - @pytest.mark.anyio - async def test_artifact_candidates_skip_non_registry_tarballs(self, temp_cache_dir): - """Test that arbitrary gzip tarballs do not trigger sidecar lookups.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with patch( - "tracecat.executor.registry_artifacts.blob.file_exists", - new_callable=AsyncMock, - ) as file_exists: - cache_key = compute_registry_artifact_cache_key( - "s3://bucket/path/custom.tar.gz" - ) - ctx = cache._context_for(cache_key) - candidates = await cache._artifact_candidates( - ctx, "s3://bucket/path/custom.tar.gz" - ) - - artifact = candidates[0] - assert len(candidates) == 1 - assert isinstance(artifact, TarballArtifact) - assert artifact.uri == "s3://bucket/path/custom.tar.gz" - assert artifact.format == RegistryArtifactFormat.TAR_GZ - file_exists.assert_not_awaited() - - def test_runtime_for_same_key(self, temp_cache_dir): - """The same cache key returns the same runtime state.""" - cache = RegistryArtifactCache(temp_cache_dir) - - runtime1 = cache._runtime_for("key1") - runtime2 = cache._runtime_for("key1") - - assert runtime1 is runtime2 - - def test_runtime_for_different_keys(self, temp_cache_dir): - """Different cache keys return different runtime state.""" - cache = RegistryArtifactCache(temp_cache_dir) - - runtime1 = cache._runtime_for("key1") - runtime2 = cache._runtime_for("key2") - - assert runtime1 is not runtime2 diff --git a/tests/unit/executor/test_registry_artifact_startup.py b/tests/unit/executor/test_registry_artifact_startup.py deleted file mode 100644 index cdc33d3a62..0000000000 --- a/tests/unit/executor/test_registry_artifact_startup.py +++ /dev/null @@ -1,533 +0,0 @@ -"""Registry artifact startup recovery tests.""" - -from __future__ import annotations - -import asyncio -import os -import threading -from pathlib import Path -from unittest.mock import AsyncMock, patch - -import pytest - -from tracecat.executor.registry_artifacts import ( - RegistryArtifactCache, - _delete_cache_path, - bundled_builtin_registry_uri, -) - -from .registry_artifact_test_helpers import ( - MAX_BYTES_CONFIG, - MAX_ENTRIES_CONFIG, - write_image_entry, - write_tarball_entry, -) - - -class TestRegistryArtifactCacheStartupSweep: - """Tests for the startup sweep that reclaims state from a dead process.""" - - @pytest.mark.anyio - async def test_sweep_uses_entry_root_mtime_for_restart_safe_lru( - self, temp_cache_dir - ): - """Startup trimming preserves a touched older tarball-only entry.""" - previous_cache = RegistryArtifactCache(temp_cache_dir) - old = write_tarball_entry(temp_cache_dir, "old") - previous_cache._touch_entry("old") - - old_mtime = previous_cache._paths_for("old").entry_dir.stat().st_mtime - new = write_tarball_entry(temp_cache_dir, "new") - new_entry_dir = previous_cache._paths_for("new").entry_dir - os.utime(new_entry_dir, (old_mtime - 1, old_mtime - 1)) - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - ): - cache = RegistryArtifactCache(temp_cache_dir) - await cache.ensure_swept() - - assert old.is_dir() - assert not new.exists() - - @pytest.mark.anyio - async def test_sweep_removes_incomplete_shell_before_lru_trimming( - self, - temp_cache_dir: Path, - ) -> None: - """A crashed cold admission cannot displace a reusable cache entry.""" - cache = RegistryArtifactCache(temp_cache_dir) - reusable_image = write_image_entry( - temp_cache_dir, - "reusable", - size=16, - mtime=100.0, - ) - incomplete = cache._paths_for("incomplete") - incomplete.entry_dir.mkdir(parents=True) - incomplete.squashfs_mount_dir.mkdir() - os.utime(incomplete.entry_dir, (200.0, 200.0)) - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - ): - await cache.ensure_swept() - - assert reusable_image.is_file() - assert not incomplete.entry_dir.exists() - assert not any(cache.trash_dir.iterdir()) - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_sweep_retires_malformed_squashfs_image( - self, temp_cache_dir: Path - ) -> None: - cache = RegistryArtifactCache(temp_cache_dir) - malformed = cache._paths_for("malformed-image") - malformed.squashfs_image_path.mkdir(parents=True) - (malformed.squashfs_image_path / "stale").write_bytes(b"stale") - - await cache.ensure_swept() - - assert not malformed.entry_dir.exists() - assert not any(cache.trash_dir.iterdir()) - - @pytest.mark.anyio - async def test_sweep_tolerates_missing_cache_dir(self, temp_cache_dir): - """A cache directory that does not exist yet is a no-op.""" - cache_dir = temp_cache_dir / "missing" - - cache = RegistryArtifactCache(cache_dir) - await cache.ensure_swept() - - assert cache.cache_dir == cache_dir - assert not cache_dir.exists() - - @pytest.mark.anyio - async def test_sweep_removes_orphaned_work(self, temp_cache_dir): - """Interrupted work is reclaimed without touching unrelated paths.""" - cache = RegistryArtifactCache(temp_cache_dir) - orphaned = cache.staging_dir / "abc123.999999.4321.squashfs" - orphaned.parent.mkdir() - orphaned.write_bytes(b"partial") - orphaned_dir = cache.trash_dir / "abc123.999999.4321" - orphaned_dir.mkdir(parents=True) - unrelated = temp_cache_dir / "unrelated" - unrelated.mkdir() - unrelated_file = unrelated / "keep.txt" - unrelated_file.write_text("keep") - entry_dir = write_tarball_entry(temp_cache_dir, "abc123") - - await cache.ensure_swept() - - assert not orphaned.exists() - assert not orphaned_dir.exists() - assert unrelated_file.read_text() == "keep" - assert entry_dir.is_dir() - - @pytest.mark.anyio - @pytest.mark.parametrize("work_dir_name", ["staging", "trash"]) - async def test_sweep_rejects_symlinked_work_directory( - self, - temp_cache_dir: Path, - work_dir_name: str, - ) -> None: - """Startup cleanup cannot follow a cache work root outside the cache.""" - cache_dir = temp_cache_dir / "cache" - cache_dir.mkdir() - cache = RegistryArtifactCache(cache_dir) - outside_dir = temp_cache_dir / f"outside-{work_dir_name}" - outside_dir.mkdir() - outside_file = outside_dir / "keep.txt" - outside_file.write_text("keep") - getattr(cache, f"{work_dir_name}_dir").symlink_to( - outside_dir, - target_is_directory=True, - ) - - with pytest.raises(OSError, match="Unsafe .*registry cache work path"): - await cache.ensure_swept() - - assert outside_file.read_text() == "keep" - - @pytest.mark.anyio - async def test_sweep_rejects_symlinked_entries_directory( - self, - temp_cache_dir: Path, - ) -> None: - """Startup inspection cannot retire entries outside the configured cache.""" - cache_dir = temp_cache_dir / "cache" - cache_dir.mkdir() - cache = RegistryArtifactCache(cache_dir) - outside_dir = temp_cache_dir / "outside-entries" - outside_entry = write_tarball_entry(outside_dir, "keep") - cache.entries_dir.symlink_to( - outside_dir / "entries", - target_is_directory=True, - ) - - with pytest.raises(OSError, match="Unsafe .*registry cache entries path"): - await cache.ensure_swept() - - assert outside_entry.is_dir() - - @pytest.mark.anyio - async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): - """A live mountpoint belongs to a running process and must survive.""" - cache = RegistryArtifactCache(temp_cache_dir) - paths = cache._paths_for("mounted") - paths.entry_dir.mkdir(parents=True) - mount_dir = paths.squashfs_mount_dir - mount_dir.mkdir() - idle_dir = write_tarball_entry(temp_cache_dir, "idle") - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path == mount_dir, - ), - ): - await cache.ensure_swept() - - assert mount_dir.is_dir() - assert not idle_dir.exists() - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_ensure_swept_runs_once(self, temp_cache_dir): - """Repeated startup-sweep requests invoke the sweep once.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with patch.object( - cache, - "_sweep_startup_state", - wraps=cache._sweep_startup_state, - ) as sweep: - await cache.ensure_swept() - await cache.ensure_swept() - - assert sweep.call_count == 1 - - @pytest.mark.anyio - async def test_concurrent_ensure_swept_runs_once(self, temp_cache_dir): - """Concurrent startup-sweep requests invoke the sweep once.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with patch.object( - cache, - "_sweep_startup_state", - wraps=cache._sweep_startup_state, - ) as sweep: - await asyncio.gather(*(cache.ensure_swept() for _ in range(4))) - - assert sweep.call_count == 1 - - @pytest.mark.anyio - async def test_cancelled_waiter_joins_same_startup_sweep(self, temp_cache_dir): - """A cancelled waiter cannot start a second concurrent startup sweep.""" - cache = RegistryArtifactCache(temp_cache_dir) - sweep_started = threading.Event() - sweep_release = threading.Event() - invocation_count = 0 - - def blocking_sweep() -> None: - nonlocal invocation_count - invocation_count += 1 - sweep_started.set() - sweep_release.wait() - - with patch.object( - cache, - "_sweep_startup_state", - side_effect=blocking_sweep, - ): - first_waiter = asyncio.create_task(cache.ensure_swept()) - assert await asyncio.to_thread(sweep_started.wait, 1) - first_waiter.cancel() - - with pytest.raises(asyncio.CancelledError): - await first_waiter - - second_waiter = asyncio.create_task(cache.ensure_swept()) - await asyncio.sleep(0) - sweep_release.set() - await second_waiter - - assert invocation_count == 1 - assert cache._swept is True - - @pytest.mark.anyio - async def test_cache_lease_triggers_startup_sweep(self, temp_cache_dir): - """Lease admission reclaims startup scratch before yielding paths.""" - cache = RegistryArtifactCache(temp_cache_dir) - orphaned_dir = cache.staging_dir / "abc123.999999.4321" - orphaned_dir.mkdir(parents=True) - - with patch.object( - cache, - "_lease_artifact", - new_callable=AsyncMock, - return_value=(None, []), - ): - async with cache.lease(["s3://bucket/registry.tar.gz"]): - assert not orphaned_dir.exists() - - @pytest.mark.anyio - @pytest.mark.parametrize( - "artifact_uris", - [None, [bundled_builtin_registry_uri("1.2.3")]], - ) - async def test_cache_free_lease_skips_failed_startup_sweep( - self, - temp_cache_dir, - artifact_uris: list[str] | None, - ): - """Unrelated cache inspection failures cannot block cache-free actions.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch.object( - cache, - "ensure_swept", - new_callable=AsyncMock, - side_effect=OSError("simulated sweep failure"), - ) as ensure_swept, - patch.object( - cache, - "_lease_artifact", - new_callable=AsyncMock, - return_value=(None, [temp_cache_dir / "builtin"]), - ), - ): - async with cache.lease(artifact_uris): - pass - - ensure_swept.assert_not_awaited() - - @pytest.mark.anyio - async def test_failed_ensure_swept_retries(self, temp_cache_dir): - """A failed startup sweep is retried by the next caller.""" - cache = RegistryArtifactCache(temp_cache_dir) - orphaned_dir = cache.staging_dir / "abc123.999999.4321" - orphaned_dir.mkdir(parents=True) - - with ( - patch.object( - cache, - "_clear_work_dir", - side_effect=OSError("simulated sweep failure"), - ), - pytest.raises(OSError), - ): - await cache.ensure_swept() - - assert orphaned_dir.is_dir() - - await cache.ensure_swept() - - assert not orphaned_dir.exists() - - @pytest.mark.parametrize("work_dir_name", ["staging", "trash"]) - def test_work_directory_inspection_errors_propagate( - self, - temp_cache_dir, - work_dir_name: str, - ): - """Unreadable work directories must trigger a later cleanup retry.""" - cache = RegistryArtifactCache(temp_cache_dir) - work_dir = getattr(cache, f"{work_dir_name}_dir") - work_dir.mkdir() - - with ( - patch.object(Path, "iterdir", side_effect=PermissionError("denied")), - pytest.raises(PermissionError, match="denied"), - ): - cache._clear_work_dir(work_dir) - - @pytest.mark.anyio - async def test_active_entry_scan_errors_propagate(self, temp_cache_dir): - """Unreadable active entries cannot be treated as an empty cache.""" - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch( - "tracecat.executor.registry_artifact_storage.os.scandir", - side_effect=PermissionError("denied"), - ), - pytest.raises(PermissionError, match="denied"), - ): - await cache._enforce_cache_budget() - - @pytest.mark.anyio - async def test_failed_trash_cleanup_skips_startup_trimming(self, temp_cache_dir): - """A failed orphan deletion must not cascade into active retirements.""" - cache = RegistryArtifactCache(temp_cache_dir) - orphan = cache.trash_dir / "orphan" - orphan.mkdir(parents=True) - oldest = write_image_entry( - temp_cache_dir, - "oldest", - size=16, - mtime=100.0, - ) - newest = write_image_entry( - temp_cache_dir, - "newest", - size=16, - mtime=200.0, - ) - real_delete = _delete_cache_path - - def fail_orphan(path: Path) -> bool: - if path == orphan: - return False - return real_delete(path) - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - side_effect=fail_orphan, - ), - ): - await cache.ensure_swept() - - assert orphan.is_dir() - assert oldest.is_file() - assert newest.is_file() - assert cache._budget_dirty is True - - @pytest.mark.anyio - async def test_failed_startup_cleanup_retries_exact_path(self, temp_cache_dir): - """Failed startup scratch cleanup retries without sweeping live staging.""" - cache = RegistryArtifactCache(temp_cache_dir) - orphaned = cache.staging_dir / "orphaned.123.456.tmp" - orphaned.parent.mkdir(parents=True) - orphaned.write_bytes(b"partial") - real_delete = _delete_cache_path - failed_once = False - - def fail_once(path: Path) -> bool: - nonlocal failed_once - if path == orphaned and not failed_once: - failed_once = True - return False - return real_delete(path) - - with patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - side_effect=fail_once, - ): - await cache.ensure_swept() - assert orphaned.is_file() - assert cache._deferred_staging_cleanup == {orphaned} - assert cache._budget_dirty is True - - assert await cache._enforce_cache_budget() is True - - assert not orphaned.exists() - assert cache._deferred_staging_cleanup == set() - - @pytest.mark.anyio - async def test_failed_startup_retirement_stays_dirty_and_retries( - self, temp_cache_dir - ): - """A startup rename failure preserves entries for later enforcement.""" - oldest = write_image_entry( - temp_cache_dir, - "oldest", - size=16, - mtime=100.0, - ) - newest = write_image_entry( - temp_cache_dir, - "newest", - size=16, - mtime=200.0, - ) - cache = RegistryArtifactCache(temp_cache_dir) - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - patch( - "tracecat.executor.registry_artifact_storage._move_entry_to_trash", - side_effect=OSError("rename failed"), - ), - ): - await cache.ensure_swept() - - assert oldest.is_file() - assert newest.is_file() - assert cache._budget_dirty is True - - with ( - patch(MAX_ENTRIES_CONFIG, 1), - patch(MAX_BYTES_CONFIG, 0), - ): - await cache._converge_cache_budget() - - assert not oldest.exists() - assert newest.is_file() - assert cache._budget_dirty is False - - @pytest.mark.anyio - async def test_failed_startup_physical_delete_retries_exact_path( - self, temp_cache_dir - ): - """Startup stops retiring entries when bytes were not reclaimed.""" - oldest = write_image_entry( - temp_cache_dir, - "oldest", - size=16, - mtime=100.0, - ) - older = write_image_entry( - temp_cache_dir, - "older", - size=16, - mtime=200.0, - ) - newest = write_image_entry( - temp_cache_dir, - "newest", - size=16, - mtime=300.0, - ) - cache = RegistryArtifactCache(temp_cache_dir) - real_delete = _delete_cache_path - failed_once = False - - def fail_once(path: Path) -> bool: - nonlocal failed_once - if path.parent == cache.trash_dir and not failed_once: - failed_once = True - return False - return real_delete(path) - - with ( - patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 16), - patch( - "tracecat.executor.registry_artifact_storage._delete_cache_path", - side_effect=fail_once, - ), - ): - await cache.ensure_swept() - assert not oldest.exists() - assert older.is_file() - assert newest.is_file() - assert len(tuple(cache.trash_dir.iterdir())) == 1 - assert cache._budget_dirty is True - - await cache._converge_cache_budget() - - assert not older.exists() - assert newest.is_file() - assert not any(cache.trash_dir.iterdir()) - assert cache._budget_dirty is False diff --git a/tests/unit/executor/test_run_python_sdk_context.py b/tests/unit/executor/test_run_python_sdk_context.py index bdef85587c..9adb24b45d 100644 --- a/tests/unit/executor/test_run_python_sdk_context.py +++ b/tests/unit/executor/test_run_python_sdk_context.py @@ -473,18 +473,13 @@ class _FakeRunPythonRegistryArtifacts: def __init__(self, paths: list[Path]) -> None: self.paths = paths self.artifact_uris: list[str] | None = None - self.paths_may_be_modified: bool | None = None self.leased = False @asynccontextmanager async def lease( - self, - artifact_uris: list[str] | None = None, - *, - paths_may_be_modified: bool = False, + self, artifact_uris: list[str] | None = None ) -> AsyncIterator[list[Path]]: self.artifact_uris = artifact_uris - self.paths_may_be_modified = paths_may_be_modified self.leased = True try: yield self.paths @@ -1237,7 +1232,6 @@ async def run_python(self, **kwargs: Any) -> dict[str, bool]: assert result.type == "success" assert leased_during_run == [True] assert fake_runner.registry_artifacts.leased is False - assert fake_runner.registry_artifacts.paths_may_be_modified is True @pytest.mark.anyio diff --git a/tests/unit/executor/test_test_backend_no_registry_action.py b/tests/unit/executor/test_test_backend_no_registry_action.py index fdd08769bb..50363e03fc 100644 --- a/tests/unit/executor/test_test_backend_no_registry_action.py +++ b/tests/unit/executor/test_test_backend_no_registry_action.py @@ -12,8 +12,8 @@ import sys import threading import uuid -from collections.abc import AsyncIterator, Awaitable -from contextlib import AsyncExitStack, asynccontextmanager +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from datetime import UTC, datetime from pathlib import Path @@ -29,7 +29,6 @@ RunContext, ) from tracecat.executor.backends.test import TestBackend -from tracecat.executor.registry_artifacts import bundled_builtin_registry_uri from tracecat.executor.schemas import ( ActionImplementation, ExecutorResult, @@ -107,31 +106,6 @@ def test_run_action_input() -> RunActionInput: class TestTestBackendNoRegistryAction: """Test that TestBackend does not query RegistryActionsService.""" - def test_backend_instances_do_not_share_loop_affine_cache( - self, - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - ) -> None: - """Function-scoped backends can run in successive pytest event loops.""" - monkeypatch.setattr( - "tracecat.executor.backends.test.config" - ".TRACECAT__EXECUTOR_REGISTRY_CACHE_DIR", - str(tmp_path), - ) - first_backend = TestBackend() - second_backend = TestBackend() - - async def sweep(backend: TestBackend) -> None: - await backend._registry_artifact_cache().ensure_swept() - - asyncio.run(sweep(first_backend)) - asyncio.run(sweep(second_backend)) - - assert ( - first_backend._registry_artifact_cache() - is not second_backend._registry_artifact_cache() - ) - @pytest.mark.anyio async def test_execute_udf_without_db_lookup( self, @@ -297,18 +271,11 @@ class FakeRegistryArtifacts: def __init__(self) -> None: self.active = 0 - async def ensure_swept(self) -> None: - pass - @asynccontextmanager async def lease( - self, - artifact_uris: list[str] | None = None, - *, - paths_may_be_modified: bool = False, + self, artifact_uris: list[str] | None = None ) -> AsyncIterator[list[Path]]: - assert paths_may_be_modified is True - if broken_uri in (artifact_uris or []): + if artifact_uris == [broken_uri]: raise RuntimeError("artifact unavailable") self.active += 1 try: @@ -336,9 +303,8 @@ async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: False, ) monkeypatch.setattr( - backend, - "_registry_artifact_cache", - lambda: fake_runner.registry_artifacts, + "tracecat.executor.backends.test.get_action_runner", + lambda: fake_runner, ) monkeypatch.setattr(backend, "_get_artifact_uris", _get_artifact_uris) monkeypatch.setattr( @@ -368,126 +334,6 @@ async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: finally: await backend.shutdown() - @pytest.mark.anyio - async def test_execute_surfaces_registry_cache_sweep_failure( - self, - test_role: Role, - test_run_action_input: RunActionInput, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - """Global cache inspection errors are not treated as one bad artifact.""" - artifact_uri = "s3://bucket/registry.tar.gz" - - class FakeRegistryArtifacts: - def __init__(self) -> None: - self.lease_attempted = False - - async def ensure_swept(self) -> None: - raise PermissionError("cannot inspect registry cache") - - @asynccontextmanager - async def lease( - self, - artifact_uris: list[str] | None = None, - *, - paths_may_be_modified: bool = False, - ) -> AsyncIterator[list[Path]]: - del artifact_uris, paths_may_be_modified - self.lease_attempted = True - yield [] - - class FakeActionRunner: - def __init__(self) -> None: - self.registry_artifacts = FakeRegistryArtifacts() - - fake_runner = FakeActionRunner() - - async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: - return [artifact_uri] - - backend = TestBackend() - monkeypatch.setattr( - "tracecat.executor.backends.test.config.TRACECAT__LOCAL_REPOSITORY_ENABLED", - False, - ) - monkeypatch.setattr( - backend, - "_registry_artifact_cache", - lambda: fake_runner.registry_artifacts, - ) - monkeypatch.setattr(backend, "_get_artifact_uris", _get_artifact_uris) - - with pytest.raises(PermissionError, match="cannot inspect registry cache"): - async with AsyncExitStack() as leases: - await backend._lease_registry_artifacts( - leases, - test_run_action_input, - test_role, - ) - - assert fake_runner.registry_artifacts.lease_attempted is False - - @pytest.mark.anyio - async def test_builtin_only_execution_skips_registry_cache_sweep( - self, - test_role: Role, - test_run_action_input: RunActionInput, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - """Cache-free builtin execution is independent of cache inspection.""" - artifact_uri = bundled_builtin_registry_uri("1.2.3") - - class FakeRegistryArtifacts: - def __init__(self) -> None: - self.sweep_attempted = False - self.lease_attempted = False - - async def ensure_swept(self) -> None: - self.sweep_attempted = True - raise PermissionError("cannot inspect registry cache") - - @asynccontextmanager - async def lease( - self, - artifact_uris: list[str] | None = None, - *, - paths_may_be_modified: bool = False, - ) -> AsyncIterator[list[Path]]: - assert artifact_uris == [artifact_uri] - assert paths_may_be_modified is False - self.lease_attempted = True - yield [] - - registry_artifacts = FakeRegistryArtifacts() - - async def _get_artifact_uris(_input: RunActionInput, _role: Role) -> list[str]: - return [artifact_uri] - - backend = TestBackend() - monkeypatch.setattr( - "tracecat.executor.backends.test.config.TRACECAT__LOCAL_REPOSITORY_ENABLED", - False, - ) - monkeypatch.setattr( - backend, - "_registry_artifact_cache", - lambda: registry_artifacts, - ) - monkeypatch.setattr(backend, "_get_artifact_uris", _get_artifact_uris) - - async with AsyncExitStack() as leases: - assert ( - await backend._lease_registry_artifacts( - leases, - test_run_action_input, - test_role, - ) - == [] - ) - - assert registry_artifacts.sweep_attempted is False - assert registry_artifacts.lease_attempted is True - @pytest.mark.anyio async def test_timed_out_sync_udf_keeps_artifact_lease_until_thread_finishes( self, @@ -503,24 +349,16 @@ async def test_timed_out_sync_udf_keeps_artifact_lease_until_thread_finishes( artifact_uri = "s3://bucket/sync-timeout.tar.gz" worker_started = threading.Event() finish_worker = threading.Event() - timeout_triggered = asyncio.Event() class FakeRegistryArtifacts: def __init__(self) -> None: self.active = 0 - async def ensure_swept(self) -> None: - pass - @asynccontextmanager async def lease( - self, - artifact_uris: list[str] | None = None, - *, - paths_may_be_modified: bool = False, + self, artifact_uris: list[str] | None = None ) -> AsyncIterator[list[Path]]: assert artifact_uris == [artifact_uri] - assert paths_may_be_modified is True self.active += 1 try: yield [artifact_path] @@ -541,21 +379,6 @@ def blocking_udf(**_kwargs: object) -> str: assert finish_worker.wait(timeout=5) return "finished" - async def wait_for_after_worker_started[T]( - awaitable: Awaitable[T], - timeout: float | None, - ) -> T: - """Drive wait_for cancellation only after the UDF thread exists.""" - del timeout - task = asyncio.ensure_future(awaitable) - assert await asyncio.to_thread(worker_started.wait, 1) - task.cancel() - timeout_triggered.set() - try: - return await task - except asyncio.CancelledError as e: - raise TimeoutError from e - backend = TestBackend() await backend.start() execution: asyncio.Task[ExecutorResult] | None = None @@ -566,9 +389,8 @@ async def wait_for_after_worker_started[T]( False, ) monkeypatch.setattr( - backend, - "_registry_artifact_cache", - lambda: fake_runner.registry_artifacts, + "tracecat.executor.backends.test.get_action_runner", + lambda: fake_runner, ) monkeypatch.setattr(backend, "_get_artifact_uris", _get_artifact_uris) monkeypatch.setattr( @@ -576,17 +398,17 @@ async def wait_for_after_worker_started[T]( "_load_udf_callable", lambda _action_impl: blocking_udf, ) - monkeypatch.setattr(asyncio, "wait_for", wait_for_after_worker_started) execution = asyncio.create_task( backend.execute( input=test_run_action_input, role=test_role, resolved_context=test_resolved_context, - timeout=30.0, + timeout=0.01, ) ) - await timeout_triggered.wait() + assert await asyncio.to_thread(worker_started.wait, 1) + await asyncio.sleep(0.05) assert not execution.done() assert fake_runner.registry_artifacts.active == 1 diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index 3ee69dd017..ac749d35e8 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -6,10 +6,8 @@ from __future__ import annotations import asyncio -import contextlib import tempfile import uuid -from collections.abc import Awaitable, Callable from datetime import UTC, datetime from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch @@ -102,15 +100,12 @@ async def communicate( *, input: bytes | None = None, # noqa: A002 timeout: float | None = None, - terminate: Callable[[asyncio.subprocess.Process], Awaitable[None]] - | None = None, ) -> tuple[bytes, bytes]: if isinstance(process, asyncio.subprocess.Process): return await real_communication( process, input=input, timeout=timeout, - terminate=terminate, ) stdout, stderr = await asyncio.wait_for( process.communicate(input=input), @@ -532,10 +527,9 @@ async def test_execute_action_disables_new_privileges_for_direct_subprocess( temp_cache_dir, mock_run_action_input, mock_role, - mock_process_group_communication: AsyncMock, monkeypatch: pytest.MonkeyPatch, ): - """Test Linux direct execution isolates and drops supervisor privileges.""" + """Test direct subprocess execution disables new Linux privileges.""" runner = ActionRunner(cache_dir=temp_cache_dir) base_dir = temp_cache_dir / "base" base_dir.mkdir() @@ -579,17 +573,8 @@ async def create_subprocess_exec_side_effect(*args, **kwargs): # noqa: ARG001 "--inh-caps=-all", "--ambient-caps=-all", ] - assert captured_args[-5] == action_runner.sys.executable - assert captured_args[-4] == "-I" - assert captured_args[-3].endswith("process_supervisor.py") assert captured_args[-2] == action_runner.sys.executable assert captured_args[-1].endswith("minimal_runner.py") - communication_call = mock_process_group_communication.await_args - assert communication_call is not None - assert ( - communication_call.kwargs["terminate"] - is action_runner.terminate_supervised_process - ) @pytest.mark.anyio async def test_execute_action_invalid_json_response( @@ -664,7 +649,6 @@ async def test_execute_action_holds_registry_lease_for_whole_subprocess( cache_key = compute_registry_artifact_cache_key(artifact_uri) entry_dir = runner.registry_artifacts._paths_for(cache_key).tarball_target_dir entry_dir.mkdir(parents=True) - await runner.registry_artifacts.ensure_swept() monkeypatch.setattr( action_runner.config, "TRACECAT__EXECUTOR_SANDBOX_ENABLED", False @@ -694,23 +678,15 @@ async def create_subprocess_exec_side_effect(*args, **kwargs): # noqa: ARG001 env = kwargs.get("env") assert isinstance(env, dict) registry_paths.append(env["PYTHONPATH"]) - (entry_dir / "action-output.bin").write_bytes(b"x" * 4096) mock_proc = AsyncMock() mock_proc.returncode = 0 mock_proc.communicate = AsyncMock(return_value=(success_response, b"")) return mock_proc - with ( - patch( - "asyncio.create_subprocess_exec", - side_effect=create_subprocess_exec_side_effect, - ), - patch.object( - runner.registry_artifacts, - "_scan_cache_entries", - wraps=runner.registry_artifacts._scan_cache_entries, - ) as scan_cache_entries, + with patch( + "asyncio.create_subprocess_exec", + side_effect=create_subprocess_exec_side_effect, ): result = await runner.execute_action( input=mock_run_action_input, @@ -724,7 +700,6 @@ async def create_subprocess_exec_side_effect(*args, **kwargs): # noqa: ARG001 assert refcounts == [1] assert registry_paths[0].startswith(str(entry_dir)) assert runner.registry_artifacts._refcount(cache_key) == 0 - assert scan_cache_entries.call_count == 1 @pytest.mark.anyio async def test_cancelled_action_reaps_child_before_releasing_mounted_artifact( @@ -772,10 +747,7 @@ async def test_cancelled_action_reaps_child_before_releasing_mounted_artifact( ) real_create_subprocess_exec = asyncio.create_subprocess_exec - real_terminate_supervised_process = action_runner.terminate_supervised_process process_started = asyncio.Event() - termination_started = asyncio.Event() - finish_termination = asyncio.Event() process: asyncio.subprocess.Process | None = None reaped_before_unmount: list[bool] = [] @@ -785,13 +757,6 @@ async def capture_subprocess(*args, **kwargs): process_started.set() return process - async def controlled_termination( - requested_process: asyncio.subprocess.Process, - ) -> None: - termination_started.set() - await finish_termination.wait() - await real_terminate_supervised_process(requested_process) - async def release_mount(mount_dir: Path) -> bool: reaped_before_unmount.append( process is not None and process.returncode is not None @@ -800,11 +765,7 @@ async def release_mount(mount_dir: Path) -> bool: return True with ( - patch.object(action_runner.sys, "platform", "linux"), - patch( - "tracecat.executor.registry_artifact_mounts.is_mount", - lambda path: path in mounted, - ), + patch.object(Path, "is_mount", lambda path: path in mounted), patch.object( action_runner, "_direct_subprocess_command", @@ -814,11 +775,6 @@ async def release_mount(mount_dir: Path) -> bool: "tracecat.executor.action_runner.asyncio.create_subprocess_exec", side_effect=capture_subprocess, ), - patch.object( - action_runner, - "terminate_supervised_process", - side_effect=controlled_termination, - ), patch.object(cache, "_unmount", side_effect=release_mount), ): execution = asyncio.create_task( @@ -831,26 +787,13 @@ async def release_mount(mount_dir: Path) -> bool: ) ) try: - await asyncio.wait_for(process_started.wait(), timeout=5) + await process_started.wait() assert cache._refcount(cache_key) == 1 execution.cancel() - await termination_started.wait() - execution.cancel() - await asyncio.sleep(0) - assert not execution.done() - assert cache._refcount(cache_key) == 1 - assert paths.squashfs_mount_dir in mounted - - finish_termination.set() with pytest.raises(asyncio.CancelledError): await execution finally: - finish_termination.set() - if not execution.done(): - execution.cancel() - with contextlib.suppress(asyncio.CancelledError, Exception): - await execution if process is not None and process.returncode is None: process.kill() await process.wait() diff --git a/tests/unit/test_executor_sandbox_nsjail.py b/tests/unit/test_executor_sandbox_nsjail.py index d8619ed137..c451035e9e 100644 --- a/tests/unit/test_executor_sandbox_nsjail.py +++ b/tests/unit/test_executor_sandbox_nsjail.py @@ -573,7 +573,7 @@ async def download_artifact(self, ctx, output_path: Path) -> float: # path instead of attempting a loopback mount. patches.append( patch( - "tracecat.executor.registry_artifact_materialization.shutil.which", + "tracecat.executor.registry_artifacts.shutil.which", return_value=None, ) ) @@ -757,15 +757,13 @@ async def capture_subprocess(*args, **kwargs): ) try: - await asyncio.wait_for(process_started.wait(), timeout=5) + await process_started.wait() for _ in range(100): - try: - descendant_pid = int(descendant_pid_path.read_text()) - except (FileNotFoundError, ValueError): - await asyncio.sleep(0.01) - else: + if descendant_pid_path.exists(): break - assert descendant_pid is not None + await asyncio.sleep(0.01) + assert descendant_pid_path.is_file() + descendant_pid = int(descendant_pid_path.read_text()) execution.cancel() @@ -784,10 +782,6 @@ async def capture_subprocess(*args, **kwargs): else: pytest.fail("nsjail descendant survived process-group cleanup") finally: - if not execution.done(): - execution.cancel() - with contextlib.suppress(asyncio.CancelledError): - await execution if process is not None and process.returncode is None: process.kill() await process.wait() diff --git a/tests/unit/executor/test_registry_artifact_tarball.py b/tests/unit/test_multitenant_registry.py similarity index 83% rename from tests/unit/executor/test_registry_artifact_tarball.py rename to tests/unit/test_multitenant_registry.py index c25937e17d..7857dd6326 100644 --- a/tests/unit/executor/test_registry_artifact_tarball.py +++ b/tests/unit/test_multitenant_registry.py @@ -1,8 +1,14 @@ -"""Tarball cache behavior through the public lease API.""" +"""Tests for tarball cache behavior in registry action runner. + +These tests verify: +1. Tarball cache behavior (concurrent downloads, cache keys) +2. Cache key isolation per tarball URI +""" from __future__ import annotations import asyncio +import tempfile from pathlib import Path from unittest.mock import patch @@ -14,10 +20,29 @@ compute_registry_artifact_cache_key, ) -from .registry_artifact_test_helpers import ( - lease_paths, - tarball_payload, -) +# ============================================================================= +# Fixtures +# ============================================================================= + + +@pytest.fixture +def temp_cache_dir(): + """Create a temporary cache directory for each test.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield Path(tmpdir) + + +async def _lease_paths( + cache: RegistryArtifactCache, + artifact_uri: str, +) -> list[Path]: + async with cache.lease([artifact_uri]) as paths: + return paths + + +# ============================================================================= +# Test Class: Tarball Cache Behavior +# ============================================================================= class TestTarballCacheBehavior: @@ -52,7 +77,7 @@ async def test_concurrent_same_uri_single_download(self, temp_cache_dir: Path): async def mock_download(self, ctx, path: Path): download_count[0] += 1 await asyncio.sleep(0.1) # Simulate network delay - path.write_bytes(tarball_payload(size=1)) + path.write_bytes(b"fake tarball") async def mock_extract(self, tarball_path: Path, target_dir: Path): (target_dir / "extracted.txt").write_text("content") @@ -63,10 +88,10 @@ async def mock_extract(self, tarball_path: Path, target_dir: Path): ): # Launch multiple concurrent requests results = await asyncio.gather( - lease_paths(cache, tarball_uri), - lease_paths(cache, tarball_uri), - lease_paths(cache, tarball_uri), - lease_paths(cache, tarball_uri), + _lease_paths(cache, tarball_uri), + _lease_paths(cache, tarball_uri), + _lease_paths(cache, tarball_uri), + _lease_paths(cache, tarball_uri), ) # All should return same path @@ -99,7 +124,7 @@ async def test_different_uris_separate_cache_entries(self, temp_cache_dir: Path) async def mock_download(self, ctx, path: Path): download_calls.append(self.uri) - path.write_bytes(tarball_payload(size=1)) + path.write_bytes(b"fake tarball") async def mock_extract(self, tarball_path: Path, target_dir: Path): (target_dir / "extracted.txt").write_text("content") @@ -110,7 +135,7 @@ async def mock_extract(self, tarball_path: Path, target_dir: Path): ): results = [] for uri in uris: - result = await lease_paths(cache, uri) + result = await _lease_paths(cache, uri) results.append(result) # All results should be different paths @@ -140,7 +165,7 @@ async def test_failed_extraction_cleans_up_temp_files(self, temp_cache_dir: Path cache_key = compute_registry_artifact_cache_key(tarball_uri) async def mock_download(self, ctx, path: Path): - path.write_bytes(tarball_payload(size=1)) + path.write_bytes(b"corrupt tarball") async def mock_extract(self, tarball_path: Path, target_dir: Path): raise RuntimeError("Extraction failed - corrupt tarball") @@ -150,7 +175,7 @@ async def mock_extract(self, tarball_path: Path, target_dir: Path): patch.object(TarballArtifact, "extract", mock_extract), ): with pytest.raises(RuntimeError, match="Extraction failed"): - await lease_paths(cache, tarball_uri) + await _lease_paths(cache, tarball_uri) # Verify no temp files remain temp_files = ( @@ -181,7 +206,7 @@ async def test_cache_reused_on_second_request(self, temp_cache_dir: Path): async def mock_download(self, ctx, path: Path): download_count[0] += 1 - path.write_bytes(tarball_payload(size=1)) + path.write_bytes(b"tarball") async def mock_extract(self, tarball_path: Path, target_dir: Path): (target_dir / "file.txt").write_text("content") @@ -191,11 +216,11 @@ async def mock_extract(self, tarball_path: Path, target_dir: Path): patch.object(TarballArtifact, "extract", mock_extract), ): # First request - result1 = await lease_paths(cache, tarball_uri) + result1 = await _lease_paths(cache, tarball_uri) assert download_count[0] == 1 # Second request (should use cache) - result2 = await lease_paths(cache, tarball_uri) + result2 = await _lease_paths(cache, tarball_uri) assert download_count[0] == 1 # No additional download assert result1 == result2 diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py new file mode 100644 index 0000000000..6f2f2e5d73 --- /dev/null +++ b/tests/unit/test_registry_artifacts.py @@ -0,0 +1,3418 @@ +"""Tests for executor registry artifact materialization.""" + +from __future__ import annotations + +import asyncio +import io +import os +import tarfile +import tempfile +import threading +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from pathlib import Path +from unittest.mock import ANY, AsyncMock, patch + +import httpx +import pytest +import tracecat_registry + +from tracecat.executor.registry_artifacts import ( + SQUASHFS_MOUNT_OPTIONS, + RegistryArtifactCache, + RegistryArtifactCacheCapacityError, + RegistryArtifactCacheLoopError, + RegistryArtifactEviction, + RegistryArtifactFormat, + RegistryArtifactMaterializationContext, + SquashfsArtifact, + SquashfsMountCommandError, + TarballArtifact, + _delete_cache_path, + _squashfs_listing_size, + bundled_builtin_registry_uri, + compute_registry_artifact_cache_key, +) +from tracecat.registry.artifact_keys import parse_s3_uri + +MAX_ENTRIES_CONFIG = ( + "tracecat.executor.registry_artifacts.config" + ".TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES" +) +MAX_BYTES_CONFIG = ( + "tracecat.executor.registry_artifacts.config" + ".TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES" +) +SQUASHFS_ENABLED_CONFIG = ( + "tracecat.executor.registry_artifacts.config" + ".TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED" +) + + +def _write_tarball_entry(cache_dir: Path, cache_key: str) -> Path: + """Create a materialized tarball cache entry on disk.""" + target_dir = cache_dir / "entries" / cache_key / "tarball" + target_dir.mkdir(parents=True) + (target_dir / "module.py").write_text("VALUE = 1") + return target_dir + + +def _write_image_entry( + cache_dir: Path, cache_key: str, *, size: int, mtime: float +) -> Path: + """Create a downloaded SquashFS image cache entry with a fixed mtime.""" + entry_dir = cache_dir / "entries" / cache_key + entry_dir.mkdir(parents=True, exist_ok=True) + image_path = entry_dir / "image.squashfs" + image_path.write_bytes(b"x" * size) + os.utime(image_path, (mtime, mtime)) + os.utime(entry_dir, (mtime, mtime)) + return image_path + + +def _tarball_payload(*, size: int) -> bytes: + """Return a gzip tarball containing one synthetic regular file.""" + payload = b"x" * size + output = io.BytesIO() + with tarfile.open(fileobj=output, mode="w:gz") as tar: + member = tarfile.TarInfo("module.py") + member.size = len(payload) + tar.addfile(member, io.BytesIO(payload)) + return output.getvalue() + + +@dataclass(slots=True) +class _SquashfsMountHarness: + """Model mount ownership without consuming host loop devices.""" + + cache: RegistryArtifactCache + failed_mount_keys: set[str] = field(default_factory=set) + mounted: set[Path] = field(default_factory=set) + mount_attempts: list[str] = field(default_factory=list) + extraction_attempts: list[str] = field(default_factory=list) + unmounts: list[Path] = field(default_factory=list) + + def seed_mount(self, cache_key: str) -> Path: + """Create one reusable image and mark its mount directory active.""" + paths = self.cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True, exist_ok=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir(exist_ok=True) + (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") + self.mounted.add(paths.squashfs_mount_dir) + return paths.squashfs_mount_dir + + async def mount( + self, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> Path: + """Record a mount attempt, failing selected keys like loop exhaustion.""" + self.mount_attempts.append(ctx.cache_key) + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + image_path.write_bytes(b"squashfs") + if ctx.cache_key in self.failed_mount_keys: + raise SquashfsMountCommandError("no free loop device") + ctx.paths.squashfs_mount_dir.mkdir(exist_ok=True) + (ctx.paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") + self.mounted.add(ctx.paths.squashfs_mount_dir) + return ctx.paths.squashfs_mount_dir + + async def extract( + self, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> Path: + """Publish an extracted fallback from the image retained by mount.""" + assert image_path.is_file() + self.extraction_attempts.append(ctx.cache_key) + ctx.paths.squashfs_extract_dir.mkdir(parents=True, exist_ok=True) + (ctx.paths.squashfs_extract_dir / "module.py").write_text("VALUE = 1") + return ctx.paths.squashfs_extract_dir + + async def unmount(self, mount_dir: Path) -> bool: + """Record final-release unmounts and release the modeled loop device.""" + self.unmounts.append(mount_dir) + self.mounted.discard(mount_dir) + return True + + +async def _materialize( + cache: RegistryArtifactCache, + cache_key: str, + artifact_uri: str, +) -> list[Path]: + """Exercise internal materialization while releasing its test-only lease.""" + await cache.ensure_swept() + ctx = cache._context_for(cache_key) + lock = cache._runtime_for(cache_key).lock + lease_acquired = False + try: + async with lock: + cache._acquire_lease(cache_key) + lease_acquired = True + candidates = await cache._artifact_candidates(ctx, artifact_uri) + if cached_paths := cache._first_cached_path(candidates, ctx): + return cached_paths + paths = await cache._materialize_candidates(ctx, candidates) + cache._touch_entry(cache_key) + await cache._enforce_cache_budget(protected_key=cache_key) + return paths + finally: + if lease_acquired: + cache._release_lease(cache_key) + + +class _BlockingSubprocess: + """Fake subprocess that blocks in communicate until it is cancelled.""" + + def __init__(self) -> None: + self.communicate_started = asyncio.Event() + self.cleanup_calls: list[str] = [] + self.returncode: int | None = None + + async def communicate(self) -> tuple[bytes, bytes]: + """Block until the task awaiting subprocess completion is cancelled.""" + self.communicate_started.set() + await asyncio.Event().wait() + return b"", b"" + + def kill(self) -> None: + """Record that the subprocess was killed.""" + self.cleanup_calls.append("kill") + self.returncode = -9 + + async def wait(self) -> int: + """Record that the killed subprocess was reaped.""" + self.cleanup_calls.append("wait") + return -9 + + +class _CapturedSubprocess: + """Capture cleanup of a real subprocess used by cancellation tests.""" + + def __init__(self, process: asyncio.subprocess.Process) -> None: + self.process = process + self.killed = False + self.reaped = False + + @property + def returncode(self) -> int | None: + """Return the wrapped subprocess exit status.""" + return self.process.returncode + + async def communicate(self) -> tuple[bytes, bytes]: + """Wait for the wrapped subprocess and collect its output.""" + return await self.process.communicate() + + def kill(self) -> None: + """Kill the wrapped subprocess and record the signal.""" + self.killed = True + self.process.kill() + + async def wait(self) -> int: + """Reap the wrapped subprocess and record completion.""" + returncode = await self.process.wait() + self.reaped = True + return returncode + + +@pytest.fixture +def temp_cache_dir(): + """Create a temporary cache directory.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield Path(tmpdir) + + +class TestParseS3Uri: + """Tests for parse_s3_uri function.""" + + def test_valid_uri(self): + """Test parsing a valid S3 URI.""" + bucket, key = parse_s3_uri("s3://my-bucket/path/to/file.tar.gz") + assert bucket == "my-bucket" + assert key == "path/to/file.tar.gz" + + def test_uri_with_nested_path(self): + """Test parsing URI with deeply nested path.""" + bucket, key = parse_s3_uri("s3://bucket/a/b/c/d/e/file.tar.gz") + assert bucket == "bucket" + assert key == "a/b/c/d/e/file.tar.gz" + + def test_invalid_uri_no_prefix(self): + """Test that non-S3 URIs raise ValueError.""" + with pytest.raises(ValueError, match="Invalid S3 URI"): + parse_s3_uri("https://bucket/key") + + def test_invalid_uri_no_key(self): + """Test that URIs without keys raise ValueError.""" + with pytest.raises(ValueError, match="Invalid S3 URI"): + parse_s3_uri("s3://bucket") + + def test_invalid_uri_empty_bucket(self): + """Test that URIs with empty bucket raise ValueError.""" + with pytest.raises(ValueError, match="Invalid S3 URI"): + parse_s3_uri("s3:///key") + + +class TestRegistryArtifactCache: + """Tests for registry artifact cache behavior.""" + + def test_compute_registry_artifact_cache_key_deterministic(self): + """Test that cache key computation is deterministic.""" + uri = "s3://bucket/path/to/registry-v1.2.3.tar.gz" + + key1 = compute_registry_artifact_cache_key(uri) + key2 = compute_registry_artifact_cache_key(uri) + + assert key1 == key2 + assert len(key1) == 16 + + def test_compute_registry_artifact_cache_key_case_sensitive(self): + """Test that cache key is case-sensitive because S3 keys are case-sensitive.""" + key1 = compute_registry_artifact_cache_key("s3://BUCKET/PATH/FILE.tar.gz") + key2 = compute_registry_artifact_cache_key("s3://bucket/path/file.tar.gz") + + assert key1 != key2 + + def test_compute_registry_artifact_cache_key_empty(self): + """Test that empty URI returns the base cache key.""" + assert compute_registry_artifact_cache_key("") == "base" + + @pytest.mark.anyio + async def test_download_artifact_uses_blob_download_file_to_path( + self, temp_cache_dir + ): + """Test that artifact downloads stay behind the blob storage helper.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key="download-test", + ) + ctx = cache._context_for(artifact.cache_key) + output_path = temp_cache_dir / "artifact.squashfs" + + async def mock_download_file_to_path( + *, + key: str, + bucket: str, + output_path: Path, + ) -> None: + output_path.write_bytes(b"squashfs") + + with patch( + "tracecat.executor.registry_artifacts.blob.download_file_to_path", + new_callable=AsyncMock, + side_effect=mock_download_file_to_path, + ) as download_file_to_path: + await artifact.download(ctx, output_path) + + download_file_to_path.assert_awaited_once() + await_args = download_file_to_path.await_args + assert await_args is not None + assert await_args.kwargs["key"] == "path/site-packages.squashfs" + assert await_args.kwargs["bucket"] == "bucket" + assert output_path.read_bytes() == b"squashfs" + + @pytest.mark.anyio + async def test_lease_uses_bundled_current_builtin( + self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch + ): + """In-tree builtin registry returns only the installed site-packages.""" + version = "1.2.3" + site_packages = temp_cache_dir / "venv" / "site-packages" + package_dir = site_packages / "tracecat_registry" + package_dir.mkdir(parents=True) + package_file = package_dir / "__init__.py" + package_file.write_text("") + + monkeypatch.setattr(tracecat_registry, "__version__", version) + monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) + monkeypatch.setattr( + "tracecat.executor.registry_artifacts.sysconfig.get_path", + lambda name: str(site_packages) if name == "purelib" else None, + ) + + cache = RegistryArtifactCache(temp_cache_dir) + async with cache.lease([bundled_builtin_registry_uri(version)]) as result: + assert result == [site_packages.resolve()] + + @pytest.mark.anyio + async def test_lease_exposes_editable_builtin_parent( + self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch + ): + """Editable builtin registry exposes the package wrapper + site-packages.""" + version = "1.2.3" + site_packages = temp_cache_dir / "venv" / "site-packages" + dependency_dir = site_packages / "orjson" + dependency_dir.mkdir(parents=True) + (dependency_dir / "__init__.py").write_text("VALUE = 1\n") + source_root = temp_cache_dir / "src" / "tracecat-registry" + package_dir = source_root / "tracecat_registry" + package_dir.mkdir(parents=True) + package_file = package_dir / "__init__.py" + package_file.write_text("__version__ = '1.2.3'\n") + + monkeypatch.setattr(tracecat_registry, "__version__", version) + monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) + monkeypatch.setattr( + "tracecat.executor.registry_artifacts.sysconfig.get_path", + lambda name: str(site_packages) if name == "purelib" else None, + ) + + cache = RegistryArtifactCache(temp_cache_dir) + async with cache.lease([bundled_builtin_registry_uri(version)]) as result: + assert result == [source_root.resolve(), site_packages.resolve()] + + @pytest.mark.anyio + async def test_lease_rejects_stale_bundled_builtin( + self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch + ): + """Bundled pseudo-URIs must match this executor's installed package.""" + monkeypatch.setattr(tracecat_registry, "__version__", "1.2.3") + + cache = RegistryArtifactCache(temp_cache_dir) + with pytest.raises(RuntimeError, match="does not match installed version"): + async with cache.lease([bundled_builtin_registry_uri("1.2.4")]): + pass + + @pytest.mark.anyio + async def test_download_artifact_normalizes_missing_objects_to_http_404( + self, temp_cache_dir + ): + """Preserve the missing-artifact error contract from presigned downloads.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact = TarballArtifact( + uri="s3://bucket/path/site-packages.tar.gz", + cache_key="missing-test", + ) + ctx = cache._context_for(artifact.cache_key) + output_path = temp_cache_dir / "artifact.tar.gz" + + with patch( + "tracecat.executor.registry_artifacts.blob.download_file_to_path", + new_callable=AsyncMock, + side_effect=FileNotFoundError, + ): + with pytest.raises(httpx.HTTPStatusError) as exc_info: + await artifact.download(ctx, output_path) + + assert exc_info.value.response.status_code == 404 + assert isinstance(exc_info.value.__cause__, FileNotFoundError) + + @pytest.mark.anyio + async def test_artifact_candidates_prefer_squashfs_sidecar(self, temp_cache_dir): + """Test that gzip tarballs prefer a sibling SquashFS sidecar.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + return_value=True, + ) as file_exists, + patch.object(cache, "_can_try_squashfs", return_value=True), + ): + cache_key = compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.tar.gz" + ) + ctx = cache._context_for(cache_key) + candidates = await cache._artifact_candidates( + ctx, "s3://bucket/path/site-packages.tar.gz" + ) + + artifact = candidates[0] + assert len(candidates) == 2 + assert isinstance(artifact, SquashfsArtifact) + assert isinstance(candidates[1], TarballArtifact) + assert artifact.uri == "s3://bucket/path/site-packages.squashfs" + assert artifact.format == RegistryArtifactFormat.SQUASHFS + file_exists.assert_awaited_once_with( + key="path/site-packages.squashfs", + bucket="bucket", + ) + + @pytest.mark.anyio + async def test_same_key_cold_fan_in_materializes_and_enforces_once( + self, temp_cache_dir + ): + """Same-key waiters share candidate lookup, materialization, and enforcement.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/path/site-packages.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + materialization_started = asyncio.Event() + finish_materialization = asyncio.Event() + + async def mock_materialize( + self: TarballArtifact, + ctx: RegistryArtifactMaterializationContext, + ) -> list[Path]: + materialization_started.set() + await finish_materialization.wait() + ctx.paths.tarball_target_dir.mkdir(parents=True) + return [ctx.paths.tarball_target_dir] + + async def take_lease() -> list[Path]: + async with cache.lease([artifact_uri]) as paths: + return paths + + with ( + patch.object( + cache, + "_sidecar_exists", + new_callable=AsyncMock, + return_value=False, + ) as sidecar_exists, + patch.object( + cache, + "_enforce_cache_budget", + new_callable=AsyncMock, + return_value=True, + ) as enforce_cache_budget, + patch.object(cache, "_converge_cache_budget", new_callable=AsyncMock), + patch.object(TarballArtifact, "materialize", mock_materialize), + ): + leases = [asyncio.create_task(take_lease()) for _ in range(5)] + await materialization_started.wait() + await asyncio.sleep(0) + finish_materialization.set() + results = await asyncio.gather(*leases) + + expected_paths = [cache._paths_for(cache_key).tarball_target_dir] + assert results == [expected_paths] * 5 + sidecar_exists.assert_awaited_once() + enforce_cache_budget.assert_awaited_once_with(protected_key=cache_key) + + @pytest.mark.anyio + async def test_different_cold_keys_serialize_materialization( + self, temp_cache_dir: Path + ) -> None: + """Distinct cold keys cannot multiply staging and extraction peaks.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + uris = ["s3://bucket/first.tar.gz", "s3://bucket/second.tar.gz"] + first_started = asyncio.Event() + release_first = asyncio.Event() + second_started = asyncio.Event() + started_keys: list[str] = [] + + async def controlled_materialize( + self: TarballArtifact, + ctx: RegistryArtifactMaterializationContext, + ) -> list[Path]: + del self + started_keys.append(ctx.cache_key) + if len(started_keys) == 1: + first_started.set() + await release_first.wait() + else: + second_started.set() + ctx.paths.tarball_target_dir.mkdir(parents=True) + return [ctx.paths.tarball_target_dir] + + async def take_lease(uri: str) -> None: + async with cache.lease([uri]): + pass + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "materialize", controlled_materialize), + patch.object(cache, "_enforce_cache_budget", new_callable=AsyncMock), + patch.object(cache, "_converge_cache_budget", new_callable=AsyncMock), + ): + first = asyncio.create_task(take_lease(uris[0])) + await first_started.wait() + second = asyncio.create_task(take_lease(uris[1])) + await asyncio.sleep(0) + assert not second_started.is_set() + + release_first.set() + await second_started.wait() + await asyncio.gather(first, second) + + assert started_keys == [ + compute_registry_artifact_cache_key(uri) for uri in uris + ] + + def test_squashfs_listing_size_sums_files_and_symlinks(self) -> None: + listing = b"\n".join( + [ + b"drwxr-xr-x 0/0 64 2026-01-01 00:00 squashfs-root", + b"-rw-r--r-- 0/0 123 2026-01-01 00:00 squashfs-root/module.py", + b"lrwxrwxrwx 0/0 9 2026-01-01 00:00 squashfs-root/current -> module.py", + ] + ) + + assert _squashfs_listing_size(listing) == 132 + + def test_squashfs_listing_size_rejects_unparseable_files(self) -> None: + with pytest.raises(ValueError, match="Could not parse SquashFS listing"): + _squashfs_listing_size(b"-rw-r--r-- malformed") + + @pytest.mark.anyio + async def test_artifact_candidates_direct_squashfs_include_gzip_fallback( + self, temp_cache_dir + ): + """Test direct SquashFS URIs fall back to sibling gzip tarballs.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with patch.object(cache, "_can_try_squashfs") as can_try_squashfs: + cache_key = compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.squashfs" + ) + ctx = cache._context_for(cache_key) + candidates = await cache._artifact_candidates( + ctx, + "s3://bucket/path/site-packages.squashfs", + ) + + assert isinstance(candidates[0], SquashfsArtifact) + assert isinstance(candidates[1], TarballArtifact) + assert [artifact.uri for artifact in candidates] == [ + "s3://bucket/path/site-packages.squashfs", + "s3://bucket/path/site-packages.tar.gz", + ] + assert [artifact.format for artifact in candidates] == [ + RegistryArtifactFormat.SQUASHFS, + RegistryArtifactFormat.TAR_GZ, + ] + can_try_squashfs.assert_not_called() + + @pytest.mark.anyio + async def test_artifact_candidates_fall_back_to_gzip(self, temp_cache_dir): + """Test that gzip tarballs are used when no sidecar exists.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + return_value=False, + ), + patch.object(cache, "_can_try_squashfs", return_value=True), + ): + cache_key = compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.tar.gz" + ) + ctx = cache._context_for(cache_key) + candidates = await cache._artifact_candidates( + ctx, "s3://bucket/path/site-packages.tar.gz" + ) + + artifact = candidates[0] + assert len(candidates) == 1 + assert isinstance(artifact, TarballArtifact) + assert artifact.uri == "s3://bucket/path/site-packages.tar.gz" + assert artifact.format == RegistryArtifactFormat.TAR_GZ + + def test_can_try_squashfs_does_not_require_mount_binary(self, temp_cache_dir): + """Prefer SquashFS whenever enabled; extraction may work without mounts.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value=None, + ), + patch( + "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", + True, + ), + ): + ctx = cache._context_for("squashfs-test") + assert cache._can_try_squashfs() is True + assert ctx.can_mount_squashfs() is False + + @pytest.mark.anyio + async def test_artifact_candidates_skip_non_registry_tarballs(self, temp_cache_dir): + """Test that arbitrary gzip tarballs do not trigger sidecar lookups.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + ) as file_exists: + cache_key = compute_registry_artifact_cache_key( + "s3://bucket/path/custom.tar.gz" + ) + ctx = cache._context_for(cache_key) + candidates = await cache._artifact_candidates( + ctx, "s3://bucket/path/custom.tar.gz" + ) + + artifact = candidates[0] + assert len(candidates) == 1 + assert isinstance(artifact, TarballArtifact) + assert artifact.uri == "s3://bucket/path/custom.tar.gz" + assert artifact.format == RegistryArtifactFormat.TAR_GZ + file_exists.assert_not_awaited() + + @pytest.mark.anyio + async def test_materialize_mounts_squashfs_sidecar(self, temp_cache_dir): + """Test that a SquashFS sidecar is mounted instead of extracting tarballs.""" + cache = RegistryArtifactCache(temp_cache_dir) + + async def mock_mount(self, ctx, image_path): + assert image_path.name.endswith(".squashfs") + target_dir = ctx.paths.squashfs_mount_dir + target_dir.mkdir(parents=True, exist_ok=True) + (target_dir / "module.py").write_text("VALUE = 1") + return target_dir + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + return_value=True, + ), + patch( + "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", + True, + ), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + patch.object( + TarballArtifact, + "materialize", + new_callable=AsyncMock, + ) as tarball_materialize, + ): + result = await _materialize( + cache, + "squashfs-key", + "s3://bucket/path/site-packages.tar.gz", + ) + + assert len(result) == 1 + assert (result[0] / "module.py").read_text() == "VALUE = 1" + tarball_materialize.assert_not_awaited() + + @pytest.mark.anyio + async def test_mount_squashfs_uses_hardened_read_only_options( + self, + temp_cache_dir, + ): + """Test that SquashFS images are mounted read-only without device/setuid bits.""" + cache_key = "cache-key" + cache = RegistryArtifactCache(temp_cache_dir) + ctx = cache._context_for(cache_key) + artifact = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key=cache_key, + ) + image_path = ctx.paths.squashfs_image_path + target_dir = ctx.paths.squashfs_mount_dir + ctx.paths.entry_dir.mkdir(parents=True) + image_path.write_bytes(b"squashfs") + target_dir.mkdir() + process = AsyncMock() + process.communicate.return_value = (b"", b"") + process.returncode = 0 + + with patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ) as create_subprocess_exec: + await artifact.mount(ctx, image_path) + + create_subprocess_exec.assert_awaited_once_with( + "mount", + "-t", + "squashfs", + "-o", + SQUASHFS_MOUNT_OPTIONS, + str(image_path), + str(target_dir), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + + @pytest.mark.anyio + async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): + """Cancellation cannot leave an orphan mount process after lock release.""" + cache_key = "cancelled-mount" + cache = RegistryArtifactCache(temp_cache_dir) + ctx = cache._context_for(cache_key) + artifact = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key=cache_key, + ) + image_path = ctx.paths.squashfs_image_path + target_dir = ctx.paths.squashfs_mount_dir + ctx.paths.entry_dir.mkdir(parents=True) + image_path.write_bytes(b"squashfs") + target_dir.mkdir() + process = _BlockingSubprocess() + + with patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ): + mounting = asyncio.create_task( + artifact._mount_image(image_path, target_dir) + ) + await process.communicate_started.wait() + mounting.cancel() + + with pytest.raises(asyncio.CancelledError): + await mounting + + assert process.cleanup_calls == ["kill", "wait"] + assert target_dir.is_dir() + assert not target_dir.is_mount() + + @pytest.mark.anyio + async def test_cancelled_squashfs_extract_kills_and_reaps_subprocess( + self, temp_cache_dir + ): + """Cancellation cannot leave unsquashfs writing into scratch.""" + cache_key = "cancelled-extract" + cache = RegistryArtifactCache(temp_cache_dir) + ctx = cache._context_for(cache_key) + artifact = SquashfsArtifact( + uri="s3://bucket/path/site-packages.squashfs", + cache_key=cache_key, + ) + image_path = ctx.paths.squashfs_image_path + target_dir = ctx.paths.squashfs_extract_dir + ctx.paths.entry_dir.mkdir(parents=True) + image_path.write_bytes(b"squashfs") + target_dir.mkdir() + real_create_subprocess_exec = asyncio.create_subprocess_exec + process_started = asyncio.Event() + captured_processes: list[_CapturedSubprocess] = [] + + async def create_sleep_subprocess( + *args: object, **kwargs: object + ) -> _CapturedSubprocess: + del args, kwargs + process = await real_create_subprocess_exec( + "/bin/sleep", + "30", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + captured = _CapturedSubprocess(process) + captured_processes.append(captured) + process_started.set() + return captured + + with ( + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/usr/bin/unsquashfs", + ), + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + side_effect=create_sleep_subprocess, + ), + ): + extracting = asyncio.create_task( + artifact._extract_image(image_path, target_dir) + ) + await process_started.wait() + captured = captured_processes[0] + extracting.cancel() + + try: + with pytest.raises(asyncio.CancelledError): + await extracting + finally: + if captured.returncode is None: + captured.kill() + await captured.wait() + + assert captured.killed is True + assert captured.reaped is True + assert captured.returncode is not None + + @pytest.mark.anyio + async def test_repeatedly_cancelled_tarball_extract_rejoins_thread( + self, temp_cache_dir + ): + """Repeated cancellation waits until the tar extractor stops writing.""" + artifact = TarballArtifact( + uri="s3://bucket/path/site-packages.tar.gz", + cache_key="cancelled-tarball-extract", + ) + tarball_path = temp_cache_dir / "artifact.tar.gz" + target_dir = temp_cache_dir / "target" + target_dir.mkdir() + with tarfile.open(tarball_path, "w:gz"): + pass + + extraction_started = threading.Event() + extraction_release = threading.Event() + extraction_finished = threading.Event() + + def blocking_extractall(*args: object, **kwargs: object) -> None: + del args, kwargs + extraction_started.set() + extraction_release.wait() + extraction_finished.set() + + with patch.object( + tarfile.TarFile, + "extractall", + new=blocking_extractall, + ): + extracting = asyncio.create_task(artifact.extract(tarball_path, target_dir)) + assert await asyncio.to_thread(extraction_started.wait, 1) + extracting.cancel() + done, _ = await asyncio.wait({extracting}, timeout=0.05) + first_cancellation_propagated_early = bool(done) + + extracting.cancel() + done, _ = await asyncio.wait({extracting}, timeout=0.05) + second_cancellation_propagated_early = bool(done) + extraction_release.set() + + with pytest.raises(asyncio.CancelledError): + await extracting + + assert extraction_finished.is_set() + assert first_cancellation_propagated_early is False + assert second_cancellation_propagated_early is False + + @pytest.mark.anyio + async def test_materialize_extracts_squashfs_when_mount_fails(self, temp_cache_dir): + """Test that SquashFS mount failures fall back to unsquashfs extraction.""" + cache = RegistryArtifactCache(temp_cache_dir) + + async def mock_mount(self, ctx, image_path): + raise SquashfsMountCommandError("operation not permitted") + + async def mock_extract(self, ctx, image_path): + target_dir = ctx.paths.squashfs_extract_dir + target_dir.mkdir(parents=True, exist_ok=True) + (target_dir / "module.py").write_text("VALUE = 1") + return target_dir + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + return_value=True, + ), + patch( + "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", + True, + ), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + patch.object(SquashfsArtifact, "extract", mock_extract), + patch.object( + TarballArtifact, + "materialize", + new_callable=AsyncMock, + ) as tarball_materialize, + ): + result = await _materialize( + cache, + "fallback-key", + "s3://bucket/path/site-packages.tar.gz", + ) + + assert len(result) == 1 + assert (result[0] / "module.py").read_text() == "VALUE = 1" + assert result[0].name == "extracted" + tarball_materialize.assert_not_awaited() + + @pytest.mark.anyio + async def test_materialize_extracts_squashfs_without_mount_binary( + self, temp_cache_dir + ): + """Test that SquashFS is still preferred when only unsquashfs is available.""" + cache = RegistryArtifactCache(temp_cache_dir) + + async def mock_extract(self, ctx, image_path): + target_dir = ctx.paths.squashfs_extract_dir + target_dir.mkdir(parents=True, exist_ok=True) + (target_dir / "module.py").write_text("VALUE = 1") + return target_dir + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + return_value=True, + ), + patch( + "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", + True, + ), + patch( + "tracecat.executor.registry_artifacts.shutil.which", return_value=None + ), + patch.object(SquashfsArtifact, "extract", mock_extract), + ): + result = await _materialize( + cache, + "extract-key", + "s3://bucket/path/site-packages.tar.gz", + ) + + assert len(result) == 1 + assert (result[0] / "module.py").read_text() == "VALUE = 1" + assert result[0].name == "extracted" + + @pytest.mark.anyio + async def test_materialize_falls_back_to_gzip_when_squashfs_extract_fails( + self, temp_cache_dir + ): + """Test that legacy gzip remains the final compatibility fallback.""" + cache = RegistryArtifactCache(temp_cache_dir) + source = temp_cache_dir / "source" + source.mkdir() + (source / "module.py").write_text("VALUE = 1") + + async def mock_tarball_download(self, ctx, path): + with tarfile.open(path, "w:gz") as tar: + tar.add(source / "module.py", arcname="module.py") + + async def mock_mount(self, ctx, image_path): + raise SquashfsMountCommandError("operation not permitted") + + async def mock_extract(self, ctx, image_path): + raise RuntimeError("unsquashfs unavailable") + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + side_effect=[True, False], + ), + patch( + "tracecat.executor.registry_artifacts.config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED", + True, + ), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + patch.object(SquashfsArtifact, "extract", mock_extract), + patch.object(TarballArtifact, "download", mock_tarball_download), + ): + result = await _materialize( + cache, + "gzip-fallback-key", + "s3://bucket/path/site-packages.tar.gz", + ) + + assert len(result) == 1 + assert (result[0] / "module.py").read_text() == "VALUE = 1" + assert result[0].name == "tarball" + + @pytest.mark.anyio + async def test_materialize_treats_unknown_suffix_as_gzip(self, temp_cache_dir): + """Test that existing gzip artifacts can use arbitrary S3 key suffixes.""" + cache = RegistryArtifactCache(temp_cache_dir) + source = temp_cache_dir / "source" + source.mkdir() + (source / "module.py").write_text("VALUE = 1") + + async def mock_download(self, ctx, path): + assert path.name.endswith(".tar.gz") + with tarfile.open(path, "w:gz") as tar: + tar.add(source / "module.py", arcname="module.py") + + with patch.object(TarballArtifact, "download", mock_download): + result = await _materialize( + cache, + "custom-key-test", + "s3://bucket/path/custom-key", + ) + + assert len(result) == 1 + assert (result[0] / "module.py").read_text() == "VALUE = 1" + + @pytest.mark.anyio + async def test_materialize_caches_result(self, temp_cache_dir): + """Test that tarball extraction is cached.""" + cache = RegistryArtifactCache(temp_cache_dir) + cache_key = "test-cache-key" + target_dir = cache._paths_for(cache_key).tarball_target_dir + target_dir.mkdir(parents=True) + + result = await _materialize( + cache, + cache_key, + "s3://bucket/test.tar.gz", + ) + + assert result == [target_dir] + + @pytest.mark.anyio + async def test_materialize_concurrent_requests(self, temp_cache_dir): + """Test that concurrent requests for same artifact do not race.""" + cache = RegistryArtifactCache(temp_cache_dir) + cache_key = "concurrent-test" + download_count = 0 + + async def mock_download(self, ctx, path): + nonlocal download_count + download_count += 1 + await asyncio.sleep(0.1) + path.write_bytes(b"fake tarball content") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "extracted.txt").write_text("extracted") + + with ( + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + results = await asyncio.gather( + _materialize(cache, cache_key, "s3://bucket/test.tar.gz"), + _materialize(cache, cache_key, "s3://bucket/test.tar.gz"), + _materialize(cache, cache_key, "s3://bucket/test.tar.gz"), + ) + + assert all(r == results[0] for r in results) + assert download_count == 1 + + def test_runtime_for_same_key(self, temp_cache_dir): + """The same cache key returns the same runtime state.""" + cache = RegistryArtifactCache(temp_cache_dir) + + runtime1 = cache._runtime_for("key1") + runtime2 = cache._runtime_for("key1") + + assert runtime1 is runtime2 + + def test_runtime_for_different_keys(self, temp_cache_dir): + """Different cache keys return different runtime state.""" + cache = RegistryArtifactCache(temp_cache_dir) + + runtime1 = cache._runtime_for("key1") + runtime2 = cache._runtime_for("key2") + + assert runtime1 is not runtime2 + + +class TestRegistryArtifactCacheLease: + """Tests for lease-based pinning of registry artifact cache entries.""" + + @pytest.mark.anyio + async def test_cache_rejects_use_from_a_second_event_loop( + self, temp_cache_dir: Path + ) -> None: + """Protect the process-wide cache from thread-local Temporal loops. + + The cache owns asyncio locks, tasks, and refcounts as one unit. A typed + ownership failure prevents a future synchronous activity from silently + sharing only part of that state across thread-local event loops, which + previously caused intermittent cross-loop failures in storage caches. + """ + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + + def use_cache_from_another_thread() -> None: + asyncio.run(cache.ensure_swept()) + + with pytest.raises(RegistryArtifactCacheLoopError): + await asyncio.to_thread(use_cache_from_another_thread) + + # A rejected caller must not poison the owning loop's cache. + async with cache.lease(None) as registry_paths: + assert registry_paths == [temp_cache_dir / "base"] + + def test_touch_entry_refreshes_tarball_root_mtime(self, temp_cache_dir): + """Touching a tarball-only entry persists its restart-safe recency.""" + cache = RegistryArtifactCache(temp_cache_dir) + cache_key = "tarball-only" + _write_tarball_entry(temp_cache_dir, cache_key) + entry_dir = cache._paths_for(cache_key).entry_dir + os.utime(entry_dir, (100.0, 100.0)) + + cache._touch_entry(cache_key) + + assert entry_dir.stat().st_mtime > 100.0 + + @pytest.mark.anyio + async def test_lease_refcounts_and_touches_image_mtime(self, temp_cache_dir): + """A lease pins its entry and refreshes the restart-safe LRU timestamp.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/leased.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + target_dir = _write_tarball_entry(temp_cache_dir, cache_key) + image_path = _write_image_entry(temp_cache_dir, cache_key, size=16, mtime=100.0) + entry_dir = cache._paths_for(cache_key).entry_dir + + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [target_dir] + assert cache._refcount(cache_key) == 1 + assert entry_dir.stat().st_mtime > 100.0 + + assert cache._refcount(cache_key) == 0 + assert image_path.is_file() + + @pytest.mark.anyio + async def test_lease_releases_refcount_when_materialization_fails( + self, temp_cache_dir + ): + """A failed materialization must not leak a pin or empty cache entry.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/broken.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def mock_download(self, ctx, path): + raise RuntimeError("download failed") + + with patch.object(TarballArtifact, "download", mock_download): + with pytest.raises(RuntimeError, match="download failed"): + async with cache.lease([artifact_uri]): + pass + + assert cache._refcount(cache_key) == 0 + assert not cache._paths_for(cache_key).entry_dir.exists() + assert cache_key not in cache._discover_cache_keys() + + @pytest.mark.anyio + async def test_failed_first_admission_converges_deposited_image( + self, temp_cache_dir: Path + ) -> None: + """A failed first artifact must not strand an over-budget image.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/failed-first.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) + + async def fail_after_deposit( + artifact: SquashfsArtifact, + ctx: RegistryArtifactMaterializationContext, + ) -> list[Path]: + del artifact + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + ctx.paths.squashfs_image_path.write_bytes(b"reusable image") + raise RuntimeError("mount failed") + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 1), + patch.object( + cache, + "_artifact_candidates", + new_callable=AsyncMock, + return_value=[artifact], + ), + patch.object(SquashfsArtifact, "materialize", fail_after_deposit), + ): + with pytest.raises(RuntimeError, match="mount failed"): + async with cache.lease([artifact_uri]): + pass + + assert cache._refcount(cache_key) == 0 + assert not cache._paths_for(cache_key).entry_dir.exists() + assert cache._discover_cache_keys() == set() + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_cancelled_lease_admission_releases_refcount(self, temp_cache_dir): + """Cancellation during candidate lookup must not leak a permanent pin.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/site-packages.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + lookup_started = asyncio.Event() + finish_lookup = asyncio.Event() + + async def blocked_sidecar_lookup(**kwargs): + lookup_started.set() + await finish_lookup.wait() + return False + + async def take_lease() -> None: + async with cache.lease([artifact_uri]): + pass + + with ( + patch(SQUASHFS_ENABLED_CONFIG, True), + patch.object(cache, "_sidecar_exists", blocked_sidecar_lookup), + ): + acquisition = asyncio.create_task(take_lease()) + await lookup_started.wait() + assert cache._refcount(cache_key) == 1 + acquisition.cancel() + with pytest.raises(asyncio.CancelledError): + await acquisition + + assert cache._refcount(cache_key) == 0 + assert not cache._paths_for(cache_key).entry_dir.exists() + + @pytest.mark.anyio + async def test_cancelled_waiter_preserves_existing_same_key_lease( + self, temp_cache_dir: Path + ) -> None: + """A waiter cancelled before admission cannot release another holder.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/path/shared.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + lock = cache._runtime_for(cache_key).lock + + async def take_lease() -> None: + async with cache.lease([artifact_uri]): + pytest.fail("cancelled waiter must not enter the lease context") + + unmount_idle_entry = AsyncMock() + with patch.object(cache, "_unmount_idle_entry", unmount_idle_entry): + async with lock: + cache._acquire_lease(cache_key) + try: + waiter = asyncio.create_task(take_lease()) + await asyncio.sleep(0) + assert not waiter.done() + + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + + assert cache._refcount(cache_key) == 1 + unmount_idle_entry.assert_not_awaited() + finally: + cache._release_lease(cache_key) + + @pytest.mark.anyio + async def test_lease_without_uris_returns_base_pythonpath_dir(self, temp_cache_dir): + """No artifact URIs still yields the base PYTHONPATH directory.""" + cache = RegistryArtifactCache(temp_cache_dir) + + async with cache.lease(None) as registry_paths: + assert registry_paths == [temp_cache_dir / "base"] + assert registry_paths[0].is_dir() + + assert cache._runtime == {} + + @pytest.mark.anyio + async def test_lease_preserves_uri_order(self, temp_cache_dir): + """Multiple artifacts keep their deterministic PYTHONPATH order.""" + cache = RegistryArtifactCache(temp_cache_dir) + uris = ["s3://bucket/first.tar.gz", "s3://bucket/second.tar.gz"] + expected = [ + _write_tarball_entry( + temp_cache_dir, compute_registry_artifact_cache_key(uri) + ) + for uri in uris + ] + + async with cache.lease(uris) as registry_paths: + assert registry_paths == expected + + @pytest.mark.anyio + async def test_overlapping_same_key_leases_share_one_mount_until_final_release( + self, temp_cache_dir: Path + ) -> None: + """Protect refcount and loop-device lifetime under heavy fan-in. + + Every holder must observe one published mount, intermediate releases + must leave that mount usable, and exactly the final release may reclaim + its loop device while retaining the downloaded image for a remount. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/shared.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + harness = _SquashfsMountHarness(cache) + holder_count = 32 + entered = 0 + all_entered = asyncio.Event() + releases = [asyncio.Event() for _ in range(holder_count)] + + async def hold_lease(index: int) -> None: + nonlocal entered + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [ + cache._paths_for(cache_key).squashfs_mount_dir + ] + entered += 1 + if entered == holder_count: + all_entered.set() + await releases[index].wait() + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", harness.mount), + patch.object(SquashfsArtifact, "extract", harness.extract), + patch.object(cache, "_unmount", harness.unmount), + ): + holders = [ + asyncio.create_task(hold_lease(index)) for index in range(holder_count) + ] + await asyncio.wait_for(all_entered.wait(), timeout=5) + + mount_dir = cache._paths_for(cache_key).squashfs_mount_dir + assert cache._refcount(cache_key) == holder_count + assert harness.mount_attempts == [cache_key] + assert harness.extraction_attempts == [] + assert mount_dir in harness.mounted + + for release in releases[:-1]: + release.set() + await asyncio.gather(*holders[:-1]) + + assert cache._refcount(cache_key) == 1 + assert mount_dir in harness.mounted + assert harness.unmounts == [] + + releases[-1].set() + await holders[-1] + + assert cache._refcount(cache_key) == 0 + assert harness.unmounts == [mount_dir] + assert mount_dir not in harness.mounted + assert cache._paths_for(cache_key).squashfs_image_path.is_file() + assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) + assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) + + @pytest.mark.anyio + async def test_cancelling_one_holder_preserves_sibling_leases( + self, temp_cache_dir: Path + ) -> None: + """Protect sibling actions from another holder's cancellation. + + Cancellation must release exactly one pin without unmounting the shared + artifact beneath surviving actions; the last surviving holder remains + solely responsible for loop-device reclamation. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/cancel-one.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + harness = _SquashfsMountHarness(cache) + entered = [asyncio.Event() for _ in range(3)] + releases = [asyncio.Event() for _ in range(3)] + + async def hold_lease(index: int) -> None: + async with cache.lease([artifact_uri]): + entered[index].set() + await releases[index].wait() + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", harness.mount), + patch.object(SquashfsArtifact, "extract", harness.extract), + patch.object(cache, "_unmount", harness.unmount), + ): + holders = [asyncio.create_task(hold_lease(index)) for index in range(3)] + await asyncio.wait_for( + asyncio.gather(*(event.wait() for event in entered)), + timeout=5, + ) + + holders[0].cancel() + with pytest.raises(asyncio.CancelledError): + await holders[0] + + mount_dir = cache._paths_for(cache_key).squashfs_mount_dir + assert cache._refcount(cache_key) == 2 + assert mount_dir in harness.mounted + assert harness.unmounts == [] + + releases[1].set() + await holders[1] + assert cache._refcount(cache_key) == 1 + assert mount_dir in harness.mounted + assert harness.unmounts == [] + + releases[2].set() + await holders[2] + + assert cache._refcount(cache_key) == 0 + assert harness.mount_attempts == [cache_key] + assert harness.unmounts == [mount_dir] + + @pytest.mark.anyio + async def test_repeated_cancellation_finishes_all_lease_cleanup( + self, temp_cache_dir: Path + ) -> None: + """Every unmount and budget pass finishes before cancellation propagates.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uris = [ + "s3://bucket/first.tar.gz", + "s3://bucket/second.tar.gz", + ] + cache_keys = [compute_registry_artifact_cache_key(uri) for uri in artifact_uris] + for cache_key in cache_keys: + _write_tarball_entry(temp_cache_dir, cache_key) + + lease_entered = asyncio.Event() + first_unmount_started = asyncio.Event() + finish_first_unmount = asyncio.Event() + cleanup_calls: list[str] = [] + + async def mock_unmount_idle_entry(cache_key: str) -> None: + cleanup_calls.append(cache_key) + if cache_key == cache_keys[0]: + first_unmount_started.set() + await finish_first_unmount.wait() + + async def mock_converge_cache_budget() -> None: + cleanup_calls.append("converge") + + async def hold_lease() -> None: + async with cache.lease(artifact_uris): + lease_entered.set() + await asyncio.Event().wait() + + with ( + patch.object(cache, "_unmount_idle_entry", mock_unmount_idle_entry), + patch.object(cache, "_converge_cache_budget", mock_converge_cache_budget), + ): + holder = asyncio.create_task(hold_lease()) + await lease_entered.wait() + holder.cancel() + await first_unmount_started.wait() + + holder.cancel() + await asyncio.sleep(0) + assert not holder.done() + + finish_first_unmount.set() + with pytest.raises(asyncio.CancelledError): + await holder + + assert cleanup_calls == [*cache_keys, "converge"] + assert all(cache._refcount(cache_key) == 0 for cache_key in cache_keys) + + @pytest.mark.anyio + async def test_new_lease_racing_final_release_prevents_stale_unmount( + self, temp_cache_dir: Path + ) -> None: + """Protect the zero-refcount-to-unmount handoff from a new acquisition. + + A lease can arrive after the old holder decrements to zero but before + unmount takes the key lock. The lock-time refcount recheck must preserve + the mount for that newcomer instead of returning a path being torn down. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/release-race.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + harness = _SquashfsMountHarness(cache) + mount_dir = harness.seed_mount(cache_key) + first_entered = asyncio.Event() + release_first = asyncio.Event() + release_reached_unmount = asyncio.Event() + allow_unmount_recheck = asyncio.Event() + newcomer_entered = asyncio.Event() + release_newcomer = asyncio.Event() + original_unmount_idle_entry = cache._unmount_idle_entry + unmount_requests = 0 + + async def pause_first_unmount_request(requested_key: str) -> None: + nonlocal unmount_requests + unmount_requests += 1 + if unmount_requests == 1: + release_reached_unmount.set() + await allow_unmount_recheck.wait() + await original_unmount_idle_entry(requested_key) + + async def old_holder() -> None: + async with cache.lease([artifact_uri]): + first_entered.set() + await release_first.wait() + + async def new_holder() -> None: + async with cache.lease([artifact_uri]): + newcomer_entered.set() + await release_newcomer.wait() + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch.object(cache, "_unmount", harness.unmount), + patch.object( + cache, + "_unmount_idle_entry", + pause_first_unmount_request, + ), + ): + old_task = asyncio.create_task(old_holder()) + await first_entered.wait() + release_first.set() + await release_reached_unmount.wait() + assert cache._refcount(cache_key) == 0 + + new_task = asyncio.create_task(new_holder()) + await newcomer_entered.wait() + assert cache._refcount(cache_key) == 1 + + allow_unmount_recheck.set() + await old_task + + assert mount_dir in harness.mounted + assert harness.unmounts == [] + + release_newcomer.set() + await new_task + + assert cache._refcount(cache_key) == 0 + assert harness.unmounts == [mount_dir] + + @pytest.mark.anyio + async def test_partial_multi_artifact_failure_rolls_back_prior_leases( + self, temp_cache_dir: Path + ) -> None: + """Protect sequential multi-artifact admission from partial pin leaks. + + If a middle artifact fails, every earlier acquisition must be released + and unmounted, the failed key must leave no shell, later artifacts must + never be acquired, and the release path must still request convergence. + """ + cache = RegistryArtifactCache(temp_cache_dir) + first_uri = "s3://bucket/first.squashfs" + failed_uri = "s3://bucket/failed.tar.gz" + untouched_uri = "s3://bucket/untouched.tar.gz" + first_key = compute_registry_artifact_cache_key(first_uri) + failed_key = compute_registry_artifact_cache_key(failed_uri) + untouched_key = compute_registry_artifact_cache_key(untouched_uri) + harness = _SquashfsMountHarness(cache) + first_mount = harness.seed_mount(first_key) + untouched_path = _write_tarball_entry(temp_cache_dir, untouched_key) + + async def fail_download( + artifact: TarballArtifact, + ctx: RegistryArtifactMaterializationContext, + output_path: Path, + ) -> None: + del artifact, ctx, output_path + raise RuntimeError("download failed") + + tracked_lease_artifact = AsyncMock(wraps=cache._lease_artifact) + converge_cache_budget = AsyncMock() + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "download", fail_download), + patch.object(cache, "_unmount", harness.unmount), + patch.object(cache, "_lease_artifact", tracked_lease_artifact), + patch.object( + cache, + "_converge_cache_budget", + converge_cache_budget, + ), + ): + with pytest.raises(RuntimeError, match="download failed"): + async with cache.lease([first_uri, failed_uri, untouched_uri]): + pass + + requested_uris = [ + await_call.args[0] for await_call in tracked_lease_artifact.await_args_list + ] + assert requested_uris == [first_uri, failed_uri] + assert cache._refcount(first_key) == 0 + assert cache._refcount(failed_key) == 0 + assert cache._refcount(untouched_key) == 0 + assert harness.unmounts == [first_mount] + assert cache._paths_for(first_key).squashfs_image_path.is_file() + assert not cache._paths_for(failed_key).entry_dir.exists() + assert untouched_path.is_dir() + converge_cache_budget.assert_awaited_once_with() + assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) + assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) + + @pytest.mark.anyio + async def test_waiter_retries_after_first_same_key_materializer_fails( + self, temp_cache_dir: Path + ) -> None: + """Protect same-key waiters from a failed cold-cache publisher. + + The first materializer may fail after writing scratch. Its waiter must + retry against a clean staging area, publish the sole valid entry, and + leave neither a leaked pin nor an empty LRU-visible cache shell. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/retry-after-failure.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + first_download_started = asyncio.Event() + fail_first_download = asyncio.Event() + retry_download_started = asyncio.Event() + allow_retry_download = asyncio.Event() + download_attempts = 0 + retry_saw_clean_staging = False + + async def controlled_download( + artifact: TarballArtifact, + ctx: RegistryArtifactMaterializationContext, + output_path: Path, + ) -> None: + del artifact, ctx + nonlocal download_attempts, retry_saw_clean_staging + download_attempts += 1 + if download_attempts == 1: + output_path.write_bytes(b"partial") + first_download_started.set() + await fail_first_download.wait() + raise RuntimeError("first publisher failed") + + retry_saw_clean_staging = not any(cache.staging_dir.iterdir()) + retry_download_started.set() + await allow_retry_download.wait() + output_path.write_bytes(b"complete") + + async def mock_extract( + artifact: TarballArtifact, + tarball_path: Path, + target_dir: Path, + ) -> None: + del artifact + assert tarball_path.read_bytes() == b"complete" + (target_dir / "module.py").write_text("VALUE = 1") + + async def take_lease() -> list[Path]: + async with cache.lease([artifact_uri]) as registry_paths: + return registry_paths + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch( + "tracecat.executor.registry_artifacts._tarball_extracted_size", + return_value=1, + ), + patch.object(TarballArtifact, "download", controlled_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + first = asyncio.create_task(take_lease()) + await first_download_started.wait() + waiter = asyncio.create_task(take_lease()) + await asyncio.sleep(0) + assert download_attempts == 1 + + fail_first_download.set() + await retry_download_started.wait() + allow_retry_download.set() + registry_paths = await waiter + with pytest.raises(RuntimeError, match="first publisher failed"): + await first + + target_dir = cache._paths_for(cache_key).tarball_target_dir + assert registry_paths == [target_dir] + assert (target_dir / "module.py").read_text() == "VALUE = 1" + assert download_attempts == 2 + assert retry_saw_clean_staging is True + assert cache._discover_cache_keys() == {cache_key} + assert cache._refcount(cache_key) == 0 + assert not any(cache.staging_dir.iterdir()) + assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) + + @pytest.mark.anyio + async def test_duplicate_uri_balances_each_acquisition_and_release( + self, temp_cache_dir: Path + ) -> None: + """Protect list bookkeeping when one lease requests the same URI twice. + + Duplicate PYTHONPATH entries intentionally acquire two pins. Both must + be released, while materialization and final unmount still occur once; + accidental deduplication on only one side would leak or underflow pins. + """ + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/duplicate.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + harness = _SquashfsMountHarness(cache) + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", harness.mount), + patch.object(SquashfsArtifact, "extract", harness.extract), + patch.object(cache, "_unmount", harness.unmount), + ): + async with cache.lease([artifact_uri, artifact_uri]) as registry_paths: + mount_dir = cache._paths_for(cache_key).squashfs_mount_dir + assert registry_paths == [mount_dir, mount_dir] + assert cache._refcount(cache_key) == 2 + assert harness.mount_attempts == [cache_key] + + assert cache._refcount(cache_key) == 0 + assert harness.unmounts == [mount_dir] + assert cache._paths_for(cache_key).squashfs_image_path.is_file() + + @pytest.mark.anyio + async def test_lease_is_never_admitted_across_an_in_flight_eviction( + self, temp_cache_dir + ): + """A lease must not return a mount an in-flight eviction is deleting.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/path/site-packages.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + mounted = {paths.squashfs_mount_dir} + umount_started = asyncio.Event() + finish_umount = asyncio.Event() + remounts: list[str] = [] + + umount_process = AsyncMock() + umount_process.communicate.return_value = (b"", b"") + umount_process.returncode = 0 + + async def mock_umount(*args, **kwargs): + umount_started.set() + await finish_umount.wait() + mounted.discard(paths.squashfs_mount_dir) + return umount_process + + async def mock_mount(self, ctx, image_path): + remounts.append(ctx.cache_key) + target_dir = ctx.paths.squashfs_mount_dir + target_dir.mkdir(parents=True, exist_ok=True) + (target_dir / "module.py").write_text("VALUE = 1") + mounted.add(target_dir) + return target_dir + + leased_paths: list[Path] = [] + leased_path_exists: list[bool] = [] + + async def take_lease() -> None: + async with cache.lease([artifact_uri]) as registry_paths: + leased_paths.extend(registry_paths) + leased_path_exists.append(registry_paths[0].is_dir()) + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/umount", + ), + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + side_effect=mock_umount, + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + ): + eviction = asyncio.create_task(cache._evict_entry(cache_key)) + await umount_started.wait() + lease = asyncio.create_task(take_lease()) + # Let the lease block on the per-key lock the eviction holds. + await asyncio.sleep(0) + finish_umount.set() + evicted, _ = await asyncio.gather(eviction, lease) + + assert evicted == RegistryArtifactEviction(retired=True, reclaimed=True) + # The lease waited for the eviction and re-materialized the entry. + assert remounts == [cache_key] + assert leased_paths == [paths.squashfs_mount_dir] + assert leased_path_exists == [True] + assert (paths.squashfs_mount_dir / "module.py").read_text() == "VALUE = 1" + + @pytest.mark.anyio + async def test_builtin_artifact_is_exempt_from_cache_accounting( + self, temp_cache_dir, monkeypatch: pytest.MonkeyPatch + ): + """The bundled builtin registry is never a cache entry.""" + version = "1.2.3" + site_packages = temp_cache_dir / "venv" / "site-packages" + package_dir = site_packages / "tracecat_registry" + package_dir.mkdir(parents=True) + package_file = package_dir / "__init__.py" + package_file.write_text("") + + monkeypatch.setattr(tracecat_registry, "__version__", version) + monkeypatch.setattr(tracecat_registry, "__file__", str(package_file)) + monkeypatch.setattr( + "tracecat.executor.registry_artifacts.sysconfig.get_path", + lambda name: str(site_packages) if name == "purelib" else None, + ) + + cache = RegistryArtifactCache(temp_cache_dir) + + with patch.object( + cache, + "_enforce_cache_budget", + new_callable=AsyncMock, + ) as enforce_cache_budget: + async with cache.lease([bundled_builtin_registry_uri(version)]) as paths: + assert paths == [site_packages.resolve()] + assert cache._runtime == {} + + enforce_cache_budget.assert_not_awaited() + + +class TestRegistryArtifactCacheEviction: + """Tests for bounded eviction of registry artifact cache entries.""" + + def test_delete_cache_path_reports_directory_failure(self, temp_cache_dir): + """Physical deletion failures are observable instead of suppressed.""" + entry_dir = temp_cache_dir / "entry" + entry_dir.mkdir() + + with ( + patch( + "tracecat.executor.registry_artifacts.shutil.rmtree", + side_effect=OSError("permission denied"), + ), + patch("tracecat.executor.registry_artifacts.logger.warning") as warning, + ): + deleted = _delete_cache_path(entry_dir) + + assert deleted is False + assert entry_dir.is_dir() + warning.assert_called_once_with( + "Failed to delete registry artifact cache path", + path=str(entry_dir), + error="permission denied", + ) + + @pytest.mark.anyio + async def test_leased_entry_survives_eviction_of_idle_entry(self, temp_cache_dir): + """Eviction must never remove an entry a live action is importing from.""" + cache = RegistryArtifactCache(temp_cache_dir) + leased_uri = "s3://bucket/leased.tar.gz" + idle_uri = "s3://bucket/idle.tar.gz" + new_uri = "s3://bucket/new.tar.gz" + leased_dir = _write_tarball_entry( + temp_cache_dir, compute_registry_artifact_cache_key(leased_uri) + ) + idle_dir = _write_tarball_entry( + temp_cache_dir, compute_registry_artifact_cache_key(idle_uri) + ) + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch(MAX_ENTRIES_CONFIG, 2), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + async with cache.lease([leased_uri]): + await _materialize( + cache, compute_registry_artifact_cache_key(new_uri), new_uri + ) + + assert leased_dir.is_dir() + assert not idle_dir.exists() + + @pytest.mark.anyio + async def test_cold_download_reserves_space_before_writing( + self, temp_cache_dir: Path + ) -> None: + """Admission evicts idle bytes before a new download enters staging.""" + cache = RegistryArtifactCache(temp_cache_dir) + idle = _write_image_entry(temp_cache_dir, "idle", size=80, mtime=100.0) + artifact_uri = "s3://bucket/new.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + payload = _tarball_payload(size=32) + max_bytes = len(payload) + 32 + capacity_checked = False + + async def download_file_to_path( + *, + key: str, + bucket: str, + output_path: Path, + max_bytes: int, + ensure_capacity: Callable[[int], Awaitable[None]], + ) -> int: + del key, bucket + nonlocal capacity_checked + assert max_bytes == len(payload) + 32 + await ensure_capacity(len(payload)) + capacity_checked = True + assert not idle.exists() + output_path.write_bytes(payload) + return len(payload) + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, max_bytes), + patch( + "tracecat.executor.registry_artifacts.blob.download_file_to_path", + side_effect=download_file_to_path, + ), + ): + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [ + cache._paths_for(cache_key).tarball_target_dir + ] + + assert capacity_checked is True + assert not idle.exists() + assert (registry_paths[0] / "module.py").read_bytes() == b"x" * 32 + + @pytest.mark.anyio + async def test_compression_heavy_tarball_is_rejected_before_extraction( + self, temp_cache_dir: Path + ) -> None: + """Compressed bytes plus declared extraction cannot exceed the cache cap.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/compression-heavy.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + payload = _tarball_payload(size=4096) + max_bytes = len(payload) + 256 + + async def download_file_to_path( + *, + key: str, + bucket: str, + output_path: Path, + max_bytes: int, + ensure_capacity: Callable[[int], Awaitable[None]], + ) -> int: + del key, bucket + assert max_bytes == len(payload) + 256 + await ensure_capacity(len(payload)) + output_path.write_bytes(payload) + return len(payload) + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, max_bytes), + patch( + "tracecat.executor.registry_artifacts.blob.download_file_to_path", + side_effect=download_file_to_path, + ), + patch.object( + TarballArtifact, + "extract", + new_callable=AsyncMock, + ) as extract, + ): + with pytest.raises(RegistryArtifactCacheCapacityError) as raised: + async with cache.lease([artifact_uri]): + pass + + assert raised.value.additional_bytes == 4096 + assert raised.value.max_bytes == max_bytes + extract.assert_not_awaited() + assert not cache._paths_for(cache_key).entry_dir.exists() + assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) + + @pytest.mark.anyio + async def test_squashfs_expansion_is_rejected_before_extraction( + self, temp_cache_dir: Path + ) -> None: + """SquashFS metadata is accounted before unsquashfs writes scratch.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/oversized.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def download( + self: SquashfsArtifact, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> float: + del self, ctx + image_path.parent.mkdir(parents=True, exist_ok=True) + image_path.write_bytes(b"image") + return 0.0 + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 100), + patch.object( + RegistryArtifactMaterializationContext, + "can_mount_squashfs", + return_value=False, + ), + patch.object(SquashfsArtifact, "download", download), + patch.object( + SquashfsArtifact, + "_squashfs_extracted_size", + new_callable=AsyncMock, + return_value=101, + ), + patch.object( + SquashfsArtifact, + "_extract_image", + new_callable=AsyncMock, + ) as extract_image, + ): + with pytest.raises(RegistryArtifactCacheCapacityError) as raised: + async with cache.lease([artifact_uri]): + pass + + assert raised.value.additional_bytes == 101 + extract_image.assert_not_awaited() + paths = cache._paths_for(cache_key) + assert paths.squashfs_image_path.read_bytes() == b"image" + assert not paths.squashfs_extract_dir.exists() + + @pytest.mark.anyio + async def test_successful_admission_enforces_actual_size_before_yield( + self, temp_cache_dir + ): + """Post-publication enforcement sees the new entry's actual size.""" + cache = RegistryArtifactCache(temp_cache_dir) + idle = _write_image_entry(temp_cache_dir, "idle", size=4096, mtime=100.0) + new_uri = "s3://bucket/new.tar.gz" + new_key = compute_registry_artifact_cache_key(new_uri) + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_bytes(b"x" * 4096) + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 6000), + patch( + "tracecat.executor.registry_artifacts._tarball_extracted_size", + return_value=4096, + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + async with cache.lease([new_uri]) as registry_paths: + assert registry_paths == [cache._paths_for(new_key).tarball_target_dir] + assert not idle.exists() + + assert not idle.exists() + assert cache._paths_for(new_key).tarball_target_dir.is_dir() + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_deletion_failure_does_not_block_materialization( + self, temp_cache_dir + ): + """Failed cleanup is reported while cache admission remains fail-open.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + _write_image_entry(temp_cache_dir, "idle", size=4096, mtime=100.0) + artifact_uri = "s3://bucket/new.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifacts._delete_cache_path", + return_value=False, + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + patch("tracecat.executor.registry_artifacts.logger.warning") as warning, + ): + registry_paths = await _materialize(cache, cache_key, artifact_uri) + + assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] + assert registry_paths[0].is_dir() + assert cache._budget_dirty is True + warning.assert_any_call( + "Registry artifact eviction remains pending physical deletion", + cache_key="idle", + trash_path=ANY, + ) + + @pytest.mark.anyio + async def test_failed_physical_delete_retries_without_extra_eviction( + self, temp_cache_dir + ): + """Failed byte reclamation stops eviction until trash cleanup succeeds.""" + cache = RegistryArtifactCache(temp_cache_dir) + oldest = _write_image_entry( + temp_cache_dir, + "oldest", + size=16, + mtime=100.0, + ) + older = _write_image_entry( + temp_cache_dir, + "older", + size=16, + mtime=200.0, + ) + newest = _write_image_entry( + temp_cache_dir, + "newest", + size=16, + mtime=300.0, + ) + real_delete = _delete_cache_path + failed_once = False + + def fail_once(path: Path) -> bool: + nonlocal failed_once + if not failed_once: + failed_once = True + return False + return real_delete(path) + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 16), + patch( + "tracecat.executor.registry_artifacts._delete_cache_path", + side_effect=fail_once, + ), + ): + assert await cache._enforce_cache_budget() is False + assert not oldest.exists() + assert older.exists() + assert newest.exists() + assert cache._discover_cache_keys() == {"older", "newest"} + assert len(tuple(cache.trash_dir.iterdir())) == 1 + + assert await cache._enforce_cache_budget() is True + + assert not older.exists() + assert newest.exists() + assert not any(cache.trash_dir.iterdir()) + + @pytest.mark.anyio + async def test_rename_failure_does_not_block_materialization(self, temp_cache_dir): + """A failed atomic retirement keeps the old entry and admits new work.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + idle = _write_image_entry(temp_cache_dir, "idle", size=16, mtime=100.0) + artifact_uri = "s3://bucket/new-after-rename-failure.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifacts._move_entry_to_trash", + side_effect=OSError("rename failed"), + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + registry_paths = await _materialize(cache, cache_key, artifact_uri) + + assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] + assert registry_paths[0].is_dir() + assert idle.is_file() + assert cache._budget_dirty is True + + @pytest.mark.anyio + async def test_cache_scan_failure_does_not_block_materialization( + self, temp_cache_dir + ): + """Maintenance errors are observable while artifact admission stays open.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/new-after-scan-failure.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch.object( + cache, + "_enforce_cache_budget", + new_callable=AsyncMock, + side_effect=PermissionError("denied"), + ), + patch( + "tracecat.executor.registry_artifacts._tarball_extracted_size", + return_value=9, + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [ + cache._paths_for(cache_key).tarball_target_dir + ] + + @pytest.mark.anyio + async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( + self, temp_cache_dir + ): + """Steady-state cache hits must not pay for a full cache scan.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/cached.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + _write_tarball_entry(temp_cache_dir, cache_key) + cache._budget_dirty = False + + with patch.object( + cache, + "_scan_cache_entries", + side_effect=AssertionError("cache hits must not scan the cache dir"), + ): + async with cache.lease([artifact_uri]): + pass + + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_failed_cold_admission_keeps_warm_lru(self, temp_cache_dir): + """A missing cold artifact must not evict an existing warm entry.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + warm_uri = "s3://bucket/warm.tar.gz" + warm_key = compute_registry_artifact_cache_key(warm_uri) + warm_dir = _write_tarball_entry(temp_cache_dir, warm_key) + missing_uri = "s3://bucket/missing.tar.gz" + missing_key = compute_registry_artifact_cache_key(missing_uri) + + async def mock_download(self, ctx, path): + raise FileNotFoundError("missing artifact") + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "download", mock_download), + ): + with pytest.raises(FileNotFoundError, match="missing artifact"): + async with cache.lease([missing_uri]): + pytest.fail("failed admission must not yield a lease") + + assert warm_dir.is_dir() + assert not cache._paths_for(missing_key).entry_dir.exists() + + @pytest.mark.anyio + async def test_failed_materialization_rearms_budget_dirty(self, temp_cache_dir): + """A failed materialization may leave a canonical image to evict.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + assert cache._budget_dirty is False + + artifact_uri = "s3://bucket/path/site-packages.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) + + async def mock_materialize(self, ctx): + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + ctx.paths.squashfs_image_path.write_bytes(b"orphaned image") + raise RuntimeError("mount failed") + + with ( + patch.object( + cache, + "_artifact_candidates", + new_callable=AsyncMock, + return_value=[artifact], + ), + patch.object(SquashfsArtifact, "materialize", mock_materialize), + ): + with pytest.raises(RuntimeError, match="mount failed"): + await _materialize(cache, cache_key, artifact_uri) + + assert cache._paths_for(cache_key).squashfs_image_path.is_file() + assert cache._budget_dirty is True + + @pytest.mark.anyio + async def test_materialization_rearms_budget_dirty_consumed_mid_flight( + self, temp_cache_dir + ): + """A convergence pass consuming the signal mid-download cannot unarm it.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + assert cache._budget_dirty is False + + artifact_uri = "s3://bucket/path/site-packages.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) + + async def mock_materialize(self, ctx): + # A concurrent lease release consumes the dirty signal and finishes + # its scan before this attempt lands its canonical bytes. + cache._budget_dirty = False + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + ctx.paths.squashfs_image_path.write_bytes(b"late image") + return [ctx.paths.squashfs_mount_dir] + + with ( + patch.object( + cache, + "_artifact_candidates", + new_callable=AsyncMock, + return_value=[artifact], + ), + patch.object(SquashfsArtifact, "materialize", mock_materialize), + ): + await _materialize(cache, cache_key, artifact_uri) + + assert cache._budget_dirty is True + + @pytest.mark.anyio + async def test_release_keeps_retrying_while_the_cache_stays_over_budget( + self, temp_cache_dir + ): + """A cache that cannot shrink yet must stay marked for re-enforcement.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/pinned.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + _write_image_entry(temp_cache_dir, cache_key, size=4096, mtime=100.0) + _write_tarball_entry(temp_cache_dir, cache_key) + cache._budget_dirty = True + # A second holder keeps the entry pinned past the inner lease. + cache._acquire_lease(cache_key) + + with patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 1): + async with cache.lease([artifact_uri]): + pass + + assert cache._budget_dirty is True + + @pytest.mark.anyio + async def test_convergence_rescans_after_concurrent_materialization( + self, temp_cache_dir + ): + """A materialization during a budget scan must schedule another scan.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/concurrent.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + scan_started = asyncio.Event() + materialized = asyncio.Event() + convergence_scans = 0 + + async def mock_enforce_cache_budget( + *, protected_key: str | None = None + ) -> bool: + nonlocal convergence_scans + if protected_key is not None: + return True + + convergence_scans += 1 + if convergence_scans == 1: + scan_started.set() + await materialized.wait() + return True + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 1") + + cache._budget_dirty = True + with ( + patch.object( + cache, + "_enforce_cache_budget", + side_effect=mock_enforce_cache_budget, + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + convergence = asyncio.create_task(cache._converge_cache_budget()) + await scan_started.wait() + registry_paths = await _materialize(cache, cache_key, artifact_uri) + materialized.set() + await convergence + + assert registry_paths == [cache._paths_for(cache_key).tarball_target_dir] + assert convergence_scans == 2 + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_cancelled_convergence_rearms_budget_dirty( + self, temp_cache_dir + ) -> None: + """Cancellation must preserve the need for a later budget pass.""" + cache = RegistryArtifactCache(temp_cache_dir) + enforcement_started = asyncio.Event() + finish_enforcement = asyncio.Event() + + async def mock_enforce_cache_budget( + *, protected_key: str | None = None + ) -> bool: + del protected_key + enforcement_started.set() + await finish_enforcement.wait() + return True + + cache._budget_dirty = True + with patch.object( + cache, + "_enforce_cache_budget", + side_effect=mock_enforce_cache_budget, + ): + convergence = asyncio.create_task(cache._converge_cache_budget()) + await enforcement_started.wait() + convergence.cancel() + + with pytest.raises(asyncio.CancelledError): + await convergence + + assert cache._budget_dirty is True + + @pytest.mark.parametrize( + "oldest_has_tarball", + [False, True], + ids=["squashfs-only", "tarball-bearing"], + ) + @pytest.mark.anyio + async def test_enforce_budget_evicts_least_recently_used_until_under_max_bytes( + self, temp_cache_dir, oldest_has_tarball: bool + ): + """Size eviction stops once the cache is within budget.""" + cache = RegistryArtifactCache(temp_cache_dir) + oldest = _write_image_entry(temp_cache_dir, "oldest", size=4096, mtime=100.0) + if oldest_has_tarball: + _write_tarball_entry(temp_cache_dir, "oldest") + oldest_entry = cache._paths_for("oldest").entry_dir + os.utime(oldest_entry, (100.0, 100.0)) + older = _write_image_entry(temp_cache_dir, "older", size=4096, mtime=200.0) + newest = _write_image_entry(temp_cache_dir, "newest", size=4096, mtime=300.0) + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 9000), + ): + within_budget = await cache._enforce_cache_budget(protected_key="pending") + + assert within_budget is True + assert not oldest.exists() + assert not cache._paths_for("oldest").tarball_target_dir.exists() + assert older.exists() + assert newest.exists() + + @pytest.mark.anyio + async def test_final_lease_release_unmounts_and_retains_image(self, temp_cache_dir): + """An idle entry releases its loop device without deleting its image.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/site-packages.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + mounted = {paths.squashfs_mount_dir} + + process = AsyncMock() + process.communicate.return_value = (b"", b"") + process.returncode = 0 + + async def mock_umount(*args, **kwargs): + mounted.discard(paths.squashfs_mount_dir) + return process + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/umount", + ), + patch.object( + asyncio, + "create_subprocess_exec", + side_effect=mock_umount, + ), + ): + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [paths.squashfs_mount_dir] + assert paths.squashfs_mount_dir in mounted + + assert paths.squashfs_mount_dir not in mounted + + assert paths.squashfs_image_path.read_bytes() == b"squashfs" + assert paths.squashfs_mount_dir.is_dir() + + @pytest.mark.anyio + async def test_concurrent_budget_passes_only_evict_once(self, temp_cache_dir): + """A waiting budget pass must re-scan after the active pass evicts.""" + cache = RegistryArtifactCache(temp_cache_dir) + oldest = _write_image_entry(temp_cache_dir, "oldest", size=16, mtime=100.0) + retained = _write_image_entry(temp_cache_dir, "retained", size=16, mtime=200.0) + original_scan = cache._scan_cache_entries + first_scan_started = threading.Event() + second_scan_started = threading.Event() + release_first_scan = threading.Event() + release_second_scan = threading.Event() + scan_count_lock = threading.Lock() + scan_count = 0 + + def controlled_scan(): + nonlocal scan_count + entries = original_scan() + with scan_count_lock: + scan_index = scan_count + scan_count += 1 + if scan_index == 0: + first_scan_started.set() + release_first_scan.wait(timeout=5) + elif scan_index == 1: + second_scan_started.set() + release_second_scan.wait(timeout=5) + return entries + + eviction_started = asyncio.Event() + finish_eviction = asyncio.Event() + extra_eviction_finished = asyncio.Event() + evicted_keys: list[str] = [] + + async def controlled_evict( + cache_key: str, + ) -> RegistryArtifactEviction: + if cache_key == "oldest": + if eviction_started.is_set(): + return RegistryArtifactEviction( + retired=False, + reclaimed=False, + ) + eviction_started.set() + await finish_eviction.wait() + _delete_cache_path(cache._paths_for("oldest").entry_dir) + else: + _delete_cache_path(cache._paths_for("retained").entry_dir) + extra_eviction_finished.set() + evicted_keys.append(cache_key) + return RegistryArtifactEviction(retired=True, reclaimed=True) + + with ( + patch.object(cache, "_scan_cache_entries", side_effect=controlled_scan), + patch.object(cache, "_evict_entry", side_effect=controlled_evict), + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + first_pass = asyncio.create_task(cache._enforce_cache_budget()) + assert await asyncio.to_thread(first_scan_started.wait, 1) + second_pass = asyncio.create_task(cache._enforce_cache_budget()) + second_scan_overlapped = await asyncio.to_thread( + second_scan_started.wait, 0.5 + ) + + release_first_scan.set() + await asyncio.wait_for(eviction_started.wait(), timeout=1) + release_second_scan.set() + try: + await asyncio.wait_for(extra_eviction_finished.wait(), timeout=0.2) + except TimeoutError: + pass + finish_eviction.set() + await asyncio.gather(first_pass, second_pass) + + assert second_scan_overlapped is False + assert scan_count == 2 + assert evicted_keys == ["oldest"] + assert not oldest.exists() + assert retained.exists() + + @pytest.mark.anyio + async def test_enforce_budget_ignores_missing_protected_key(self, temp_cache_dir): + """Only successfully published entries count against the entry budget.""" + cache = RegistryArtifactCache(temp_cache_dir) + existing = _write_image_entry(temp_cache_dir, "existing", size=16, mtime=100.0) + + with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): + within_budget = await cache._enforce_cache_budget(protected_key="missing") + + assert within_budget is True + assert existing.exists() + + @pytest.mark.anyio + async def test_enforce_budget_never_evicts_the_protected_key(self, temp_cache_dir): + """The newly materialized key is exempt even when it is the LRU entry.""" + cache = RegistryArtifactCache(temp_cache_dir) + protected = _write_image_entry( + temp_cache_dir, "protected", size=4096, mtime=100.0 + ) + other = _write_image_entry(temp_cache_dir, "other", size=4096, mtime=300.0) + + with patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0): + await cache._enforce_cache_budget(protected_key="protected") + + assert protected.exists() + assert not other.exists() + + @pytest.mark.anyio + async def test_enforce_budget_proceeds_over_budget_when_everything_is_leased( + self, temp_cache_dir + ): + """An over-budget cache must degrade, not fail the action.""" + cache = RegistryArtifactCache(temp_cache_dir) + leased = _write_image_entry(temp_cache_dir, "leased", size=4096, mtime=100.0) + cache._acquire_lease("leased") + + with patch(MAX_ENTRIES_CONFIG, 0), patch(MAX_BYTES_CONFIG, 1): + within_budget = await cache._enforce_cache_budget(protected_key="missing") + + assert within_budget is False + assert leased.exists() + + @pytest.mark.anyio + async def test_eviction_unmounts_before_deleting_the_image(self, temp_cache_dir): + """Unlinking a mounted image would strand an open-file zombie.""" + cache = RegistryArtifactCache(temp_cache_dir) + paths = cache._paths_for("mounted") + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + mounted = {paths.squashfs_mount_dir} + image_present_at_umount: list[bool] = [] + + process = AsyncMock() + process.communicate.return_value = (b"", b"") + process.returncode = 0 + + async def mock_umount(*args, **kwargs): + image_present_at_umount.append(paths.squashfs_image_path.exists()) + mounted.discard(paths.squashfs_mount_dir) + return process + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/umount", + ), + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + side_effect=mock_umount, + ) as create_subprocess_exec, + ): + evicted = await cache._evict_entry("mounted") + + assert evicted == RegistryArtifactEviction(retired=True, reclaimed=True) + assert image_present_at_umount == [True] + assert not paths.squashfs_image_path.exists() + assert not paths.squashfs_mount_dir.exists() + create_subprocess_exec.assert_called_once() + assert create_subprocess_exec.call_args.args == ( + "/sbin/umount", + str(paths.squashfs_mount_dir), + ) + + @pytest.mark.anyio + async def test_cancelled_unmount_kills_and_reaps_before_releasing_key_lock( + self, temp_cache_dir + ): + """Cancellation leaves a consistent entry for the next admission.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/cancelled-unmount.squashfs" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") + mounted = {paths.squashfs_mount_dir} + blocked_process = _BlockingSubprocess() + released_process = AsyncMock() + released_process.communicate.return_value = (b"", b"") + released_process.returncode = 0 + unmount_attempts = 0 + + async def mock_umount(*args, **kwargs): + nonlocal unmount_attempts + unmount_attempts += 1 + if unmount_attempts == 1: + return blocked_process + mounted.discard(paths.squashfs_mount_dir) + return released_process + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/umount", + ), + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + side_effect=mock_umount, + ), + ): + eviction = asyncio.create_task(cache._evict_entry(cache_key)) + await blocked_process.communicate_started.wait() + eviction.cancel() + + with pytest.raises(asyncio.CancelledError): + await eviction + + assert blocked_process.cleanup_calls == ["kill", "wait"] + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [paths.squashfs_mount_dir] + assert registry_paths[0].is_dir() + assert (registry_paths[0] / "module.py").read_text() == "VALUE = 1" + + assert paths.squashfs_image_path.is_file() + assert paths.squashfs_mount_dir.is_dir() + assert paths.squashfs_mount_dir not in mounted + + @pytest.mark.anyio + async def test_cancelled_background_deletion_leaves_a_clean_miss( + self, temp_cache_dir + ): + """Cancellation cannot expose live paths that deletion still owns.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/cancelled-eviction.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + original_target = _write_tarball_entry(temp_cache_dir, cache_key) + delete_started = threading.Event() + finish_delete = threading.Event() + delete_finished = threading.Event() + doomed: list[Path] = [] + + def blocked_delete(path: Path) -> bool: + doomed.append(path) + delete_started.set() + finish_delete.wait(timeout=5) + deleted = _delete_cache_path(path) + delete_finished.set() + return deleted + + async def mock_download(self, ctx, path): + path.write_bytes(b"fake tarball") + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_text("VALUE = 2") + + with ( + patch( + "tracecat.executor.registry_artifacts._delete_cache_path", + side_effect=blocked_delete, + ), + patch( + "tracecat.executor.registry_artifacts._tarball_extracted_size", + return_value=9, + ), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + eviction = asyncio.create_task(cache._evict_entry(cache_key)) + assert await asyncio.to_thread(delete_started.wait, 1) + eviction.cancel() + try: + await asyncio.sleep(0) + eviction.cancel() + await asyncio.sleep(0) + assert not eviction.done() + assert not original_target.exists() + assert (doomed[0] / "tarball").is_dir() + assert cache._discover_cache_keys() == set() + finally: + finish_delete.set() + + with pytest.raises(asyncio.CancelledError): + await eviction + assert await asyncio.to_thread(delete_finished.wait, 1) + assert not doomed[0].exists() + + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [original_target] + assert (original_target / "module.py").read_text() == "VALUE = 2" + + assert original_target.is_dir() + + @pytest.mark.anyio + async def test_doomed_eviction_names_are_startup_scratch(self, temp_cache_dir): + """Every renamed entry root is invisible and reclaimed on startup.""" + cache = RegistryArtifactCache(temp_cache_dir) + cache_key = "squashfs-doomed" + paths = cache._paths_for(cache_key) + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + paths.squashfs_extract_dir.mkdir() + paths.tarball_target_dir.mkdir() + + with patch( + "tracecat.executor.registry_artifacts._delete_cache_path", + return_value=True, + ) as delete_cache_path: + assert await cache._evict_entry(cache_key) == RegistryArtifactEviction( + retired=True, reclaimed=True + ) + + delete_cache_path.assert_called_once() + trash_path = delete_cache_path.call_args.args[0] + assert trash_path.parent == cache.trash_dir + assert trash_path.exists() + assert cache._discover_cache_keys() == set() + + cache._sweep_startup_state() + + assert not trash_path.exists() + + @pytest.mark.anyio + async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): + """A failed unmount skips the key instead of forcing a lazy detach.""" + cache = RegistryArtifactCache(temp_cache_dir) + stuck = cache._paths_for("stuck") + stuck.entry_dir.mkdir(parents=True) + stuck.squashfs_image_path.write_bytes(b"squashfs") + stuck.squashfs_mount_dir.mkdir() + os.utime(stuck.squashfs_image_path, (100.0, 100.0)) + idle = _write_image_entry(temp_cache_dir, "idle", size=16, mtime=300.0) + mounted = {stuck.squashfs_mount_dir} + + process = AsyncMock() + process.communicate.return_value = (b"", b"target is busy") + process.returncode = 32 + + with ( + patch.object(Path, "is_mount", lambda self: self in mounted), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/umount", + ), + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ), + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + await cache._enforce_cache_budget(protected_key="pending") + + assert stuck.squashfs_image_path.exists() + assert stuck.squashfs_mount_dir.exists() + assert not idle.exists() + + @pytest.mark.anyio + async def test_eviction_keeps_stable_runtime_state(self, temp_cache_dir): + """A key keeps one lock and zeroed lease state for the process lifetime.""" + cache = RegistryArtifactCache(temp_cache_dir) + _write_tarball_entry(temp_cache_dir, "bookkeeping") + cache._acquire_lease("bookkeeping") + cache._release_lease("bookkeeping") + lock = cache._runtime_for("bookkeeping").lock + + assert await cache._evict_entry("bookkeeping") == RegistryArtifactEviction( + retired=True, reclaimed=True + ) + runtime = cache._runtime["bookkeeping"] + assert runtime.lock is lock + assert runtime.refcount == 0 + + @pytest.mark.anyio + async def test_eviction_skips_busy_key(self, temp_cache_dir): + """A key another task is materializing is never evicted underneath it.""" + cache = RegistryArtifactCache(temp_cache_dir) + target_dir = _write_tarball_entry(temp_cache_dir, "busy") + lock = cache._runtime_for("busy").lock + + async with lock: + assert await cache._evict_entry("busy") == RegistryArtifactEviction( + retired=False, reclaimed=False + ) + + assert target_dir.is_dir() + + +class TestRegistryArtifactCacheStartupSweep: + """Tests for the startup sweep that reclaims state from a dead process.""" + + @pytest.mark.anyio + async def test_sweep_uses_entry_root_mtime_for_restart_safe_lru( + self, temp_cache_dir + ): + """Startup trimming preserves a touched older tarball-only entry.""" + previous_cache = RegistryArtifactCache(temp_cache_dir) + old = _write_tarball_entry(temp_cache_dir, "old") + previous_cache._touch_entry("old") + + old_mtime = previous_cache._paths_for("old").entry_dir.stat().st_mtime + new = _write_tarball_entry(temp_cache_dir, "new") + new_entry_dir = previous_cache._paths_for("new").entry_dir + os.utime(new_entry_dir, (old_mtime - 1, old_mtime - 1)) + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + + assert old.is_dir() + assert not new.exists() + + @pytest.mark.anyio + async def test_sweep_tolerates_missing_cache_dir(self, temp_cache_dir): + """A cache directory that does not exist yet is a no-op.""" + cache_dir = temp_cache_dir / "missing" + + cache = RegistryArtifactCache(cache_dir) + await cache.ensure_swept() + + assert cache.cache_dir == cache_dir + assert not cache_dir.exists() + + @pytest.mark.anyio + async def test_sweep_removes_orphaned_work(self, temp_cache_dir): + """Interrupted work is reclaimed without touching unrelated paths.""" + cache = RegistryArtifactCache(temp_cache_dir) + orphaned = cache.staging_dir / "abc123.999999.4321.squashfs" + orphaned.parent.mkdir() + orphaned.write_bytes(b"partial") + orphaned_dir = cache.trash_dir / "abc123.999999.4321" + orphaned_dir.mkdir(parents=True) + unrelated = temp_cache_dir / "unrelated" + unrelated.mkdir() + unrelated_file = unrelated / "keep.txt" + unrelated_file.write_text("keep") + entry_dir = _write_tarball_entry(temp_cache_dir, "abc123") + + await cache.ensure_swept() + + assert not orphaned.exists() + assert not orphaned_dir.exists() + assert unrelated_file.read_text() == "keep" + assert entry_dir.is_dir() + + @pytest.mark.anyio + async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): + """A live mountpoint belongs to a running process and must survive.""" + cache = RegistryArtifactCache(temp_cache_dir) + paths = cache._paths_for("abc123") + paths.entry_dir.mkdir(parents=True) + mount_dir = paths.squashfs_mount_dir + mount_dir.mkdir() + + with patch.object(Path, "is_mount", lambda self: self == mount_dir): + await cache.ensure_swept() + + assert mount_dir.is_dir() + + @pytest.mark.anyio + async def test_ensure_swept_runs_once(self, temp_cache_dir): + """Repeated startup-sweep requests invoke the sweep once.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with patch.object( + cache, + "_sweep_startup_state", + wraps=cache._sweep_startup_state, + ) as sweep: + await cache.ensure_swept() + await cache.ensure_swept() + + assert sweep.call_count == 1 + + @pytest.mark.anyio + async def test_concurrent_ensure_swept_runs_once(self, temp_cache_dir): + """Concurrent startup-sweep requests invoke the sweep once.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with patch.object( + cache, + "_sweep_startup_state", + wraps=cache._sweep_startup_state, + ) as sweep: + await asyncio.gather(*(cache.ensure_swept() for _ in range(4))) + + assert sweep.call_count == 1 + + @pytest.mark.anyio + async def test_cancelled_waiter_joins_same_startup_sweep(self, temp_cache_dir): + """A cancelled waiter cannot start a second concurrent startup sweep.""" + cache = RegistryArtifactCache(temp_cache_dir) + sweep_started = threading.Event() + sweep_release = threading.Event() + invocation_count = 0 + + def blocking_sweep() -> None: + nonlocal invocation_count + invocation_count += 1 + sweep_started.set() + sweep_release.wait() + + with patch.object( + cache, + "_sweep_startup_state", + side_effect=blocking_sweep, + ): + first_waiter = asyncio.create_task(cache.ensure_swept()) + assert await asyncio.to_thread(sweep_started.wait, 1) + first_waiter.cancel() + + with pytest.raises(asyncio.CancelledError): + await first_waiter + + second_waiter = asyncio.create_task(cache.ensure_swept()) + await asyncio.sleep(0) + sweep_release.set() + await second_waiter + + assert invocation_count == 1 + assert cache._swept is True + + @pytest.mark.anyio + async def test_lease_triggers_startup_sweep(self, temp_cache_dir): + """Lease admission reclaims startup scratch before yielding paths.""" + cache = RegistryArtifactCache(temp_cache_dir) + orphaned_dir = cache.staging_dir / "abc123.999999.4321" + orphaned_dir.mkdir(parents=True) + + async with cache.lease(None): + assert not orphaned_dir.exists() + + @pytest.mark.anyio + async def test_failed_ensure_swept_retries(self, temp_cache_dir): + """A failed startup sweep is retried by the next caller.""" + cache = RegistryArtifactCache(temp_cache_dir) + orphaned_dir = cache.staging_dir / "abc123.999999.4321" + orphaned_dir.mkdir(parents=True) + + with ( + patch.object( + cache, + "_clear_work_dir", + side_effect=OSError("simulated sweep failure"), + ), + pytest.raises(OSError), + ): + await cache.ensure_swept() + + assert orphaned_dir.is_dir() + + await cache.ensure_swept() + + assert not orphaned_dir.exists() + + @pytest.mark.parametrize("work_dir_name", ["staging", "trash"]) + def test_work_directory_inspection_errors_propagate( + self, + temp_cache_dir, + work_dir_name: str, + ): + """Unreadable work directories must trigger a later cleanup retry.""" + cache = RegistryArtifactCache(temp_cache_dir) + work_dir = getattr(cache, f"{work_dir_name}_dir") + work_dir.mkdir() + + with ( + patch.object(Path, "iterdir", side_effect=PermissionError("denied")), + pytest.raises(PermissionError, match="denied"), + ): + cache._clear_work_dir(work_dir) + + @pytest.mark.anyio + async def test_active_entry_scan_errors_propagate(self, temp_cache_dir): + """Unreadable active entries cannot be treated as an empty cache.""" + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch( + "tracecat.executor.registry_artifacts.os.scandir", + side_effect=PermissionError("denied"), + ), + pytest.raises(PermissionError, match="denied"), + ): + await cache._enforce_cache_budget() + + @pytest.mark.anyio + async def test_failed_trash_cleanup_skips_startup_trimming(self, temp_cache_dir): + """A failed orphan deletion must not cascade into active retirements.""" + cache = RegistryArtifactCache(temp_cache_dir) + orphan = cache.trash_dir / "orphan" + orphan.mkdir(parents=True) + oldest = _write_image_entry( + temp_cache_dir, + "oldest", + size=16, + mtime=100.0, + ) + newest = _write_image_entry( + temp_cache_dir, + "newest", + size=16, + mtime=200.0, + ) + real_delete = _delete_cache_path + + def fail_orphan(path: Path) -> bool: + if path == orphan: + return False + return real_delete(path) + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifacts._delete_cache_path", + side_effect=fail_orphan, + ), + ): + await cache.ensure_swept() + + assert orphan.is_dir() + assert oldest.is_file() + assert newest.is_file() + assert cache._budget_dirty is True + + @pytest.mark.anyio + async def test_failed_startup_cleanup_retries_exact_path(self, temp_cache_dir): + """Failed startup scratch cleanup retries without sweeping live staging.""" + cache = RegistryArtifactCache(temp_cache_dir) + orphaned = cache.staging_dir / "orphaned.123.456.tmp" + orphaned.parent.mkdir(parents=True) + orphaned.write_bytes(b"partial") + real_delete = _delete_cache_path + failed_once = False + + def fail_once(path: Path) -> bool: + nonlocal failed_once + if path == orphaned and not failed_once: + failed_once = True + return False + return real_delete(path) + + with patch( + "tracecat.executor.registry_artifacts._delete_cache_path", + side_effect=fail_once, + ): + await cache.ensure_swept() + assert orphaned.is_file() + assert cache._failed_startup_cleanup == {orphaned} + assert cache._budget_dirty is True + + assert await cache._enforce_cache_budget() is True + + assert not orphaned.exists() + assert cache._failed_startup_cleanup == set() + + @pytest.mark.anyio + async def test_failed_startup_retirement_stays_dirty_and_retries( + self, temp_cache_dir + ): + """A startup rename failure preserves entries for later enforcement.""" + oldest = _write_image_entry( + temp_cache_dir, + "oldest", + size=16, + mtime=100.0, + ) + newest = _write_image_entry( + temp_cache_dir, + "newest", + size=16, + mtime=200.0, + ) + cache = RegistryArtifactCache(temp_cache_dir) + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + patch( + "tracecat.executor.registry_artifacts._move_entry_to_trash", + side_effect=OSError("rename failed"), + ), + ): + await cache.ensure_swept() + + assert oldest.is_file() + assert newest.is_file() + assert cache._budget_dirty is True + + with ( + patch(MAX_ENTRIES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, 0), + ): + await cache._converge_cache_budget() + + assert not oldest.exists() + assert newest.is_file() + assert cache._budget_dirty is False + + @pytest.mark.anyio + async def test_failed_startup_physical_delete_retries_exact_path( + self, temp_cache_dir + ): + """Startup stops retiring entries when bytes were not reclaimed.""" + oldest = _write_image_entry( + temp_cache_dir, + "oldest", + size=16, + mtime=100.0, + ) + older = _write_image_entry( + temp_cache_dir, + "older", + size=16, + mtime=200.0, + ) + newest = _write_image_entry( + temp_cache_dir, + "newest", + size=16, + mtime=300.0, + ) + cache = RegistryArtifactCache(temp_cache_dir) + real_delete = _delete_cache_path + failed_once = False + + def fail_once(path: Path) -> bool: + nonlocal failed_once + if path.parent == cache.trash_dir and not failed_once: + failed_once = True + return False + return real_delete(path) + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, 16), + patch( + "tracecat.executor.registry_artifacts._delete_cache_path", + side_effect=fail_once, + ), + ): + await cache.ensure_swept() + assert not oldest.exists() + assert older.is_file() + assert newest.is_file() + assert len(tuple(cache.trash_dir.iterdir())) == 1 + assert cache._budget_dirty is True + + await cache._converge_cache_budget() + + assert not older.exists() + assert newest.is_file() + assert not any(cache.trash_dir.iterdir()) + assert cache._budget_dirty is False + + +class TestSquashfsMountPolicy: + """Tests for per-artifact SquashFS fallback.""" + + @pytest.mark.anyio + async def test_loop_device_exhaustion_isolated_sticky_extraction_fallback( + self, temp_cache_dir: Path + ) -> None: + """Protect the cache's fail-open policy when loop devices are saturated. + + A mount-command failure must extract only the affected cold artifact, + preserve already-leased mounts, reuse that extraction on later leases, + and still let unrelated artifacts attempt mounting. This models loop + exhaustion deterministically without consuming host-global devices. + """ + cache = RegistryArtifactCache(temp_cache_dir) + held_uri = "s3://bucket/already-mounted.squashfs" + saturated_uri = "s3://bucket/no-loop-available.squashfs" + later_uri = "s3://bucket/later-artifact.squashfs" + held_key = compute_registry_artifact_cache_key(held_uri) + saturated_key = compute_registry_artifact_cache_key(saturated_uri) + later_key = compute_registry_artifact_cache_key(later_uri) + harness = _SquashfsMountHarness( + cache, + failed_mount_keys={saturated_key}, + ) + + with ( + patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", harness.mount), + patch.object(SquashfsArtifact, "extract", harness.extract), + patch.object(cache, "_unmount", harness.unmount), + ): + async with cache.lease([held_uri]) as held_paths: + held_mount = cache._paths_for(held_key).squashfs_mount_dir + assert held_paths == [held_mount] + + async with cache.lease([saturated_uri]) as saturated_paths: + extracted = cache._paths_for(saturated_key).squashfs_extract_dir + assert saturated_paths == [extracted] + assert held_mount in harness.mounted + assert cache._refcount(held_key) == 1 + assert harness.unmounts == [] + + # The extracted directory is the cache entry's sticky path; + # freeing a loop elsewhere does not trigger a remount attempt. + async with cache.lease([saturated_uri]) as cached_paths: + assert cached_paths == [extracted] + + async with cache.lease([later_uri]) as later_paths: + later_mount = cache._paths_for(later_key).squashfs_mount_dir + assert later_paths == [later_mount] + assert held_mount in harness.mounted + + assert held_mount in harness.mounted + + assert harness.mount_attempts == [held_key, saturated_key, later_key] + assert harness.extraction_attempts == [saturated_key] + assert harness.unmounts == [later_mount, held_mount] + assert cache._paths_for(saturated_key).squashfs_image_path.is_file() + assert cache._paths_for(saturated_key).squashfs_extract_dir.is_dir() + assert cache._refcount(held_key) == 0 + assert cache._refcount(saturated_key) == 0 + assert cache._refcount(later_key) == 0 + + @pytest.mark.anyio + async def test_mount_failure_does_not_disable_later_artifacts( + self, temp_cache_dir + ) -> None: + """One failed mount falls back without changing later mount attempts.""" + cache = RegistryArtifactCache(temp_cache_dir) + first_ctx = cache._context_for("first") + second_ctx = cache._context_for("second") + first = SquashfsArtifact( + uri="s3://bucket/first.squashfs", + cache_key="first", + ) + second = SquashfsArtifact( + uri="s3://bucket/second.squashfs", + cache_key="second", + ) + mount_attempts: list[str] = [] + + async def mock_mount(self, ctx, image_path): + del image_path + mount_attempts.append(ctx.cache_key) + if ctx.cache_key == "first": + raise SquashfsMountCommandError("operation not permitted") + ctx.paths.squashfs_mount_dir.mkdir(parents=True) + return ctx.paths.squashfs_mount_dir + + async def mock_extract(self, ctx, image_path): + del image_path + ctx.paths.squashfs_extract_dir.mkdir(parents=True) + return ctx.paths.squashfs_extract_dir + + with ( + patch(SQUASHFS_ENABLED_CONFIG, True), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/mount", + ), + patch.object(SquashfsArtifact, "mount", mock_mount), + patch.object(SquashfsArtifact, "extract", mock_extract), + ): + assert await first.materialize(first_ctx) == [ + first_ctx.paths.squashfs_extract_dir + ] + assert await second.materialize(second_ctx) == [ + second_ctx.paths.squashfs_mount_dir + ] + + assert mount_attempts == ["first", "second"] diff --git a/tests/unit/test_sandbox_utils.py b/tests/unit/test_sandbox_utils.py deleted file mode 100644 index d2f1fbdc05..0000000000 --- a/tests/unit/test_sandbox_utils.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Tests for shared sandbox process utilities.""" - -from __future__ import annotations - -import asyncio -from typing import cast -from unittest.mock import patch - -import pytest - -from tracecat.sandbox.utils import communicate_process_group - - -class _BlockingProcess: - """Minimal process double whose communication blocks until cancellation.""" - - def __init__(self) -> None: - self.returncode: int | None = None - self.communicate_started = asyncio.Event() - self.communicate_finished = asyncio.Event() - - async def communicate( - self, - input: bytes | None = None, # noqa: A002 - ) -> tuple[bytes, bytes]: - del input - self.communicate_started.set() - try: - await asyncio.Event().wait() - finally: - self.communicate_finished.set() - return b"", b"" - - -@pytest.mark.anyio -async def test_repeated_cancellation_rejoins_process_group_cleanup() -> None: - """A second cancellation cannot return while group termination is live.""" - fake_process = _BlockingProcess() - process = cast(asyncio.subprocess.Process, fake_process) - termination_started = asyncio.Event() - finish_termination = asyncio.Event() - termination_finished = asyncio.Event() - - async def blocking_termination( - requested_process: asyncio.subprocess.Process, - ) -> None: - assert requested_process is process - termination_started.set() - await finish_termination.wait() - termination_finished.set() - - with patch( - "tracecat.sandbox.utils.terminate_process_group", - side_effect=blocking_termination, - ): - communication = asyncio.create_task(communicate_process_group(process)) - await fake_process.communicate_started.wait() - communication.cancel() - await termination_started.wait() - - communication.cancel() - await asyncio.sleep(0) - assert not communication.done() - - finish_termination.set() - with pytest.raises(asyncio.CancelledError): - await communication - - assert termination_finished.is_set() - assert fake_process.communicate_finished.is_set() - - -@pytest.mark.anyio -async def test_cleanup_failure_preserves_caller_cancellation() -> None: - """A failed process cleanup remains context for the caller cancellation.""" - fake_process = _BlockingProcess() - process = cast(asyncio.subprocess.Process, fake_process) - termination_started = asyncio.Event() - finish_termination = asyncio.Event() - - async def failing_termination( - requested_process: asyncio.subprocess.Process, - ) -> None: - assert requested_process is process - termination_started.set() - await finish_termination.wait() - raise RuntimeError("process cleanup failed") - - with patch( - "tracecat.sandbox.utils.terminate_process_group", - side_effect=failing_termination, - ): - communication = asyncio.create_task(communicate_process_group(process)) - await fake_process.communicate_started.wait() - communication.cancel() - await termination_started.wait() - finish_termination.set() - - with pytest.raises(asyncio.CancelledError) as raised: - await communication - - assert isinstance(raised.value.__cause__, RuntimeError) - assert fake_process.communicate_finished.is_set() diff --git a/tests/unit/test_storage_blob.py b/tests/unit/test_storage_blob.py index 2885151cb4..c960a2cf5b 100644 --- a/tests/unit/test_storage_blob.py +++ b/tests/unit/test_storage_blob.py @@ -2,14 +2,13 @@ import asyncio import hashlib -import threading from contextlib import asynccontextmanager from pathlib import Path from unittest.mock import AsyncMock, patch from urllib.parse import urlparse import pytest -from botocore.exceptions import ClientError, EndpointConnectionError +from botocore.exceptions import ClientError from tracecat.storage import blob as blob_module from tracecat.storage.blob import ( @@ -705,81 +704,6 @@ async def test_open_download_stream_yields_stream_and_length(self, mock_get_clie assert stream is mock_body assert length == 123 - @pytest.mark.anyio - @patch("tracecat.storage.blob.get_storage_client") - async def test_open_download_stream_redacts_sensitive_client_error( - self, mock_get_client - ): - mock_client = AsyncMock() - mock_get_client.return_value.__aenter__.return_value = mock_client - mock_client.get_object.side_effect = ClientError( - error_response={ - "Error": { - "Code": "AccessDenied", - "Message": "denied tenant-key/org/repository", - } - }, - operation_name="get_object", - ) - - with ( - patch("tracecat.storage.blob.logger.error") as log_error, - pytest.raises(blob_module.StorageDownloadError) as exc_info, - ): - async with open_download_stream( - key="tenant-key/org/repository", - bucket="tenant-bucket", - redact_log_identifiers=True, - ): - pass - - assert str(exc_info.value) == "Storage download failed" - assert exc_info.value.error_code == "AccessDenied" - assert exc_info.value.__cause__ is None - log_error.assert_called_once_with( - "Failed to open download stream", - key="", - bucket="", - error_code="AccessDenied", - error_type="ClientError", - ) - - @pytest.mark.anyio - @patch("tracecat.storage.blob.get_storage_client") - async def test_open_download_stream_redacts_sensitive_transport_error( - self, - mock_get_client, - ) -> None: - mock_client = AsyncMock() - mock_get_client.return_value.__aenter__.return_value = mock_client - sensitive_endpoint = "https://tenant-bucket.invalid/tenant-key/org/repository" - mock_client.get_object.side_effect = EndpointConnectionError( - endpoint_url=sensitive_endpoint - ) - - with ( - patch("tracecat.storage.blob.logger.error") as log_error, - pytest.raises(blob_module.StorageDownloadError) as exc_info, - ): - async with open_download_stream( - key="tenant-key/org/repository", - bucket="tenant-bucket", - redact_log_identifiers=True, - ): - pass - - assert str(exc_info.value) == "Storage download failed" - assert exc_info.value.error_code is None - assert exc_info.value.__cause__ is None - log_error.assert_called_once_with( - "Failed to open download stream", - key="", - bucket="", - error_code=None, - error_type="EndpointConnectionError", - ) - assert sensitive_endpoint not in repr(log_error.call_args) - @pytest.mark.anyio async def test_download_file_to_path_writes_bytes( self, tmp_path: Path, monkeypatch @@ -816,85 +740,6 @@ async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 assert bytes_written == 11 assert out.read_bytes() == b"hello world" - @pytest.mark.anyio - async def test_download_file_to_path_redacts_sensitive_success_log( - self, tmp_path: Path, monkeypatch - ): - class DummyStream: - async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 - yield b"payload" - - @asynccontextmanager - async def _fake_open_download_stream( - *, - key: str, - bucket: str, - redact_log_identifiers: bool, - ): - assert key == "tenant-key/org/repository" - assert bucket == "tenant-bucket" - assert redact_log_identifiers is True - yield DummyStream(), 7 - - monkeypatch.setattr( - "tracecat.storage.blob.open_download_stream", - _fake_open_download_stream, - ) - out = tmp_path / "out.bin" - - with patch("tracecat.storage.blob.logger.debug") as log_debug: - await download_file_to_path( - key="tenant-key/org/repository", - bucket="tenant-bucket", - output_path=out, - redact_log_identifiers=True, - ) - - log_debug.assert_called_once_with( - "File streamed to disk successfully", - key="", - bucket="", - output_path=str(out), - size=7, - ) - - @pytest.mark.anyio - async def test_download_file_to_path_redacts_sensitive_capacity_error( - self, tmp_path: Path, monkeypatch - ): - class DummyStream: - async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 - yield b"payload" - - @asynccontextmanager - async def _fake_open_download_stream( - *, - key: str, # noqa: ARG001 - bucket: str, # noqa: ARG001 - redact_log_identifiers: bool, - ): - assert redact_log_identifiers is True - yield DummyStream(), 7 - - monkeypatch.setattr( - "tracecat.storage.blob.open_download_stream", - _fake_open_download_stream, - ) - - with pytest.raises(ValueError) as exc_info: - await download_file_to_path( - key="tenant-key/org/repository", - bucket="tenant-bucket", - output_path=tmp_path / "out.bin", - max_bytes=5, - redact_log_identifiers=True, - ) - - message = str(exc_info.value) - assert "tenant-key" not in message - assert "tenant-bucket" not in message - assert "/" in message - @pytest.mark.anyio async def test_download_file_to_path_max_bytes_refuses( self, tmp_path: Path, monkeypatch @@ -963,49 +808,6 @@ async def ensure_capacity(size_bytes: int) -> None: assert reservations == [7] assert out.read_bytes() == b"payload" - @pytest.mark.anyio - async def test_download_file_to_path_grows_unknown_length_reservation( - self, tmp_path: Path, monkeypatch - ): - """Unknown-length downloads reserve each chunk above current usage.""" - - class DummyStream: - async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 - yield b"abc" - yield b"defg" - - @asynccontextmanager - async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 - yield DummyStream(), None - - monkeypatch.setattr( - "tracecat.storage.blob.open_download_stream", - _fake_open_download_stream, - ) - - out = tmp_path / "out.bin" - temp_path = tmp_path / "out.bin.part" - reservations: list[int] = [] - partial_sizes: list[int] = [] - - async def ensure_capacity(size_bytes: int) -> None: - partial_size = temp_path.stat().st_size - assert 2 + partial_size + size_bytes <= 10 - reservations.append(size_bytes) - partial_sizes.append(partial_size) - - await download_file_to_path( - key="k", - bucket="b", - output_path=out, - max_bytes=10, - ensure_capacity=ensure_capacity, - ) - - assert reservations == [3, 4] - assert partial_sizes == [0, 3] - assert out.read_bytes() == b"abcdefg" - @pytest.mark.anyio async def test_download_file_to_path_never_exceeds_reserved_length( self, tmp_path: Path, monkeypatch @@ -1076,46 +878,6 @@ async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 assert not out.exists() assert not (tmp_path / "out.bin.part").exists() - @pytest.mark.anyio - async def test_download_file_to_path_defers_failed_partial_cleanup( - self, - tmp_path: Path, - monkeypatch, - ): - """A failed partial unlink remains available for runtime cleanup.""" - - class DummyStream: - async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 - yield b"hello" - - @asynccontextmanager - async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 - yield DummyStream(), 5 - - monkeypatch.setattr( - "tracecat.storage.blob.open_download_stream", - _fake_open_download_stream, - ) - - out = tmp_path / "out.bin" - temp_path = tmp_path / "out.bin.part" - deferred_paths: list[Path] = [] - expected_sha256 = hashlib.sha256(b"hello").hexdigest() - with ( - patch.object(Path, "unlink", side_effect=PermissionError("busy")), - pytest.raises(ValueError, match="Integrity check failed"), - ): - await download_file_to_path( - key="k", - bucket="b", - output_path=out, - expected_sha256=expected_sha256 + "bad", - defer_cleanup=deferred_paths.append, - ) - - assert temp_path.exists() - assert deferred_paths == [temp_path] - @pytest.mark.anyio async def test_download_file_to_path_cancellation_cleans_partial( self, tmp_path: Path, monkeypatch @@ -1156,153 +918,6 @@ async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 assert not out.exists() assert not (tmp_path / "out.bin.part").exists() - @pytest.mark.anyio - async def test_download_file_to_path_cancellation_rejoins_active_write( - self, tmp_path: Path, monkeypatch - ): - """Cancellation cannot unlink a partial file under an active writer.""" - write_started = threading.Event() - write_release = threading.Event() - write_finished = threading.Event() - temp_path = tmp_path / "out.bin.part" - - class DummyStream: - async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 - yield b"partial" - - class BlockingFile: - async def __aenter__(self): - temp_path.touch() - return self - - async def __aexit__(self, exc_type, exc, traceback): - del exc_type, exc, traceback - - async def write(self, chunk: bytes) -> int: - def blocking_write() -> int: - write_started.set() - write_release.wait() - temp_path.write_bytes(chunk) - write_finished.set() - return len(chunk) - - return await asyncio.to_thread(blocking_write) - - @asynccontextmanager - async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 - yield DummyStream(), len(b"partial") - - monkeypatch.setattr( - "tracecat.storage.blob.open_download_stream", - _fake_open_download_stream, - ) - monkeypatch.setattr( - "tracecat.storage.blob.aiofiles.open", - lambda *args, **kwargs: BlockingFile(), - ) - - out = tmp_path / "out.bin" - download = asyncio.create_task( - download_file_to_path( - key="k", - bucket="b", - output_path=out, - ) - ) - try: - assert await asyncio.to_thread(write_started.wait, 1.0) - - download.cancel() - await asyncio.sleep(0) - assert not download.done() - assert temp_path.exists() - - download.cancel() - await asyncio.sleep(0) - assert not download.done() - assert temp_path.exists() - finally: - write_release.set() - - with pytest.raises(asyncio.CancelledError): - await download - - assert write_finished.is_set() - assert not out.exists() - assert not temp_path.exists() - - @pytest.mark.anyio - async def test_download_file_to_path_cancellation_rejoins_active_open( - self, tmp_path: Path, monkeypatch - ): - """Cancellation cannot race cleanup against the aiofiles open worker.""" - open_started = threading.Event() - open_release = threading.Event() - open_finished = threading.Event() - close_finished = threading.Event() - temp_path = tmp_path / "out.bin.part" - - class DummyStream: - async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 - yield b"unused" - - class BlockingFile: - async def __aenter__(self): - def blocking_open(): - open_started.set() - open_release.wait() - temp_path.touch() - open_finished.set() - return self - - return await asyncio.to_thread(blocking_open) - - async def __aexit__(self, exc_type, exc, traceback): - del exc_type, exc, traceback - close_finished.set() - - @asynccontextmanager - async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 - yield DummyStream(), len(b"unused") - - monkeypatch.setattr( - "tracecat.storage.blob.open_download_stream", - _fake_open_download_stream, - ) - monkeypatch.setattr( - "tracecat.storage.blob.aiofiles.open", - lambda *args, **kwargs: BlockingFile(), - ) - - out = tmp_path / "out.bin" - download = asyncio.create_task( - download_file_to_path( - key="k", - bucket="b", - output_path=out, - ) - ) - try: - assert await asyncio.to_thread(open_started.wait, 1.0) - - download.cancel() - await asyncio.sleep(0) - assert not download.done() - - download.cancel() - await asyncio.sleep(0) - assert not download.done() - finally: - open_release.set() - - with pytest.raises(asyncio.CancelledError): - await download - - assert open_finished.is_set() - assert close_finished.is_set() - assert not out.exists() - assert not temp_path.exists() - @pytest.mark.anyio @patch("tracecat.storage.blob.get_storage_client") async def test_ensure_bucket_exists_create_error_propagates(self, mock_get_client): diff --git a/tests/unit/test_unsafe_pid_executor.py b/tests/unit/test_unsafe_pid_executor.py index 96ab11578f..6667d749b0 100644 --- a/tests/unit/test_unsafe_pid_executor.py +++ b/tests/unit/test_unsafe_pid_executor.py @@ -5,7 +5,6 @@ import logging import os import signal -import sys from pathlib import Path import pytest @@ -23,12 +22,8 @@ def _process_is_running(pid: int) -> bool: return False stat_path = Path(f"/proc/{pid}/stat") - try: - if not stat_path.exists(): - return True - except (FileNotFoundError, ProcessLookupError): - # Procfs can report ESRCH while resolving a process that just exited. - return False + if not stat_path.exists(): + return True try: stat_fields = stat_path.read_text().split() except (FileNotFoundError, ProcessLookupError): @@ -74,25 +69,6 @@ def main(pid_file, wait): """ -def _detached_background_process_script() -> str: - return """ -import subprocess -import sys -from pathlib import Path - -def main(pid_file): - child = subprocess.Popen( - [sys.executable, "-c", "import time; time.sleep(30)"], - stdin=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - start_new_session=True, - ) - Path(pid_file).write_text(str(child.pid)) - return child.pid -""" - - class TestUnsafePidExecutor: @pytest.fixture def executor(self, tmp_path) -> UnsafePidExecutor: @@ -115,18 +91,6 @@ def process_disappeared( assert not _process_is_running(123) - def test_process_probe_handles_procfs_stat_exit_race( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - def process_disappeared(path: Path) -> bool: - del path - raise ProcessLookupError - - monkeypatch.setattr(os, "kill", lambda *_: None) - monkeypatch.setattr(Path, "exists", process_disappeared) - - assert not _process_is_running(123) - @pytest.mark.anyio async def test_build_execution_cmd_with_pid_namespace( self, executor: UnsafePidExecutor, monkeypatch: pytest.MonkeyPatch @@ -138,35 +102,10 @@ async def pid_namespace_available() -> bool: "tracecat.sandbox.unsafe_pid_executor.pid_namespace_available", pid_namespace_available, ) - command = await executor._build_execution_cmd( + cmd = await executor._build_execution_cmd( "python3", executor.cache_dir / "wrapper.py" ) - assert command.argv[:4] == ["unshare", "--pid", "--fork", "--kill-child"] - assert command.supervised is False - - @pytest.mark.anyio - async def test_build_execution_cmd_without_pid_namespace_uses_linux_supervisor( - self, - executor: UnsafePidExecutor, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - async def pid_namespace_unavailable() -> bool: - return False - - monkeypatch.setattr( - unsafe_pid_executor, - "pid_namespace_available", - pid_namespace_unavailable, - ) - monkeypatch.setattr(unsafe_pid_executor.sys, "platform", "linux") - - wrapper_path = executor.cache_dir / "wrapper.py" - command = await executor._build_execution_cmd("python3", wrapper_path) - - assert command.argv[:2] == [sys.executable, "-I"] - assert Path(command.argv[2]).name == "process_supervisor.py" - assert command.argv[-2:] == ["python3", str(wrapper_path)] - assert command.supervised is True + assert cmd[:4] == ["unshare", "--pid", "--fork", "--kill-child"] @pytest.mark.anyio async def test_pid_isolation_warning_logged_once( @@ -250,90 +189,6 @@ async def fake_wait_for(awaitable, *args, **kwargs): unsafe_pid_executor.pid_namespace_probe_error() == "unshare probe timed out" ) - @pytest.mark.parametrize("operation", ["create-venv", "install-packages"]) - @pytest.mark.anyio - async def test_cancelled_dependency_setup_kills_process_group( - self, - executor: UnsafePidExecutor, - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - operation: str, - ) -> None: - """Cancellation cannot leave package-build descendants running.""" - pid_file = tmp_path / f"{operation}-child.pid" - real_create_subprocess_exec = asyncio.create_subprocess_exec - created_processes: list[asyncio.subprocess.Process] = [] - setup_script = """ -import subprocess -import sys -import time -from pathlib import Path - -child = subprocess.Popen( - [sys.executable, "-c", "import time; time.sleep(30)"], - stdin=subprocess.DEVNULL, -) -Path(sys.argv[1]).write_text(str(child.pid)) -time.sleep(30) -""" - - async def create_dependency_process(*args, **kwargs): - del args - assert kwargs["start_new_session"] is True - process = await real_create_subprocess_exec( - sys.executable, - "-c", - setup_script, - str(pid_file), - stdout=kwargs["stdout"], - stderr=kwargs["stderr"], - start_new_session=True, - ) - created_processes.append(process) - return process - - monkeypatch.setattr( - asyncio, - "create_subprocess_exec", - create_dependency_process, - ) - - if operation == "create-venv": - setup = executor._create_venv(tmp_path / "venv") - else: - setup = executor._install_packages( - tmp_path / "venv", - ["synthetic-package"], - ) - - task = asyncio.create_task(setup) - child_pid: int | None = None - try: - await _wait_for_file(pid_file) - child_pid = int(pid_file.read_text()) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - assert len(created_processes) == 1 - assert created_processes[0].returncode is not None - await _wait_for_process_exit(child_pid) - finally: - if not task.done(): - task.cancel() - with contextlib.suppress(asyncio.CancelledError, Exception): - await task - for process in created_processes: - with contextlib.suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) - if process.returncode is None: - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - if child_pid is not None: - with contextlib.suppress(ProcessLookupError): - os.kill(child_pid, signal.SIGKILL) - @pytest.mark.anyio async def test_execute_basic_script(self, executor: UnsafePidExecutor) -> None: script = """ @@ -404,42 +259,6 @@ async def test_execute_kills_background_descendants_holding_output_pipes( with contextlib.suppress(ProcessLookupError): os.kill(child_pid, signal.SIGKILL) - @pytest.mark.skipif( - sys.platform != "linux", - reason="Detached fallback containment uses the Linux subreaper supervisor", - ) - @pytest.mark.anyio - async def test_execute_without_pid_namespace_reaps_detached_descendant( - self, - executor: UnsafePidExecutor, - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - async def pid_namespace_unavailable() -> bool: - return False - - monkeypatch.setattr( - unsafe_pid_executor, - "pid_namespace_available", - pid_namespace_unavailable, - ) - pid_file = tmp_path / "detached-child.pid" - result = await asyncio.wait_for( - executor.execute( - script=_detached_background_process_script(), - inputs={"pid_file": str(pid_file)}, - ), - timeout=5, - ) - - assert result.success - child_pid = int(pid_file.read_text()) - try: - await _wait_for_process_exit(child_pid) - finally: - with contextlib.suppress(ProcessLookupError): - os.kill(child_pid, signal.SIGKILL) - @pytest.mark.anyio async def test_execute_kills_descendants_on_timeout( self, executor: UnsafePidExecutor, tmp_path: Path diff --git a/tracecat/executor/action_runner.py b/tracecat/executor/action_runner.py index 63876bf19a..63b9c3606d 100644 --- a/tracecat/executor/action_runner.py +++ b/tracecat/executor/action_runner.py @@ -45,10 +45,7 @@ from tracecat.logger import logger from tracecat.sandbox.executor import ActionSandboxConfig, NsjailExecutor from tracecat.sandbox.types import ResourceLimits -from tracecat.sandbox.utils import ( - communicate_process_group, - terminate_supervised_process, -) +from tracecat.sandbox.utils import communicate_process_group from tracecat.secrets.common import apply_masks, apply_masks_object if TYPE_CHECKING: @@ -121,17 +118,11 @@ def _direct_subprocess_command(minimal_runner_path: Path) -> list[str]: if setpriv is None: raise RuntimeError("setpriv is required for direct action subprocess isolation") - supervisor_path = Path(__file__).with_name("process_supervisor.py") return [ setpriv, "--no-new-privs", "--inh-caps=-all", "--ambient-caps=-all", - sys.executable, - # Keep registry-controlled PYTHONPATH out of the supervisor interpreter. - # Isolated mode leaves the environment intact for the nested action. - "-I", - str(supervisor_path), *runner_command, ] @@ -181,19 +172,15 @@ async def execute_action( """ timeout = timeout or config.TRACECAT__EXECUTOR_CLIENT_TIMEOUT - # Direct subprocesses receive host paths and can modify extracted - # artifacts. NsJail exposes the same paths through read-only bind mounts. - use_sandbox = force_sandbox or ( - config.TRACECAT__EXECUTOR_SANDBOX_ENABLED and _is_sandbox_available() - ) - # Materialize each registry artifact, collect paths in deterministic order. # The lease is held for the whole subprocess execution so cache eviction # cannot delete a directory the subprocess is still importing from. - async with self.registry_artifacts.lease( - artifact_uris, - paths_may_be_modified=not use_sandbox, - ) as registry_paths: + async with self.registry_artifacts.lease(artifact_uris) as registry_paths: + # Check if sandbox execution is enabled and available + # force_sandbox=True overrides config (used by ephemeral backend) + use_sandbox = force_sandbox or ( + config.TRACECAT__EXECUTOR_SANDBOX_ENABLED and _is_sandbox_available() + ) logger.debug( "Using sandbox execution", use_sandbox=use_sandbox, @@ -456,9 +443,6 @@ async def _execute_direct( proc, input=input_json, timeout=timeout, - terminate=( - terminate_supervised_process if sys.platform == "linux" else None - ), ) elapsed_ms = (time.monotonic() - start_time) * 1000 logger.info( diff --git a/tracecat/executor/backends/base.py b/tracecat/executor/backends/base.py index 0507882b9d..082d4dfe7e 100644 --- a/tracecat/executor/backends/base.py +++ b/tracecat/executor/backends/base.py @@ -38,7 +38,6 @@ if TYPE_CHECKING: from tracecat.auth.types import Role from tracecat.dsl.schemas import RunActionInput - from tracecat.executor.registry_artifacts import RegistryArtifactCache from tracecat.executor.schemas import ResolvedContext @@ -149,14 +148,9 @@ async def _execute_run_python( ) # The lease is held for the whole sandbox run so cache eviction cannot - # delete a directory the script is still importing from. SandboxService - # may select UnsafePidExecutor, which exposes these host paths writable, - # so conservatively rescan their footprint after every successful run. - registry_artifacts = self._registry_artifact_cache() - async with registry_artifacts.lease( - artifact_uris, - paths_may_be_modified=True, - ) as registry_paths: + # delete a directory the script is still importing from. + registry_artifacts = get_action_runner().registry_artifacts + async with registry_artifacts.lease(artifact_uris) as registry_paths: return await self._run_python_in_sandbox( script=script, args=args, @@ -165,10 +159,6 @@ async def _execute_run_python( resolved_context=resolved_context, ) - def _registry_artifact_cache(self) -> RegistryArtifactCache: - """Return the cache owned by this backend's executor process.""" - return get_action_runner().registry_artifacts - async def _run_python_in_sandbox( self, *, diff --git a/tracecat/executor/backends/test.py b/tracecat/executor/backends/test.py index 2086a7b261..696141180c 100644 --- a/tracecat/executor/backends/test.py +++ b/tracecat/executor/backends/test.py @@ -18,10 +18,10 @@ from __future__ import annotations import asyncio +import contextlib import sys import threading from contextlib import AsyncExitStack, contextmanager -from pathlib import Path from typing import TYPE_CHECKING, Any from tracecat_registry import secrets as registry_secrets @@ -37,14 +37,9 @@ ctx_session_id, ) from tracecat.executor.action_gateway.config import action_gateway_socket_path +from tracecat.executor.action_runner import get_action_runner from tracecat.executor.backends.base import ExecutorBackend from tracecat.executor.backends.registry_helpers import get_registry_artifact_uris -from tracecat.executor.registry_artifact_materialization import ( - _artifact_uri_for_logging, - _is_cache_entry_uri, - _run_blocking_rejoin_on_cancel, -) -from tracecat.executor.registry_artifacts import RegistryArtifactCache from tracecat.executor.schemas import ( ActionImplementation, ExecutorActionErrorInfo, @@ -112,21 +107,6 @@ def _temporary_sys_path(paths: list[str]) -> Iterator[None]: class TestBackend(ExecutorBackend): """In-process execution backend for tests only.""" - __test__ = False - - def __init__(self) -> None: - # Pytest creates a backend inside each function-scoped event loop. Keep - # its loop-affine cache on the same lifecycle instead of reusing the - # process-global ActionRunner cache across closed test loops. - self._owned_registry_artifacts: RegistryArtifactCache | None = None - - def _registry_artifact_cache(self) -> RegistryArtifactCache: - if self._owned_registry_artifacts is None: - self._owned_registry_artifacts = RegistryArtifactCache( - Path(config.TRACECAT__EXECUTOR_REGISTRY_CACHE_DIR) - ) - return self._owned_registry_artifacts - async def _execute( self, input: RunActionInput, @@ -268,7 +248,21 @@ async def _run_sync_udf( leases, temporary ``sys.path`` entries, and secret contexts alive until the function actually stops. """ - return await _run_blocking_rejoin_on_cancel(lambda: fn(**args)) + worker = asyncio.ensure_future(asyncio.to_thread(fn, **args)) + try: + return await asyncio.shield(worker) + except asyncio.CancelledError: + while not worker.done(): + try: + await asyncio.shield(worker) + except asyncio.CancelledError: + continue + except Exception: + break + if not worker.cancelled(): + with contextlib.suppress(Exception): + worker.result() + raise def _load_udf_callable(self, action_impl: ActionImplementation): """Load the UDF callable from action_impl metadata.""" @@ -314,32 +308,21 @@ async def _lease_registry_artifacts( logger.debug("No artifact URIs found, using empty paths") return [] - registry_artifacts = self._registry_artifact_cache() - if any(_is_cache_entry_uri(uri) for uri in artifact_uris): - await registry_artifacts.ensure_swept() + registry_artifacts = get_action_runner().registry_artifacts extracted_paths: list[str] = [] - mutable_rescan_registered = False for artifact_uri in artifact_uris: - rescan_on_release = not mutable_rescan_registered and _is_cache_entry_uri( - artifact_uri - ) try: artifact_paths = await leases.enter_async_context( - registry_artifacts.lease( - [artifact_uri], - paths_may_be_modified=rescan_on_release, - ) + registry_artifacts.lease([artifact_uri]) ) except Exception as e: logger.warning( "Failed to materialize artifact for test execution", - artifact_uri=_artifact_uri_for_logging(artifact_uri), + artifact_uri=artifact_uri, error=str(e), ) continue - if rescan_on_release: - mutable_rescan_registered = True extracted_paths.extend(str(path) for path in artifact_paths) logger.debug( diff --git a/tracecat/executor/process_supervisor.py b/tracecat/executor/process_supervisor.py deleted file mode 100644 index 2820815eaa..0000000000 --- a/tracecat/executor/process_supervisor.py +++ /dev/null @@ -1,276 +0,0 @@ -"""Contain descendants of one Linux subprocess. - -The outer process remains the child observed by its caller. A detached -subreaper monitor owns the actual action process and all descendants orphaned by -it. Closing the control pipe asks that monitor to kill and reap its complete -child tree before the outer process exits. - -This module intentionally uses only the Python standard library so invoking it -does not import the Tracecat application in an untrusted action process. -""" - -from __future__ import annotations - -import ctypes -import os -import select -import signal -import sys -from collections.abc import Sequence -from contextlib import suppress -from pathlib import Path -from types import FrameType - -_PR_SET_CHILD_SUBREAPER = 36 -_PR_SET_PDEATHSIG = 1 -_PARENT_POLL_INTERVAL_MS = 10 - - -def _exit_code(wait_status: int) -> int: - """Convert a wait status into a shell-compatible non-negative exit code.""" - code = os.waitstatus_to_exitcode(wait_status) - return code if code >= 0 else 128 - code - - -def _set_child_subreaper() -> None: - """Make this process the reparenting boundary for orphaned descendants.""" - libc = ctypes.CDLL(None, use_errno=True) - prctl = libc.prctl - prctl.argtypes = [ - ctypes.c_int, - ctypes.c_ulong, - ctypes.c_ulong, - ctypes.c_ulong, - ctypes.c_ulong, - ] - prctl.restype = ctypes.c_int - if prctl(_PR_SET_CHILD_SUBREAPER, 1, 0, 0, 0) != 0: - error_number = ctypes.get_errno() - raise OSError(error_number, os.strerror(error_number)) - - -def _set_parent_death_signal(signal_number: int) -> None: - """Ask Linux to signal this process when its current parent exits.""" - libc = ctypes.CDLL(None, use_errno=True) - prctl = libc.prctl - prctl.argtypes = [ - ctypes.c_int, - ctypes.c_ulong, - ctypes.c_ulong, - ctypes.c_ulong, - ctypes.c_ulong, - ] - prctl.restype = ctypes.c_int - if prctl(_PR_SET_PDEATHSIG, signal_number, 0, 0, 0) != 0: - error_number = ctypes.get_errno() - raise OSError(error_number, os.strerror(error_number)) - - -def _direct_child_pids_from_children_file() -> list[int]: - """Return direct child PIDs using procfs's optional children file.""" - children_path = Path(f"/proc/self/task/{os.getpid()}/children") - contents = children_path.read_text().strip() - return [int(pid) for pid in contents.split()] if contents else [] - - -def _proc_parent_pid(stat_path: Path) -> int: - """Read one process's parent PID from its procfs stat record.""" - contents = stat_path.read_text() - closing_paren = contents.rfind(")") - fields = contents[closing_paren + 1 :].split() - if closing_paren < 0 or len(fields) < 2: - raise RuntimeError("Malformed procfs stat record") - return int(fields[1]) - - -def _direct_child_pids_from_proc_stat() -> list[int]: - """Return direct child PIDs by scanning portable procfs stat records.""" - parent_pid = os.getpid() - _proc_parent_pid(Path("/proc/self/stat")) # Verify procfs is readable. - child_pids: list[int] = [] - - with os.scandir("/proc") as entries: - for entry in entries: - if not entry.name.isdecimal(): - continue - try: - candidate_pid = int(entry.name) - candidate_parent_pid = _proc_parent_pid(Path(entry.path) / "stat") - except (FileNotFoundError, PermissionError, ProcessLookupError): - continue - if candidate_parent_pid == parent_pid: - child_pids.append(candidate_pid) - - return child_pids - - -def _direct_child_pids() -> list[int]: - """Return direct child PIDs from either available procfs interface.""" - try: - return _direct_child_pids_from_children_file() - except OSError: - return _direct_child_pids_from_proc_stat() - - -def _waitpid(pid: int, options: int = 0) -> tuple[int, int]: - """Wait for a child while tolerating signal interruptions.""" - while True: - try: - return os.waitpid(pid, options) - except InterruptedError: - continue - - -def _kill_and_reap_children() -> None: - """Kill every adopted child, including descendants orphaned while reaping.""" - while child_pids := _direct_child_pids(): - for child_pid in child_pids: - with suppress(ProcessLookupError): - os.kill(child_pid, signal.SIGKILL) - for child_pid in child_pids: - with suppress(ChildProcessError): - _waitpid(child_pid) - - -def _exec(command: Sequence[str], control_fd: int) -> None: - """Replace this child with the actual action command.""" - os.close(control_fd) - try: - os.execvpe(command[0], list(command), os.environ) - except OSError: - with suppress(OSError): - os.write(2, b"Failed to execute supervised process\n") - os._exit(127) - - -def _run_monitor( - control_fd: int, - command: Sequence[str], - *, - supervisor_pid: int, -) -> int: - """Run the action below a detached subreaper and contain its descendants.""" - try: - parent_cleanup_requested = False - - def request_parent_cleanup( - _signal: int, - _frame: FrameType | None, - ) -> None: - nonlocal parent_cleanup_requested - parent_cleanup_requested = True - - # The action can stop this same-UID monitor and kill its outer parent. - # SIGCONT both resumes the monitor and records that it must clean up. - signal.signal(signal.SIGCONT, request_parent_cleanup) - _set_parent_death_signal(signal.SIGCONT) - # Close the fork-to-prctl race before any untrusted command is started. - if os.getppid() != supervisor_pid or parent_cleanup_requested: - return 128 + signal.SIGTERM - - os.setsid() - _set_child_subreaper() - _direct_child_pids() # Fail before execution when procfs tracking is absent. - - action_pid = os.fork() - if action_pid == 0: - _exec(command, control_fd) - - poller = select.poll() - poller.register(control_fd, select.POLLIN | select.POLLHUP | select.POLLERR) - action_status: int | None = None - parent_closed = False - - while ( - action_status is None and not parent_closed and not parent_cleanup_requested - ): - waited_pid, wait_status = _waitpid(action_pid, os.WNOHANG) - if waited_pid == action_pid: - action_status = wait_status - break - if poller.poll(_PARENT_POLL_INTERVAL_MS): - parent_closed = os.read(control_fd, 1) == b"" - - _kill_and_reap_children() - if parent_closed or parent_cleanup_requested: - return 128 + signal.SIGTERM - if action_status is None: - raise RuntimeError("Supervised action exited without a wait status") - return _exit_code(action_status) - except BaseException: - with suppress(BaseException): - _kill_and_reap_children() - with suppress(OSError): - os.write(2, b"Process supervisor failed\n") - return 1 - finally: - with suppress(OSError): - os.close(control_fd) - - -def supervise(command: Sequence[str]) -> int: - """Run one command and return only after all of its descendants are reaped.""" - if sys.platform != "linux": - raise RuntimeError("The process supervisor requires Linux") - if not command: - raise ValueError("A supervised command is required") - - # The monitor is intentionally visible as the action's parent. Make this - # outer process a second subreaper so killing that monitor only reparents - # the action tree here, where it is still killed and reaped before return. - _set_child_subreaper() - _direct_child_pids() # Fail before execution when procfs tracking is absent. - - control_read_fd, control_write_fd = os.pipe() - writer_open = True - monitor_pid: int | None = None - - def close_control_pipe() -> None: - nonlocal writer_open - if writer_open: - with suppress(OSError): - os.close(control_write_fd) - writer_open = False - - def request_cleanup(_signal: int, _frame: FrameType | None) -> None: - close_control_pipe() - # The action shares the monitor's UID and can stop it. SIGKILL the - # monitor session so the outer subreaper adopts and reaps every member, - # including descendants that detached into their own sessions. - if monitor_pid is not None: - with suppress(ProcessLookupError): - os.killpg(monitor_pid, signal.SIGKILL) - - signal.signal(signal.SIGTERM, request_cleanup) - signal.signal(signal.SIGINT, request_cleanup) - - supervisor_pid = os.getpid() - monitor_pid = os.fork() - if monitor_pid == 0: - signal.signal(signal.SIGTERM, signal.SIG_DFL) - signal.signal(signal.SIGINT, signal.SIG_DFL) - os.close(control_write_fd) - os._exit( - _run_monitor( - control_read_fd, - command, - supervisor_pid=supervisor_pid, - ) - ) - - os.close(control_read_fd) - try: - _, monitor_status = _waitpid(monitor_pid) - return _exit_code(monitor_status) - finally: - close_control_pipe() - _kill_and_reap_children() - - -def main() -> int: - """CLI entry point.""" - return supervise(sys.argv[1:]) - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tracecat/executor/registry_artifact_cache_state.py b/tracecat/executor/registry_artifact_cache_state.py deleted file mode 100644 index b1daac941f..0000000000 --- a/tracecat/executor/registry_artifact_cache_state.py +++ /dev/null @@ -1,218 +0,0 @@ -"""Process-local state and lease bookkeeping for registry artifact caches.""" - -from __future__ import annotations - -import asyncio -import os -import threading -import time -from collections.abc import AsyncIterator -from contextlib import asynccontextmanager -from dataclasses import dataclass, field -from pathlib import Path - -from tracecat.executor.registry_artifact_materialization import ( - RegistryArtifactAdmission, - RegistryArtifactMaterializationContext, - RegistryArtifactPaths, -) -from tracecat.logger import logger - -BASE_PYTHONPATH_DIR_NAME = "base" -"""Cache subdirectory used as the PYTHONPATH entry when no artifact is requested.""" - -CACHE_ENTRIES_DIR_NAME = "entries" -"""Directory containing one atomic subdirectory per cache key.""" - -CACHE_STAGING_DIR_NAME = "staging" -"""Directory containing in-progress materialization scratch.""" - -CACHE_TRASH_DIR_NAME = "trash" -"""Directory containing atomically retired entries pending physical deletion.""" - - -class RegistryArtifactCacheLoopError(RuntimeError): - """A registry artifact cache was used outside its owning event loop.""" - - -@dataclass(slots=True) -class RegistryArtifactRuntimeState: - """Process-local synchronization and lease state for one cache key.""" - - lock: asyncio.Lock = field(default_factory=asyncio.Lock) - refcount: int = 0 - last_used: float = 0.0 - users: int = 0 - - -class _RegistryArtifactCacheState: - """Owns cache paths, event-loop affinity, and per-key lease state.""" - - def __init__(self, cache_dir: Path): - self.cache_dir = cache_dir - self.entries_dir = cache_dir / CACHE_ENTRIES_DIR_NAME - self.staging_dir = cache_dir / CACHE_STAGING_DIR_NAME - self.trash_dir = cache_dir / CACHE_TRASH_DIR_NAME - # Runtime states remain stable while their entry, leases, or lock users - # exist. Failed and evicted keys are discarded once no waiter can still - # hold the old lock identity. - self._runtime: dict[str, RegistryArtifactRuntimeState] = {} - # The cache contains asyncio locks, tasks, and multi-step lease state. - # Bind the public API to one loop/thread so a future synchronous - # Temporal activity fails immediately instead of corrupting that state - # through a thread-local event loop. - self._owner_binding_lock = threading.Lock() - self._owner_loop: asyncio.AbstractEventLoop | None = None - self._owner_thread_id: int | None = None - # Cold materializations and budget passes share this outer lock. It - # keeps byte reservations stable while downloads and extraction write. - self._admission_lock = asyncio.Lock() - self._budget_lock = asyncio.Lock() - # Guard the off-loop startup sweep independently from cache operations. - self._swept: bool = False - self._sweep_task: asyncio.Task[None] | None = None - self._sweep_lock = asyncio.Lock() - # Startup is the only time the whole staging directory is swept. Exact - # startup or runtime paths that could not be removed are safe to retry. - self._deferred_staging_cleanup: set[Path] = set() - # Final-release unmount failures are retried by later lease cleanup so - # transient errors cannot accumulate idle loop devices indefinitely. - self._failed_unmounts: set[str] = set() - # Whether the on-disk cache may exceed its budget. Set when a new entry - # is materialized and cleared once enforcement measures a cache that - # fits, so steady-state cache hits never pay for a disk scan. - self._budget_dirty = True - - def _assert_owner_loop(self) -> None: - """Bind to the current loop or reject use from another loop/thread.""" - current_loop = asyncio.get_running_loop() - current_thread_id = threading.get_ident() - with self._owner_binding_lock: - if self._owner_loop is None: - self._owner_loop = current_loop - self._owner_thread_id = current_thread_id - return - if ( - self._owner_loop is current_loop - and self._owner_thread_id == current_thread_id - ): - return - - raise RegistryArtifactCacheLoopError( - "RegistryArtifactCache is bound to one event loop and thread " - f"(owner_loop={id(self._owner_loop)}, " - f"owner_thread={self._owner_thread_id}, " - f"current_loop={id(current_loop)}, " - f"current_thread={current_thread_id})" - ) - - def _runtime_for(self, cache_key: str) -> RegistryArtifactRuntimeState: - """Return the process-local state for one cache key.""" - if runtime := self._runtime.get(cache_key): - return runtime - runtime = RegistryArtifactRuntimeState() - self._runtime[cache_key] = runtime - return runtime - - @asynccontextmanager - async def _runtime_lock( - self, cache_key: str - ) -> AsyncIterator[RegistryArtifactRuntimeState]: - """Lock one key while keeping its runtime identity pinned for waiters.""" - runtime = self._runtime_for(cache_key) - runtime.users += 1 - try: - async with runtime.lock: - yield runtime - finally: - runtime.users -= 1 - self._discard_idle_runtime(cache_key, runtime) - - def _discard_idle_runtime( - self, - cache_key: str, - runtime: RegistryArtifactRuntimeState, - ) -> None: - """Discard state only after its entry and every possible user are gone.""" - if runtime.users > 0 or runtime.refcount > 0 or runtime.lock.locked(): - return - try: - if self._paths_for(cache_key).entry_dir.exists(): - return - except OSError: - # Retaining a small state object is safer than splitting lock identity - # when the entry cannot be inspected. - return - if self._runtime.get(cache_key) is runtime: - del self._runtime[cache_key] - - def _context_for( - self, - cache_key: str, - *, - admission: RegistryArtifactAdmission | None = None, - ) -> RegistryArtifactMaterializationContext: - """Return a materialization context for a registry artifact key.""" - return RegistryArtifactMaterializationContext( - cache_key=cache_key, - staging_dir=self.staging_dir, - paths=self._paths_for(cache_key), - defer_cleanup=self._deferred_staging_cleanup.add, - admission=admission, - ) - - def _base_pythonpath_dir(self) -> Path: - """Return the base PYTHONPATH directory used when no artifact is requested.""" - base_dir = self.cache_dir / BASE_PYTHONPATH_DIR_NAME - base_dir.mkdir(parents=True, exist_ok=True) - return base_dir - - def _acquire_lease(self, cache_key: str) -> None: - """Pin a cache entry against eviction and mark it as recently used. - - Callers must hold the per-key lock so the increment is ordered against - in-flight eviction of the same key. - """ - runtime = self._runtime_for(cache_key) - runtime.refcount += 1 - runtime.last_used = time.time() - self._touch_entry(cache_key) - - def _release_lease(self, cache_key: str) -> bool: - """Release one pin and return whether the entry became idle.""" - runtime = self._runtime.get(cache_key) - if runtime is None or runtime.refcount == 0: - return False - runtime.refcount -= 1 - runtime.last_used = time.time() - became_idle = runtime.refcount == 0 - if became_idle: - self._touch_entry(cache_key) - return became_idle - - def _refcount(self, cache_key: str) -> int: - """Return the number of live leases on a cache entry.""" - runtime = self._runtime.get(cache_key) - return 0 if runtime is None else runtime.refcount - - def _touch_entry(self, cache_key: str) -> None: - """Best-effort refresh of the entry-root mtime for restart-safe LRU.""" - entry_dir = self._paths_for(cache_key).entry_dir - try: - os.utime(entry_dir) - except OSError: - logger.debug( - "Could not refresh registry artifact entry mtime", - cache_key=cache_key, - ) - - def _paths_for(self, cache_key: str) -> RegistryArtifactPaths: - """Return local cache paths for a registry artifact key.""" - entry_dir = self.entries_dir / cache_key - return RegistryArtifactPaths( - entry_dir=entry_dir, - squashfs_image_path=entry_dir / "image.squashfs", - squashfs_mount_dir=entry_dir / "mount", - squashfs_extract_dir=entry_dir / "extracted", - tarball_target_dir=entry_dir / "tarball", - ) diff --git a/tracecat/executor/registry_artifact_materialization.py b/tracecat/executor/registry_artifact_materialization.py deleted file mode 100644 index 111064205c..0000000000 --- a/tracecat/executor/registry_artifact_materialization.py +++ /dev/null @@ -1,1141 +0,0 @@ -"""Registry artifact formats and local materialization primitives.""" - -from __future__ import annotations - -import asyncio -import contextlib -import hashlib -import os -import secrets -import shutil -import sysconfig -import tarfile -import time -from abc import ABC, abstractmethod -from collections.abc import Awaitable, Callable -from dataclasses import dataclass -from enum import StrEnum -from pathlib import Path, PurePosixPath -from urllib.parse import urlsplit, urlunsplit - -import httpx -import tracecat_registry - -from tracecat import config -from tracecat.executor import registry_artifact_mounts -from tracecat.logger import logger -from tracecat.registry.artifact_keys import parse_s3_uri -from tracecat.registry.constants import DEFAULT_REGISTRY_ORIGIN -from tracecat.sandbox.utils import communicate_process_group -from tracecat.storage import blob - -__all__ = [ - "_artifact_uri_for_logging", - "_is_cache_entry_uri", - "_squashfs_sidecar_uri", - "_tarball_uri_for_squashfs", -] - - -def _artifact_uri_for_logging(artifact_uri: str) -> str: - """Retain only the non-sensitive scheme of an artifact URI.""" - try: - parsed = urlsplit(artifact_uri) - except ValueError: - return "" - if not parsed.scheme or not parsed.hostname: - return "" - return urlunsplit((parsed.scheme, "", "", "", "")) - - -class RegistryArtifactFormat(StrEnum): - """Executor-supported registry artifact encodings.""" - - BUILTIN = "builtin" - SQUASHFS = "squashfs" - TAR_GZ = "tar.gz" - - -SQUASHFS_MOUNT_OPTIONS = "loop,ro,nodev,nosuid" -"""Mount options for executor-managed SquashFS registry artifacts. - -The image must stay read-only and should not expose device nodes or setuid bits -from registry package contents. Avoid noexec because Python packages may include -native extension modules that need to be loaded from the mounted artifact. -""" - -BUNDLED_BUILTIN_REGISTRY_URI_PREFIX = f"tracecat-builtin://{DEFAULT_REGISTRY_ORIGIN}/" -"""Pseudo-URI for the builtin registry already installed in the executor image.""" - - -@dataclass(frozen=True, slots=True) -class RegistryArtifactPaths: - """Executor-local cache paths for one registry artifact key.""" - - entry_dir: Path - squashfs_image_path: Path - squashfs_mount_dir: Path - squashfs_extract_dir: Path - tarball_target_dir: Path - - -class SquashfsMountCommandError(RuntimeError): - """The ``mount`` command itself failed for a SquashFS registry artifact. - - Only this error drives SquashFS mount policy. Download, mkdir, and other - preparation failures must not be mistaken for a missing mount capability or - for loop-device exhaustion. - """ - - -class RegistryArtifactUriError(ValueError): - """A registry artifact URI is malformed, with identifiers suppressed.""" - - -class RegistryArtifactExtractionError(RuntimeError): - """A registry archive could not be inspected or extracted safely.""" - - def __init__(self) -> None: - super().__init__("Registry artifact extraction failed") - - -@dataclass(frozen=True, slots=True) -class RegistryArtifactAdmission: - """Byte-bound admission hook shared by one cold materialization.""" - - max_bytes: int - allocation_unit: int - ensure_capacity: Callable[[int], Awaitable[None]] - - -@dataclass(slots=True) -class RegistryArtifactMaterializationContext: - """Shared runtime state for artifact materialization.""" - - cache_key: str - staging_dir: Path - paths: RegistryArtifactPaths - defer_cleanup: Callable[[Path], None] - admission: RegistryArtifactAdmission | None = None - - def can_mount_squashfs(self) -> bool: - return config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED and ( - shutil.which("mount") is not None - ) - - -@dataclass(frozen=True, slots=True) -class RegistryArtifact(ABC): - """An executor-local materializable registry artifact.""" - - uri: str - cache_key: str - - @property - @abstractmethod - def format(self) -> RegistryArtifactFormat: - """Artifact format used for logging and dispatch.""" - - @abstractmethod - def cached_path( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path] | None: - """Return already-materialized import paths for this artifact, if present.""" - - @abstractmethod - async def materialize( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path]: - """Return importable Python paths, materializing the artifact if needed.""" - - def discard_failed_materialization( - self, ctx: RegistryArtifactMaterializationContext - ) -> None: - """Discard canonical bytes that cannot serve a fallback candidate.""" - del ctx - - def _temp_path( - self, - ctx: RegistryArtifactMaterializationContext, - suffix: str, - ) -> Path: - _validate_cache_work_directory(ctx.staging_dir) - ctx.staging_dir.mkdir(parents=True, exist_ok=True) - _validate_cache_work_directory(ctx.staging_dir) - while True: - attempt_id = secrets.token_hex(8) - candidate = ( - ctx.staging_dir / f"{self.cache_key}.{os.getpid()}.{attempt_id}{suffix}" - ) - if not os.path.lexists(candidate): - return candidate - - -async def _rejoin_future_on_cancel[T](future: asyncio.Future[T]) -> T: - """Shield a future and rejoin it through repeated caller cancellation.""" - try: - return await asyncio.shield(future) - except asyncio.CancelledError: - while not future.done(): - try: - await asyncio.shield(future) - except asyncio.CancelledError: - continue - except Exception: - break - if not future.cancelled(): - with contextlib.suppress(Exception): - future.result() - raise - - -async def _run_blocking_rejoin_on_cancel[T](operation: Callable[[], T]) -> T: - """Run blocking work without abandoning its thread on cancellation.""" - worker = asyncio.ensure_future(asyncio.to_thread(operation)) - return await _rejoin_future_on_cancel(worker) - - -async def _remove_tree_rejoin_on_cancel( - path: Path, - *, - defer_cleanup: Callable[[Path], None], -) -> None: - """Remove a tree off-loop and retain failed paths for a later retry.""" - if not path.exists(): - return - try: - await _run_blocking_rejoin_on_cancel(lambda: shutil.rmtree(path)) - except FileNotFoundError: - return - except asyncio.CancelledError: - if path.exists(): - defer_cleanup(path) - logger.warning( - "Deferred cancelled registry artifact staging cleanup", - path=str(path), - ) - raise - except OSError as e: - defer_cleanup(path) - logger.warning( - "Deferred failed registry artifact staging cleanup", - path=str(path), - error=str(e), - ) - - -def _remove_file_or_defer( - path: Path, - *, - defer_cleanup: Callable[[Path], None], -) -> None: - """Remove one artifact file without masking the materialization outcome.""" - try: - path.unlink(missing_ok=True) - except OSError as e: - defer_cleanup(path) - logger.warning( - "Deferred failed registry artifact file cleanup", - path=str(path), - error=str(e), - ) - - -def _is_reusable_cache_file(path: Path) -> bool: - """Return whether a canonical cache file is regular and not a symlink.""" - return path.is_file() and not path.is_symlink() - - -def _is_reusable_cache_directory(path: Path) -> bool: - """Return whether a canonical cache directory is real and not a symlink.""" - return path.is_dir() and not path.is_symlink() - - -def _validate_cache_work_directory(path: Path) -> None: - """Reject a staging or trash root redirected outside its cache root.""" - if os.path.lexists(path) and not _is_reusable_cache_directory(path): - raise OSError("Unsafe symlink or non-directory registry cache work path") - if not path.exists(): - return - - resolved_cache_dir = path.parent.resolve(strict=True) - resolved_work_dir = path.resolve(strict=True) - if resolved_work_dir.parent != resolved_cache_dir: - raise OSError("Unsafe registry cache work path outside the cache directory") - - -def _validate_cache_entries_directory(entries_dir: Path) -> None: - """Reject an entries root redirected outside its configured cache root.""" - if os.path.lexists(entries_dir) and not _is_reusable_cache_directory(entries_dir): - raise OSError("Unsafe symlink or non-directory registry cache entries path") - if not entries_dir.exists(): - return - - resolved_cache_dir = entries_dir.parent.resolve(strict=True) - resolved_entries_dir = entries_dir.resolve(strict=True) - if resolved_entries_dir.parent != resolved_cache_dir: - raise OSError("Unsafe registry cache entries path outside the cache directory") - - -def _validate_cache_entry_path(paths: RegistryArtifactPaths) -> None: - """Reject an entry root redirected outside the configured entries directory.""" - entry_dir = paths.entry_dir - entries_dir = entry_dir.parent - _validate_cache_entries_directory(entries_dir) - if os.path.lexists(entry_dir) and not _is_reusable_cache_directory(entry_dir): - raise OSError("Unsafe symlink or non-directory registry cache entry path") - if not entry_dir.exists(): - return - - resolved_entries_dir = entries_dir.resolve(strict=True) - resolved_entry_dir = entry_dir.resolve(strict=True) - if resolved_entry_dir.parent != resolved_entries_dir: - raise OSError("Unsafe registry cache entry path outside the entries directory") - - -def _validate_squashfs_mount_paths(paths: RegistryArtifactPaths) -> None: - """Reject mount paths that can escape their canonical cache entry.""" - entry_dir = paths.entry_dir - mount_dir = paths.squashfs_mount_dir - _validate_cache_entry_path(paths) - if mount_dir.parent != entry_dir: - raise OSError("Unsafe SquashFS mount path outside its cache entry") - - if os.path.lexists(mount_dir) and not _is_reusable_cache_directory(mount_dir): - raise OSError("Unsafe symlink or non-directory SquashFS mount path") - - if not entry_dir.exists() or not mount_dir.exists(): - return - - resolved_entry_dir = entry_dir.resolve(strict=True) - resolved_mount_dir = mount_dir.resolve(strict=True) - if resolved_mount_dir.parent != resolved_entry_dir: - raise OSError("Unsafe SquashFS mount path outside its cache entry") - - -async def _reuse_or_reclaim_squashfs_image( - path: Path, - *, - defer_cleanup: Callable[[Path], None], -) -> bool: - """Reuse a safe image or reclaim a malformed target before downloading.""" - if _is_reusable_cache_file(path): - return True - if not os.path.lexists(path): - return False - - if _is_reusable_cache_directory(path): - await _remove_tree_rejoin_on_cancel(path, defer_cleanup=defer_cleanup) - else: - _remove_file_or_defer(path, defer_cleanup=defer_cleanup) - - if os.path.lexists(path): - raise OSError("Failed to reclaim malformed SquashFS image target") - return False - - -def _is_reusable_extraction_dir( - path: Path, - *, - defer_cleanup: Callable[[Path], None], -) -> bool: - """Accept canonical directories and reclaim malformed file or symlink targets.""" - if _is_reusable_cache_directory(path): - return True - if os.path.lexists(path): - _remove_file_or_defer(path, defer_cleanup=defer_cleanup) - return False - - -@dataclass(frozen=True, slots=True) -class BuiltinArtifact(RegistryArtifact): - """Current builtin registry package already installed in the executor image.""" - - version: str - - @property - def format(self) -> RegistryArtifactFormat: - return RegistryArtifactFormat.BUILTIN - - def cached_path( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path] | None: - return None - - async def materialize( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path]: - del ctx - import_paths = _bundled_builtin_registry_import_paths(self.version) - logger.info( - "Using bundled builtin registry environment", - registry_version=self.version, - paths=[str(p) for p in import_paths], - ) - return import_paths - - -@dataclass(frozen=True, slots=True) -class SquashfsArtifact(RegistryArtifact): - """SquashFS registry environment image.""" - - @property - def format(self) -> RegistryArtifactFormat: - return RegistryArtifactFormat.SQUASHFS - - def cached_path( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path] | None: - _validate_squashfs_mount_paths(ctx.paths) - if registry_artifact_mounts.is_mount(ctx.paths.squashfs_mount_dir): - logger.debug( - "Using cached SquashFS registry mount", - cache_key=ctx.cache_key, - ) - return [ctx.paths.squashfs_mount_dir] - if _is_reusable_extraction_dir( - ctx.paths.squashfs_extract_dir, - defer_cleanup=ctx.defer_cleanup, - ): - logger.debug( - "Using cached SquashFS registry extraction", - cache_key=ctx.cache_key, - ) - return [ctx.paths.squashfs_extract_dir] - return None - - def discard_failed_materialization( - self, ctx: RegistryArtifactMaterializationContext - ) -> None: - """Remove an unusable image before admitting a tarball fallback.""" - try: - if self.cached_path(ctx) is not None: - return - except OSError as e: - logger.warning( - "Cannot determine whether failed SquashFS candidate is reusable", - cache_key=ctx.cache_key, - artifact_uri=_artifact_uri_for_logging(self.uri), - error=str(e), - ) - return - - _remove_file_or_defer( - ctx.paths.squashfs_image_path, - defer_cleanup=ctx.defer_cleanup, - ) - - for directory in ( - ctx.paths.squashfs_extract_dir, - ctx.paths.squashfs_mount_dir, - ctx.paths.entry_dir, - ): - with contextlib.suppress(OSError): - directory.rmdir() - - async def materialize( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path]: - image_path = ctx.paths.squashfs_image_path - if ctx.can_mount_squashfs(): - try: - return [await self.mount(ctx, image_path)] - except SquashfsMountCommandError as e: - logger.warning( - "Failed to mount SquashFS registry artifact, trying extraction", - cache_key=ctx.cache_key, - artifact_uri=_artifact_uri_for_logging(self.uri), - artifact_format=self.format.value, - error=str(e), - ) - - return [await self.extract(ctx, image_path)] - - async def download( - self, - ctx: RegistryArtifactMaterializationContext, - image_path: Path, - ) -> float: - """Ensure the SquashFS image exists locally and return download time.""" - _validate_cache_entry_path(ctx.paths) - if await _reuse_or_reclaim_squashfs_image( - image_path, - defer_cleanup=ctx.defer_cleanup, - ): - return 0.0 - - image_path.parent.mkdir(parents=True, exist_ok=True) - temp_image = self._temp_path(ctx, ".squashfs") - try: - download_start = time.monotonic() - await _download_s3_artifact( - self.uri, - temp_image, - admission=ctx.admission, - defer_cleanup=ctx.defer_cleanup, - ) - _validate_cache_entry_path(ctx.paths) - try: - temp_image.rename(image_path) - except OSError: - _validate_cache_entry_path(ctx.paths) - if not await _reuse_or_reclaim_squashfs_image( - image_path, - defer_cleanup=ctx.defer_cleanup, - ): - temp_image.rename(image_path) - return (time.monotonic() - download_start) * 1000 - finally: - _remove_file_or_defer( - temp_image, - defer_cleanup=ctx.defer_cleanup, - ) - - async def mount( - self, - ctx: RegistryArtifactMaterializationContext, - image_path: Path, - ) -> Path: - """Download the image if needed and mount it read-only. - - Args: - ctx: Materialization context for the artifact being mounted. - image_path: Local path of the SquashFS image. - - Returns: - The mount directory. - - Raises: - SquashfsMountCommandError: The ``mount`` command failed. - Exception: The image could not be downloaded or prepared. - """ - target_dir = ctx.paths.squashfs_mount_dir - _validate_squashfs_mount_paths(ctx.paths) - if registry_artifact_mounts.is_mount(target_dir): - return target_dir - - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - target_dir.mkdir(parents=True, exist_ok=True) - _validate_squashfs_mount_paths(ctx.paths) - - logger.info( - "Materializing SquashFS registry artifact", - cache_key=ctx.cache_key, - artifact_uri=_artifact_uri_for_logging(self.uri), - artifact_format=self.format.value, - ) - start_time = time.monotonic() - download_elapsed = await self.download(ctx, image_path) - - mount_start = time.monotonic() - _validate_squashfs_mount_paths(ctx.paths) - await self._mount_image(image_path, target_dir) - mount_elapsed = (time.monotonic() - mount_start) * 1000 - total_elapsed = (time.monotonic() - start_time) * 1000 - - logger.info( - "SquashFS registry artifact mounted", - cache_key=ctx.cache_key, - artifact_uri=_artifact_uri_for_logging(self.uri), - artifact_format=self.format.value, - download_ms=f"{download_elapsed:.1f}", - mount_ms=f"{mount_elapsed:.1f}", - total_ms=f"{total_elapsed:.1f}", - ) - return target_dir - - async def extract( - self, - ctx: RegistryArtifactMaterializationContext, - image_path: Path, - ) -> Path: - target_dir = ctx.paths.squashfs_extract_dir - _validate_cache_entry_path(ctx.paths) - if _is_reusable_extraction_dir( - target_dir, - defer_cleanup=ctx.defer_cleanup, - ): - return target_dir - - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - _validate_cache_entry_path(ctx.paths) - - logger.info( - "Extracting SquashFS registry artifact", - cache_key=ctx.cache_key, - artifact_uri=_artifact_uri_for_logging(self.uri), - artifact_format=self.format.value, - ) - start_time = time.monotonic() - download_elapsed = await self.download(ctx, image_path) - - temp_dir = self._temp_path(ctx, ".unsquashfs") - try: - if ctx.admission is not None: - extracted_size = await self._squashfs_extracted_size( - image_path, - allocation_unit=ctx.admission.allocation_unit, - ) - await ctx.admission.ensure_capacity(extracted_size) - extract_start = time.monotonic() - temp_dir.mkdir(parents=True, exist_ok=True) - await self._extract_image(image_path, temp_dir) - extract_elapsed = (time.monotonic() - extract_start) * 1000 - - _validate_cache_entry_path(ctx.paths) - try: - temp_dir.rename(target_dir) - total_elapsed = (time.monotonic() - start_time) * 1000 - logger.info( - "SquashFS registry artifact extracted", - cache_key=ctx.cache_key, - artifact_uri=_artifact_uri_for_logging(self.uri), - artifact_format=self.format.value, - download_ms=f"{download_elapsed:.1f}", - extract_ms=f"{extract_elapsed:.1f}", - total_ms=f"{total_elapsed:.1f}", - ) - except OSError: - _validate_cache_entry_path(ctx.paths) - if _is_reusable_extraction_dir( - target_dir, - defer_cleanup=ctx.defer_cleanup, - ): - logger.debug( - "SquashFS already extracted by another process", - cache_key=ctx.cache_key, - artifact_uri=_artifact_uri_for_logging(self.uri), - artifact_format=self.format.value, - ) - else: - raise - finally: - await _remove_tree_rejoin_on_cancel( - temp_dir, - defer_cleanup=ctx.defer_cleanup, - ) - - return target_dir - - async def _mount_image(self, image_path: Path, target_dir: Path) -> None: - """Mount a SquashFS image read-only at target_dir. - - Cancellation kills and reaps the mount subprocess before propagating, - so the caller's per-key lock covers the complete mount lifecycle. The - target therefore remains an unmounted cache miss if mount never took - effect, or is already mounted and reusable by the next admission. - - Args: - image_path: Local path of the SquashFS image. - target_dir: Existing directory to mount the image at. - - Raises: - SquashfsMountCommandError: The ``mount`` command failed. - """ - if registry_artifact_mounts.is_mount(target_dir): - return - - proc = await asyncio.create_subprocess_exec( - "mount", - "-t", - "squashfs", - "-o", - SQUASHFS_MOUNT_OPTIONS, - str(image_path), - str(target_dir), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - start_new_session=True, - ) - stdout, stderr = await communicate_process_group(proc) - - if proc.returncode == 0 or registry_artifact_mounts.is_mount(target_dir): - return - - output = (stderr or stdout).decode(errors="replace").strip() - raise SquashfsMountCommandError(output or "mount command failed") - - async def _extract_image(self, image_path: Path, target_dir: Path) -> None: - """Extract a SquashFS image to target_dir using unsquashfs. - - Cancellation kills and reaps the extractor before the caller removes - its scratch directory. Otherwise a live extractor could recreate - scratch after cleanup, outside startup-sweep discovery. - """ - unsquashfs = shutil.which("unsquashfs") - if unsquashfs is None: - raise RuntimeError("unsquashfs command is not installed") - - proc = await asyncio.create_subprocess_exec( - unsquashfs, - "-f", - "-d", - str(target_dir), - str(image_path), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - start_new_session=True, - ) - await communicate_process_group(proc) - - if proc.returncode == 0: - return - - raise RegistryArtifactExtractionError() - - async def _squashfs_extracted_size( - self, - image_path: Path, - *, - allocation_unit: int = 1, - ) -> int: - """Return a conservative allocated size for a SquashFS extraction.""" - unsquashfs = shutil.which("unsquashfs") - if unsquashfs is None: - raise RuntimeError("unsquashfs command is not installed") - - proc = await asyncio.create_subprocess_exec( - unsquashfs, - "-lln", - str(image_path), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - start_new_session=True, - ) - stdout, _ = await communicate_process_group(proc) - - if proc.returncode != 0: - raise RegistryArtifactExtractionError() - try: - return _squashfs_listing_size(stdout, allocation_unit=allocation_unit) - except Exception: - raise RegistryArtifactExtractionError() from None - - -@dataclass(frozen=True, slots=True) -class TarballArtifact(RegistryArtifact): - """Legacy gzip tarball registry environment.""" - - @property - def format(self) -> RegistryArtifactFormat: - return RegistryArtifactFormat.TAR_GZ - - def cached_path( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path] | None: - _validate_cache_entry_path(ctx.paths) - if _is_reusable_extraction_dir( - ctx.paths.tarball_target_dir, - defer_cleanup=ctx.defer_cleanup, - ): - logger.debug( - "Using cached tarball extraction", - cache_key=ctx.cache_key, - ) - return [ctx.paths.tarball_target_dir] - return None - - async def materialize( - self, ctx: RegistryArtifactMaterializationContext - ) -> list[Path]: - target_dir = ctx.paths.tarball_target_dir - _validate_cache_entry_path(ctx.paths) - logger.info( - "Materializing tarball registry artifact", - cache_key=ctx.cache_key, - artifact_uri=_artifact_uri_for_logging(self.uri), - artifact_format=self.format.value, - ) - start_time = time.monotonic() - - temp_tarball = self._temp_path(ctx, ".tar.gz") - temp_dir = self._temp_path(ctx, ".tmp") - - try: - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - _validate_cache_entry_path(ctx.paths) - - download_start = time.monotonic() - await self.download(ctx, temp_tarball) - download_elapsed = (time.monotonic() - download_start) * 1000 - - admission = ctx.admission - if admission is not None: - try: - extracted_size = await _run_blocking_rejoin_on_cancel( - lambda: _tarball_extracted_size( - temp_tarball, - allocation_unit=admission.allocation_unit, - ) - ) - except Exception: - raise RegistryArtifactExtractionError() from None - await admission.ensure_capacity(extracted_size) - - extract_start = time.monotonic() - temp_dir.mkdir(parents=True, exist_ok=True) - await self.extract(temp_tarball, temp_dir) - extract_elapsed = (time.monotonic() - extract_start) * 1000 - - _validate_cache_entry_path(ctx.paths) - try: - temp_dir.rename(target_dir) - total_elapsed = (time.monotonic() - start_time) * 1000 - logger.info( - "Tarball extracted and cached", - cache_key=ctx.cache_key, - artifact_uri=_artifact_uri_for_logging(self.uri), - artifact_format=self.format.value, - download_ms=f"{download_elapsed:.1f}", - extract_ms=f"{extract_elapsed:.1f}", - total_ms=f"{total_elapsed:.1f}", - ) - except OSError: - _validate_cache_entry_path(ctx.paths) - if _is_reusable_extraction_dir( - target_dir, - defer_cleanup=ctx.defer_cleanup, - ): - logger.debug( - "Tarball already extracted by another process", - cache_key=ctx.cache_key, - artifact_uri=_artifact_uri_for_logging(self.uri), - artifact_format=self.format.value, - ) - else: - raise - finally: - try: - await _remove_tree_rejoin_on_cancel( - temp_dir, - defer_cleanup=ctx.defer_cleanup, - ) - finally: - _remove_file_or_defer( - temp_tarball, - defer_cleanup=ctx.defer_cleanup, - ) - - return [target_dir] - - async def download( - self, - ctx: RegistryArtifactMaterializationContext, - output_path: Path, - ) -> None: - await _download_s3_artifact( - self.uri, - output_path, - admission=ctx.admission, - defer_cleanup=ctx.defer_cleanup, - ) - - async def extract(self, tarball_path: Path, target_dir: Path) -> None: - """Extract a supported registry tarball to target directory. - - A cancelled caller rejoins the non-interruptible extraction thread - before propagating cancellation so scratch cleanup cannot race a live - writer and leave undiscoverable ephemeral storage behind. - """ - - def _do_extract() -> None: - if tarball_path.name.endswith(".tar.gz"): - with tarfile.open(tarball_path, "r:gz") as tar: - tar.extractall(path=target_dir, filter="data") - return - - raise ValueError(f"Unsupported tarball format: {tarball_path}") - - try: - await _run_blocking_rejoin_on_cancel(_do_extract) - except Exception: - raise RegistryArtifactExtractionError() from None - - logger.debug( - "Tarball extracted", - target=str(target_dir), - artifact_format=_artifact_format(str(tarball_path)).value, - ) - - -async def _download_s3_artifact( - artifact_uri: str, - output_path: Path, - *, - admission: RegistryArtifactAdmission | None = None, - defer_cleanup: Callable[[Path], None], -) -> None: - """Download an S3 registry artifact to a local path.""" - try: - bucket, key = parse_s3_uri(artifact_uri) - except ValueError: - raise RegistryArtifactUriError("Invalid registry artifact URI") from None - try: - if admission is None: - await blob.download_file_to_path( - key=key, - bucket=bucket, - output_path=output_path, - defer_cleanup=defer_cleanup, - redact_log_identifiers=True, - ) - else: - await blob.download_file_to_path( - key=key, - bucket=bucket, - output_path=output_path, - max_bytes=admission.max_bytes, - ensure_capacity=admission.ensure_capacity, - defer_cleanup=defer_cleanup, - redact_log_identifiers=True, - ) - except FileNotFoundError as e: - request = httpx.Request("GET", artifact_uri) - response = httpx.Response(status_code=404, request=request) - raise httpx.HTTPStatusError( - f"Registry artifact not found: {_artifact_uri_for_logging(artifact_uri)}", - request=request, - response=response, - ) from e - - -def compute_registry_artifact_cache_key(artifact_uri: str) -> str: - """Compute the local cache key for a registry artifact URI.""" - if not artifact_uri: - return "base" - # S3 keys are byte-sensitive, so hash the exact URI used for retrieval. - return hashlib.sha256(artifact_uri.encode()).hexdigest()[:16] - - -def bundled_builtin_registry_uri(version: str) -> str: - """Return the pseudo-URI for the installed builtin registry package.""" - return f"{BUNDLED_BUILTIN_REGISTRY_URI_PREFIX}{version}" - - -def _bundled_builtin_registry_version(artifact_uri: str) -> str | None: - """Return the builtin registry version encoded in a bundled pseudo-URI.""" - if not artifact_uri.startswith(BUNDLED_BUILTIN_REGISTRY_URI_PREFIX): - return None - version = artifact_uri.removeprefix(BUNDLED_BUILTIN_REGISTRY_URI_PREFIX) - return version or None - - -def _bundled_builtin_registry_import_paths(version: str) -> list[Path]: - """Return import paths for the current builtin registry and its dependencies. - - Dependencies always live in the executor's site-packages. For editable - installs the parent of ``package_dir`` (the package wrapper, e.g. - ``packages/tracecat-registry/``) is exposed first so its ``tracecat_registry/`` - shadows any stale copy in site-packages. - """ - installed_version = tracecat_registry.__version__ - if version != installed_version: - raise RuntimeError( - "Bundled builtin registry version does not match installed version: " - f"requested={version!r}, installed={installed_version!r}" - ) - - package_file = tracecat_registry.__file__ - if package_file is None: - raise RuntimeError("Installed tracecat_registry package has no __file__") - - site_packages_path = sysconfig.get_path("purelib") - if site_packages_path is None: - raise RuntimeError("Could not resolve installed Python site-packages path") - - site_packages = Path(site_packages_path).resolve() - if not site_packages.exists(): - raise RuntimeError( - f"Installed Python site-packages path does not exist: {site_packages}" - ) - - package_dir = Path(package_file).resolve().parent - if package_dir.is_relative_to(site_packages): - return [site_packages] - - return [package_dir.parent, site_packages] - - -def _squashfs_sidecar_uri(tarball_uri: str) -> str | None: - """Return the sibling SquashFS URI for registry site-packages tarballs.""" - if not tarball_uri.endswith("site-packages.tar.gz"): - return None - return tarball_uri.removesuffix(".tar.gz") + ".squashfs" - - -def _tarball_uri_for_squashfs(squashfs_uri: str) -> str | None: - """Return the sibling gzip tarball URI for registry SquashFS artifacts.""" - if not squashfs_uri.endswith("site-packages.squashfs"): - return None - return squashfs_uri.removesuffix(".squashfs") + ".tar.gz" - - -def _artifact_format(artifact_uri: str) -> RegistryArtifactFormat: - """Return the materialization format for an artifact URI.""" - if artifact_uri.endswith(".squashfs"): - return RegistryArtifactFormat.SQUASHFS - return RegistryArtifactFormat.TAR_GZ - - -def _is_cache_entry_uri(artifact_uri: str) -> bool: - """Return whether an artifact URI materializes into an evictable cache entry. - - The bundled builtin registry is served from the executor image and never - writes into the cache directory, so it is exempt from eviction accounting. - """ - return _bundled_builtin_registry_version(artifact_uri) is None - - -def _allocated_size_bound(size_bytes: int, *, allocation_unit: int) -> int: - """Round one filesystem object up to its minimum allocated footprint.""" - if size_bytes < 0: - raise ValueError("size_bytes must be non-negative") - if allocation_unit <= 0: - raise ValueError("allocation_unit must be positive") - return ( - max(1, (size_bytes + allocation_unit - 1) // allocation_unit) * allocation_unit - ) - - -_DIRECTORY_ENTRY_OVERHEAD_BYTES = 32 -"""Conservative per-child filesystem directory-record overhead.""" - - -def _directory_entry_size_bound( - child_name: str, - *, - allocation_unit: int, -) -> int: - """Reserve directory storage for one unique child name. - - Each child receives at least one allocation unit. This intentionally - overbounds filesystem-specific records and index blocks without assuming a - particular executor filesystem layout. - """ - return _allocated_size_bound( - _DIRECTORY_ENTRY_OVERHEAD_BYTES + len(os.fsencode(child_name)), - allocation_unit=allocation_unit, - ) - - -def _tarball_extracted_size( - tarball_path: Path, - *, - allocation_unit: int = 1, -) -> int: - """Return a conservative allocated size bound for a tarball extraction.""" - total_bytes = 0 - root_path = PurePosixPath(".") - has_explicit_root_directory = False - required_parent_dirs: set[PurePosixPath] = set() - explicit_dirs: set[PurePosixPath] = set() - directory_children: dict[PurePosixPath, set[str]] = {} - - def record_directory_child(path: PurePosixPath) -> None: - if path == root_path: - return - directory_children.setdefault(path.parent, set()).add(path.name) - - with tarfile.open(tarball_path, "r:gz") as tar: - for member in tar: - if member.size < 0: - raise ValueError( - f"Registry tarball member has a negative size: {member.name}" - ) - total_bytes += _allocated_size_bound( - member.size, - allocation_unit=allocation_unit, - ) - member_path = PurePosixPath(member.name) - record_directory_child(member_path) - if member.isdir(): - explicit_dirs.add(member_path) - has_explicit_root_directory |= member_path == root_path - for parent in member_path.parents: - if parent == root_path: - break - required_parent_dirs.add(parent) - record_directory_child(parent) - # Extraction creates a target root even when the tar manifest omits it. - if not has_explicit_root_directory: - total_bytes += _allocated_size_bound(0, allocation_unit=allocation_unit) - implicit_parent_dirs = required_parent_dirs - explicit_dirs - total_bytes += len(implicit_parent_dirs) * allocation_unit - total_bytes += sum( - _directory_entry_size_bound( - child_name, - allocation_unit=allocation_unit, - ) - for child_names in directory_children.values() - for child_name in child_names - ) - return total_bytes - - -def _squashfs_listing_size(output: bytes, *, allocation_unit: int = 1) -> int: - """Bound allocated bytes from ``unsquashfs -lln`` output, failing closed.""" - total_bytes = 0 - parsed_entries = 0 - root_path = PurePosixPath(".") - listing_root: str | None = None - required_dirs: set[PurePosixPath] = {root_path} - listed_dirs: set[PurePosixPath] = set() - directory_children: dict[PurePosixPath, set[str]] = {} - - def record_directory_child(path: PurePosixPath) -> None: - if path == root_path: - return - directory_children.setdefault(path.parent, set()).add(path.name) - - for raw_line in output.decode(errors="replace").splitlines(): - line = raw_line.strip() - if not line: - continue - fields = line.split(maxsplit=5) - mode = fields[0] - if len(mode) != 10 or mode[0] not in "bcdlps-": - continue - if len(fields) < 6 or "/" not in fields[1] or not fields[2].isdigit(): - raise ValueError("Could not parse SquashFS listing line") - - listed_path_text = fields[5] - if mode[0] == "l": - listed_path_text = listed_path_text.split(" -> ", maxsplit=1)[0] - listed_path = PurePosixPath(listed_path_text) - if listed_path.is_absolute() or not listed_path.parts: - raise ValueError("Could not parse SquashFS listing path") - if listing_root is None: - listing_root = listed_path.parts[0] - elif listed_path.parts[0] != listing_root: - raise ValueError("Inconsistent SquashFS listing root") - - relative_parts = listed_path.parts[1:] - entry_path = PurePosixPath(*relative_parts) if relative_parts else root_path - if entry_path == root_path and mode[0] != "d": - raise ValueError("Could not parse SquashFS listing root") - - parsed_entries += 1 - total_bytes += _allocated_size_bound( - int(fields[2]), - allocation_unit=allocation_unit, - ) - record_directory_child(entry_path) - if mode[0] == "d": - listed_dirs.add(entry_path) - for parent in entry_path.parents: - required_dirs.add(parent) - if parent == root_path: - break - record_directory_child(parent) - if parsed_entries == 0: - raise ValueError("Could not parse any SquashFS listing entries") - total_bytes += len(required_dirs - listed_dirs) * allocation_unit - total_bytes += sum( - _directory_entry_size_bound( - child_name, - allocation_unit=allocation_unit, - ) - for child_names in directory_children.values() - for child_name in child_names - ) - return total_bytes diff --git a/tracecat/executor/registry_artifact_mounts.py b/tracecat/executor/registry_artifact_mounts.py deleted file mode 100644 index 98db4250bc..0000000000 --- a/tracecat/executor/registry_artifact_mounts.py +++ /dev/null @@ -1,30 +0,0 @@ -"""Fail-closed mount-state inspection for registry artifact caches.""" - -from __future__ import annotations - -import os -import stat -from pathlib import Path - - -def is_mount(path: Path) -> bool: - """Return whether a path is mounted without hiding inspection failures. - - ``Path.is_mount()`` delegates to ``os.path.ismount()``, which converts every - ``OSError`` from ``lstat`` into ``False``. Cache cleanup must distinguish a - missing mount directory from an unreadable one so it never deletes the - backing image while mount state is unknown. - """ - try: - path_stat = path.lstat() - except FileNotFoundError: - return False - - if stat.S_ISLNK(path_stat.st_mode): - return False - - parent = Path(os.path.realpath(path / "..", strict=True)) - parent_stat = parent.lstat() - return ( - path_stat.st_dev != parent_stat.st_dev or path_stat.st_ino == parent_stat.st_ino - ) diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py deleted file mode 100644 index 6e4360689b..0000000000 --- a/tracecat/executor/registry_artifact_storage.py +++ /dev/null @@ -1,1022 +0,0 @@ -"""Disk-budget enforcement and cleanup for registry artifact caches.""" - -from __future__ import annotations - -import asyncio -import os -import shutil -import stat -import time -from collections.abc import Iterable -from dataclasses import dataclass -from pathlib import Path - -from tracecat import config -from tracecat.executor import registry_artifact_mounts -from tracecat.executor.registry_artifact_cache_state import ( - _RegistryArtifactCacheState, -) -from tracecat.executor.registry_artifact_materialization import ( - RegistryArtifactAdmission, - _allocated_size_bound, - _is_reusable_cache_directory, - _is_reusable_cache_file, - _rejoin_future_on_cancel, - _run_blocking_rejoin_on_cancel, - _validate_cache_entries_directory, - _validate_cache_entry_path, - _validate_cache_work_directory, -) -from tracecat.logger import logger -from tracecat.sandbox.utils import communicate_process_group - - -class RegistryArtifactCacheCapacityError(RuntimeError): - """A cold artifact cannot fit within the configured cache byte budget.""" - - def __init__( - self, - *, - current_bytes: int, - additional_bytes: int, - max_bytes: int, - ) -> None: - super().__init__( - "Registry artifact admission exceeds the cache byte budget: " - f"current_bytes={current_bytes}, additional_bytes={additional_bytes}, " - f"max_bytes={max_bytes}" - ) - self.current_bytes = current_bytes - self.additional_bytes = additional_bytes - self.max_bytes = max_bytes - - -@dataclass(frozen=True, slots=True) -class RegistryArtifactEviction: - """Outcome of atomically retiring and physically deleting one entry.""" - - retired: bool - reclaimed: bool - - -@dataclass(frozen=True, slots=True) -class _RegistryArtifactEvictionPass: - """Result of applying one cache limit policy to evictable entries.""" - - total_bytes: int - fits: bool - exhausted_candidates: bool - - -@dataclass(frozen=True, slots=True) -class RegistryArtifactCacheEntry: - """Measured on-disk footprint and recency for one registry artifact key.""" - - cache_key: str - size_bytes: int - last_used: float - - -def _allocated_stat_size( - file_stat: os.stat_result, - *, - allocation_unit: int, -) -> int: - """Return allocated bytes while charging at least one unit per inode.""" - if allocation_unit <= 0: - raise ValueError("allocation_unit must be positive") - blocks = getattr(file_stat, "st_blocks", None) - if blocks is None: - allocated_bytes = file_stat.st_size - else: - allocated_bytes = blocks * 512 - return max(allocation_unit, allocated_bytes) - - -def _filesystem_allocation_unit(path: Path) -> int: - """Return the allocation unit for a path or its nearest existing parent.""" - candidate = path - while True: - try: - filesystem = os.statvfs(candidate) - break - except FileNotFoundError: - parent = candidate.parent - if parent == candidate: - raise - candidate = parent - - return filesystem.f_frsize or filesystem.f_bsize or 1 - - -def _directory_footprint( - directory: Path, - *, - allocation_unit: int | None = None, - pruned_directories: Iterable[Path] = (), - include_root: bool = True, -) -> int: - """Return the allocated footprint of a cache directory tree. - - Args: - directory: Cache directory to measure. - pruned_directories: Directories whose own inodes are counted without - descending into their contents. - include_root: Whether to count the root directory's inode. Disable this - when a separate structural scan already owns that inode. - - Returns: - Total allocated bytes of unique contained inodes, or zero when the - directory is missing. - """ - - def raise_walk_error(error: OSError) -> None: - raise error - - if allocation_unit is None: - allocation_unit = _filesystem_allocation_unit(directory) - - total_bytes = 0 - seen_inodes: set[tuple[int, int]] = set() - pruned_paths = frozenset(pruned_directories) - - def allocated_inode_size(file_stat: os.stat_result) -> int: - inode_key = (file_stat.st_dev, file_stat.st_ino) - if inode_key in seen_inodes: - return 0 - seen_inodes.add(inode_key) - return _allocated_stat_size( - file_stat, - allocation_unit=allocation_unit, - ) - - try: - walker = os.walk(directory, onerror=raise_walk_error) - for root, dirs, files in walker: - root_path = Path(root) - if include_root or root_path != directory: - try: - total_bytes += allocated_inode_size(os.lstat(root)) - except FileNotFoundError: - continue - for file_name in files: - try: - total_bytes += allocated_inode_size( - os.lstat(os.path.join(root, file_name)) - ) - except FileNotFoundError: - continue - traversed_directories: list[str] = [] - for directory_name in dirs: - child_path = root_path / directory_name - try: - directory_stat = child_path.lstat() - except FileNotFoundError: - continue - if stat.S_ISLNK(directory_stat.st_mode) or child_path in pruned_paths: - total_bytes += allocated_inode_size(directory_stat) - continue - traversed_directories.append(directory_name) - dirs[:] = traversed_directories - except FileNotFoundError: - return 0 - return total_bytes - - -def _delete_cache_path(path: Path) -> bool: - """Best-effort delete one cache path while reporting filesystem failures.""" - try: - if path.is_dir(): - shutil.rmtree(path) - else: - path.unlink(missing_ok=True) - except OSError as e: - logger.warning( - "Failed to delete registry artifact cache path", - path=str(path), - error=str(e), - ) - return False - return True - - -async def _delete_cache_path_off_loop(path: Path) -> bool: - """Delete one path without abandoning its worker thread on cancellation.""" - return await _run_blocking_rejoin_on_cancel(lambda: _delete_cache_path(path)) - - -def _unique_work_path(root: Path, cache_key: str) -> Path: - """Return a unique path beneath a cache work directory.""" - _validate_cache_work_directory(root) - root.mkdir(parents=True, exist_ok=True) - _validate_cache_work_directory(root) - unique_id = time.time_ns() - while True: - path = root / f"{cache_key}.{os.getpid()}.{unique_id}" - if not path.exists(): - return path - unique_id += 1 - - -def _move_entry_to_trash(entry_dir: Path, trash_dir: Path, cache_key: str) -> Path: - """Atomically retire one cache entry and return its trash path.""" - trash_path = _unique_work_path(trash_dir, cache_key) - entry_dir.rename(trash_path) - return trash_path - - -class _RegistryArtifactCacheStorage(_RegistryArtifactCacheState): - """Adds startup recovery, eviction, and byte-budget enforcement.""" - - def _cache_structural_footprint(self) -> int: - """Measure cache roots and non-entry data exactly once. - - Entry, staging, and trash contents are measured separately. Pruning - those trees here retains their root-directory blocks, including growth - caused by child names, while avoiding double-counting their contents. - """ - return _directory_footprint( - self.cache_dir, - pruned_directories=( - self.entries_dir, - self.staging_dir, - self.trash_dir, - ), - ) - - async def ensure_swept(self) -> None: - """Run the startup sweep exactly once successfully, off the event loop. - - Idempotent and cancellation-safe under concurrency: every caller joins - one stored sweep task, and cancelling a waiter never abandons or - restarts its live sweep. The lock is deliberately held while awaiting - that shared task so queued callers observe its result before proceeding. - Failures clear the task so the next caller retries. The sweep runs - before the first lease or materialization, so it never observes - in-flight cache entries. - """ - self._assert_owner_loop() - if self._swept: - return - async with self._sweep_lock: - if self._swept: - return - sweep_task = self._sweep_task - if sweep_task is None or ( - sweep_task.done() - and (sweep_task.cancelled() or sweep_task.exception() is not None) - ): - sweep_task = asyncio.ensure_future( - asyncio.to_thread(self._sweep_startup_state) - ) - self._sweep_task = sweep_task - try: - await asyncio.shield(sweep_task) - except asyncio.CancelledError: - if sweep_task.cancelled() and self._sweep_task is sweep_task: - self._sweep_task = None - raise - except Exception: - if self._sweep_task is sweep_task: - self._sweep_task = None - raise - self._swept = True - - def _admission_for(self, cache_key: str) -> RegistryArtifactAdmission | None: - """Return byte-bound admission controls for one cold cache key.""" - max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES - if max_bytes <= 0: - return None - - allocation_unit = _filesystem_allocation_unit(self.cache_dir) - verified_headroom_bytes = 0 - - async def ensure_capacity(additional_bytes: int) -> None: - nonlocal verified_headroom_bytes - allocated_bytes = _allocated_size_bound( - additional_bytes, - allocation_unit=allocation_unit, - ) - if allocated_bytes <= verified_headroom_bytes: - verified_headroom_bytes -= allocated_bytes - return - - verified_headroom_bytes = await self._ensure_cache_capacity( - additional_bytes=allocated_bytes, - protected_key=cache_key, - max_bytes=max_bytes, - ) - - return RegistryArtifactAdmission( - max_bytes=max_bytes, - allocation_unit=allocation_unit, - ensure_capacity=ensure_capacity, - ) - - async def _unmount_idle_entry(self, cache_key: str) -> None: - """Best-effort unmount an entry after its final lease is released.""" - try: - unmounted = await self._unmount_entry(cache_key) - except OSError as e: - self._failed_unmounts.add(cache_key) - logger.warning( - "Failed to release idle registry artifact mount", - cache_key=cache_key, - error=str(e), - ) - return - - if unmounted: - self._failed_unmounts.discard(cache_key) - return - - mount_dir = self._paths_for(cache_key).squashfs_mount_dir - try: - retry = self._refcount( - cache_key - ) == 0 and registry_artifact_mounts.is_mount(mount_dir) - except OSError as e: - retry = True - logger.warning( - "Failed to inspect idle registry artifact mount", - cache_key=cache_key, - mount_dir=str(mount_dir), - error=str(e), - ) - - if retry: - self._failed_unmounts.add(cache_key) - else: - self._failed_unmounts.discard(cache_key) - - async def _retry_failed_unmounts(self, *, excluded: set[str]) -> None: - """Retry prior unmount failures once on a later lease cleanup.""" - for cache_key in sorted(self._failed_unmounts - excluded): - await self._unmount_idle_entry(cache_key) - - async def _converge_cache_budget(self) -> None: - """Bring an idle cache back under budget after a lease is released. - - Successful materialization enforces the budget after publication while - protecting the new entry. The cache can still sit over budget while - entries are leased. This runs on release, when every newly idle entry is - evictable. The scan is skipped entirely unless a materialization attempt - has occurred since the last successful enforcement. - - Each successful pass consumes the dirty signal before its awaited scan. - A follow-up pass therefore occurs only when a concurrent materialization - sets the flag again. Without new materializations the loop terminates, - while an over-budget or failed scan restores the flag and breaks so it - cannot spin while entries remain leased. Cancellation also restores the - consumed flag before propagating. - """ - while self._budget_dirty: - self._budget_dirty = False - try: - within_budget = await self._enforce_cache_budget() - except OSError as e: - logger.warning( - "Failed to converge registry artifact cache to budget", - cache_dir=str(self.cache_dir), - error=str(e), - ) - self._budget_dirty = True - break - except BaseException: - # Cancellation must re-arm the consumed dirty signal before propagating. - self._budget_dirty = True - raise - - if not within_budget: - self._budget_dirty = True - break - - async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bool: - """Evict least-recently-used idle entries until the cache fits its budget. - - The admission lock excludes cold writers before the budget lock begins - a scan/select/evict pass. Callers invoke enforcement without holding a - per-key lock. Cold writers already hold the admission lock and use - ``_ensure_cache_capacity`` for their staged reservations instead. - - Args: - protected_key: Newly materialized cache key. It is counted against - the budget when present but never evicted. None when enforcing - against idle entries after leases are released. - - Returns: - Whether the cache is within budget once eviction has finished. - """ - async with self._admission_lock: - async with self._budget_lock: - within_budget = await self._enforce_cache_budget_locked( - protected_key=protected_key - ) - if within_budget: - # Clear while cold writers remain excluded so a later - # materialization cannot have its dirty signal erased. - self._budget_dirty = False - return within_budget - - async def _enforce_cache_budget_locked( - self, - *, - protected_key: str | None, - ) -> bool: - """Enforce entry and byte limits while both cache-wide locks are held.""" - if not await self._cleanup_cache_work_dirs(): - return False - - max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES - max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES - if max_entries <= 0 and max_bytes <= 0: - return True - - entries, structural_bytes, staging_bytes, trash_bytes = await asyncio.gather( - asyncio.to_thread(self._scan_cache_entries), - asyncio.to_thread(self._cache_structural_footprint), - asyncio.to_thread( - _directory_footprint, - self.staging_dir, - include_root=False, - ), - asyncio.to_thread( - _directory_footprint, - self.trash_dir, - include_root=False, - ), - ) - total_bytes = ( - structural_bytes - + staging_bytes - + trash_bytes - + sum(entry.size_bytes for entry in entries.values()) - ) - protected = set() if protected_key is None else {protected_key} - eviction_pass = await self._evict_until_fits( - entries, - total_bytes=total_bytes, - excluded=protected, - max_entries=max_entries, - max_bytes=max_bytes, - ) - if not eviction_pass.fits: - if eviction_pass.exhausted_candidates: - logger.warning( - "Registry artifact cache is over budget but every entry is in use", - cache_dir=str(self.cache_dir), - entries=len(entries), - max_entries=max_entries, - total_bytes=eviction_pass.total_bytes, - max_bytes=max_bytes, - ) - return False - - return True - - async def _cleanup_cache_work_dirs(self) -> bool: - """Rejoin cache-wide cleanup workers and report complete reclamation.""" - cleanup = asyncio.gather( - asyncio.to_thread(self._clear_work_dir, self.trash_dir), - asyncio.to_thread(self._retry_deferred_staging_cleanup), - ) - trash_clean, staging_clean = await _rejoin_future_on_cancel(cleanup) - return trash_clean and staging_clean - - async def _ensure_cache_capacity( - self, - *, - additional_bytes: int, - protected_key: str, - max_bytes: int, - ) -> int: - """Reserve peak bytes for a cold writer without exceeding the cap. - - The caller holds the admission lock and its key lock. Every normal - budget pass takes the admission lock first, so acquiring the budget - lock here cannot deadlock with eviction of the protected key. - - Returns the additional byte headroom proven by the same scan after the - requested allocation. The admission callback consumes that headroom - before rescanning, which keeps unknown-length chunked downloads cheap. - """ - if additional_bytes < 0: - raise ValueError("additional_bytes must be non-negative") - - async with self._budget_lock: - cleanup_complete = await self._cleanup_cache_work_dirs() - ( - entries, - structural_bytes, - staging_bytes, - trash_bytes, - ) = await asyncio.gather( - asyncio.to_thread(self._scan_cache_entries), - asyncio.to_thread(self._cache_structural_footprint), - asyncio.to_thread( - _directory_footprint, - self.staging_dir, - include_root=False, - ), - asyncio.to_thread( - _directory_footprint, - self.trash_dir, - include_root=False, - ), - ) - total_bytes = ( - structural_bytes - + staging_bytes - + trash_bytes - + sum(entry.size_bytes for entry in entries.values()) - ) - non_evictable_bytes = ( - structural_bytes - + staging_bytes - + trash_bytes - + sum( - entry.size_bytes - for entry in entries.values() - if entry.cache_key == protected_key - or self._refcount(entry.cache_key) > 0 - or ( - (runtime := self._runtime.get(entry.cache_key)) is not None - and runtime.lock.locked() - ) - ) - ) - if non_evictable_bytes + additional_bytes > max_bytes: - raise RegistryArtifactCacheCapacityError( - current_bytes=non_evictable_bytes, - additional_bytes=additional_bytes, - max_bytes=max_bytes, - ) - if not cleanup_complete and total_bytes + additional_bytes > max_bytes: - raise RegistryArtifactCacheCapacityError( - current_bytes=total_bytes, - additional_bytes=additional_bytes, - max_bytes=max_bytes, - ) - eviction_pass = await self._evict_until_fits( - entries, - total_bytes=total_bytes, - excluded={protected_key}, - max_entries=0, - max_bytes=max_bytes, - additional_bytes=additional_bytes, - ) - if not eviction_pass.fits: - raise RegistryArtifactCacheCapacityError( - current_bytes=eviction_pass.total_bytes, - additional_bytes=additional_bytes, - max_bytes=max_bytes, - ) - - return max_bytes - eviction_pass.total_bytes - additional_bytes - - async def _evict_until_fits( - self, - entries: dict[str, RegistryArtifactCacheEntry], - *, - total_bytes: int, - excluded: set[str], - max_entries: int, - max_bytes: int, - additional_bytes: int = 0, - ) -> _RegistryArtifactEvictionPass: - """Apply shared LRU retirement mechanics until the given limits fit.""" - skipped = set(excluded) - - while (max_entries > 0 and len(entries) > max_entries) or ( - max_bytes > 0 and total_bytes + additional_bytes > max_bytes - ): - candidate = self._least_recently_used( - entries.values(), - excluded=skipped, - ) - if candidate is None: - return _RegistryArtifactEvictionPass( - total_bytes=total_bytes, - fits=False, - exhausted_candidates=True, - ) - - eviction = await self._evict_entry(candidate.cache_key) - if eviction.retired: - del entries[candidate.cache_key] - if not eviction.reclaimed: - return _RegistryArtifactEvictionPass( - total_bytes=total_bytes, - fits=False, - exhausted_candidates=False, - ) - total_bytes -= candidate.size_bytes - else: - skipped.add(candidate.cache_key) - - return _RegistryArtifactEvictionPass( - total_bytes=total_bytes, - fits=True, - exhausted_candidates=False, - ) - - def _least_recently_used( - self, - entries: Iterable[RegistryArtifactCacheEntry], - *, - excluded: set[str], - ) -> RegistryArtifactCacheEntry | None: - """Return the least recently used idle entry eligible for eviction.""" - eligible = [ - entry - for entry in entries - if entry.cache_key not in excluded and self._refcount(entry.cache_key) == 0 - ] - if not eligible: - return None - return min(eligible, key=self._recency) - - def _recency(self, entry: RegistryArtifactCacheEntry) -> float: - """Return the most recent known use time for a cache entry.""" - runtime = self._runtime.get(entry.cache_key) - if runtime is None: - return entry.last_used - return max(entry.last_used, runtime.last_used) - - async def _unmount_entry(self, cache_key: str) -> bool: - """Unmount one idle cache entry while retaining its reusable image. - - Loop-device reclamation is independent from disk-budget eviction. The - per-key lock and lease recheck prevent an entry from being unmounted - while an action is importing from it. The image and empty mount - directory remain cached so a later admission can remount without - downloading the artifact again. - - Args: - cache_key: Cache key whose mounted artifact should be released. - - Returns: - Whether a mounted entry was unmounted. - """ - runtime = self._runtime.get(cache_key) - if runtime is not None and runtime.lock.locked(): - logger.debug( - "Skipping unmount of busy registry artifact", - cache_key=cache_key, - ) - return False - - async with self._runtime_lock(cache_key): - if self._refcount(cache_key) > 0: - return False - - mount_dir = self._paths_for(cache_key).squashfs_mount_dir - if not registry_artifact_mounts.is_mount(mount_dir): - return False - if not await self._unmount(mount_dir): - logger.warning( - "Failed to unmount registry artifact for loop-device reclamation", - cache_key=cache_key, - mount_dir=str(mount_dir), - ) - return False - - logger.info( - "Unmounted idle registry artifact", - cache_key=cache_key, - mount_dir=str(mount_dir), - ) - return True - - async def _evict_entry(self, cache_key: str) -> RegistryArtifactEviction: - """Remove one cache entry from disk, unmounting it first. - - The entry is skipped rather than forced when it is leased, busy, or - cannot be unmounted: deleting the image file behind a live mount would - leave an open-file zombie holding the loop device. - - After unmounting, the entry root is atomically renamed into ``trash`` - under the per-key lock. The lock is then released before physical - deletion runs in a worker thread. - - Args: - cache_key: Cache key to evict. - - Returns: - Whether the entry was retired and its bytes were reclaimed. - """ - runtime = self._runtime.get(cache_key) - if runtime is not None and runtime.lock.locked(): - logger.debug( - "Skipping eviction of busy registry artifact", - cache_key=cache_key, - ) - return RegistryArtifactEviction(retired=False, reclaimed=False) - - async with self._runtime_lock(cache_key): - if self._refcount(cache_key) > 0: - return RegistryArtifactEviction(retired=False, reclaimed=False) - - paths = self._paths_for(cache_key) - _validate_cache_entry_path(paths) - if not paths.entry_dir.exists(): - return RegistryArtifactEviction(retired=True, reclaimed=True) - if registry_artifact_mounts.is_mount( - paths.squashfs_mount_dir - ) and not await self._unmount(paths.squashfs_mount_dir): - logger.warning( - "Failed to unmount registry artifact, skipping eviction", - cache_key=cache_key, - mount_dir=str(paths.squashfs_mount_dir), - ) - return RegistryArtifactEviction(retired=False, reclaimed=False) - - _validate_cache_entry_path(paths) - try: - trash_path = _move_entry_to_trash( - paths.entry_dir, - self.trash_dir, - cache_key, - ) - except OSError as e: - self._budget_dirty = True - logger.warning( - "Failed to retire registry artifact cache entry", - cache_key=cache_key, - entry_dir=str(paths.entry_dir), - error=str(e), - ) - return RegistryArtifactEviction(retired=False, reclaimed=False) - - try: - deleted = await _delete_cache_path_off_loop(trash_path) - except BaseException: - self._budget_dirty = True - raise - if not deleted: - self._budget_dirty = True - logger.warning( - "Registry artifact eviction remains pending physical deletion", - cache_key=cache_key, - trash_path=str(trash_path), - ) - return RegistryArtifactEviction(retired=True, reclaimed=False) - - logger.info("Evicted registry artifact from cache", cache_key=cache_key) - return RegistryArtifactEviction(retired=True, reclaimed=True) - - async def _unmount(self, mount_dir: Path) -> bool: - """Unmount a SquashFS artifact directory, releasing its loop device. - - Cancellation kills and reaps the umount subprocess before propagating, - so the caller's per-key lock covers the complete unmount lifecycle. If - umount never took effect, the mounted entry stays consistent and can be - reused; if it already took effect, the missing extraction directory - makes the entry a plain cache miss on the next admission. - """ - umount = shutil.which("umount") - if umount is None: - return False - - proc = await asyncio.create_subprocess_exec( - umount, - str(mount_dir), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - start_new_session=True, - ) - stdout, stderr = await communicate_process_group(proc) - if proc.returncode == 0 or not registry_artifact_mounts.is_mount(mount_dir): - return True - - logger.warning( - "umount command failed", - mount_dir=str(mount_dir), - output=(stderr or stdout).decode(errors="replace").strip(), - ) - return False - - def _scan_cache_entries(self) -> dict[str, RegistryArtifactCacheEntry]: - """Measure every registry artifact entry currently on disk.""" - return { - cache_key: self._measure_entry(cache_key) - for cache_key in self._discover_cache_keys() - } - - def _discover_cache_keys(self) -> set[str]: - """Return cache keys represented by atomic entry directories.""" - _validate_cache_entries_directory(self.entries_dir) - try: - entries = list(os.scandir(self.entries_dir)) - except FileNotFoundError: - return set() - - return { - entry.name - for entry in entries - if entry.name and entry.is_dir(follow_symlinks=False) - } - - def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: - """Measure the on-disk footprint and recency of one cache entry. - - Every entry-owned inode is included, including paths created by mutable - consumers. Active mounted contents are pruned because the image already - accounts for their backing bytes. - """ - paths = self._paths_for(cache_key) - allocation_unit = _filesystem_allocation_unit(self.cache_dir) - try: - mount_is_active = registry_artifact_mounts.is_mount( - paths.squashfs_mount_dir - ) - except FileNotFoundError: - mount_is_active = False - pruned_directories = (paths.squashfs_mount_dir,) if mount_is_active else () - size_bytes = _directory_footprint( - paths.entry_dir, - allocation_unit=allocation_unit, - pruned_directories=pruned_directories, - ) - - try: - last_used = paths.entry_dir.stat().st_mtime - except FileNotFoundError: - last_used = 0.0 - - return RegistryArtifactCacheEntry( - cache_key=cache_key, - size_bytes=size_bytes, - last_used=last_used, - ) - - def _sweep_startup_state(self) -> None: - """Reclaim orphaned cache state left behind by a previous process. - - Scratch and trash paths from interrupted work are removed, and active - entries are trimmed to budget using entry-root mtimes as LRU order. - - The worker warms this sweep before activities can run; lazy first-use - sweeping remains a safe fallback. - """ - if not self.cache_dir.is_dir(): - self._budget_dirty = False - return - - try: - staging_clean = self._clear_work_dir( - self.staging_dir, - remember_failures=True, - ) - trash_clean = self._clear_work_dir(self.trash_dir) - cleanup_complete = staging_clean and trash_clean - within_budget = cleanup_complete and self._trim_startup_cache() - self._budget_dirty = not (cleanup_complete and within_budget) - except OSError as e: - logger.warning( - "Failed to sweep registry artifact cache", - cache_dir=str(self.cache_dir), - error=str(e), - ) - raise - - def _clear_work_dir( - self, - work_dir: Path, - *, - remember_failures: bool = False, - ) -> bool: - """Best-effort remove every child of a staging or trash directory.""" - _validate_cache_work_directory(work_dir) - try: - paths = list(work_dir.iterdir()) - except FileNotFoundError: - return True - except OSError as e: - logger.warning( - "Failed to inspect registry artifact work directory", - path=str(work_dir), - error=str(e), - ) - raise - - deleted = True - for path in paths: - if _delete_cache_path(path): - if remember_failures: - self._deferred_staging_cleanup.discard(path) - logger.info( - "Removed registry artifact work path", - path=str(path), - ) - else: - deleted = False - if remember_failures: - self._deferred_staging_cleanup.add(path) - return deleted - - def _retry_deferred_staging_cleanup(self) -> bool: - """Retry exact failed paths without sweeping live staging work.""" - _validate_cache_work_directory(self.staging_dir) - for path in tuple(self._deferred_staging_cleanup): - if _delete_cache_path(path): - self._deferred_staging_cleanup.discard(path) - return not self._deferred_staging_cleanup - - def _trim_startup_cache(self) -> bool: - """Trim the cache to budget before any artifact is leased. - - Returns whether active entries and pending physical deletion fit within - the configured budget. - """ - max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES - max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES - entries = self._scan_cache_entries() - for entry in tuple(entries.values()): - paths = self._paths_for(entry.cache_key) - if ( - registry_artifact_mounts.is_mount(paths.squashfs_mount_dir) - or _is_reusable_cache_file(paths.squashfs_image_path) - or _is_reusable_cache_directory(paths.squashfs_extract_dir) - or _is_reusable_cache_directory(paths.tarball_target_dir) - ): - continue - - try: - trash_path = _move_entry_to_trash( - paths.entry_dir, - self.trash_dir, - entry.cache_key, - ) - except OSError as e: - logger.warning( - "Failed to retire incomplete registry artifact during startup sweep", - cache_key=entry.cache_key, - entry_dir=str(paths.entry_dir), - error=str(e), - ) - return False - - del entries[entry.cache_key] - if not _delete_cache_path(trash_path): - return False - logger.info( - "Removed incomplete registry artifact during startup sweep", - cache_key=entry.cache_key, - size_bytes=entry.size_bytes, - ) - - if max_entries <= 0 and max_bytes <= 0: - return True - - total_bytes = self._cache_structural_footprint() + sum( - entry.size_bytes for entry in entries.values() - ) - # Mounted entries belong to a live process sharing this cache directory. - candidates = sorted( - ( - entry - for entry in entries.values() - if not registry_artifact_mounts.is_mount( - self._paths_for(entry.cache_key).squashfs_mount_dir - ) - ), - key=lambda entry: entry.last_used, - ) - - def within_budget() -> bool: - return (max_entries <= 0 or len(entries) <= max_entries) and ( - max_bytes <= 0 or total_bytes <= max_bytes - ) - - for entry in candidates: - if within_budget(): - break - paths = self._paths_for(entry.cache_key) - try: - trash_path = _move_entry_to_trash( - paths.entry_dir, - self.trash_dir, - entry.cache_key, - ) - except OSError as e: - logger.warning( - "Failed to retire stale registry artifact during startup sweep", - cache_key=entry.cache_key, - entry_dir=str(paths.entry_dir), - error=str(e), - ) - return False - - del entries[entry.cache_key] - if _delete_cache_path(trash_path): - total_bytes -= entry.size_bytes - else: - return False - logger.info( - "Evicted stale registry artifact during startup sweep", - cache_key=entry.cache_key, - size_bytes=entry.size_bytes, - ) - - return within_budget() diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 4156cd1dc5..58e8ad2129 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -4,117 +4,971 @@ import asyncio import contextlib -from collections.abc import AsyncIterator +import hashlib +import os +import shutil +import sysconfig +import tarfile +import threading +import time +from abc import ABC, abstractmethod +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable from contextlib import asynccontextmanager +from dataclasses import dataclass, field +from enum import StrEnum from pathlib import Path +import httpx +import tracecat_registry + from tracecat import config -from tracecat.executor import registry_artifact_mounts -from tracecat.executor.registry_artifact_cache_state import ( - BASE_PYTHONPATH_DIR_NAME, - CACHE_ENTRIES_DIR_NAME, - CACHE_STAGING_DIR_NAME, - CACHE_TRASH_DIR_NAME, - RegistryArtifactCacheLoopError, - RegistryArtifactRuntimeState, -) -from tracecat.executor.registry_artifact_materialization import ( - BUNDLED_BUILTIN_REGISTRY_URI_PREFIX, - SQUASHFS_MOUNT_OPTIONS, - BuiltinArtifact, - RegistryArtifact, - RegistryArtifactAdmission, - RegistryArtifactExtractionError, - RegistryArtifactFormat, - RegistryArtifactMaterializationContext, - RegistryArtifactPaths, - SquashfsArtifact, - SquashfsMountCommandError, - TarballArtifact, - _artifact_format, - _artifact_uri_for_logging, - _bundled_builtin_registry_import_paths, - _bundled_builtin_registry_version, - _download_s3_artifact, - _is_cache_entry_uri, - _is_reusable_cache_file, - _squashfs_listing_size, - _squashfs_sidecar_uri, - _tarball_extracted_size, - _tarball_uri_for_squashfs, - bundled_builtin_registry_uri, - compute_registry_artifact_cache_key, -) -from tracecat.executor.registry_artifact_storage import ( - RegistryArtifactCacheCapacityError, - RegistryArtifactCacheEntry, - RegistryArtifactEviction, - _allocated_stat_size, - _delete_cache_path, - _delete_cache_path_off_loop, - _directory_footprint, - _move_entry_to_trash, - _RegistryArtifactCacheStorage, - _unique_work_path, -) from tracecat.logger import logger from tracecat.registry.artifact_keys import parse_s3_uri +from tracecat.registry.constants import DEFAULT_REGISTRY_ORIGIN from tracecat.storage import blob -__all__ = [ - "BASE_PYTHONPATH_DIR_NAME", - "BUNDLED_BUILTIN_REGISTRY_URI_PREFIX", - "CACHE_ENTRIES_DIR_NAME", - "CACHE_STAGING_DIR_NAME", - "CACHE_TRASH_DIR_NAME", - "SQUASHFS_MOUNT_OPTIONS", - "BuiltinArtifact", - "RegistryArtifact", - "RegistryArtifactAdmission", - "RegistryArtifactCache", - "RegistryArtifactCacheCapacityError", - "RegistryArtifactCacheEntry", - "RegistryArtifactCacheLoopError", - "RegistryArtifactEviction", - "RegistryArtifactExtractionError", - "RegistryArtifactFormat", - "RegistryArtifactMaterializationContext", - "RegistryArtifactPaths", - "RegistryArtifactRuntimeState", - "SquashfsArtifact", - "SquashfsMountCommandError", - "TarballArtifact", - "_artifact_format", - "_artifact_uri_for_logging", - "_allocated_stat_size", - "_bundled_builtin_registry_import_paths", - "_bundled_builtin_registry_version", - "_delete_cache_path", - "_delete_cache_path_off_loop", - "_directory_footprint", - "_download_s3_artifact", - "_is_cache_entry_uri", - "_move_entry_to_trash", - "_squashfs_listing_size", - "_squashfs_sidecar_uri", - "_tarball_extracted_size", - "_tarball_uri_for_squashfs", - "_unique_work_path", - "bundled_builtin_registry_uri", - "compute_registry_artifact_cache_key", -] - - -class RegistryArtifactCache(_RegistryArtifactCacheStorage): - """Materializes and leases executor-local registry artifact paths.""" - @asynccontextmanager - async def lease( +class RegistryArtifactFormat(StrEnum): + """Executor-supported registry artifact encodings.""" + + BUILTIN = "builtin" + SQUASHFS = "squashfs" + TAR_GZ = "tar.gz" + + +SQUASHFS_MOUNT_OPTIONS = "loop,ro,nodev,nosuid" +"""Mount options for executor-managed SquashFS registry artifacts. + +The image must stay read-only and should not expose device nodes or setuid bits +from registry package contents. Avoid noexec because Python packages may include +native extension modules that need to be loaded from the mounted artifact. +""" + +BUNDLED_BUILTIN_REGISTRY_URI_PREFIX = f"tracecat-builtin://{DEFAULT_REGISTRY_ORIGIN}/" +"""Pseudo-URI for the builtin registry already installed in the executor image.""" + +BASE_PYTHONPATH_DIR_NAME = "base" +"""Cache subdirectory used as the PYTHONPATH entry when no artifact is requested.""" + +CACHE_ENTRIES_DIR_NAME = "entries" +"""Directory containing one atomic subdirectory per cache key.""" + +CACHE_STAGING_DIR_NAME = "staging" +"""Directory containing in-progress materialization scratch.""" + +CACHE_TRASH_DIR_NAME = "trash" +"""Directory containing atomically retired entries pending physical deletion.""" + + +class SquashfsMountCommandError(RuntimeError): + """The ``mount`` command itself failed for a SquashFS registry artifact. + + Only this error drives SquashFS mount policy. Download, mkdir, and other + preparation failures must not be mistaken for a missing mount capability or + for loop-device exhaustion. + """ + + +class RegistryArtifactCacheLoopError(RuntimeError): + """A registry artifact cache was used outside its owning event loop.""" + + +class RegistryArtifactCacheCapacityError(RuntimeError): + """A cold artifact cannot fit within the configured cache byte budget.""" + + def __init__( self, - artifact_uris: list[str] | None, *, - paths_may_be_modified: bool = False, - ) -> AsyncIterator[list[Path]]: + current_bytes: int, + additional_bytes: int, + max_bytes: int, + ) -> None: + super().__init__( + "Registry artifact admission exceeds the cache byte budget: " + f"current_bytes={current_bytes}, additional_bytes={additional_bytes}, " + f"max_bytes={max_bytes}" + ) + self.current_bytes = current_bytes + self.additional_bytes = additional_bytes + self.max_bytes = max_bytes + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactPaths: + """Executor-local cache paths for one registry artifact key.""" + + entry_dir: Path + squashfs_image_path: Path + squashfs_mount_dir: Path + squashfs_extract_dir: Path + tarball_target_dir: Path + + +@dataclass(slots=True) +class RegistryArtifactRuntimeState: + """Process-local synchronization and lease state for one cache key.""" + + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + refcount: int = 0 + last_used: float = 0.0 + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactEviction: + """Outcome of atomically retiring and physically deleting one entry.""" + + retired: bool + reclaimed: bool + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactCacheEntry: + """Measured on-disk footprint and recency for one registry artifact key.""" + + cache_key: str + size_bytes: int + last_used: float + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactAdmission: + """Byte-bound admission hook shared by one cold materialization.""" + + max_bytes: int + ensure_capacity: Callable[[int], Awaitable[None]] + + +@dataclass(slots=True) +class RegistryArtifactMaterializationContext: + """Shared runtime state for artifact materialization.""" + + cache_key: str + staging_dir: Path + paths: RegistryArtifactPaths + admission: RegistryArtifactAdmission | None = None + + def can_mount_squashfs(self) -> bool: + return config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED and ( + shutil.which("mount") is not None + ) + + +@dataclass(frozen=True, slots=True) +class RegistryArtifact(ABC): + """An executor-local materializable registry artifact.""" + + uri: str + cache_key: str + + @property + @abstractmethod + def format(self) -> RegistryArtifactFormat: + """Artifact format used for logging and dispatch.""" + + @abstractmethod + def cached_path( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path] | None: + """Return already-materialized import paths for this artifact, if present.""" + + @abstractmethod + async def materialize( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path]: + """Return importable Python paths, materializing the artifact if needed.""" + + def _temp_path( + self, + ctx: RegistryArtifactMaterializationContext, + suffix: str, + ) -> Path: + unique_id = id(asyncio.current_task()) + ctx.staging_dir.mkdir(parents=True, exist_ok=True) + return ctx.staging_dir / f"{self.cache_key}.{os.getpid()}.{unique_id}{suffix}" + + +@dataclass(frozen=True, slots=True) +class BuiltinArtifact(RegistryArtifact): + """Current builtin registry package already installed in the executor image.""" + + version: str + + @property + def format(self) -> RegistryArtifactFormat: + return RegistryArtifactFormat.BUILTIN + + def cached_path( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path] | None: + return None + + async def materialize( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path]: + del ctx + import_paths = _bundled_builtin_registry_import_paths(self.version) + logger.info( + "Using bundled builtin registry environment", + registry_version=self.version, + paths=[str(p) for p in import_paths], + ) + return import_paths + + +@dataclass(frozen=True, slots=True) +class SquashfsArtifact(RegistryArtifact): + """SquashFS registry environment image.""" + + @property + def format(self) -> RegistryArtifactFormat: + return RegistryArtifactFormat.SQUASHFS + + def cached_path( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path] | None: + if ctx.paths.squashfs_mount_dir.is_mount(): + logger.debug( + "Using cached SquashFS registry mount", + cache_key=ctx.cache_key, + ) + return [ctx.paths.squashfs_mount_dir] + if ctx.paths.squashfs_extract_dir.exists(): + logger.debug( + "Using cached SquashFS registry extraction", + cache_key=ctx.cache_key, + ) + return [ctx.paths.squashfs_extract_dir] + return None + + async def materialize( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path]: + image_path = ctx.paths.squashfs_image_path + if ctx.can_mount_squashfs(): + try: + return [await self.mount(ctx, image_path)] + except SquashfsMountCommandError as e: + logger.warning( + "Failed to mount SquashFS registry artifact, trying extraction", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + error=str(e), + ) + + return [await self.extract(ctx, image_path)] + + async def download( + self, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> float: + """Ensure the SquashFS image exists locally and return download time.""" + if image_path.exists(): + return 0.0 + + image_path.parent.mkdir(parents=True, exist_ok=True) + temp_image = self._temp_path(ctx, ".squashfs") + try: + download_start = time.monotonic() + await _download_s3_artifact( + self.uri, + temp_image, + admission=ctx.admission, + ) + try: + temp_image.rename(image_path) + except OSError: + if not image_path.exists(): + raise + return (time.monotonic() - download_start) * 1000 + finally: + temp_image.unlink(missing_ok=True) + + async def mount( + self, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> Path: + """Download the image if needed and mount it read-only. + + Args: + ctx: Materialization context for the artifact being mounted. + image_path: Local path of the SquashFS image. + + Returns: + The mount directory. + + Raises: + SquashfsMountCommandError: The ``mount`` command failed. + Exception: The image could not be downloaded or prepared. + """ + target_dir = ctx.paths.squashfs_mount_dir + if target_dir.is_mount(): + return target_dir + + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + target_dir.mkdir(parents=True, exist_ok=True) + + logger.info( + "Materializing SquashFS registry artifact", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + ) + start_time = time.monotonic() + download_elapsed = await self.download(ctx, image_path) + + mount_start = time.monotonic() + await self._mount_image(image_path, target_dir) + mount_elapsed = (time.monotonic() - mount_start) * 1000 + total_elapsed = (time.monotonic() - start_time) * 1000 + + logger.info( + "SquashFS registry artifact mounted", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + download_ms=f"{download_elapsed:.1f}", + mount_ms=f"{mount_elapsed:.1f}", + total_ms=f"{total_elapsed:.1f}", + ) + return target_dir + + async def extract( + self, + ctx: RegistryArtifactMaterializationContext, + image_path: Path, + ) -> Path: + target_dir = ctx.paths.squashfs_extract_dir + if target_dir.exists(): + return target_dir + + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + + logger.info( + "Extracting SquashFS registry artifact", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + ) + start_time = time.monotonic() + download_elapsed = await self.download(ctx, image_path) + + temp_dir = self._temp_path(ctx, ".unsquashfs") + try: + if ctx.admission is not None: + extracted_size = await self._squashfs_extracted_size(image_path) + await ctx.admission.ensure_capacity(extracted_size) + extract_start = time.monotonic() + temp_dir.mkdir(parents=True, exist_ok=True) + await self._extract_image(image_path, temp_dir) + extract_elapsed = (time.monotonic() - extract_start) * 1000 + + try: + temp_dir.rename(target_dir) + total_elapsed = (time.monotonic() - start_time) * 1000 + logger.info( + "SquashFS registry artifact extracted", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + download_ms=f"{download_elapsed:.1f}", + extract_ms=f"{extract_elapsed:.1f}", + total_ms=f"{total_elapsed:.1f}", + ) + except OSError: + if target_dir.exists(): + logger.debug( + "SquashFS already extracted by another process", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + ) + else: + raise + finally: + if temp_dir.exists(): + shutil.rmtree(temp_dir, ignore_errors=True) + + return target_dir + + async def _mount_image(self, image_path: Path, target_dir: Path) -> None: + """Mount a SquashFS image read-only at target_dir. + + Cancellation kills and reaps the mount subprocess before propagating, + so the caller's per-key lock covers the complete mount lifecycle. The + target therefore remains an unmounted cache miss if mount never took + effect, or is already mounted and reusable by the next admission. + + Args: + image_path: Local path of the SquashFS image. + target_dir: Existing directory to mount the image at. + + Raises: + SquashfsMountCommandError: The ``mount`` command failed. + """ + if target_dir.is_mount(): + return + + proc = await asyncio.create_subprocess_exec( + "mount", + "-t", + "squashfs", + "-o", + SQUASHFS_MOUNT_OPTIONS, + str(image_path), + str(target_dir), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + proc.kill() + await proc.wait() + raise + + if proc.returncode == 0 or target_dir.is_mount(): + return + + output = (stderr or stdout).decode(errors="replace").strip() + raise SquashfsMountCommandError(output or "mount command failed") + + async def _extract_image(self, image_path: Path, target_dir: Path) -> None: + """Extract a SquashFS image to target_dir using unsquashfs. + + Cancellation kills and reaps the extractor before the caller removes + its scratch directory. Otherwise a live extractor could recreate + scratch after cleanup, outside startup-sweep discovery. + """ + unsquashfs = shutil.which("unsquashfs") + if unsquashfs is None: + raise RuntimeError("unsquashfs command is not installed") + + proc = await asyncio.create_subprocess_exec( + unsquashfs, + "-f", + "-d", + str(target_dir), + str(image_path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + proc.kill() + await proc.wait() + raise + + if proc.returncode == 0: + return + + output = (stderr or stdout).decode(errors="replace").strip() + raise RuntimeError(output or "unsquashfs command failed") + + async def _squashfs_extracted_size(self, image_path: Path) -> int: + """Return a conservative logical size for an extracted SquashFS image.""" + unsquashfs = shutil.which("unsquashfs") + if unsquashfs is None: + raise RuntimeError("unsquashfs command is not installed") + + proc = await asyncio.create_subprocess_exec( + unsquashfs, + "-lln", + str(image_path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + proc.kill() + await proc.wait() + raise + + if proc.returncode != 0: + output = (stderr or stdout).decode(errors="replace").strip() + raise RuntimeError(output or "unsquashfs listing failed") + return _squashfs_listing_size(stdout) + + +@dataclass(frozen=True, slots=True) +class TarballArtifact(RegistryArtifact): + """Legacy gzip tarball registry environment.""" + + @property + def format(self) -> RegistryArtifactFormat: + return RegistryArtifactFormat.TAR_GZ + + def cached_path( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path] | None: + if ctx.paths.tarball_target_dir.exists(): + logger.debug( + "Using cached tarball extraction", + cache_key=ctx.cache_key, + ) + return [ctx.paths.tarball_target_dir] + return None + + async def materialize( + self, ctx: RegistryArtifactMaterializationContext + ) -> list[Path]: + target_dir = ctx.paths.tarball_target_dir + logger.info( + "Materializing tarball registry artifact", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + ) + start_time = time.monotonic() + + temp_tarball = self._temp_path(ctx, ".tar.gz") + temp_dir = self._temp_path(ctx, ".tmp") + + try: + ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + + download_start = time.monotonic() + await self.download(ctx, temp_tarball) + download_elapsed = (time.monotonic() - download_start) * 1000 + + if ctx.admission is not None: + extracted_size = await asyncio.to_thread( + _tarball_extracted_size, + temp_tarball, + ) + await ctx.admission.ensure_capacity(extracted_size) + + extract_start = time.monotonic() + temp_dir.mkdir(parents=True, exist_ok=True) + await self.extract(temp_tarball, temp_dir) + extract_elapsed = (time.monotonic() - extract_start) * 1000 + + try: + temp_dir.rename(target_dir) + total_elapsed = (time.monotonic() - start_time) * 1000 + logger.info( + "Tarball extracted and cached", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + download_ms=f"{download_elapsed:.1f}", + extract_ms=f"{extract_elapsed:.1f}", + total_ms=f"{total_elapsed:.1f}", + ) + except OSError: + if target_dir.exists(): + logger.debug( + "Tarball already extracted by another process", + cache_key=ctx.cache_key, + artifact_uri=self.uri, + artifact_format=self.format.value, + ) + else: + raise + finally: + if temp_dir.exists(): + shutil.rmtree(temp_dir, ignore_errors=True) + if temp_tarball.exists(): + temp_tarball.unlink(missing_ok=True) + + return [target_dir] + + async def download( + self, + ctx: RegistryArtifactMaterializationContext, + output_path: Path, + ) -> None: + await _download_s3_artifact( + self.uri, + output_path, + admission=ctx.admission, + ) + + async def extract(self, tarball_path: Path, target_dir: Path) -> None: + """Extract a supported registry tarball to target directory. + + A cancelled caller rejoins the non-interruptible extraction thread + before propagating cancellation so scratch cleanup cannot race a live + writer and leave undiscoverable ephemeral storage behind. + """ + + def _do_extract() -> None: + if tarball_path.name.endswith(".tar.gz"): + with tarfile.open(tarball_path, "r:gz") as tar: + tar.extractall(path=target_dir, filter="data") + return + + raise ValueError(f"Unsupported tarball format: {tarball_path}") + + extraction = asyncio.ensure_future(asyncio.to_thread(_do_extract)) + try: + await asyncio.shield(extraction) + except asyncio.CancelledError: + # A thread cannot be killed. Rejoin it before materialize removes + # scratch. Each cancellation can interrupt shield without stopping + # the thread, so keep waiting until extraction reaches a terminal + # state before propagating the original cancellation. + while not extraction.done(): + try: + await asyncio.shield(extraction) + except asyncio.CancelledError: + continue + except Exception: + break + if not extraction.cancelled(): + with contextlib.suppress(Exception): + extraction.result() + raise + + logger.debug( + "Tarball extracted", + target=str(target_dir), + artifact_format=_artifact_format(str(tarball_path)).value, + ) + + +async def _download_s3_artifact( + artifact_uri: str, + output_path: Path, + *, + admission: RegistryArtifactAdmission | None = None, +) -> None: + """Download an S3 registry artifact to a local path.""" + bucket, key = parse_s3_uri(artifact_uri) + try: + if admission is None: + await blob.download_file_to_path( + key=key, + bucket=bucket, + output_path=output_path, + ) + else: + await blob.download_file_to_path( + key=key, + bucket=bucket, + output_path=output_path, + max_bytes=admission.max_bytes, + ensure_capacity=admission.ensure_capacity, + ) + except FileNotFoundError as e: + request = httpx.Request("GET", artifact_uri) + response = httpx.Response(status_code=404, request=request) + raise httpx.HTTPStatusError( + f"Registry artifact not found: {artifact_uri}", + request=request, + response=response, + ) from e + + +def compute_registry_artifact_cache_key(artifact_uri: str) -> str: + """Compute the local cache key for a registry artifact URI.""" + if not artifact_uri: + return "base" + # S3 keys are case-sensitive, so preserve URI case when hashing. + content = artifact_uri.strip() + return hashlib.sha256(content.encode()).hexdigest()[:16] + + +def bundled_builtin_registry_uri(version: str) -> str: + """Return the pseudo-URI for the installed builtin registry package.""" + return f"{BUNDLED_BUILTIN_REGISTRY_URI_PREFIX}{version}" + + +def _bundled_builtin_registry_version(artifact_uri: str) -> str | None: + """Return the builtin registry version encoded in a bundled pseudo-URI.""" + if not artifact_uri.startswith(BUNDLED_BUILTIN_REGISTRY_URI_PREFIX): + return None + version = artifact_uri.removeprefix(BUNDLED_BUILTIN_REGISTRY_URI_PREFIX) + return version or None + + +def _bundled_builtin_registry_import_paths(version: str) -> list[Path]: + """Return import paths for the current builtin registry and its dependencies. + + Dependencies always live in the executor's site-packages. For editable + installs the parent of ``package_dir`` (the package wrapper, e.g. + ``packages/tracecat-registry/``) is exposed first so its ``tracecat_registry/`` + shadows any stale copy in site-packages. + """ + installed_version = tracecat_registry.__version__ + if version != installed_version: + raise RuntimeError( + "Bundled builtin registry version does not match installed version: " + f"requested={version!r}, installed={installed_version!r}" + ) + + package_file = tracecat_registry.__file__ + if package_file is None: + raise RuntimeError("Installed tracecat_registry package has no __file__") + + site_packages_path = sysconfig.get_path("purelib") + if site_packages_path is None: + raise RuntimeError("Could not resolve installed Python site-packages path") + + site_packages = Path(site_packages_path).resolve() + if not site_packages.exists(): + raise RuntimeError( + f"Installed Python site-packages path does not exist: {site_packages}" + ) + + package_dir = Path(package_file).resolve().parent + if package_dir.is_relative_to(site_packages): + return [site_packages] + + return [package_dir.parent, site_packages] + + +def _squashfs_sidecar_uri(tarball_uri: str) -> str | None: + """Return the sibling SquashFS URI for registry site-packages tarballs.""" + if not tarball_uri.endswith("site-packages.tar.gz"): + return None + return tarball_uri.removesuffix(".tar.gz") + ".squashfs" + + +def _tarball_uri_for_squashfs(squashfs_uri: str) -> str | None: + """Return the sibling gzip tarball URI for registry SquashFS artifacts.""" + if not squashfs_uri.endswith("site-packages.squashfs"): + return None + return squashfs_uri.removesuffix(".squashfs") + ".tar.gz" + + +def _artifact_format(artifact_uri: str) -> RegistryArtifactFormat: + """Return the materialization format for an artifact URI.""" + if artifact_uri.endswith(".squashfs"): + return RegistryArtifactFormat.SQUASHFS + return RegistryArtifactFormat.TAR_GZ + + +def _is_cache_entry_uri(artifact_uri: str) -> bool: + """Return whether an artifact URI materializes into an evictable cache entry. + + The bundled builtin registry is served from the executor image and never + writes into the cache directory, so it is exempt from eviction accounting. + """ + return _bundled_builtin_registry_version(artifact_uri) is None + + +def _tarball_extracted_size(tarball_path: Path) -> int: + """Return a conservative logical size for all tarball members.""" + total_bytes = 0 + with tarfile.open(tarball_path, "r:gz") as tar: + for member in tar: + if member.size < 0: + raise ValueError( + f"Registry tarball member has a negative size: {member.name}" + ) + total_bytes += member.size + return total_bytes + + +def _squashfs_listing_size(output: bytes) -> int: + """Sum file sizes from ``unsquashfs -lln`` output, failing closed.""" + total_bytes = 0 + for raw_line in output.decode(errors="strict").splitlines(): + line = raw_line.strip() + if not line: + continue + fields = line.split(maxsplit=4) + mode = fields[0] + if len(mode) != 10 or mode[0] not in "bcdlps-": + continue + if mode[0] not in "-l": + continue + if len(fields) < 5 or "/" not in fields[1] or not fields[2].isdigit(): + raise ValueError(f"Could not parse SquashFS listing line: {line}") + total_bytes += int(fields[2]) + return total_bytes + + +def _directory_footprint(directory: Path) -> int: + """Return the total file size of a cache directory. + + Args: + directory: Cache directory to measure. + + Returns: + Total byte size of contained files, or zero when the directory is + missing. + """ + + def raise_walk_error(error: OSError) -> None: + raise error + + total_bytes = 0 + try: + walker = os.walk(directory, onerror=raise_walk_error) + for root, _dirs, files in walker: + for file_name in files: + try: + total_bytes += os.lstat(os.path.join(root, file_name)).st_size + except FileNotFoundError: + continue + except FileNotFoundError: + return 0 + return total_bytes + + +def _delete_cache_path(path: Path) -> bool: + """Best-effort delete one cache path while reporting filesystem failures.""" + try: + if path.is_dir(): + shutil.rmtree(path) + else: + path.unlink(missing_ok=True) + except OSError as e: + logger.warning( + "Failed to delete registry artifact cache path", + path=str(path), + error=str(e), + ) + return False + return True + + +async def _delete_cache_path_off_loop(path: Path) -> bool: + """Delete one path without abandoning its worker thread on cancellation.""" + deletion = asyncio.ensure_future(asyncio.to_thread(_delete_cache_path, path)) + try: + return await asyncio.shield(deletion) + except asyncio.CancelledError: + # A worker thread cannot be killed. Rejoin it so no live deletion can + # race a later trash-directory scan. Repeated cancellation can interrupt + # shield without stopping the thread, so keep waiting for termination. + while not deletion.done(): + try: + await asyncio.shield(deletion) + except asyncio.CancelledError: + continue + except Exception: + break + if not deletion.cancelled(): + with contextlib.suppress(Exception): + deletion.result() + raise + + +def _unique_work_path(root: Path, cache_key: str) -> Path: + """Return a unique path beneath a cache work directory.""" + root.mkdir(parents=True, exist_ok=True) + unique_id = time.time_ns() + while True: + path = root / f"{cache_key}.{os.getpid()}.{unique_id}" + if not path.exists(): + return path + unique_id += 1 + + +def _move_entry_to_trash(entry_dir: Path, trash_dir: Path, cache_key: str) -> Path: + """Atomically retire one cache entry and return its trash path.""" + trash_path = _unique_work_path(trash_dir, cache_key) + entry_dir.rename(trash_path) + return trash_path + + +class RegistryArtifactCache: + """Materializes registry artifacts into executor-local Python paths.""" + + def __init__(self, cache_dir: Path): + self.cache_dir = cache_dir + self.entries_dir = cache_dir / CACHE_ENTRIES_DIR_NAME + self.staging_dir = cache_dir / CACHE_STAGING_DIR_NAME + self.trash_dir = cache_dir / CACHE_TRASH_DIR_NAME + # Runtime states live for the process lifetime so every operation for a + # key always serializes on the same lock. + self._runtime: dict[str, RegistryArtifactRuntimeState] = {} + # The cache contains asyncio locks, tasks, and multi-step lease state. + # Bind the public API to one loop/thread so a future synchronous + # Temporal activity fails immediately instead of corrupting that state + # through a thread-local event loop. + self._owner_binding_lock = threading.Lock() + self._owner_loop: asyncio.AbstractEventLoop | None = None + self._owner_thread_id: int | None = None + # Cold materializations and budget passes share this outer lock. It + # keeps byte reservations stable while downloads and extraction write. + self._admission_lock = asyncio.Lock() + self._budget_lock = asyncio.Lock() + # Guard the off-loop startup sweep independently from cache operations. + self._swept: bool = False + self._sweep_task: asyncio.Task[None] | None = None + self._sweep_lock = asyncio.Lock() + # Startup is the only time the whole staging directory is swept. Exact + # paths that could not be removed are safe to retry later. + self._failed_startup_cleanup: set[Path] = set() + # Whether the on-disk cache may exceed its budget. Set when a new entry + # is materialized and cleared once enforcement measures a cache that + # fits, so steady-state cache hits never pay for a disk scan. + self._budget_dirty = True + + async def ensure_swept(self) -> None: + """Run the startup sweep exactly once successfully, off the event loop. + + Idempotent and cancellation-safe under concurrency: every caller joins + one stored sweep task, and cancelling a waiter never abandons or + restarts its live sweep. The lock is deliberately held while awaiting + that shared task so queued callers observe its result before proceeding. + Failures clear the task so the next caller retries. The sweep runs + before the first lease or materialization, so it never observes + in-flight cache entries. + """ + self._assert_owner_loop() + if self._swept: + return + async with self._sweep_lock: + if self._swept: + return + sweep_task = self._sweep_task + if sweep_task is None or ( + sweep_task.done() + and (sweep_task.cancelled() or sweep_task.exception() is not None) + ): + sweep_task = asyncio.ensure_future( + asyncio.to_thread(self._sweep_startup_state) + ) + self._sweep_task = sweep_task + try: + await asyncio.shield(sweep_task) + except asyncio.CancelledError: + if sweep_task.cancelled() and self._sweep_task is sweep_task: + self._sweep_task = None + raise + except Exception: + if self._sweep_task is sweep_task: + self._sweep_task = None + raise + self._swept = True + + def _assert_owner_loop(self) -> None: + """Bind to the current loop or reject use from another loop/thread.""" + current_loop = asyncio.get_running_loop() + current_thread_id = threading.get_ident() + with self._owner_binding_lock: + if self._owner_loop is None: + self._owner_loop = current_loop + self._owner_thread_id = current_thread_id + return + if ( + self._owner_loop is current_loop + and self._owner_thread_id == current_thread_id + ): + return + + raise RegistryArtifactCacheLoopError( + "RegistryArtifactCache is bound to one event loop and thread " + f"(owner_loop={id(self._owner_loop)}, " + f"owner_thread={self._owner_thread_id}, " + f"current_loop={id(current_loop)}, " + f"current_thread={current_thread_id})" + ) + + @asynccontextmanager + async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Path]]: """Materialize registry artifacts and pin them for the life of the context. Leased cache entries are never evicted, so callers may keep importing @@ -123,34 +977,18 @@ async def lease( Args: artifact_uris: Registry artifact URIs in deterministic PYTHONPATH order, or None to use the base PYTHONPATH directory. - paths_may_be_modified: Whether the consumer can write to returned - cache paths. Mutable leases re-arm byte-budget convergence when - execution ends so post-admission growth is measured. Yields: Importable Python paths for the requested artifacts. """ + await self.ensure_swept() + if not artifact_uris: logger.info("No registry artifact URIs provided, using base PYTHONPATH") yield [self._base_pythonpath_dir()] return - if not any(_is_cache_entry_uri(uri) for uri in artifact_uris): - cache_free_paths: list[Path] = [] - for artifact_uri in artifact_uris: - _, artifact_paths = await self._lease_artifact(artifact_uri) - cache_free_paths.extend(artifact_paths) - logger.info( - "Using cache-free registry artifact environments", - count=len(cache_free_paths), - ) - yield cache_free_paths - return - - await self.ensure_swept() - leased_keys: list[str] = [] - lease_setup_complete = False try: registry_paths: list[Path] = [] for artifact_uri in artifact_uris: @@ -162,20 +1000,12 @@ async def lease( "Using registry artifact environments", count=len(registry_paths), ) - lease_setup_complete = True yield registry_paths finally: - if paths_may_be_modified and lease_setup_complete and leased_keys: - self._budget_dirty = True idle_keys = [ cache_key for cache_key in leased_keys if self._release_lease(cache_key) ] - cleanup_task = asyncio.ensure_future( - self._finish_lease_cleanup( - idle_keys, - converge=not lease_setup_complete or bool(idle_keys), - ) - ) + cleanup_task = asyncio.ensure_future(self._finish_lease_cleanup(idle_keys)) pending_cancellation: asyncio.CancelledError | None = None while True: try: @@ -185,29 +1015,15 @@ async def lease( if cleanup_task.cancelled(): raise pending_cancellation = e - except Exception as e: - logger.error( - "Registry artifact lease cleanup failed; preserving caller outcome", - cache_dir=str(self.cache_dir), - error_type=type(e).__name__, - ) - break if pending_cancellation is not None: raise pending_cancellation - async def _finish_lease_cleanup( - self, - idle_keys: list[str], - *, - converge: bool, - ) -> None: - """Unmount newly idle entries and converge after meaningful changes.""" + async def _finish_lease_cleanup(self, idle_keys: list[str]) -> None: + """Unmount every newly idle entry and converge the cache budget.""" for cache_key in idle_keys: await self._unmount_idle_entry(cache_key) - await self._retry_failed_unmounts(excluded=set(idle_keys)) - if converge: - await self._converge_cache_budget() + await self._converge_cache_budget() async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Path]]: """Pin and materialize one artifact, returning its releasable cache key.""" @@ -217,16 +1033,17 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat candidates = await self._artifact_candidates(ctx, artifact_uri) return None, await self._materialize_candidates(ctx, candidates) + lock = self._runtime_for(cache_key).lock lease_acquired = False try: - async with self._runtime_lock(cache_key): + async with lock: self._acquire_lease(cache_key) lease_acquired = True if cached_paths := self._locally_cached_path(ctx, artifact_uri): return cache_key, cached_paths async with self._admission_lock: - async with self._runtime_lock(cache_key): + async with lock: if cached_paths := self._locally_cached_path(ctx, artifact_uri): return cache_key, cached_paths ctx = self._context_for( @@ -255,19 +1072,7 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat return cache_key, paths except BaseException: if lease_acquired and self._release_lease(cache_key): - rollback_task = asyncio.ensure_future( - self._unmount_idle_entry(cache_key) - ) - while not rollback_task.done(): - try: - await asyncio.shield(rollback_task) - except asyncio.CancelledError: - continue - except Exception: - break - if not rollback_task.cancelled(): - with contextlib.suppress(Exception): - rollback_task.result() + await self._unmount_idle_entry(cache_key) raise async def _materialize_candidates( @@ -288,7 +1093,7 @@ async def _materialize_candidates( logger.info( "Trying registry artifact candidate", cache_key=cache_key, - artifact_uri=_artifact_uri_for_logging(artifact.uri), + artifact_uri=artifact.uri, artifact_format=artifact.format.value, candidate=index + 1, candidates=len(candidates), @@ -308,17 +1113,120 @@ async def _materialize_candidates( except Exception as e: if index == len(candidates) - 1: raise - artifact.discard_failed_materialization(ctx) logger.warning( "Failed to materialize registry artifact candidate, trying fallback", cache_key=cache_key, - artifact_uri=_artifact_uri_for_logging(artifact.uri), + artifact_uri=artifact.uri, artifact_format=artifact.format.value, - error_type=type(e).__name__, + error=str(e), ) raise RuntimeError(f"No registry artifact candidates for {ctx.cache_key}") + def _runtime_for(self, cache_key: str) -> RegistryArtifactRuntimeState: + """Return the process-local state for one cache key.""" + if runtime := self._runtime.get(cache_key): + return runtime + runtime = RegistryArtifactRuntimeState() + self._runtime[cache_key] = runtime + return runtime + + def _context_for( + self, + cache_key: str, + *, + admission: RegistryArtifactAdmission | None = None, + ) -> RegistryArtifactMaterializationContext: + """Return a materialization context for a registry artifact key.""" + return RegistryArtifactMaterializationContext( + cache_key=cache_key, + staging_dir=self.staging_dir, + paths=self._paths_for(cache_key), + admission=admission, + ) + + def _admission_for(self, cache_key: str) -> RegistryArtifactAdmission | None: + """Return byte-bound admission controls for one cold cache key.""" + max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES + if max_bytes <= 0: + return None + + async def ensure_capacity(additional_bytes: int) -> None: + await self._ensure_cache_capacity( + additional_bytes=additional_bytes, + protected_key=cache_key, + max_bytes=max_bytes, + ) + + return RegistryArtifactAdmission( + max_bytes=max_bytes, + ensure_capacity=ensure_capacity, + ) + + def _base_pythonpath_dir(self) -> Path: + """Return the base PYTHONPATH directory used when no artifact is requested.""" + base_dir = self.cache_dir / BASE_PYTHONPATH_DIR_NAME + base_dir.mkdir(parents=True, exist_ok=True) + return base_dir + + def _acquire_lease(self, cache_key: str) -> None: + """Pin a cache entry against eviction and mark it as recently used. + + Callers must hold the per-key lock so the increment is ordered against + in-flight eviction of the same key. + """ + runtime = self._runtime_for(cache_key) + runtime.refcount += 1 + runtime.last_used = time.time() + self._touch_entry(cache_key) + + def _release_lease(self, cache_key: str) -> bool: + """Release one pin and return whether the entry became idle.""" + runtime = self._runtime.get(cache_key) + if runtime is None or runtime.refcount == 0: + return False + runtime.refcount -= 1 + runtime.last_used = time.time() + return runtime.refcount == 0 + + async def _unmount_idle_entry(self, cache_key: str) -> None: + """Best-effort unmount an entry after its final lease is released.""" + try: + await self._unmount_entry(cache_key) + except OSError as e: + logger.warning( + "Failed to release idle registry artifact mount", + cache_key=cache_key, + error=str(e), + ) + + def _refcount(self, cache_key: str) -> int: + """Return the number of live leases on a cache entry.""" + runtime = self._runtime.get(cache_key) + return 0 if runtime is None else runtime.refcount + + def _touch_entry(self, cache_key: str) -> None: + """Best-effort refresh of the entry-root mtime for restart-safe LRU.""" + entry_dir = self._paths_for(cache_key).entry_dir + try: + os.utime(entry_dir) + except OSError: + logger.debug( + "Could not refresh registry artifact entry mtime", + cache_key=cache_key, + ) + + def _paths_for(self, cache_key: str) -> RegistryArtifactPaths: + """Return local cache paths for a registry artifact key.""" + entry_dir = self.entries_dir / cache_key + return RegistryArtifactPaths( + entry_dir=entry_dir, + squashfs_image_path=entry_dir / "image.squashfs", + squashfs_mount_dir=entry_dir / "mount", + squashfs_extract_dir=entry_dir / "extracted", + tarball_target_dir=entry_dir / "tarball", + ) + def _first_cached_path( self, candidates: list[RegistryArtifact], @@ -369,7 +1277,7 @@ def _remove_unpublished_entry( """ paths = ctx.paths try: - if registry_artifact_mounts.is_mount(paths.squashfs_mount_dir): + if paths.squashfs_mount_dir.is_mount(): return except OSError: return @@ -414,7 +1322,7 @@ async def _artifact_candidates( if self._can_try_squashfs(): squashfs_uri = _squashfs_sidecar_uri(artifact_uri) if squashfs_uri: - if _is_reusable_cache_file(ctx.paths.squashfs_image_path): + if ctx.paths.squashfs_image_path.exists(): candidates.append( SquashfsArtifact( uri=squashfs_uri, @@ -449,23 +1357,23 @@ async def _sidecar_exists( artifact_format: RegistryArtifactFormat, ) -> bool: """Return whether a registry sidecar exists, logging lookup failures.""" + bucket, key = parse_s3_uri(sidecar_uri) try: - bucket, key = parse_s3_uri(sidecar_uri) if await blob.file_exists(key=key, bucket=bucket): logger.debug( "Using registry artifact sidecar", - artifact_uri=_artifact_uri_for_logging(base_uri), - sidecar_uri=_artifact_uri_for_logging(sidecar_uri), + artifact_uri=base_uri, + sidecar_uri=sidecar_uri, artifact_format=artifact_format.value, ) return True except Exception as e: logger.warning( "Failed to check for registry artifact sidecar, falling back", - artifact_uri=_artifact_uri_for_logging(base_uri), - sidecar_uri=_artifact_uri_for_logging(sidecar_uri), + artifact_uri=base_uri, + sidecar_uri=sidecar_uri, artifact_format=artifact_format.value, - error_type=type(e).__name__, + error=str(e), ) return False @@ -473,3 +1381,534 @@ async def _sidecar_exists( def _can_try_squashfs(self) -> bool: """Return whether this process should prefer SquashFS artifacts.""" return config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED + + async def _converge_cache_budget(self) -> None: + """Bring an idle cache back under budget after a lease is released. + + Successful materialization enforces the budget after publication while + protecting the new entry. The cache can still sit over budget while + entries are leased. This runs on release, when every newly idle entry is + evictable. The scan is skipped entirely unless a materialization attempt + has occurred since the last successful enforcement. + + Each successful pass consumes the dirty signal before its awaited scan. + A follow-up pass therefore occurs only when a concurrent materialization + sets the flag again. Without new materializations the loop terminates, + while an over-budget or failed scan restores the flag and breaks so it + cannot spin while entries remain leased. Cancellation also restores the + consumed flag before propagating. + """ + while self._budget_dirty: + self._budget_dirty = False + try: + within_budget = await self._enforce_cache_budget() + except OSError as e: + logger.warning( + "Failed to converge registry artifact cache to budget", + cache_dir=str(self.cache_dir), + error=str(e), + ) + self._budget_dirty = True + break + except BaseException: + # Cancellation must re-arm the consumed dirty signal before propagating. + self._budget_dirty = True + raise + + if not within_budget: + self._budget_dirty = True + break + + async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bool: + """Evict least-recently-used idle entries until the cache fits its budget. + + The admission lock excludes cold writers before the budget lock begins + a scan/select/evict pass. Callers invoke enforcement without holding a + per-key lock. Cold writers already hold the admission lock and use + ``_ensure_cache_capacity`` for their staged reservations instead. + + Args: + protected_key: Newly materialized cache key. It is counted against + the budget when present but never evicted. None when enforcing + against idle entries after leases are released. + + Returns: + Whether the cache is within budget once eviction has finished. + """ + async with self._admission_lock: + async with self._budget_lock: + return await self._enforce_cache_budget_locked( + protected_key=protected_key + ) + + async def _enforce_cache_budget_locked( + self, + *, + protected_key: str | None, + ) -> bool: + """Enforce entry and byte limits while both cache-wide locks are held.""" + trash_clean, startup_clean = await asyncio.gather( + asyncio.to_thread(self._clear_work_dir, self.trash_dir), + asyncio.to_thread(self._retry_failed_startup_cleanup), + ) + cleanup_complete = trash_clean and startup_clean + if not cleanup_complete: + return False + + max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES + max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES + if max_entries <= 0 and max_bytes <= 0: + return True + + entries = await asyncio.to_thread(self._scan_cache_entries) + total_bytes = sum(entry.size_bytes for entry in entries.values()) + protected = set() if protected_key is None else {protected_key} + skipped: set[str] = set() + + while (max_entries > 0 and len(entries) > max_entries) or ( + max_bytes > 0 and total_bytes > max_bytes + ): + candidate = self._least_recently_used( + entries.values(), + excluded=skipped | protected, + ) + if candidate is None: + logger.warning( + "Registry artifact cache is over budget but every entry is in use", + cache_dir=str(self.cache_dir), + entries=len(entries), + max_entries=max_entries, + total_bytes=total_bytes, + max_bytes=max_bytes, + ) + return False + + eviction = await self._evict_entry(candidate.cache_key) + if eviction.retired: + del entries[candidate.cache_key] + if not eviction.reclaimed: + return False + total_bytes -= candidate.size_bytes + else: + skipped.add(candidate.cache_key) + + return True + + async def _ensure_cache_capacity( + self, + *, + additional_bytes: int, + protected_key: str, + max_bytes: int, + ) -> None: + """Reserve peak bytes for a cold writer without exceeding the cap. + + The caller holds the admission lock and its key lock. Every normal + budget pass takes the admission lock first, so acquiring the budget + lock here cannot deadlock with eviction of the protected key. + """ + if additional_bytes < 0: + raise ValueError("additional_bytes must be non-negative") + + async with self._budget_lock: + await asyncio.gather( + asyncio.to_thread(self._clear_work_dir, self.trash_dir), + asyncio.to_thread(self._retry_failed_startup_cleanup), + ) + entries = await asyncio.to_thread(self._scan_cache_entries) + staging_bytes, trash_bytes = await asyncio.gather( + asyncio.to_thread(_directory_footprint, self.staging_dir), + asyncio.to_thread(_directory_footprint, self.trash_dir), + ) + total_bytes = ( + sum(entry.size_bytes for entry in entries.values()) + + staging_bytes + + trash_bytes + ) + skipped = {protected_key} + + while total_bytes + additional_bytes > max_bytes: + candidate = self._least_recently_used( + entries.values(), + excluded=skipped, + ) + if candidate is None: + raise RegistryArtifactCacheCapacityError( + current_bytes=total_bytes, + additional_bytes=additional_bytes, + max_bytes=max_bytes, + ) + + eviction = await self._evict_entry(candidate.cache_key) + if eviction.retired: + del entries[candidate.cache_key] + if not eviction.reclaimed: + raise RegistryArtifactCacheCapacityError( + current_bytes=total_bytes, + additional_bytes=additional_bytes, + max_bytes=max_bytes, + ) + total_bytes -= candidate.size_bytes + else: + skipped.add(candidate.cache_key) + + def _least_recently_used( + self, + entries: Iterable[RegistryArtifactCacheEntry], + *, + excluded: set[str], + ) -> RegistryArtifactCacheEntry | None: + """Return the least recently used idle entry eligible for eviction.""" + eligible = [ + entry + for entry in entries + if entry.cache_key not in excluded and self._refcount(entry.cache_key) == 0 + ] + if not eligible: + return None + return min(eligible, key=self._recency) + + def _recency(self, entry: RegistryArtifactCacheEntry) -> float: + """Return the most recent known use time for a cache entry.""" + runtime = self._runtime.get(entry.cache_key) + if runtime is None: + return entry.last_used + return max(entry.last_used, runtime.last_used) + + async def _unmount_entry(self, cache_key: str) -> bool: + """Unmount one idle cache entry while retaining its reusable image. + + Loop-device reclamation is independent from disk-budget eviction. The + per-key lock and lease recheck prevent an entry from being unmounted + while an action is importing from it. The image and empty mount + directory remain cached so a later admission can remount without + downloading the artifact again. + + Args: + cache_key: Cache key whose mounted artifact should be released. + + Returns: + Whether a mounted entry was unmounted. + """ + lock = self._runtime_for(cache_key).lock + if lock.locked(): + logger.debug( + "Skipping unmount of busy registry artifact", + cache_key=cache_key, + ) + return False + + async with lock: + if self._refcount(cache_key) > 0: + return False + + mount_dir = self._paths_for(cache_key).squashfs_mount_dir + if not mount_dir.is_mount(): + return False + if not await self._unmount(mount_dir): + logger.warning( + "Failed to unmount registry artifact for loop-device reclamation", + cache_key=cache_key, + mount_dir=str(mount_dir), + ) + return False + + logger.info( + "Unmounted idle registry artifact", + cache_key=cache_key, + mount_dir=str(mount_dir), + ) + return True + + async def _evict_entry(self, cache_key: str) -> RegistryArtifactEviction: + """Remove one cache entry from disk, unmounting it first. + + The entry is skipped rather than forced when it is leased, busy, or + cannot be unmounted: deleting the image file behind a live mount would + leave an open-file zombie holding the loop device. + + After unmounting, the entry root is atomically renamed into ``trash`` + under the per-key lock. The lock is then released before physical + deletion runs in a worker thread. + + Args: + cache_key: Cache key to evict. + + Returns: + Whether the entry was retired and its bytes were reclaimed. + """ + lock = self._runtime_for(cache_key).lock + if lock.locked(): + logger.debug( + "Skipping eviction of busy registry artifact", + cache_key=cache_key, + ) + return RegistryArtifactEviction(retired=False, reclaimed=False) + + async with lock: + if self._refcount(cache_key) > 0: + return RegistryArtifactEviction(retired=False, reclaimed=False) + + paths = self._paths_for(cache_key) + if not paths.entry_dir.exists(): + return RegistryArtifactEviction(retired=True, reclaimed=True) + if paths.squashfs_mount_dir.is_mount() and not await self._unmount( + paths.squashfs_mount_dir + ): + logger.warning( + "Failed to unmount registry artifact, skipping eviction", + cache_key=cache_key, + mount_dir=str(paths.squashfs_mount_dir), + ) + return RegistryArtifactEviction(retired=False, reclaimed=False) + + try: + trash_path = _move_entry_to_trash( + paths.entry_dir, + self.trash_dir, + cache_key, + ) + except OSError as e: + self._budget_dirty = True + logger.warning( + "Failed to retire registry artifact cache entry", + cache_key=cache_key, + entry_dir=str(paths.entry_dir), + error=str(e), + ) + return RegistryArtifactEviction(retired=False, reclaimed=False) + + try: + deleted = await _delete_cache_path_off_loop(trash_path) + except BaseException: + self._budget_dirty = True + raise + if not deleted: + self._budget_dirty = True + logger.warning( + "Registry artifact eviction remains pending physical deletion", + cache_key=cache_key, + trash_path=str(trash_path), + ) + return RegistryArtifactEviction(retired=True, reclaimed=False) + + logger.info("Evicted registry artifact from cache", cache_key=cache_key) + return RegistryArtifactEviction(retired=True, reclaimed=True) + + async def _unmount(self, mount_dir: Path) -> bool: + """Unmount a SquashFS artifact directory, releasing its loop device. + + Cancellation kills and reaps the umount subprocess before propagating, + so the caller's per-key lock covers the complete unmount lifecycle. If + umount never took effect, the mounted entry stays consistent and can be + reused; if it already took effect, the missing extraction directory + makes the entry a plain cache miss on the next admission. + """ + umount = shutil.which("umount") + if umount is None: + return False + + proc = await asyncio.create_subprocess_exec( + umount, + str(mount_dir), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + proc.kill() + await proc.wait() + raise + if proc.returncode == 0 or not mount_dir.is_mount(): + return True + + logger.warning( + "umount command failed", + mount_dir=str(mount_dir), + output=(stderr or stdout).decode(errors="replace").strip(), + ) + return False + + def _scan_cache_entries(self) -> dict[str, RegistryArtifactCacheEntry]: + """Measure every registry artifact entry currently on disk.""" + return { + cache_key: self._measure_entry(cache_key) + for cache_key in self._discover_cache_keys() + } + + def _discover_cache_keys(self) -> set[str]: + """Return cache keys represented by atomic entry directories.""" + try: + entries = list(os.scandir(self.entries_dir)) + except FileNotFoundError: + return set() + + return { + entry.name + for entry in entries + if entry.name and entry.is_dir(follow_symlinks=False) + } + + def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: + """Measure the on-disk footprint and recency of one cache entry. + + The mount directory is excluded because a mounted view only costs the + image file that backs it. The image is measured with a single ``stat`` + so a concurrent eviction deleting it cannot fail the scan. + """ + paths = self._paths_for(cache_key) + size_bytes = 0 + + try: + image_stat = paths.squashfs_image_path.stat() + except FileNotFoundError: + pass + else: + size_bytes += image_stat.st_size + + for directory in (paths.squashfs_extract_dir, paths.tarball_target_dir): + size_bytes += _directory_footprint(directory) + + try: + last_used = paths.entry_dir.stat().st_mtime + except FileNotFoundError: + last_used = 0.0 + + return RegistryArtifactCacheEntry( + cache_key=cache_key, + size_bytes=size_bytes, + last_used=last_used, + ) + + def _sweep_startup_state(self) -> None: + """Reclaim orphaned cache state left behind by a previous process. + + Scratch and trash paths from interrupted work are removed, and active + entries are trimmed to budget using entry-root mtimes as LRU order. + + The worker warms this sweep before activities can run; lazy first-use + sweeping remains a safe fallback. + """ + if not self.cache_dir.is_dir(): + self._budget_dirty = False + return + + try: + staging_clean = self._clear_work_dir( + self.staging_dir, + remember_failures=True, + ) + trash_clean = self._clear_work_dir(self.trash_dir) + cleanup_complete = staging_clean and trash_clean + within_budget = cleanup_complete and self._trim_startup_cache() + self._budget_dirty = not (cleanup_complete and within_budget) + except OSError as e: + logger.warning( + "Failed to sweep registry artifact cache", + cache_dir=str(self.cache_dir), + error=str(e), + ) + raise + + def _clear_work_dir( + self, + work_dir: Path, + *, + remember_failures: bool = False, + ) -> bool: + """Best-effort remove every child of a staging or trash directory.""" + try: + paths = list(work_dir.iterdir()) + except FileNotFoundError: + return True + except OSError as e: + logger.warning( + "Failed to inspect registry artifact work directory", + path=str(work_dir), + error=str(e), + ) + raise + + deleted = True + for path in paths: + if _delete_cache_path(path): + if remember_failures: + self._failed_startup_cleanup.discard(path) + logger.info( + "Removed registry artifact work path", + path=str(path), + ) + else: + deleted = False + if remember_failures: + self._failed_startup_cleanup.add(path) + return deleted + + def _retry_failed_startup_cleanup(self) -> bool: + """Retry exact startup paths without sweeping live staging work.""" + for path in tuple(self._failed_startup_cleanup): + if _delete_cache_path(path): + self._failed_startup_cleanup.discard(path) + return not self._failed_startup_cleanup + + def _trim_startup_cache(self) -> bool: + """Trim the cache to budget before any artifact is leased. + + Returns whether active entries and pending physical deletion fit within + the configured budget. + """ + max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES + max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES + if max_entries <= 0 and max_bytes <= 0: + return True + + entries = self._scan_cache_entries() + total_bytes = sum(entry.size_bytes for entry in entries.values()) + # Mounted entries belong to a live process sharing this cache directory. + candidates = sorted( + ( + entry + for entry in entries.values() + if not self._paths_for(entry.cache_key).squashfs_mount_dir.is_mount() + ), + key=lambda entry: entry.last_used, + ) + + def within_budget() -> bool: + return (max_entries <= 0 or len(entries) <= max_entries) and ( + max_bytes <= 0 or total_bytes <= max_bytes + ) + + for entry in candidates: + if within_budget(): + break + paths = self._paths_for(entry.cache_key) + try: + trash_path = _move_entry_to_trash( + paths.entry_dir, + self.trash_dir, + entry.cache_key, + ) + except OSError as e: + logger.warning( + "Failed to retire stale registry artifact during startup sweep", + cache_key=entry.cache_key, + entry_dir=str(paths.entry_dir), + error=str(e), + ) + return False + + del entries[entry.cache_key] + if _delete_cache_path(trash_path): + total_bytes -= entry.size_bytes + else: + return False + logger.info( + "Evicted stale registry artifact during startup sweep", + cache_key=entry.cache_key, + size_bytes=entry.size_bytes, + ) + + return within_budget() diff --git a/tracecat/sandbox/unsafe_pid_executor.py b/tracecat/sandbox/unsafe_pid_executor.py index c308db0cbf..7c8ebdc400 100644 --- a/tracecat/sandbox/unsafe_pid_executor.py +++ b/tracecat/sandbox/unsafe_pid_executor.py @@ -6,15 +6,14 @@ """ import asyncio +import contextlib import hashlib import json import logging import os import shutil -import sys import tempfile import time -from dataclasses import dataclass from pathlib import Path from typing import Any @@ -35,7 +34,6 @@ communicate_process_group, pid_namespace_available, pid_namespace_probe_error, - terminate_supervised_process, ) module_logger = logging.getLogger(__name__) @@ -243,14 +241,6 @@ def main(): ''' -@dataclass(frozen=True, slots=True) -class _ExecutionCommand: - """Command metadata for one unsafe executor subprocess.""" - - argv: list[str] - supervised: bool - - class UnsafePidExecutor: """Executor for Python scripts without nsjail, using subprocess isolation.""" @@ -318,31 +308,17 @@ def _with_python_paths( async def _build_execution_cmd( self, python_path: str, wrapper_path: Path - ) -> _ExecutionCommand: + ) -> list[str]: base_cmd = [python_path, str(wrapper_path)] if await pid_namespace_available(): - return _ExecutionCommand( - argv=["unshare", "--pid", "--fork", "--kill-child", *base_cmd], - supervised=False, - ) + return ["unshare", "--pid", "--fork", "--kill-child", *base_cmd] if not self._pid_isolation_warning_emitted: message = "PID namespace isolation unavailable; running script without PID isolation" logger.warning(message, reason=pid_namespace_probe_error()) module_logger.warning(message) self._pid_isolation_warning_emitted = True - - if sys.platform == "linux": - supervisor_path = ( - Path(__file__).resolve().parents[1] - / "executor" - / "process_supervisor.py" - ) - return _ExecutionCommand( - argv=[sys.executable, "-I", str(supervisor_path), *base_cmd], - supervised=True, - ) - return _ExecutionCommand(argv=base_cmd, supervised=False) + return base_cmd async def _create_venv(self, venv_path: Path) -> None: create_cmd = ["uv", "venv", str(venv_path), "--python", "3.12"] @@ -355,11 +331,17 @@ async def _create_venv(self, venv_path: Path) -> None: "HOME": os.environ.get("HOME", "/tmp"), "UV_CACHE_DIR": str(self.uv_cache), }, - start_new_session=True, ) try: - _, stderr = await communicate_process_group(process, timeout=60) + _, stderr = await asyncio.wait_for(process.communicate(), timeout=60) + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + raise except TimeoutError as e: + process.kill() + await process.wait() raise PackageInstallError("Virtual environment creation timed out") from e if process.returncode != 0: raise PackageInstallError( @@ -395,15 +377,21 @@ async def _install_packages( "HOME": os.environ.get("HOME", "/tmp"), "UV_CACHE_DIR": str(self.uv_cache), }, - start_new_session=True, ) try: - _, stderr = await communicate_process_group( - process, + _, stderr = await asyncio.wait_for( + process.communicate(), timeout=timeout_seconds, ) + except asyncio.CancelledError: + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + raise except TimeoutError as e: + process.kill() + await process.wait() raise PackageInstallError( f"Package installation timed out after {timeout_seconds}s" ) from e @@ -489,12 +477,9 @@ async def execute( if execution_env_vars: exec_env.update(execution_env_vars) - execution_command = await self._build_execution_cmd( - python_path, - wrapper_path, - ) + cmd = await self._build_execution_cmd(python_path, wrapper_path) process = await asyncio.create_subprocess_exec( - *execution_command.argv, + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=str(work_dir), @@ -505,11 +490,6 @@ async def execute( stdout_bytes, stderr_bytes = await communicate_process_group( process, timeout=timeout_seconds, - terminate=( - terminate_supervised_process - if execution_command.supervised - else None - ), ) except TimeoutError as e: raise SandboxTimeoutError( diff --git a/tracecat/sandbox/utils.py b/tracecat/sandbox/utils.py index f75fd0773a..360b1349e3 100644 --- a/tracecat/sandbox/utils.py +++ b/tracecat/sandbox/utils.py @@ -11,7 +11,6 @@ import shutil import signal import subprocess -from collections.abc import Awaitable, Callable from contextlib import suppress from pathlib import Path @@ -42,69 +41,11 @@ async def terminate_process_group(process: asyncio.subprocess.Process) -> None: await process.wait() -async def terminate_supervised_process(process: asyncio.subprocess.Process) -> None: - """Request descendant cleanup from a Linux process supervisor.""" - if process.returncode is None: - with suppress(ProcessLookupError): - os.kill(process.pid, signal.SIGTERM) - # An untrusted action runs under the supervisor's UID and can stop the - # outer supervisor. SIGTERM remains pending for a stopped process, so - # resume it before waiting for its cleanup handler to run. - with suppress(ProcessLookupError): - os.kill(process.pid, signal.SIGCONT) - # The supervisor exits only after its detached subreaper has killed and - # reaped the action's complete descendant tree. Do not impose a shorter - # wait that could release registry leases while cleanup is still active. - await process.wait() - - -async def _finish_process_group_cleanup( - process: asyncio.subprocess.Process, - communicate_task: asyncio.Task[tuple[bytes | None, bytes | None]], - termination_task: asyncio.Future[None] | None, - terminate: Callable[[asyncio.subprocess.Process], Awaitable[None]], -) -> None: - """Finish process termination and consume the communication task.""" - if termination_task is None: - termination_task = asyncio.ensure_future(terminate(process)) - try: - await termination_task - finally: - if not communicate_task.done(): - communicate_task.cancel() - with suppress(asyncio.CancelledError): - await communicate_task - - -async def _rejoin_cleanup_through_cancellation( - cleanup_task: asyncio.Task[None], -) -> None: - """Wait for cleanup despite repeated caller cancellation.""" - pending_cancellation: asyncio.CancelledError | None = None - while not cleanup_task.done(): - try: - await asyncio.shield(cleanup_task) - except asyncio.CancelledError as e: - if cleanup_task.cancelled(): - raise - pending_cancellation = e - - try: - cleanup_task.result() - except BaseException as cleanup_error: - if pending_cancellation is not None: - raise pending_cancellation from cleanup_error - raise - if pending_cancellation is not None: - raise pending_cancellation - - async def communicate_process_group( process: asyncio.subprocess.Process, *, input: bytes | None = None, # noqa: A002 timeout: float | None = None, - terminate: Callable[[asyncio.subprocess.Process], Awaitable[None]] | None = None, ) -> tuple[bytes, bytes]: """Communicate with a process while containing its process group. @@ -112,39 +53,24 @@ async def communicate_process_group( ``communicate()`` and asyncio's ``wait()`` to wait after the leader exits. Polling ``returncode`` observes the leader exit independently, letting us terminate the group immediately and close those pipes. Cancellation also - terminates the group before it propagates. Cleanup runs in an independent - task and is rejoined through repeated cancellation so callers cannot release - resources while the process group is still alive. + terminates the group before it propagates. """ - terminator = terminate or terminate_process_group communicate_task = asyncio.create_task(process.communicate(input=input)) - termination_task: asyncio.Future[None] | None = None - operation_error: BaseException | None = None + group_terminated = False try: async with asyncio.timeout(timeout): while process.returncode is None: await asyncio.sleep(_PROCESS_EXIT_POLL_INTERVAL_SECONDS) - termination_task = asyncio.ensure_future(terminator(process)) - await asyncio.shield(termination_task) + await terminate_process_group(process) + group_terminated = True stdout, stderr = await communicate_task - except BaseException as e: - operation_error = e - raise finally: - cleanup_task = asyncio.create_task( - _finish_process_group_cleanup( - process, - communicate_task, - termination_task, - terminator, - ) - ) - try: - await _rejoin_cleanup_through_cancellation(cleanup_task) - except BaseException as cleanup_error: - if operation_error is not None: - raise operation_error from cleanup_error - raise + if not group_terminated: + await terminate_process_group(process) + if not communicate_task.done(): + communicate_task.cancel() + with suppress(asyncio.CancelledError): + await communicate_task if stdout is None or stderr is None: raise RuntimeError("Captured stdout and stderr are required") diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index 7c2cd09fab..79fb6e5e22 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -11,13 +11,13 @@ from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass, field from pathlib import Path -from typing import TYPE_CHECKING, Protocol +from typing import TYPE_CHECKING import aioboto3 import aiofiles from aiobotocore.config import AioConfig from boto3.s3.transfer import TransferConfig -from botocore.exceptions import BotoCoreError, ClientError +from botocore.exceptions import ClientError from tracecat import config from tracecat.logger import logger @@ -35,32 +35,6 @@ DEFAULT_UPLOAD_CHUNK_SIZE_BYTES = 8 * 1024 * 1024 # 8MB DEFAULT_UPLOAD_MAX_CONCURRENCY = 4 DEFAULT_UPLOAD_MAX_IO_QUEUE_SIZE = 2 -_REDACTED_STORAGE_IDENTIFIER = "" - - -class _AsyncWritableFile(Protocol): - async def write(self, data: bytes, /) -> int: ... - - -class StorageDownloadError(RuntimeError): - """A storage download failed without exposing sensitive object identifiers.""" - - def __init__(self, *, error_code: str | None) -> None: - super().__init__("Storage download failed") - self.error_code = error_code - - -def _download_log_identifiers( - key: str, - bucket: str, - *, - redact: bool, -) -> tuple[str, str]: - """Return storage identifiers that are safe for logs and error messages.""" - if redact: - return _REDACTED_STORAGE_IDENTIFIER, _REDACTED_STORAGE_IDENTIFIER - return key, bucket - # Shared S3/MinIO client config: explicit standard-mode retries so transient # failures (throttling, 5xx, connection resets) are retried with backoff instead @@ -729,8 +703,6 @@ async def download_file_range( async def open_download_stream( key: str, bucket: str, - *, - redact_log_identifiers: bool = False, ) -> AsyncIterator[tuple[StreamingBody, int | None]]: """Open a streaming download for an S3/MinIO object. @@ -746,22 +718,14 @@ async def open_download_stream( Args: key: The S3 object key. bucket: Bucket name (required). - redact_log_identifiers: Replace the key and bucket in logs and suppress - raw client-error text for sensitive internal objects. Yields: Tuple of (streaming body, content_length). Raises: ClientError: If the download fails. - StorageDownloadError: If a redacted download fails. FileNotFoundError: If the file doesn't exist. """ - log_key, log_bucket = _download_log_identifiers( - key, - bucket, - redact=redact_log_identifiers, - ) try: async with get_storage_client() as s3_client: response = await s3_client.get_object(Bucket=bucket, Key=key) @@ -770,43 +734,20 @@ async def open_download_stream( async with body: yield body, content_length except ClientError as e: - error_code = e.response.get("Error", {}).get("Code") - if error_code == "NoSuchKey": + if e.response.get("Error", {}).get("Code") == "NoSuchKey": logger.warning( "File not found in storage", - key=log_key, - bucket=log_bucket, + key=key, + bucket=bucket, ) - if redact_log_identifiers: - raise FileNotFoundError from None raise FileNotFoundError from e - if redact_log_identifiers: - logger.error( - "Failed to open download stream", - key=log_key, - bucket=log_bucket, - error_code=error_code, - error_type=type(e).__name__, - ) - raise StorageDownloadError(error_code=error_code) from None logger.error( "Failed to open download stream", - key=log_key, - bucket=log_bucket, + key=key, + bucket=bucket, error=str(e), ) raise - except BotoCoreError as e: - if not redact_log_identifiers: - raise - logger.error( - "Failed to open download stream", - key=log_key, - bucket=log_bucket, - error_code=None, - error_type=type(e).__name__, - ) - raise StorageDownloadError(error_code=None) from None async def download_file_to_path( @@ -818,8 +759,6 @@ async def download_file_to_path( max_bytes: int | None = None, expected_sha256: str | None = None, ensure_capacity: Callable[[int], Awaitable[None]] | None = None, - defer_cleanup: Callable[[Path], None] | None = None, - redact_log_identifiers: bool = False, ) -> int: """Stream an S3/MinIO object to a local file. @@ -835,12 +774,7 @@ async def download_file_to_path( expected_sha256: Optional integrity check; raise if computed SHA-256 differs. ensure_capacity: Optional callback invoked before the first disk write with the maximum number of bytes the download may occupy. When the server - omits ContentLength, max_bytes is required and capacity is checked - incrementally before each chunk is written. - defer_cleanup: Optional callback that retains a partial-file path for a - later cleanup retry when immediate deletion fails. - redact_log_identifiers: Replace the key and bucket in logs and generated - error messages for sensitive internal objects. + omits ContentLength, max_bytes is required to provide that bound. Returns: Total bytes written. @@ -850,97 +784,11 @@ async def download_file_to_path( hasher = hashlib.sha256() if expected_sha256 is not None else None bytes_written = 0 - log_key, log_bucket = _download_log_identifiers( - key, - bucket, - redact=redact_log_identifiers, - ) - - @asynccontextmanager - async def open_file_rejoin_on_cancel() -> AsyncIterator[_AsyncWritableFile]: - """Keep the aiofiles open/close workers joined through cancellation.""" - opened: asyncio.Future[_AsyncWritableFile] = ( - asyncio.get_running_loop().create_future() - ) - close_file = asyncio.Event() - - async def file_lifecycle() -> None: - try: - async with aiofiles.open(temp_path, "wb", buffering=0) as file: - opened.set_result(file) - await close_file.wait() - except BaseException as e: - if not opened.done(): - opened.set_exception(e) - return - raise - - lifecycle = asyncio.create_task(file_lifecycle()) - operation_error: BaseException | None = None - try: - file = await asyncio.shield(opened) - yield file - except BaseException as e: - operation_error = e - raise - finally: - close_file.set() - pending_cancellation: asyncio.CancelledError | None = None - while not lifecycle.done(): - try: - await asyncio.shield(lifecycle) - except asyncio.CancelledError as e: - if lifecycle.cancelled(): - raise - pending_cancellation = e - - try: - lifecycle.result() - except BaseException as cleanup_error: - if operation_error is not None: - raise operation_error from cleanup_error - raise - if pending_cancellation is not None: - raise pending_cancellation - - async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: - """Write one chunk without abandoning the aiofiles worker thread.""" - writer = asyncio.ensure_future(file.write(chunk)) - try: - await asyncio.shield(writer) - except asyncio.CancelledError: - # aiofiles delegates writes to a thread that cannot be killed. Keep - # the partial file live until the worker stops touching it, even if - # the caller is cancelled repeatedly while cleanup is in progress. - while not writer.done(): - try: - await asyncio.shield(writer) - except asyncio.CancelledError: - continue - except Exception: - break - if not writer.cancelled(): - try: - writer.result() - except Exception: - pass - raise try: - download_stream = ( - open_download_stream( - key=key, - bucket=bucket, - redact_log_identifiers=True, - ) - if redact_log_identifiers - else open_download_stream(key=key, bucket=bucket) - ) - async with ( - download_stream as ( - stream, - content_length, - ) + async with open_download_stream(key=key, bucket=bucket) as ( + stream, + content_length, ): if ( max_bytes is not None @@ -948,53 +796,47 @@ async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: and content_length > max_bytes ): raise ValueError( - f"Refusing to download {log_bucket}/{log_key} to disk: " + f"Refusing to download {bucket}/{key} to disk: " f"ContentLength={content_length} exceeds max_bytes={max_bytes}" ) download_limit = max_bytes - grow_reservation_by_chunk = False if ensure_capacity is not None: reserved_bytes = content_length if reserved_bytes is None: if max_bytes is None: raise ValueError( "Cannot reserve disk capacity for a download without " - f"ContentLength or max_bytes: {log_bucket}/{log_key}" + f"ContentLength or max_bytes: {bucket}/{key}" ) - grow_reservation_by_chunk = True - else: - await ensure_capacity(reserved_bytes) - download_limit = ( - reserved_bytes - if download_limit is None - else min(download_limit, reserved_bytes) - ) + reserved_bytes = max_bytes + await ensure_capacity(reserved_bytes) + download_limit = ( + reserved_bytes + if download_limit is None + else min(download_limit, reserved_bytes) + ) - # Unbuffered writes keep the partial file's allocated size visible - # to incremental capacity scans between unknown-length chunks. - async with open_file_rejoin_on_cancel() as f: + async with aiofiles.open(temp_path, "wb") as f: async for chunk in stream.iter_chunks(chunk_size=chunk_size): if not chunk: continue bytes_written += len(chunk) if download_limit is not None and bytes_written > download_limit: raise ValueError( - f"Refusing to download {log_bucket}/{log_key} to disk: " + f"Refusing to download {bucket}/{key} to disk: " f"bytes_written={bytes_written} exceeds " f"max_bytes={download_limit}" ) - if grow_reservation_by_chunk and ensure_capacity is not None: - await ensure_capacity(len(chunk)) if hasher is not None: hasher.update(chunk) - await write_chunk_rejoin_on_cancel(f, chunk) + await f.write(chunk) if hasher is not None: actual_sha256 = hasher.hexdigest() if actual_sha256 != expected_sha256: raise ValueError( - f"Integrity check failed for {log_bucket}/{log_key}: " + f"Integrity check failed for {bucket}/{key}: " f"expected {expected_sha256}, got {actual_sha256}" ) @@ -1007,14 +849,12 @@ async def write_chunk_rejoin_on_cancel(file, chunk: bytes) -> None: "Failed to cleanup partial download", temp_path=str(temp_path), ) - if defer_cleanup is not None: - defer_cleanup(temp_path) raise logger.debug( "File streamed to disk successfully", - key=log_key, - bucket=log_bucket, + key=key, + bucket=bucket, output_path=str(output_path), size=bytes_written, ) From 0fa03c038defa91c4d76f75e020a9bcd74be3960 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:52:08 -0400 Subject: [PATCH 147/161] fix(executor): preserve cache process lifetimes --- tests/unit/test_action_runner.py | 32 ++++- tests/unit/test_registry_artifacts.py | 174 +++++++++++++++++++++++- tests/unit/test_unsafe_pid_executor.py | 44 ++++++ tracecat/executor/registry_artifacts.py | 131 ++++++++++-------- tracecat/sandbox/unsafe_pid_executor.py | 23 +--- tracecat/sandbox/utils.py | 69 ++++++++-- 6 files changed, 381 insertions(+), 92 deletions(-) diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index ac749d35e8..536ceadb56 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio +import contextlib import tempfile import uuid from datetime import UTC, datetime @@ -31,6 +32,7 @@ from tracecat.executor.secret_preprocessors import SecretEnvProjection from tracecat.identifiers.workflow import WorkflowUUID from tracecat.registry.lock.types import RegistryLock +from tracecat.sandbox import utils as sandbox_utils from tracecat.sandbox.types import SandboxResult @@ -747,7 +749,10 @@ async def test_cancelled_action_reaps_child_before_releasing_mounted_artifact( ) real_create_subprocess_exec = asyncio.create_subprocess_exec + real_terminate_process_group = sandbox_utils.terminate_process_group process_started = asyncio.Event() + termination_started = asyncio.Event() + finish_termination = asyncio.Event() process: asyncio.subprocess.Process | None = None reaped_before_unmount: list[bool] = [] @@ -757,6 +762,13 @@ async def capture_subprocess(*args, **kwargs): process_started.set() return process + async def controlled_termination( + requested_process: asyncio.subprocess.Process, + ) -> None: + termination_started.set() + await finish_termination.wait() + await real_terminate_process_group(requested_process) + async def release_mount(mount_dir: Path) -> bool: reaped_before_unmount.append( process is not None and process.returncode is not None @@ -775,6 +787,11 @@ async def release_mount(mount_dir: Path) -> bool: "tracecat.executor.action_runner.asyncio.create_subprocess_exec", side_effect=capture_subprocess, ), + patch.object( + sandbox_utils, + "terminate_process_group", + side_effect=controlled_termination, + ), patch.object(cache, "_unmount", side_effect=release_mount), ): execution = asyncio.create_task( @@ -787,13 +804,26 @@ async def release_mount(mount_dir: Path) -> bool: ) ) try: - await process_started.wait() + await asyncio.wait_for(process_started.wait(), timeout=5) assert cache._refcount(cache_key) == 1 execution.cancel() + await termination_started.wait() + execution.cancel() + await asyncio.sleep(0) + assert not execution.done() + assert cache._refcount(cache_key) == 1 + assert paths.squashfs_mount_dir in mounted + + finish_termination.set() with pytest.raises(asyncio.CancelledError): await execution finally: + finish_termination.set() + if not execution.done(): + execution.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await execution if process is not None and process.returncode is None: process.kill() await process.wait() diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 6f2f2e5d73..d11accd294 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -166,10 +166,13 @@ async def _materialize( class _BlockingSubprocess: """Fake subprocess that blocks in communicate until it is cancelled.""" - def __init__(self) -> None: + def __init__(self, *, block_wait: bool = False) -> None: self.communicate_started = asyncio.Event() + self.wait_started = asyncio.Event() + self.release_wait = asyncio.Event() self.cleanup_calls: list[str] = [] self.returncode: int | None = None + self._block_wait = block_wait async def communicate(self) -> tuple[bytes, bytes]: """Block until the task awaiting subprocess completion is cancelled.""" @@ -185,6 +188,9 @@ def kill(self) -> None: async def wait(self) -> int: """Record that the killed subprocess was reaped.""" self.cleanup_calls.append("wait") + self.wait_started.set() + if self._block_wait: + await self.release_wait.wait() return -9 @@ -733,8 +739,8 @@ async def test_mount_squashfs_uses_hardened_read_only_options( ) @pytest.mark.anyio - async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): - """Cancellation cannot leave an orphan mount process after lock release.""" + async def test_repeatedly_cancelled_mount_reaps_subprocess(self, temp_cache_dir): + """Repeated cancellation cannot abandon a killed mount process.""" cache_key = "cancelled-mount" cache = RegistryArtifactCache(temp_cache_dir) ctx = cache._context_for(cache_key) @@ -747,7 +753,7 @@ async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): ctx.paths.entry_dir.mkdir(parents=True) image_path.write_bytes(b"squashfs") target_dir.mkdir() - process = _BlockingSubprocess() + process = _BlockingSubprocess(block_wait=True) with patch( "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", @@ -759,7 +765,12 @@ async def test_cancelled_mount_kills_and_reaps_subprocess(self, temp_cache_dir): ) await process.communicate_started.wait() mounting.cancel() + await process.wait_started.wait() + mounting.cancel() + done, _ = await asyncio.wait({mounting}, timeout=0.05) + assert not done + process.release_wait.set() with pytest.raises(asyncio.CancelledError): await mounting @@ -880,6 +891,55 @@ def blocking_extractall(*args: object, **kwargs: object) -> None: assert first_cancellation_propagated_early is False assert second_cancellation_propagated_early is False + @pytest.mark.anyio + async def test_repeatedly_cancelled_tarball_size_scan_rejoins_thread( + self, temp_cache_dir: Path + ) -> None: + """Cancellation cannot unlink a tarball under an active size scan.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/path/slow-size-scan.tar.gz" + scan_started = threading.Event() + scan_release = threading.Event() + input_present_at_finish: list[bool] = [] + downloaded_paths: list[Path] = [] + + async def mock_download(self, ctx, path): + del self, ctx + path.write_bytes(_tarball_payload(size=1)) + downloaded_paths.append(path) + + def blocking_size_scan(path: Path) -> int: + scan_started.set() + scan_release.wait() + input_present_at_finish.append(path.exists()) + return 1 + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "download", mock_download), + patch( + "tracecat.executor.registry_artifacts._tarball_extracted_size", + side_effect=blocking_size_scan, + ), + patch.object(TarballArtifact, "extract", new_callable=AsyncMock) as extract, + ): + materializing = asyncio.create_task(cache._lease_artifact(artifact_uri)) + assert await asyncio.to_thread(scan_started.wait, 1) + materializing.cancel() + await asyncio.sleep(0) + materializing.cancel() + await asyncio.sleep(0) + assert not materializing.done() + assert downloaded_paths[0].exists() + scan_release.set() + + with pytest.raises(asyncio.CancelledError): + await materializing + + assert input_present_at_finish == [True] + assert not downloaded_paths[0].exists() + extract.assert_not_awaited() + @pytest.mark.anyio async def test_materialize_extracts_squashfs_when_mount_fails(self, temp_cache_dir): """Test that SquashFS mount failures fall back to unsquashfs extraction.""" @@ -1250,6 +1310,49 @@ async def take_lease() -> None: assert cache._refcount(cache_key) == 0 assert not cache._paths_for(cache_key).entry_dir.exists() + @pytest.mark.anyio + async def test_repeated_cancellation_finishes_acquisition_rollback( + self, temp_cache_dir: Path + ) -> None: + """A second cancellation cannot abandon final rollback cleanup.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/path/site-packages.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + lookup_started = asyncio.Event() + rollback_started = asyncio.Event() + finish_rollback = asyncio.Event() + + async def blocked_sidecar_lookup(**kwargs): + del kwargs + lookup_started.set() + await asyncio.Event().wait() + + async def blocked_unmount(requested_key: str) -> None: + assert requested_key == cache_key + rollback_started.set() + await finish_rollback.wait() + + with ( + patch(SQUASHFS_ENABLED_CONFIG, True), + patch.object(cache, "_sidecar_exists", blocked_sidecar_lookup), + patch.object(cache, "_unmount_idle_entry", blocked_unmount), + ): + acquisition = asyncio.create_task(cache._lease_artifact(artifact_uri)) + await lookup_started.wait() + acquisition.cancel() + await rollback_started.wait() + + acquisition.cancel() + await asyncio.sleep(0) + assert not acquisition.done() + finish_rollback.set() + + with pytest.raises(asyncio.CancelledError): + await acquisition + + assert cache._refcount(cache_key) == 0 + @pytest.mark.anyio async def test_cancelled_waiter_preserves_existing_same_key_lease( self, temp_cache_dir: Path @@ -2709,10 +2812,10 @@ async def mock_umount(*args, **kwargs): ) @pytest.mark.anyio - async def test_cancelled_unmount_kills_and_reaps_before_releasing_key_lock( + async def test_repeatedly_cancelled_unmount_reaps_before_releasing_key_lock( self, temp_cache_dir ): - """Cancellation leaves a consistent entry for the next admission.""" + """Repeated cancellation leaves a consistent entry for next admission.""" cache = RegistryArtifactCache(temp_cache_dir) artifact_uri = "s3://bucket/path/cancelled-unmount.squashfs" cache_key = compute_registry_artifact_cache_key(artifact_uri) @@ -2722,7 +2825,7 @@ async def test_cancelled_unmount_kills_and_reaps_before_releasing_key_lock( paths.squashfs_mount_dir.mkdir() (paths.squashfs_mount_dir / "module.py").write_text("VALUE = 1") mounted = {paths.squashfs_mount_dir} - blocked_process = _BlockingSubprocess() + blocked_process = _BlockingSubprocess(block_wait=True) released_process = AsyncMock() released_process.communicate.return_value = (b"", b"") released_process.returncode = 0 @@ -2750,7 +2853,12 @@ async def mock_umount(*args, **kwargs): eviction = asyncio.create_task(cache._evict_entry(cache_key)) await blocked_process.communicate_started.wait() eviction.cancel() + await blocked_process.wait_started.wait() + eviction.cancel() + done, _ = await asyncio.wait({eviction}, timeout=0.05) + assert not done + blocked_process.release_wait.set() with pytest.raises(asyncio.CancelledError): await eviction @@ -2830,6 +2938,58 @@ async def mock_extract(self, tarball_path, target_dir): assert original_target.is_dir() + @pytest.mark.parametrize("operation", ["budget", "admission"]) + @pytest.mark.anyio + async def test_cancelled_cleanup_rejoins_workers_before_releasing_locks( + self, temp_cache_dir: Path, operation: str + ) -> None: + """Cache locks outlive repeatedly cancelled cleanup workers.""" + cache = RegistryArtifactCache(temp_cache_dir) + cleanup_started = threading.Event() + cleanup_release = threading.Event() + cleanup_finished = threading.Event() + + def blocking_clear(work_dir: Path) -> bool: + assert work_dir == cache.trash_dir + cleanup_started.set() + cleanup_release.wait(timeout=5) + cleanup_finished.set() + return True + + async def run_operation() -> None: + if operation == "budget": + await cache._enforce_cache_budget() + else: + await cache._ensure_cache_capacity( + additional_bytes=0, + protected_key="pending", + max_bytes=1, + ) + + with ( + patch.object(cache, "_clear_work_dir", side_effect=blocking_clear), + patch.object(cache, "_retry_failed_startup_cleanup", return_value=True), + ): + running = asyncio.create_task(run_operation()) + try: + assert await asyncio.to_thread(cleanup_started.wait, 1) + running.cancel() + await asyncio.sleep(0) + running.cancel() + await asyncio.sleep(0) + assert not running.done() + assert cache._budget_lock.locked() + assert cache._admission_lock.locked() is (operation == "budget") + finally: + cleanup_release.set() + + with pytest.raises(asyncio.CancelledError): + await running + + assert cleanup_finished.is_set() + assert not cache._budget_lock.locked() + assert not cache._admission_lock.locked() + @pytest.mark.anyio async def test_doomed_eviction_names_are_startup_scratch(self, temp_cache_dir): """Every renamed entry root is invisible and reclaimed on startup.""" diff --git a/tests/unit/test_unsafe_pid_executor.py b/tests/unit/test_unsafe_pid_executor.py index 6667d749b0..00d5f46834 100644 --- a/tests/unit/test_unsafe_pid_executor.py +++ b/tests/unit/test_unsafe_pid_executor.py @@ -6,6 +6,7 @@ import os import signal from pathlib import Path +from unittest.mock import AsyncMock, patch import pytest @@ -74,6 +75,49 @@ class TestUnsafePidExecutor: def executor(self, tmp_path) -> UnsafePidExecutor: return UnsafePidExecutor(cache_dir=str(tmp_path)) + @pytest.mark.parametrize( + ("operation", "timeout"), + [("create-venv", 60), ("install-packages", 300)], + ) + @pytest.mark.anyio + async def test_dependency_setup_uses_process_group_cleanup( + self, + executor: UnsafePidExecutor, + tmp_path: Path, + operation: str, + timeout: int, + ) -> None: + """Dependency setup descendants stay inside a terminated process group.""" + process = AsyncMock() + process.returncode = 0 + create_process = AsyncMock(return_value=process) + communicate = AsyncMock(return_value=(b"", b"")) + + with ( + patch.object( + unsafe_pid_executor.asyncio, + "create_subprocess_exec", + create_process, + ), + patch.object( + unsafe_pid_executor, + "communicate_process_group", + communicate, + ), + ): + if operation == "create-venv": + await executor._create_venv(tmp_path / "venv") + else: + await executor._install_packages( + tmp_path / "venv", + ["synthetic-package"], + ) + + await_args = create_process.await_args + assert await_args is not None + assert await_args.kwargs["start_new_session"] is True + communicate.assert_awaited_once_with(process, timeout=timeout) + def test_process_probe_handles_procfs_exit_race( self, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 58e8ad2129..e0003c3bff 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -187,6 +187,61 @@ def _temp_path( return ctx.staging_dir / f"{self.cache_key}.{os.getpid()}.{unique_id}{suffix}" +async def _drain_future_through_cancellation[T]( + future: asyncio.Future[T], +) -> None: + """Wait for a future to finish despite repeated caller cancellation.""" + while not future.done(): + try: + await asyncio.shield(future) + except asyncio.CancelledError: + continue + except Exception: + break + if not future.cancelled(): + with contextlib.suppress(Exception): + future.result() + + +async def _rejoin_future_on_cancel[T](future: asyncio.Future[T]) -> T: + """Shield a future and rejoin it before propagating cancellation.""" + try: + return await asyncio.shield(future) + except asyncio.CancelledError: + await _drain_future_through_cancellation(future) + raise + + +async def _run_blocking_rejoin_on_cancel[T](operation: Callable[[], T]) -> T: + """Run blocking work without abandoning its worker thread.""" + return await _rejoin_future_on_cancel( + asyncio.ensure_future(asyncio.to_thread(operation)) + ) + + +async def _kill_and_reap_subprocess(process: asyncio.subprocess.Process) -> None: + """Kill a subprocess and wait until its child state is reaped.""" + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + + +async def _communicate_rejoin_on_cancel( + process: asyncio.subprocess.Process, +) -> tuple[bytes, bytes]: + """Communicate without allowing cancellation to abandon child cleanup.""" + try: + stdout, stderr = await process.communicate() + except asyncio.CancelledError: + reaper = asyncio.ensure_future(_kill_and_reap_subprocess(process)) + await _drain_future_through_cancellation(reaper) + raise + + if stdout is None or stderr is None: + raise RuntimeError("Captured subprocess output is required") + return stdout, stderr + + @dataclass(frozen=True, slots=True) class BuiltinArtifact(RegistryArtifact): """Current builtin registry package already installed in the executor image.""" @@ -422,13 +477,7 @@ async def _mount_image(self, image_path: Path, target_dir: Path) -> None: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - try: - stdout, stderr = await proc.communicate() - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - proc.kill() - await proc.wait() - raise + stdout, stderr = await _communicate_rejoin_on_cancel(proc) if proc.returncode == 0 or target_dir.is_mount(): return @@ -456,13 +505,7 @@ async def _extract_image(self, image_path: Path, target_dir: Path) -> None: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - try: - stdout, stderr = await proc.communicate() - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - proc.kill() - await proc.wait() - raise + stdout, stderr = await _communicate_rejoin_on_cancel(proc) if proc.returncode == 0: return @@ -483,13 +526,7 @@ async def _squashfs_extracted_size(self, image_path: Path) -> int: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - try: - stdout, stderr = await proc.communicate() - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - proc.kill() - await proc.wait() - raise + stdout, stderr = await _communicate_rejoin_on_cancel(proc) if proc.returncode != 0: output = (stderr or stdout).decode(errors="replace").strip() @@ -539,9 +576,8 @@ async def materialize( download_elapsed = (time.monotonic() - download_start) * 1000 if ctx.admission is not None: - extracted_size = await asyncio.to_thread( - _tarball_extracted_size, - temp_tarball, + extracted_size = await _run_blocking_rejoin_on_cancel( + lambda: _tarball_extracted_size(temp_tarball) ) await ctx.admission.ensure_capacity(extracted_size) @@ -607,25 +643,7 @@ def _do_extract() -> None: raise ValueError(f"Unsupported tarball format: {tarball_path}") - extraction = asyncio.ensure_future(asyncio.to_thread(_do_extract)) - try: - await asyncio.shield(extraction) - except asyncio.CancelledError: - # A thread cannot be killed. Rejoin it before materialize removes - # scratch. Each cancellation can interrupt shield without stopping - # the thread, so keep waiting until extraction reaches a terminal - # state before propagating the original cancellation. - while not extraction.done(): - try: - await asyncio.shield(extraction) - except asyncio.CancelledError: - continue - except Exception: - break - if not extraction.cancelled(): - with contextlib.suppress(Exception): - extraction.result() - raise + await _run_blocking_rejoin_on_cancel(_do_extract) logger.debug( "Tarball extracted", @@ -1072,7 +1090,8 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat return cache_key, paths except BaseException: if lease_acquired and self._release_lease(cache_key): - await self._unmount_idle_entry(cache_key) + rollback = asyncio.ensure_future(self._unmount_idle_entry(cache_key)) + await _drain_future_through_cancellation(rollback) raise async def _materialize_candidates( @@ -1447,9 +1466,11 @@ async def _enforce_cache_budget_locked( protected_key: str | None, ) -> bool: """Enforce entry and byte limits while both cache-wide locks are held.""" - trash_clean, startup_clean = await asyncio.gather( - asyncio.to_thread(self._clear_work_dir, self.trash_dir), - asyncio.to_thread(self._retry_failed_startup_cleanup), + trash_clean, startup_clean = await _rejoin_future_on_cancel( + asyncio.gather( + asyncio.to_thread(self._clear_work_dir, self.trash_dir), + asyncio.to_thread(self._retry_failed_startup_cleanup), + ) ) cleanup_complete = trash_clean and startup_clean if not cleanup_complete: @@ -1511,9 +1532,11 @@ async def _ensure_cache_capacity( raise ValueError("additional_bytes must be non-negative") async with self._budget_lock: - await asyncio.gather( - asyncio.to_thread(self._clear_work_dir, self.trash_dir), - asyncio.to_thread(self._retry_failed_startup_cleanup), + await _rejoin_future_on_cancel( + asyncio.gather( + asyncio.to_thread(self._clear_work_dir, self.trash_dir), + asyncio.to_thread(self._retry_failed_startup_cleanup), + ) ) entries = await asyncio.to_thread(self._scan_cache_entries) staging_bytes, trash_bytes = await asyncio.gather( @@ -1714,13 +1737,7 @@ async def _unmount(self, mount_dir: Path) -> bool: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - try: - stdout, stderr = await proc.communicate() - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - proc.kill() - await proc.wait() - raise + stdout, stderr = await _communicate_rejoin_on_cancel(proc) if proc.returncode == 0 or not mount_dir.is_mount(): return True diff --git a/tracecat/sandbox/unsafe_pid_executor.py b/tracecat/sandbox/unsafe_pid_executor.py index 7c8ebdc400..2be02133cd 100644 --- a/tracecat/sandbox/unsafe_pid_executor.py +++ b/tracecat/sandbox/unsafe_pid_executor.py @@ -6,7 +6,6 @@ """ import asyncio -import contextlib import hashlib import json import logging @@ -331,17 +330,11 @@ async def _create_venv(self, venv_path: Path) -> None: "HOME": os.environ.get("HOME", "/tmp"), "UV_CACHE_DIR": str(self.uv_cache), }, + start_new_session=True, ) try: - _, stderr = await asyncio.wait_for(process.communicate(), timeout=60) - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - raise + _, stderr = await communicate_process_group(process, timeout=60) except TimeoutError as e: - process.kill() - await process.wait() raise PackageInstallError("Virtual environment creation timed out") from e if process.returncode != 0: raise PackageInstallError( @@ -377,21 +370,15 @@ async def _install_packages( "HOME": os.environ.get("HOME", "/tmp"), "UV_CACHE_DIR": str(self.uv_cache), }, + start_new_session=True, ) try: - _, stderr = await asyncio.wait_for( - process.communicate(), + _, stderr = await communicate_process_group( + process, timeout=timeout_seconds, ) - except asyncio.CancelledError: - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - raise except TimeoutError as e: - process.kill() - await process.wait() raise PackageInstallError( f"Package installation timed out after {timeout_seconds}s" ) from e diff --git a/tracecat/sandbox/utils.py b/tracecat/sandbox/utils.py index 360b1349e3..0904235a48 100644 --- a/tracecat/sandbox/utils.py +++ b/tracecat/sandbox/utils.py @@ -41,6 +41,46 @@ async def terminate_process_group(process: asyncio.subprocess.Process) -> None: await process.wait() +async def _finish_process_group_cleanup( + process: asyncio.subprocess.Process, + communicate_task: asyncio.Task[tuple[bytes | None, bytes | None]], + termination_task: asyncio.Task[None] | None, +) -> None: + """Finish process termination and consume the communication task.""" + if termination_task is None: + termination_task = asyncio.create_task(terminate_process_group(process)) + try: + await termination_task + finally: + if not communicate_task.done(): + communicate_task.cancel() + with suppress(asyncio.CancelledError): + await communicate_task + + +async def _rejoin_cleanup_through_cancellation( + cleanup_task: asyncio.Task[None], +) -> None: + """Wait for cleanup despite repeated caller cancellation.""" + pending_cancellation: asyncio.CancelledError | None = None + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError as e: + if cleanup_task.cancelled(): + raise + pending_cancellation = e + + try: + cleanup_task.result() + except BaseException as cleanup_error: + if pending_cancellation is not None: + raise pending_cancellation from cleanup_error + raise + if pending_cancellation is not None: + raise pending_cancellation + + async def communicate_process_group( process: asyncio.subprocess.Process, *, @@ -56,21 +96,32 @@ async def communicate_process_group( terminates the group before it propagates. """ communicate_task = asyncio.create_task(process.communicate(input=input)) - group_terminated = False + termination_task: asyncio.Task[None] | None = None + operation_error: BaseException | None = None try: async with asyncio.timeout(timeout): while process.returncode is None: await asyncio.sleep(_PROCESS_EXIT_POLL_INTERVAL_SECONDS) - await terminate_process_group(process) - group_terminated = True + termination_task = asyncio.create_task(terminate_process_group(process)) + await asyncio.shield(termination_task) stdout, stderr = await communicate_task + except BaseException as e: + operation_error = e + raise finally: - if not group_terminated: - await terminate_process_group(process) - if not communicate_task.done(): - communicate_task.cancel() - with suppress(asyncio.CancelledError): - await communicate_task + cleanup_task = asyncio.create_task( + _finish_process_group_cleanup( + process, + communicate_task, + termination_task, + ) + ) + try: + await _rejoin_cleanup_through_cancellation(cleanup_task) + except BaseException as cleanup_error: + if operation_error is not None: + raise operation_error from cleanup_error + raise if stdout is None or stderr is None: raise RuntimeError("Captured stdout and stderr are required") From 8c41cf8e70486eb97a9686468b5ec3c41b4fbb2e Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:53:30 -0400 Subject: [PATCH 148/161] fix(executor): reject destructive cache admissions --- tests/unit/test_registry_artifacts.py | 43 +++++++++++++++++++++++-- tracecat/executor/registry_artifacts.py | 31 +++++++++++++++++- 2 files changed, 71 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index d11accd294..835b4d8ecf 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -2064,11 +2064,12 @@ async def download_file_to_path( assert (registry_paths[0] / "module.py").read_bytes() == b"x" * 32 @pytest.mark.anyio - async def test_compression_heavy_tarball_is_rejected_before_extraction( + async def test_impossible_tarball_reservation_preserves_warm_entry( self, temp_cache_dir: Path ) -> None: - """Compressed bytes plus declared extraction cannot exceed the cache cap.""" + """Impossible extraction cannot evict warm entries before rejection.""" cache = RegistryArtifactCache(temp_cache_dir) + warm = _write_image_entry(temp_cache_dir, "warm", size=64, mtime=100.0) artifact_uri = "s3://bucket/compression-heavy.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) payload = _tarball_payload(size=4096) @@ -2101,6 +2102,9 @@ async def download_file_to_path( "extract", new_callable=AsyncMock, ) as extract, + patch.object( + cache, "_evict_entry", wraps=cache._evict_entry + ) as evict_entry, ): with pytest.raises(RegistryArtifactCacheCapacityError) as raised: async with cache.lease([artifact_uri]): @@ -2109,6 +2113,8 @@ async def download_file_to_path( assert raised.value.additional_bytes == 4096 assert raised.value.max_bytes == max_bytes extract.assert_not_awaited() + evict_entry.assert_not_awaited() + assert warm.is_file() assert not cache._paths_for(cache_key).entry_dir.exists() assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) @@ -2290,6 +2296,39 @@ def fail_once(path: Path) -> bool: assert newest.exists() assert not any(cache.trash_dir.iterdir()) + @pytest.mark.anyio + async def test_undeletable_trash_does_not_evict_warm_entries_for_admission( + self, temp_cache_dir: Path + ) -> None: + """Unreclaimed trash blocks a write without cascading eviction.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + warm = _write_image_entry(temp_cache_dir, "warm", size=32, mtime=100.0) + stale_trash = cache.trash_dir / "stale" + stale_trash.mkdir(parents=True) + (stale_trash / "image.squashfs").write_bytes(b"x" * 32) + real_delete = _delete_cache_path + + def fail_stale_trash(path: Path) -> bool: + if path == stale_trash: + return False + return real_delete(path) + + with patch( + "tracecat.executor.registry_artifacts._delete_cache_path", + side_effect=fail_stale_trash, + ): + with pytest.raises(RegistryArtifactCacheCapacityError) as raised: + await cache._ensure_cache_capacity( + additional_bytes=16, + protected_key="new", + max_bytes=64, + ) + + assert raised.value.current_bytes == 64 + assert warm.exists() + assert stale_trash.exists() + @pytest.mark.anyio async def test_rename_failure_does_not_block_materialization(self, temp_cache_dir): """A failed atomic retirement keeps the old entry and admits new work.""" diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index e0003c3bff..08c9963b09 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -1532,7 +1532,7 @@ async def _ensure_cache_capacity( raise ValueError("additional_bytes must be non-negative") async with self._budget_lock: - await _rejoin_future_on_cancel( + trash_clean, startup_clean = await _rejoin_future_on_cancel( asyncio.gather( asyncio.to_thread(self._clear_work_dir, self.trash_dir), asyncio.to_thread(self._retry_failed_startup_cleanup), @@ -1548,6 +1548,35 @@ async def _ensure_cache_capacity( + staging_bytes + trash_bytes ) + non_evictable_bytes = ( + staging_bytes + + trash_bytes + + sum( + entry.size_bytes + for entry in entries.values() + if entry.cache_key == protected_key + or self._refcount(entry.cache_key) > 0 + or ( + (runtime := self._runtime.get(entry.cache_key)) is not None + and runtime.lock.locked() + ) + ) + ) + if non_evictable_bytes + additional_bytes > max_bytes: + raise RegistryArtifactCacheCapacityError( + current_bytes=non_evictable_bytes, + additional_bytes=additional_bytes, + max_bytes=max_bytes, + ) + if ( + not (trash_clean and startup_clean) + and total_bytes + additional_bytes > max_bytes + ): + raise RegistryArtifactCacheCapacityError( + current_bytes=total_bytes, + additional_bytes=additional_bytes, + max_bytes=max_bytes, + ) skipped = {protected_key} while total_bytes + additional_bytes > max_bytes: From 35faa532769d98661e83bf14dc23df2a0f65b418 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 3 Aug 2026 01:48:10 -0400 Subject: [PATCH 149/161] test(executor): use valid tarball cache fixtures --- tests/unit/test_multitenant_registry.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_multitenant_registry.py b/tests/unit/test_multitenant_registry.py index 7857dd6326..f8bd6b1b04 100644 --- a/tests/unit/test_multitenant_registry.py +++ b/tests/unit/test_multitenant_registry.py @@ -8,6 +8,7 @@ from __future__ import annotations import asyncio +import tarfile import tempfile from pathlib import Path from unittest.mock import patch @@ -40,6 +41,12 @@ async def _lease_paths( return paths +def _write_empty_tarball(path: Path) -> None: + """Write a valid archive for tests that mock the extraction contents.""" + with tarfile.open(path, "w:gz"): + pass + + # ============================================================================= # Test Class: Tarball Cache Behavior # ============================================================================= @@ -77,7 +84,7 @@ async def test_concurrent_same_uri_single_download(self, temp_cache_dir: Path): async def mock_download(self, ctx, path: Path): download_count[0] += 1 await asyncio.sleep(0.1) # Simulate network delay - path.write_bytes(b"fake tarball") + _write_empty_tarball(path) async def mock_extract(self, tarball_path: Path, target_dir: Path): (target_dir / "extracted.txt").write_text("content") @@ -124,7 +131,7 @@ async def test_different_uris_separate_cache_entries(self, temp_cache_dir: Path) async def mock_download(self, ctx, path: Path): download_calls.append(self.uri) - path.write_bytes(b"fake tarball") + _write_empty_tarball(path) async def mock_extract(self, tarball_path: Path, target_dir: Path): (target_dir / "extracted.txt").write_text("content") @@ -165,7 +172,7 @@ async def test_failed_extraction_cleans_up_temp_files(self, temp_cache_dir: Path cache_key = compute_registry_artifact_cache_key(tarball_uri) async def mock_download(self, ctx, path: Path): - path.write_bytes(b"corrupt tarball") + _write_empty_tarball(path) async def mock_extract(self, tarball_path: Path, target_dir: Path): raise RuntimeError("Extraction failed - corrupt tarball") @@ -206,7 +213,7 @@ async def test_cache_reused_on_second_request(self, temp_cache_dir: Path): async def mock_download(self, ctx, path: Path): download_count[0] += 1 - path.write_bytes(b"tarball") + _write_empty_tarball(path) async def mock_extract(self, tarball_path: Path, target_dir: Path): (target_dir / "file.txt").write_text("content") From 7c2e17f1989dd1c981307f1b2a58ccd3e7036061 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 3 Aug 2026 01:03:47 -0700 Subject: [PATCH 150/161] refactor(executor): deduplicate registry cache helpers --- tests/unit/test_registry_artifacts.py | 30 +- tracecat/concurrency.py | 33 ++ tracecat/executor/backends/test.py | 19 +- tracecat/executor/registry_artifacts.py | 517 +++++++++++------------- tracecat/storage/blob.py | 6 +- 5 files changed, 283 insertions(+), 322 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 835b4d8ecf..e816edb2e3 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -303,7 +303,11 @@ async def mock_download_file_to_path( key: str, bucket: str, output_path: Path, + max_bytes: int | None = None, + ensure_capacity: Callable[[int], Awaitable[None]] | None = None, ) -> None: + assert max_bytes is None + assert ensure_capacity is None output_path.write_bytes(b"squashfs") with patch( @@ -2319,11 +2323,12 @@ def fail_stale_trash(path: Path) -> bool: side_effect=fail_stale_trash, ): with pytest.raises(RegistryArtifactCacheCapacityError) as raised: - await cache._ensure_cache_capacity( - additional_bytes=16, - protected_key="new", - max_bytes=64, - ) + async with cache._admission_lock: + await cache._ensure_cache_capacity( + additional_bytes=16, + protected_key="new", + max_bytes=64, + ) assert raised.value.current_bytes == 64 assert warm.exists() @@ -2999,11 +3004,12 @@ async def run_operation() -> None: if operation == "budget": await cache._enforce_cache_budget() else: - await cache._ensure_cache_capacity( - additional_bytes=0, - protected_key="pending", - max_bytes=1, - ) + async with cache._admission_lock: + await cache._ensure_cache_capacity( + additional_bytes=0, + protected_key="pending", + max_bytes=1, + ) with ( patch.object(cache, "_clear_work_dir", side_effect=blocking_clear), @@ -3017,8 +3023,7 @@ async def run_operation() -> None: running.cancel() await asyncio.sleep(0) assert not running.done() - assert cache._budget_lock.locked() - assert cache._admission_lock.locked() is (operation == "budget") + assert cache._admission_lock.locked() finally: cleanup_release.set() @@ -3026,7 +3031,6 @@ async def run_operation() -> None: await running assert cleanup_finished.is_set() - assert not cache._budget_lock.locked() assert not cache._admission_lock.locked() @pytest.mark.anyio diff --git a/tracecat/concurrency.py b/tracecat/concurrency.py index e0dbd5d8a5..3878f75c92 100644 --- a/tracecat/concurrency.py +++ b/tracecat/concurrency.py @@ -1,4 +1,5 @@ import asyncio +import contextlib from collections.abc import AsyncGenerator, Awaitable, Callable, Iterable from concurrent.futures import Future, ProcessPoolExecutor from typing import Any, TypeVar, override @@ -10,6 +11,38 @@ T = TypeVar("T") +async def drain_future_through_cancellation[T]( + future: asyncio.Future[T], +) -> None: + """Wait for a future to finish despite repeated caller cancellation.""" + while not future.done(): + try: + await asyncio.shield(future) + except asyncio.CancelledError: + continue + except Exception: + break + if not future.cancelled(): + with contextlib.suppress(Exception): + future.result() + + +async def rejoin_future_on_cancel[T](future: asyncio.Future[T]) -> T: + """Shield a future and rejoin it before propagating cancellation.""" + try: + return await asyncio.shield(future) + except asyncio.CancelledError: + await drain_future_through_cancellation(future) + raise + + +async def run_blocking_rejoin_on_cancel[T](operation: Callable[[], T]) -> T: + """Run blocking work without abandoning its worker thread.""" + return await rejoin_future_on_cancel( + asyncio.ensure_future(asyncio.to_thread(operation)) + ) + + def apartial[T](coro: Callable[..., Awaitable[T]], /, *bind_args, **bind_kwargs): async def wrapped(*args, **kwargs): keywords = {**bind_kwargs, **kwargs} diff --git a/tracecat/executor/backends/test.py b/tracecat/executor/backends/test.py index 696141180c..ce7b2025ae 100644 --- a/tracecat/executor/backends/test.py +++ b/tracecat/executor/backends/test.py @@ -18,7 +18,7 @@ from __future__ import annotations import asyncio -import contextlib +import functools import sys import threading from contextlib import AsyncExitStack, contextmanager @@ -29,6 +29,7 @@ from tracecat_registry.sdk.client import TracecatClient from tracecat import config +from tracecat.concurrency import run_blocking_rejoin_on_cancel from tracecat.contexts import ( ctx_interaction, ctx_logger, @@ -248,21 +249,7 @@ async def _run_sync_udf( leases, temporary ``sys.path`` entries, and secret contexts alive until the function actually stops. """ - worker = asyncio.ensure_future(asyncio.to_thread(fn, **args)) - try: - return await asyncio.shield(worker) - except asyncio.CancelledError: - while not worker.done(): - try: - await asyncio.shield(worker) - except asyncio.CancelledError: - continue - except Exception: - break - if not worker.cancelled(): - with contextlib.suppress(Exception): - worker.result() - raise + return await run_blocking_rejoin_on_cancel(functools.partial(fn, **args)) def _load_udf_callable(self, action_impl: ActionImplementation): """Load the UDF callable from action_impl metadata.""" diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 08c9963b09..c43eeb2013 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -4,6 +4,7 @@ import asyncio import contextlib +import functools import hashlib import os import shutil @@ -22,6 +23,11 @@ import tracecat_registry from tracecat import config +from tracecat.concurrency import ( + drain_future_through_cancellation, + rejoin_future_on_cancel, + run_blocking_rejoin_on_cancel, +) from tracecat.logger import logger from tracecat.registry.artifact_keys import parse_s3_uri from tracecat.registry.constants import DEFAULT_REGISTRY_ORIGIN @@ -182,41 +188,7 @@ def _temp_path( ctx: RegistryArtifactMaterializationContext, suffix: str, ) -> Path: - unique_id = id(asyncio.current_task()) - ctx.staging_dir.mkdir(parents=True, exist_ok=True) - return ctx.staging_dir / f"{self.cache_key}.{os.getpid()}.{unique_id}{suffix}" - - -async def _drain_future_through_cancellation[T]( - future: asyncio.Future[T], -) -> None: - """Wait for a future to finish despite repeated caller cancellation.""" - while not future.done(): - try: - await asyncio.shield(future) - except asyncio.CancelledError: - continue - except Exception: - break - if not future.cancelled(): - with contextlib.suppress(Exception): - future.result() - - -async def _rejoin_future_on_cancel[T](future: asyncio.Future[T]) -> T: - """Shield a future and rejoin it before propagating cancellation.""" - try: - return await asyncio.shield(future) - except asyncio.CancelledError: - await _drain_future_through_cancellation(future) - raise - - -async def _run_blocking_rejoin_on_cancel[T](operation: Callable[[], T]) -> T: - """Run blocking work without abandoning its worker thread.""" - return await _rejoin_future_on_cancel( - asyncio.ensure_future(asyncio.to_thread(operation)) - ) + return _unique_work_path(ctx.staging_dir, self.cache_key, suffix=suffix) async def _kill_and_reap_subprocess(process: asyncio.subprocess.Process) -> None: @@ -234,7 +206,7 @@ async def _communicate_rejoin_on_cancel( stdout, stderr = await process.communicate() except asyncio.CancelledError: reaper = asyncio.ensure_future(_kill_and_reap_subprocess(process)) - await _drain_future_through_cancellation(reaper) + await drain_future_through_cancellation(reaper) raise if stdout is None or stderr is None: @@ -576,7 +548,7 @@ async def materialize( download_elapsed = (time.monotonic() - download_start) * 1000 if ctx.admission is not None: - extracted_size = await _run_blocking_rejoin_on_cancel( + extracted_size = await run_blocking_rejoin_on_cancel( lambda: _tarball_extracted_size(temp_tarball) ) await ctx.admission.ensure_capacity(extracted_size) @@ -643,7 +615,7 @@ def _do_extract() -> None: raise ValueError(f"Unsupported tarball format: {tarball_path}") - await _run_blocking_rejoin_on_cancel(_do_extract) + await run_blocking_rejoin_on_cancel(_do_extract) logger.debug( "Tarball extracted", @@ -661,20 +633,13 @@ async def _download_s3_artifact( """Download an S3 registry artifact to a local path.""" bucket, key = parse_s3_uri(artifact_uri) try: - if admission is None: - await blob.download_file_to_path( - key=key, - bucket=bucket, - output_path=output_path, - ) - else: - await blob.download_file_to_path( - key=key, - bucket=bucket, - output_path=output_path, - max_bytes=admission.max_bytes, - ensure_capacity=admission.ensure_capacity, - ) + await blob.download_file_to_path( + key=key, + bucket=bucket, + output_path=output_path, + max_bytes=None if admission is None else admission.max_bytes, + ensure_capacity=None if admission is None else admission.ensure_capacity, + ) except FileNotFoundError as e: request = httpx.Request("GET", artifact_uri) response = httpx.Response(status_code=404, request=request) @@ -795,9 +760,7 @@ def _squashfs_listing_size(output: bytes) -> int: continue fields = line.split(maxsplit=4) mode = fields[0] - if len(mode) != 10 or mode[0] not in "bcdlps-": - continue - if mode[0] not in "-l": + if len(mode) != 10 or mode[0] not in "-l": continue if len(fields) < 5 or "/" not in fields[1] or not fields[2].isdigit(): raise ValueError(f"Could not parse SquashFS listing line: {line}") @@ -852,32 +815,19 @@ def _delete_cache_path(path: Path) -> bool: async def _delete_cache_path_off_loop(path: Path) -> bool: """Delete one path without abandoning its worker thread on cancellation.""" - deletion = asyncio.ensure_future(asyncio.to_thread(_delete_cache_path, path)) - try: - return await asyncio.shield(deletion) - except asyncio.CancelledError: - # A worker thread cannot be killed. Rejoin it so no live deletion can - # race a later trash-directory scan. Repeated cancellation can interrupt - # shield without stopping the thread, so keep waiting for termination. - while not deletion.done(): - try: - await asyncio.shield(deletion) - except asyncio.CancelledError: - continue - except Exception: - break - if not deletion.cancelled(): - with contextlib.suppress(Exception): - deletion.result() - raise + # A worker thread cannot be killed. Rejoin it so no live deletion can race a + # later trash-directory scan, even through repeated caller cancellation. + return await run_blocking_rejoin_on_cancel( + functools.partial(_delete_cache_path, path) + ) -def _unique_work_path(root: Path, cache_key: str) -> Path: +def _unique_work_path(root: Path, cache_key: str, *, suffix: str = "") -> Path: """Return a unique path beneath a cache work directory.""" root.mkdir(parents=True, exist_ok=True) unique_id = time.time_ns() while True: - path = root / f"{cache_key}.{os.getpid()}.{unique_id}" + path = root / f"{cache_key}.{os.getpid()}.{unique_id}{suffix}" if not path.exists(): return path unique_id += 1 @@ -911,7 +861,6 @@ def __init__(self, cache_dir: Path): # Cold materializations and budget passes share this outer lock. It # keeps byte reservations stable while downloads and extraction write. self._admission_lock = asyncio.Lock() - self._budget_lock = asyncio.Lock() # Guard the off-loop startup sweep independently from cache operations. self._swept: bool = False self._sweep_task: asyncio.Task[None] | None = None @@ -1023,19 +972,22 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat idle_keys = [ cache_key for cache_key in leased_keys if self._release_lease(cache_key) ] - cleanup_task = asyncio.ensure_future(self._finish_lease_cleanup(idle_keys)) - pending_cancellation: asyncio.CancelledError | None = None - while True: - try: - await asyncio.shield(cleanup_task) - break - except asyncio.CancelledError as e: - if cleanup_task.cancelled(): - raise - pending_cancellation = e - - if pending_cancellation is not None: - raise pending_cancellation + if idle_keys or self._budget_dirty: + cleanup_task = asyncio.ensure_future( + self._finish_lease_cleanup(idle_keys) + ) + pending_cancellation: asyncio.CancelledError | None = None + while True: + try: + await asyncio.shield(cleanup_task) + break + except asyncio.CancelledError as e: + if cleanup_task.cancelled(): + raise + pending_cancellation = e + + if pending_cancellation is not None: + raise pending_cancellation async def _finish_lease_cleanup(self, idle_keys: list[str]) -> None: """Unmount every newly idle entry and converge the cache budget.""" @@ -1069,6 +1021,7 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat admission=self._admission_for(cache_key), ) candidates = await self._artifact_candidates(ctx, artifact_uri) + # Recheck after acquiring the lock. if cached_paths := self._first_cached_path(candidates, ctx): return cache_key, cached_paths paths = await self._materialize_candidates(ctx, candidates) @@ -1091,7 +1044,7 @@ async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Pat except BaseException: if lease_acquired and self._release_lease(cache_key): rollback = asyncio.ensure_future(self._unmount_idle_entry(cache_key)) - await _drain_future_through_cancellation(rollback) + await drain_future_through_cancellation(rollback) raise async def _materialize_candidates( @@ -1103,9 +1056,6 @@ async def _materialize_candidates( Callers hold the cache key's lock for evictable entries. """ - if cached_paths := self._first_cached_path(candidates, ctx): - return cached_paths - cache_key = ctx.cache_key for index, artifact in enumerate(candidates): try: @@ -1209,9 +1159,46 @@ def _release_lease(self, cache_key: str) -> bool: return runtime.refcount == 0 async def _unmount_idle_entry(self, cache_key: str) -> None: - """Best-effort unmount an entry after its final lease is released.""" + """Best-effort unmount one idle entry while retaining its reusable image. + + Loop-device reclamation is independent from disk-budget eviction. The + per-key lock and lease recheck prevent an entry from being unmounted + while an action is importing from it. The image and empty mount + directory remain cached so a later admission can remount without + downloading the artifact again. + + Args: + cache_key: Cache key whose mounted artifact should be released. + """ try: - await self._unmount_entry(cache_key) + lock = self._runtime_for(cache_key).lock + if lock.locked(): + logger.debug( + "Skipping unmount of busy registry artifact", + cache_key=cache_key, + ) + return + + async with lock: + if self._refcount(cache_key) > 0: + return + + mount_dir = self._paths_for(cache_key).squashfs_mount_dir + if not mount_dir.is_mount(): + return + if not await self._unmount(mount_dir): + logger.warning( + "Failed to unmount registry artifact for loop-device reclamation", + cache_key=cache_key, + mount_dir=str(mount_dir), + ) + return + + logger.info( + "Unmounted idle registry artifact", + cache_key=cache_key, + mount_dir=str(mount_dir), + ) except OSError as e: logger.warning( "Failed to release idle registry artifact mount", @@ -1263,27 +1250,69 @@ def _locally_cached_path( artifact_uri: str, ) -> list[Path] | None: """Return a reusable local candidate without probing remote sidecars.""" + include_unverified_sidecar = ( + _bundled_builtin_registry_version(artifact_uri) is None + and _artifact_format(artifact_uri) == RegistryArtifactFormat.TAR_GZ + and self._can_try_squashfs() + ) + candidates = self._candidate_artifacts( + ctx, + artifact_uri, + include_unverified_sidecar=include_unverified_sidecar, + ) + return self._first_cached_path(candidates, ctx) + + def _candidate_artifacts( + self, + ctx: RegistryArtifactMaterializationContext, + artifact_uri: str, + *, + include_unverified_sidecar: bool, + ) -> list[RegistryArtifact]: + """Build artifact candidates in executor preference order.""" + if version := _bundled_builtin_registry_version(artifact_uri): + return [ + BuiltinArtifact( + uri=artifact_uri, + cache_key=ctx.cache_key, + version=version, + ) + ] + artifact_format = _artifact_format(artifact_uri) - candidates: list[RegistryArtifact] = [] if artifact_format == RegistryArtifactFormat.SQUASHFS: - candidates.append( - SquashfsArtifact(uri=artifact_uri, cache_key=ctx.cache_key) - ) - if tarball_uri := _tarball_uri_for_squashfs(artifact_uri): - candidates.append( - TarballArtifact(uri=tarball_uri, cache_key=ctx.cache_key) + candidates: list[RegistryArtifact] = [ + SquashfsArtifact( + uri=artifact_uri, + cache_key=ctx.cache_key, ) - else: - if self._can_try_squashfs() and ( - squashfs_uri := _squashfs_sidecar_uri(artifact_uri) - ): + ] + if tarball_uri := _tarball_uri_for_squashfs(artifact_uri): candidates.append( - SquashfsArtifact(uri=squashfs_uri, cache_key=ctx.cache_key) + TarballArtifact( + uri=tarball_uri, + cache_key=ctx.cache_key, + ) ) + return candidates + + candidates = [] + if include_unverified_sidecar and ( + squashfs_uri := _squashfs_sidecar_uri(artifact_uri) + ): candidates.append( - TarballArtifact(uri=artifact_uri, cache_key=ctx.cache_key) + SquashfsArtifact( + uri=squashfs_uri, + cache_key=ctx.cache_key, + ) ) - return self._first_cached_path(candidates, ctx) + candidates.append( + TarballArtifact( + uri=artifact_uri, + cache_key=ctx.cache_key, + ) + ) + return candidates def _remove_unpublished_entry( self, @@ -1311,62 +1340,34 @@ async def _artifact_candidates( artifact_uri: str, ) -> list[RegistryArtifact]: """Return artifact candidates in executor preference order.""" - if version := _bundled_builtin_registry_version(artifact_uri): - return [ - BuiltinArtifact( - uri=artifact_uri, - cache_key=ctx.cache_key, - version=version, - ) - ] + if _bundled_builtin_registry_version(artifact_uri) is not None: + return self._candidate_artifacts( + ctx, + artifact_uri, + include_unverified_sidecar=False, + ) artifact_format = _artifact_format(artifact_uri) - if artifact_format == RegistryArtifactFormat.SQUASHFS: - candidates = [ - SquashfsArtifact( - uri=artifact_uri, - cache_key=ctx.cache_key, - ) - ] - if tarball_uri := _tarball_uri_for_squashfs(artifact_uri): - candidates.append( - TarballArtifact( - uri=tarball_uri, - cache_key=ctx.cache_key, - ) - ) - return candidates - - candidates: list[RegistryArtifact] = [] - if self._can_try_squashfs(): - squashfs_uri = _squashfs_sidecar_uri(artifact_uri) - if squashfs_uri: - if ctx.paths.squashfs_image_path.exists(): - candidates.append( - SquashfsArtifact( - uri=squashfs_uri, - cache_key=ctx.cache_key, - ) - ) - elif await self._sidecar_exists( + include_unverified_sidecar = False + if ( + artifact_format == RegistryArtifactFormat.TAR_GZ + and self._can_try_squashfs() + and (squashfs_uri := _squashfs_sidecar_uri(artifact_uri)) + ): + include_unverified_sidecar = ( + ctx.paths.squashfs_image_path.exists() + or await self._sidecar_exists( base_uri=artifact_uri, sidecar_uri=squashfs_uri, artifact_format=RegistryArtifactFormat.SQUASHFS, - ): - candidates.append( - SquashfsArtifact( - uri=squashfs_uri, - cache_key=ctx.cache_key, - ) - ) - - candidates.append( - TarballArtifact( - uri=artifact_uri, - cache_key=ctx.cache_key, + ) ) + + return self._candidate_artifacts( + ctx, + artifact_uri, + include_unverified_sidecar=include_unverified_sidecar, ) - return candidates async def _sidecar_exists( self, @@ -1441,10 +1442,10 @@ async def _converge_cache_budget(self) -> None: async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bool: """Evict least-recently-used idle entries until the cache fits its budget. - The admission lock excludes cold writers before the budget lock begins - a scan/select/evict pass. Callers invoke enforcement without holding a - per-key lock. Cold writers already hold the admission lock and use - ``_ensure_cache_capacity`` for their staged reservations instead. + The admission lock excludes cold writers during the scan/select/evict + pass. Callers invoke enforcement without holding a per-key lock. Cold + writers already hold the admission lock and use ``_ensure_cache_capacity`` + for their staged reservations instead. Args: protected_key: Newly materialized cache key. It is counted against @@ -1455,23 +1456,24 @@ async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bo Whether the cache is within budget once eviction has finished. """ async with self._admission_lock: - async with self._budget_lock: - return await self._enforce_cache_budget_locked( - protected_key=protected_key - ) + return await self._enforce_cache_budget_locked(protected_key=protected_key) - async def _enforce_cache_budget_locked( - self, - *, - protected_key: str | None, - ) -> bool: - """Enforce entry and byte limits while both cache-wide locks are held.""" - trash_clean, startup_clean = await _rejoin_future_on_cancel( + async def _reclaim_pending_work(self) -> tuple[bool, bool]: + """Retry pending trash and startup cleanup off the event loop.""" + return await rejoin_future_on_cancel( asyncio.gather( asyncio.to_thread(self._clear_work_dir, self.trash_dir), asyncio.to_thread(self._retry_failed_startup_cleanup), ) ) + + async def _enforce_cache_budget_locked( + self, + *, + protected_key: str | None, + ) -> bool: + """Enforce entry and byte limits while the admission lock is held.""" + trash_clean, startup_clean = await self._reclaim_pending_work() cleanup_complete = trash_clean and startup_clean if not cleanup_complete: return False @@ -1524,85 +1526,69 @@ async def _ensure_cache_capacity( ) -> None: """Reserve peak bytes for a cold writer without exceeding the cap. - The caller holds the admission lock and its key lock. Every normal - budget pass takes the admission lock first, so acquiring the budget - lock here cannot deadlock with eviction of the protected key. + The caller must hold the admission lock and its key lock so reservations + stay serialized with normal budget passes and other cold writers. """ if additional_bytes < 0: raise ValueError("additional_bytes must be non-negative") - async with self._budget_lock: - trash_clean, startup_clean = await _rejoin_future_on_cancel( - asyncio.gather( - asyncio.to_thread(self._clear_work_dir, self.trash_dir), - asyncio.to_thread(self._retry_failed_startup_cleanup), - ) - ) - entries = await asyncio.to_thread(self._scan_cache_entries) - staging_bytes, trash_bytes = await asyncio.gather( - asyncio.to_thread(_directory_footprint, self.staging_dir), - asyncio.to_thread(_directory_footprint, self.trash_dir), - ) - total_bytes = ( - sum(entry.size_bytes for entry in entries.values()) - + staging_bytes - + trash_bytes + def capacity_error(current_bytes: int) -> RegistryArtifactCacheCapacityError: + return RegistryArtifactCacheCapacityError( + current_bytes=current_bytes, + additional_bytes=additional_bytes, + max_bytes=max_bytes, ) - non_evictable_bytes = ( - staging_bytes - + trash_bytes - + sum( - entry.size_bytes - for entry in entries.values() - if entry.cache_key == protected_key - or self._refcount(entry.cache_key) > 0 - or ( - (runtime := self._runtime.get(entry.cache_key)) is not None - and runtime.lock.locked() - ) + + trash_clean, startup_clean = await self._reclaim_pending_work() + entries = await asyncio.to_thread(self._scan_cache_entries) + staging_bytes, trash_bytes = await asyncio.gather( + asyncio.to_thread(_directory_footprint, self.staging_dir), + asyncio.to_thread(_directory_footprint, self.trash_dir), + ) + total_bytes = ( + sum(entry.size_bytes for entry in entries.values()) + + staging_bytes + + trash_bytes + ) + non_evictable_bytes = ( + staging_bytes + + trash_bytes + + sum( + entry.size_bytes + for entry in entries.values() + if entry.cache_key == protected_key + or self._refcount(entry.cache_key) > 0 + or ( + (runtime := self._runtime.get(entry.cache_key)) is not None + and runtime.lock.locked() ) ) - if non_evictable_bytes + additional_bytes > max_bytes: - raise RegistryArtifactCacheCapacityError( - current_bytes=non_evictable_bytes, - additional_bytes=additional_bytes, - max_bytes=max_bytes, - ) - if ( - not (trash_clean and startup_clean) - and total_bytes + additional_bytes > max_bytes - ): - raise RegistryArtifactCacheCapacityError( - current_bytes=total_bytes, - additional_bytes=additional_bytes, - max_bytes=max_bytes, - ) - skipped = {protected_key} + ) + if non_evictable_bytes + additional_bytes > max_bytes: + raise capacity_error(non_evictable_bytes) + if ( + not (trash_clean and startup_clean) + and total_bytes + additional_bytes > max_bytes + ): + raise capacity_error(total_bytes) + skipped = {protected_key} - while total_bytes + additional_bytes > max_bytes: - candidate = self._least_recently_used( - entries.values(), - excluded=skipped, - ) - if candidate is None: - raise RegistryArtifactCacheCapacityError( - current_bytes=total_bytes, - additional_bytes=additional_bytes, - max_bytes=max_bytes, - ) + while total_bytes + additional_bytes > max_bytes: + candidate = self._least_recently_used( + entries.values(), + excluded=skipped, + ) + if candidate is None: + raise capacity_error(total_bytes) - eviction = await self._evict_entry(candidate.cache_key) - if eviction.retired: - del entries[candidate.cache_key] - if not eviction.reclaimed: - raise RegistryArtifactCacheCapacityError( - current_bytes=total_bytes, - additional_bytes=additional_bytes, - max_bytes=max_bytes, - ) - total_bytes -= candidate.size_bytes - else: - skipped.add(candidate.cache_key) + eviction = await self._evict_entry(candidate.cache_key) + if eviction.retired: + del entries[candidate.cache_key] + if not eviction.reclaimed: + raise capacity_error(total_bytes) + total_bytes -= candidate.size_bytes + else: + skipped.add(candidate.cache_key) def _least_recently_used( self, @@ -1627,51 +1613,6 @@ def _recency(self, entry: RegistryArtifactCacheEntry) -> float: return entry.last_used return max(entry.last_used, runtime.last_used) - async def _unmount_entry(self, cache_key: str) -> bool: - """Unmount one idle cache entry while retaining its reusable image. - - Loop-device reclamation is independent from disk-budget eviction. The - per-key lock and lease recheck prevent an entry from being unmounted - while an action is importing from it. The image and empty mount - directory remain cached so a later admission can remount without - downloading the artifact again. - - Args: - cache_key: Cache key whose mounted artifact should be released. - - Returns: - Whether a mounted entry was unmounted. - """ - lock = self._runtime_for(cache_key).lock - if lock.locked(): - logger.debug( - "Skipping unmount of busy registry artifact", - cache_key=cache_key, - ) - return False - - async with lock: - if self._refcount(cache_key) > 0: - return False - - mount_dir = self._paths_for(cache_key).squashfs_mount_dir - if not mount_dir.is_mount(): - return False - if not await self._unmount(mount_dir): - logger.warning( - "Failed to unmount registry artifact for loop-device reclamation", - cache_key=cache_key, - mount_dir=str(mount_dir), - ) - return False - - logger.info( - "Unmounted idle registry artifact", - cache_key=cache_key, - mount_dir=str(mount_dir), - ) - return True - async def _evict_entry(self, cache_key: str) -> RegistryArtifactEviction: """Remove one cache entry from disk, unmounting it first. diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index 79fb6e5e22..1fccec7362 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -811,11 +811,7 @@ async def download_file_to_path( ) reserved_bytes = max_bytes await ensure_capacity(reserved_bytes) - download_limit = ( - reserved_bytes - if download_limit is None - else min(download_limit, reserved_bytes) - ) + download_limit = reserved_bytes async with aiofiles.open(temp_path, "wb") as f: async for chunk in stream.iter_chunks(chunk_size=chunk_size): From a78918607f200a83b82cbe540caeb23132cb84d0 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 3 Aug 2026 11:37:43 -0700 Subject: [PATCH 151/161] refactor(executor): modularize registry cache and action supervision --- .../unit/executor/test_process_supervisor.py | 303 ++++ .../executor/test_registry_artifact_budget.py | 60 + .../executor/test_run_python_sdk_context.py | 8 +- .../test_test_backend_no_registry_action.py | 12 +- tests/unit/test_action_runner.py | 43 +- tests/unit/test_concurrency.py | 55 + tests/unit/test_executor_sandbox_nsjail.py | 8 + tests/unit/test_registry_artifacts.py | 487 +++++- tests/unit/test_storage_blob.py | 131 +- tracecat/concurrency.py | 29 + tracecat/executor/action_runner.py | 34 +- tracecat/executor/backends/base.py | 9 +- tracecat/executor/backends/test.py | 12 +- tracecat/executor/process_supervisor.py | 276 +++ tracecat/executor/registry_artifact_budget.py | 85 + .../executor/registry_artifact_storage.py | 1192 +++++++++++++ tracecat/executor/registry_artifacts.py | 1558 ++++++----------- tracecat/sandbox/utils.py | 58 +- tracecat/storage/blob.py | 114 +- 19 files changed, 3246 insertions(+), 1228 deletions(-) create mode 100644 tests/unit/executor/test_process_supervisor.py create mode 100644 tests/unit/executor/test_registry_artifact_budget.py create mode 100644 tracecat/executor/process_supervisor.py create mode 100644 tracecat/executor/registry_artifact_budget.py create mode 100644 tracecat/executor/registry_artifact_storage.py diff --git a/tests/unit/executor/test_process_supervisor.py b/tests/unit/executor/test_process_supervisor.py new file mode 100644 index 0000000000..2ba7f9a36d --- /dev/null +++ b/tests/unit/executor/test_process_supervisor.py @@ -0,0 +1,303 @@ +"""Linux process-tree containment for direct registry actions.""" + +from __future__ import annotations + +import asyncio +import contextlib +import os +import signal +import sys +from pathlib import Path + +import pytest + +from tracecat.executor import process_supervisor +from tracecat.sandbox.utils import ( + communicate_process_group, + terminate_supervised_process, +) + +pytestmark = pytest.mark.skipif( + sys.platform != "linux", + reason="The direct action supervisor uses Linux prctl and procfs", +) + + +def _process_is_running(pid: int) -> bool: + """Return whether a process is alive, treating zombies as terminated.""" + try: + os.kill(pid, 0) + stat_path = Path(f"/proc/{pid}/stat") + return not stat_path.exists() or stat_path.read_text().split()[2] != "Z" + except (FileNotFoundError, ProcessLookupError): + return False + + +async def _wait_for_file(path: Path) -> None: + for _ in range(500): + if path.exists(): + return + await asyncio.sleep(0.01) + raise AssertionError(f"Timed out waiting for {path}") + + +def _write_action_script(path: Path) -> None: + path.write_text( + """ +import os +import signal +import subprocess +import sys +import time +from pathlib import Path + +pid_file = Path(sys.argv[1]) +mode = sys.argv[2] +child = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(30)"], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, +) +monitor_pid = os.getppid() +outer_pid = int( + next( + line.removeprefix("PPid:").strip() + for line in Path(f"/proc/{monitor_pid}/status").read_text().splitlines() + if line.startswith("PPid:") + ) +) +if mode == "stop-monitor-kill-supervisor": + pid_file.write_text(f"{os.getpid()} {child.pid} {monitor_pid} {outer_pid}") +else: + pid_file.write_text(f"{os.getpid()} {child.pid}") +if mode == "failure": + raise SystemExit(23) +if mode == "kill-monitor": + os.kill(monitor_pid, signal.SIGKILL) +if mode in {"stop-monitor", "stop-supervisors"}: + os.kill(monitor_pid, signal.SIGSTOP) + if mode == "stop-supervisors": + os.kill(outer_pid, signal.SIGSTOP) +if mode == "stop-monitor-kill-supervisor": + os.kill(monitor_pid, signal.SIGSTOP) + os.kill(outer_pid, signal.SIGKILL) +if mode in { + "block", + "kill-monitor", + "stop-monitor", + "stop-supervisors", + "stop-monitor-kill-supervisor", +}: + time.sleep(30) +""".lstrip() + ) + + +async def _spawn_supervised_action( + tmp_path: Path, + *, + mode: str, +) -> tuple[asyncio.subprocess.Process, Path]: + action_script = tmp_path / f"action-{mode}.py" + pid_file = tmp_path / f"action-{mode}.pid" + _write_action_script(action_script) + supervisor_path = Path(process_supervisor.__file__) + process = await asyncio.create_subprocess_exec( + sys.executable, + str(supervisor_path), + sys.executable, + str(action_script), + str(pid_file), + mode, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + start_new_session=True, + ) + return process, pid_file + + +@pytest.mark.anyio +@pytest.mark.parametrize( + ("mode", "expected_returncode"), + [("success", 0), ("failure", 23)], +) +async def test_supervisor_propagates_exit_and_reaps_detached_descendant( + tmp_path: Path, + mode: str, + expected_returncode: int, +) -> None: + process, pid_file = await _spawn_supervised_action(tmp_path, mode=mode) + detached_pid: int | None = None + try: + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=5) + + assert process.returncode == expected_returncode, stderr.decode() + assert stdout == b"" + _, detached_pid = (int(pid) for pid in pid_file.read_text().split()) + assert not _process_is_running(detached_pid) + finally: + if detached_pid is None and pid_file.exists(): + _, detached_pid = (int(pid) for pid in pid_file.read_text().split()) + if detached_pid is not None: + with contextlib.suppress(ProcessLookupError): + os.kill(detached_pid, signal.SIGKILL) + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() + + +@pytest.mark.anyio +async def test_supervisor_reaps_detached_descendant_before_cancellation_returns( + tmp_path: Path, +) -> None: + process, pid_file = await _spawn_supervised_action(tmp_path, mode="block") + await _wait_for_file(pid_file) + action_pid, detached_pid = (int(pid) for pid in pid_file.read_text().split()) + communication = asyncio.create_task( + communicate_process_group( + process, + timeout=30, + terminate=terminate_supervised_process, + ) + ) + + try: + await asyncio.sleep(0) + communication.cancel() + communication.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(communication, timeout=5) + assert not _process_is_running(action_pid) + assert not _process_is_running(detached_pid) + finally: + for pid in (action_pid, detached_pid): + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() + + +@pytest.mark.anyio +async def test_supervisor_reaps_action_that_kills_its_monitor( + tmp_path: Path, +) -> None: + process, pid_file = await _spawn_supervised_action( + tmp_path, + mode="kill-monitor", + ) + tracked_pids: tuple[int, ...] = () + try: + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=5) + + assert process.returncode == 128 + signal.SIGKILL, stderr.decode() + assert stdout == b"" + tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) + action_pid, detached_pid = tracked_pids + assert not _process_is_running(action_pid) + assert not _process_is_running(detached_pid) + finally: + if not tracked_pids and pid_file.exists(): + tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) + for pid in tracked_pids: + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() + + +@pytest.mark.anyio +@pytest.mark.parametrize("mode", ["stop-monitor", "stop-supervisors"]) +async def test_termination_recovers_stopped_supervisor_tree( + tmp_path: Path, + mode: str, +) -> None: + process, pid_file = await _spawn_supervised_action(tmp_path, mode=mode) + await _wait_for_file(pid_file) + tracked_pids = tuple(int(pid) for pid in pid_file.read_text().split()) + action_pid, detached_pid = tracked_pids + + try: + with pytest.raises(TimeoutError): + await asyncio.wait_for( + communicate_process_group( + process, + timeout=0.05, + terminate=terminate_supervised_process, + ), + timeout=5, + ) + + assert process.returncode is not None + assert not _process_is_running(action_pid) + assert not _process_is_running(detached_pid) + finally: + for pid in tracked_pids: + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() + + +@pytest.mark.anyio +async def test_parent_death_resumes_monitor_and_reaps_action_tree( + tmp_path: Path, +) -> None: + process, pid_file = await _spawn_supervised_action( + tmp_path, + mode="stop-monitor-kill-supervisor", + ) + await _wait_for_file(pid_file) + action_pid, detached_pid, monitor_pid, outer_pid = ( + int(pid) for pid in pid_file.read_text().split() + ) + assert outer_pid == process.pid + + try: + stdout, stderr = await asyncio.wait_for( + communicate_process_group( + process, + timeout=2, + terminate=terminate_supervised_process, + ), + timeout=5, + ) + + assert process.returncode == -signal.SIGKILL, stderr.decode() + assert stdout == b"" + assert not _process_is_running(monitor_pid) + assert not _process_is_running(action_pid) + assert not _process_is_running(detached_pid) + finally: + for pid in (monitor_pid, action_pid, detached_pid): + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() + + +def test_direct_child_discovery_falls_back_to_proc_stat(monkeypatch) -> None: + child_pid = os.fork() + if child_pid == 0: + signal.pause() + os._exit(0) + + def missing_children_file() -> list[int]: + raise FileNotFoundError + + monkeypatch.setattr( + process_supervisor, + "_direct_child_pids_from_children_file", + missing_children_file, + ) + try: + assert child_pid in process_supervisor._direct_child_pids() + finally: + with contextlib.suppress(ProcessLookupError): + os.kill(child_pid, signal.SIGKILL) + os.waitpid(child_pid, 0) diff --git a/tests/unit/executor/test_registry_artifact_budget.py b/tests/unit/executor/test_registry_artifact_budget.py new file mode 100644 index 0000000000..a7ee09326b --- /dev/null +++ b/tests/unit/executor/test_registry_artifact_budget.py @@ -0,0 +1,60 @@ +"""Pure registry artifact cache budget policy tests.""" + +from tracecat.executor.registry_artifact_budget import ( + RegistryArtifactCacheBudget, + RegistryArtifactCacheEntry, + plan_registry_artifact_evictions, +) + + +def _entry( + cache_key: str, *, size: int, last_used: float +) -> RegistryArtifactCacheEntry: + return RegistryArtifactCacheEntry( + cache_key=cache_key, + size_bytes=size, + last_used=last_used, + ) + + +def test_plan_orders_eligible_entries_by_effective_lru() -> None: + entries = { + "old-on-disk": _entry("old-on-disk", size=40, last_used=1.0), + "recent-in-process": _entry("recent-in-process", size=40, last_used=2.0), + "protected": _entry("protected", size=40, last_used=0.0), + } + + plan = plan_registry_artifact_evictions( + entries, + total_bytes=160, + budget=RegistryArtifactCacheBudget(max_entries=1, max_bytes=50), + excluded={"protected"}, + effective_last_used={"old-on-disk": 20.0, "recent-in-process": 10.0}, + ) + + assert [entry.cache_key for entry in plan.candidates] == [ + "recent-in-process", + "old-on-disk", + ] + assert plan.can_fit is False + + +def test_plan_reports_when_eligible_evictions_can_satisfy_budget() -> None: + entries = { + "old": _entry("old", size=60, last_used=1.0), + "new": _entry("new", size=30, last_used=2.0), + } + + plan = plan_registry_artifact_evictions( + entries, + total_bytes=100, + budget=RegistryArtifactCacheBudget( + max_entries=1, + max_bytes=80, + additional_bytes=10, + ), + excluded=set(), + ) + + assert [entry.cache_key for entry in plan.candidates] == ["old", "new"] + assert plan.can_fit is True diff --git a/tests/unit/executor/test_run_python_sdk_context.py b/tests/unit/executor/test_run_python_sdk_context.py index 9adb24b45d..a096ffbbe8 100644 --- a/tests/unit/executor/test_run_python_sdk_context.py +++ b/tests/unit/executor/test_run_python_sdk_context.py @@ -473,13 +473,18 @@ class _FakeRunPythonRegistryArtifacts: def __init__(self, paths: list[Path]) -> None: self.paths = paths self.artifact_uris: list[str] | None = None + self.paths_may_be_modified = False self.leased = False @asynccontextmanager async def lease( - self, artifact_uris: list[str] | None = None + self, + artifact_uris: list[str] | None = None, + *, + paths_may_be_modified: bool = False, ) -> AsyncIterator[list[Path]]: self.artifact_uris = artifact_uris + self.paths_may_be_modified = paths_may_be_modified self.leased = True try: yield self.paths @@ -1232,6 +1237,7 @@ async def run_python(self, **kwargs: Any) -> dict[str, bool]: assert result.type == "success" assert leased_during_run == [True] assert fake_runner.registry_artifacts.leased is False + assert fake_runner.registry_artifacts.paths_may_be_modified is True @pytest.mark.anyio diff --git a/tests/unit/executor/test_test_backend_no_registry_action.py b/tests/unit/executor/test_test_backend_no_registry_action.py index 50363e03fc..1e346061bc 100644 --- a/tests/unit/executor/test_test_backend_no_registry_action.py +++ b/tests/unit/executor/test_test_backend_no_registry_action.py @@ -273,8 +273,12 @@ def __init__(self) -> None: @asynccontextmanager async def lease( - self, artifact_uris: list[str] | None = None + self, + artifact_uris: list[str] | None = None, + *, + paths_may_be_modified: bool = False, ) -> AsyncIterator[list[Path]]: + assert paths_may_be_modified is True if artifact_uris == [broken_uri]: raise RuntimeError("artifact unavailable") self.active += 1 @@ -356,9 +360,13 @@ def __init__(self) -> None: @asynccontextmanager async def lease( - self, artifact_uris: list[str] | None = None + self, + artifact_uris: list[str] | None = None, + *, + paths_may_be_modified: bool = False, ) -> AsyncIterator[list[Path]]: assert artifact_uris == [artifact_uri] + assert paths_may_be_modified is True self.active += 1 try: yield [artifact_path] diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index 536ceadb56..eafa81af77 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -7,8 +7,10 @@ import asyncio import contextlib +import shutil import tempfile import uuid +from collections.abc import Awaitable, Callable from datetime import UTC, datetime from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch @@ -90,6 +92,23 @@ def temp_cache_dir(): class TestActionRunner: """Tests for ActionRunner class.""" + @pytest.fixture(autouse=True) + def resolve_setpriv(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Resolve setpriv on non-Linux developer machines. + + The direct runner command is Linux-only in production. These tests mock + the spawn, so only the lookup has to succeed for the command to build. + """ + real_which = shutil.which + + def which(cmd: str, *args, **kwargs) -> str | None: + resolved = real_which(cmd, *args, **kwargs) + if resolved is None and cmd == "setpriv": + return "/usr/bin/setpriv" + return resolved + + monkeypatch.setattr(action_runner.shutil, "which", which) + @pytest.fixture(autouse=True) def mock_process_group_communication( self, monkeypatch: pytest.MonkeyPatch @@ -102,12 +121,15 @@ async def communicate( *, input: bytes | None = None, # noqa: A002 timeout: float | None = None, + terminate: Callable[[asyncio.subprocess.Process], Awaitable[None]] + | None = None, ) -> tuple[bytes, bytes]: if isinstance(process, asyncio.subprocess.Process): return await real_communication( process, input=input, timeout=timeout, + terminate=terminate, ) stdout, stderr = await asyncio.wait_for( process.communicate(input=input), @@ -529,9 +551,10 @@ async def test_execute_action_disables_new_privileges_for_direct_subprocess( temp_cache_dir, mock_run_action_input, mock_role, + mock_process_group_communication: AsyncMock, monkeypatch: pytest.MonkeyPatch, ): - """Test direct subprocess execution disables new Linux privileges.""" + """Test direct execution supervises the action and drops privileges.""" runner = ActionRunner(cache_dir=temp_cache_dir) base_dir = temp_cache_dir / "base" base_dir.mkdir() @@ -549,7 +572,6 @@ async def create_subprocess_exec_side_effect(*args, **kwargs): # noqa: ARG001 mock_proc.communicate = AsyncMock(return_value=(success_response, b"")) return mock_proc - monkeypatch.setattr(action_runner.sys, "platform", "linux") monkeypatch.setattr( action_runner.shutil, "which", @@ -575,8 +597,17 @@ async def create_subprocess_exec_side_effect(*args, **kwargs): # noqa: ARG001 "--inh-caps=-all", "--ambient-caps=-all", ] + assert captured_args[-5] == action_runner.sys.executable + assert captured_args[-4] == "-I" + assert captured_args[-3].endswith("process_supervisor.py") assert captured_args[-2] == action_runner.sys.executable assert captured_args[-1].endswith("minimal_runner.py") + communication_call = mock_process_group_communication.await_args + assert communication_call is not None + assert ( + communication_call.kwargs["terminate"] + is action_runner.terminate_supervised_process + ) @pytest.mark.anyio async def test_execute_action_invalid_json_response( @@ -749,7 +780,7 @@ async def test_cancelled_action_reaps_child_before_releasing_mounted_artifact( ) real_create_subprocess_exec = asyncio.create_subprocess_exec - real_terminate_process_group = sandbox_utils.terminate_process_group + real_terminate = sandbox_utils.terminate_supervised_process process_started = asyncio.Event() termination_started = asyncio.Event() finish_termination = asyncio.Event() @@ -767,7 +798,7 @@ async def controlled_termination( ) -> None: termination_started.set() await finish_termination.wait() - await real_terminate_process_group(requested_process) + await real_terminate(requested_process) async def release_mount(mount_dir: Path) -> bool: reaped_before_unmount.append( @@ -788,8 +819,8 @@ async def release_mount(mount_dir: Path) -> bool: side_effect=capture_subprocess, ), patch.object( - sandbox_utils, - "terminate_process_group", + action_runner, + "terminate_supervised_process", side_effect=controlled_termination, ), patch.object(cache, "_unmount", side_effect=release_mount), diff --git a/tests/unit/test_concurrency.py b/tests/unit/test_concurrency.py index 0f2b776efc..75891d7b29 100644 --- a/tests/unit/test_concurrency.py +++ b/tests/unit/test_concurrency.py @@ -10,6 +10,7 @@ apartial, cooperative, cooperative_every, + rejoin_future_through_cancellation, ) @@ -585,3 +586,57 @@ async def failing_sleep(delay): assert len(result) >= 1 assert result[0] == 1 assert len(result) <= 2 # At most 2 items before exception + + +@pytest.mark.anyio +async def test_rejoin_future_finishes_through_repeated_cancellation() -> None: + started = asyncio.Event() + release = asyncio.Event() + finished = asyncio.Event() + + async def operation() -> None: + started.set() + await release.wait() + finished.set() + + operation_task = asyncio.create_task(operation()) + joining_task = asyncio.create_task( + rejoin_future_through_cancellation(operation_task) + ) + await started.wait() + + joining_task.cancel() + await asyncio.sleep(0) + joining_task.cancel() + await asyncio.sleep(0) + assert not joining_task.done() + + release.set() + with pytest.raises(asyncio.CancelledError): + await joining_task + assert finished.is_set() + + +@pytest.mark.anyio +async def test_rejoin_future_chains_cleanup_failure_from_cancellation() -> None: + started = asyncio.Event() + release = asyncio.Event() + cleanup_error = RuntimeError("cleanup failed") + + async def operation() -> None: + started.set() + await release.wait() + raise cleanup_error + + operation_task = asyncio.create_task(operation()) + joining_task = asyncio.create_task( + rejoin_future_through_cancellation(operation_task) + ) + await started.wait() + + joining_task.cancel() + await asyncio.sleep(0) + release.set() + with pytest.raises(asyncio.CancelledError) as raised: + await joining_task + assert raised.value.__cause__ is cleanup_error diff --git a/tests/unit/test_executor_sandbox_nsjail.py b/tests/unit/test_executor_sandbox_nsjail.py index c451035e9e..f59a3ba1c1 100644 --- a/tests/unit/test_executor_sandbox_nsjail.py +++ b/tests/unit/test_executor_sandbox_nsjail.py @@ -113,6 +113,14 @@ def _skip_smoke(reason: str) -> NoReturn: def _missing_prerequisite(smoke_case: SmokeCase) -> str | None: + if not smoke_case.force_sandbox: + if sys.platform != "linux": + return ( + "direct action subprocesses require Linux (setpriv + subreaper " + "process supervisor)" + ) + if shutil.which("setpriv") is None: + return "setpriv is unavailable" if smoke_case.force_sandbox and not _executor_nsjail_available(): return "executor nsjail unavailable" if smoke_case.force_sandbox and not Path("/dev/net/tun").exists(): diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index e816edb2e3..6d05782f68 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -11,12 +11,20 @@ from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from pathlib import Path -from unittest.mock import ANY, AsyncMock, patch +from unittest.mock import ANY, AsyncMock, call, patch import httpx import pytest import tracecat_registry +from tracecat.executor.registry_artifact_storage import ( + RegistryArtifactMaterializationContext, + _allocated_stat_size, + _delete_cache_path, + _directory_footprint, + _filesystem_allocation_unit, + allocated_size_bound, +) from tracecat.executor.registry_artifacts import ( SQUASHFS_MOUNT_OPTIONS, RegistryArtifactCache, @@ -24,11 +32,10 @@ RegistryArtifactCacheLoopError, RegistryArtifactEviction, RegistryArtifactFormat, - RegistryArtifactMaterializationContext, + RegistryArtifactUriError, SquashfsArtifact, SquashfsMountCommandError, TarballArtifact, - _delete_cache_path, _squashfs_listing_size, bundled_builtin_registry_uri, compute_registry_artifact_cache_key, @@ -142,25 +149,19 @@ async def _materialize( cache_key: str, artifact_uri: str, ) -> list[Path]: - """Exercise internal materialization while releasing its test-only lease.""" - await cache.ensure_swept() - ctx = cache._context_for(cache_key) - lock = cache._runtime_for(cache_key).lock - lease_acquired = False - try: - async with lock: - cache._acquire_lease(cache_key) - lease_acquired = True - candidates = await cache._artifact_candidates(ctx, artifact_uri) - if cached_paths := cache._first_cached_path(candidates, ctx): - return cached_paths - paths = await cache._materialize_candidates(ctx, candidates) - cache._touch_entry(cache_key) - await cache._enforce_cache_budget(protected_key=cache_key) + """Materialize through the same public lifecycle used by executors.""" + assert cache_key == compute_registry_artifact_cache_key(artifact_uri) + async with cache.lease([artifact_uri]) as paths: return paths - finally: - if lease_acquired: - cache._release_lease(cache_key) + + +async def _lease_and_release( + cache: RegistryArtifactCache, + artifact_uri: str, +) -> None: + """Exercise one artifact through the public lease lifecycle.""" + async with cache.lease([artifact_uri]): + pass class _BlockingSubprocess: @@ -305,9 +306,13 @@ async def mock_download_file_to_path( output_path: Path, max_bytes: int | None = None, ensure_capacity: Callable[[int], Awaitable[None]] | None = None, + defer_cleanup: Callable[[Path], None] | None = None, + redact_log_identifiers: bool = False, ) -> None: assert max_bytes is None assert ensure_capacity is None + assert defer_cleanup is not None + assert redact_log_identifiers is True output_path.write_bytes(b"squashfs") with patch( @@ -410,6 +415,26 @@ async def test_download_artifact_normalizes_missing_objects_to_http_404( assert exc_info.value.response.status_code == 404 assert isinstance(exc_info.value.__cause__, FileNotFoundError) + @pytest.mark.anyio + async def test_invalid_artifact_uri_suppresses_identifiers( + self, temp_cache_dir: Path + ) -> None: + sensitive_uri = "https://affected.example/tenant/secret-artifact.tar.gz" + artifact = TarballArtifact( + uri=sensitive_uri, + cache_key="invalid-uri", + ) + cache = RegistryArtifactCache(temp_cache_dir) + + with pytest.raises(RegistryArtifactUriError) as raised: + await artifact.download( + cache._context_for(artifact.cache_key), + temp_cache_dir / "artifact.tar.gz", + ) + + assert sensitive_uri not in str(raised.value) + assert str(raised.value) == "Invalid registry artifact URI" + @pytest.mark.anyio async def test_artifact_candidates_prefer_squashfs_sidecar(self, temp_cache_dir): """Test that gzip tarballs prefer a sibling SquashFS sidecar.""" @@ -554,12 +579,60 @@ def test_squashfs_listing_size_sums_files_and_symlinks(self) -> None: ] ) - assert _squashfs_listing_size(listing) == 132 + # Includes every inode plus directory-entry overhead, not just payload + # bytes (123-byte file + 9-byte symlink). + assert _squashfs_listing_size(listing) == 276 def test_squashfs_listing_size_rejects_unparseable_files(self) -> None: with pytest.raises(ValueError, match="Could not parse SquashFS listing"): _squashfs_listing_size(b"-rw-r--r-- malformed") + def test_directory_footprint_does_not_follow_symlinked_root( + self, temp_cache_dir: Path + ) -> None: + outside = temp_cache_dir / "outside" + outside.mkdir() + (outside / "large.bin").write_bytes(b"x" * (1024 * 1024)) + symlinked_root = temp_cache_dir / "cache-link" + symlinked_root.symlink_to(outside, target_is_directory=True) + allocation_unit = _filesystem_allocation_unit(temp_cache_dir) + + assert _directory_footprint(symlinked_root) == _allocated_stat_size( + symlinked_root.lstat(), + allocation_unit=allocation_unit, + ) + + def test_budget_scan_rejects_symlinked_cache_root( + self, temp_cache_dir: Path + ) -> None: + outside = temp_cache_dir / "outside" + outside.mkdir() + symlinked_root = temp_cache_dir / "cache-link" + symlinked_root.symlink_to(outside, target_is_directory=True) + cache = RegistryArtifactCache(symlinked_root) + + with pytest.raises(OSError, match="Unsafe registry artifact cache directory"): + cache._scan_cache_snapshot() + + @pytest.mark.anyio + async def test_materialization_rejects_symlinked_cache_root( + self, temp_cache_dir: Path + ) -> None: + outside = temp_cache_dir / "outside" + outside.mkdir() + symlinked_root = temp_cache_dir / "cache-link" + symlinked_root.symlink_to(outside, target_is_directory=True) + cache = RegistryArtifactCache(symlinked_root) + artifact = TarballArtifact( + uri="s3://bucket/path/site-packages.tar.gz", + cache_key="symlinked-root", + ) + + with pytest.raises(OSError, match="Unsafe registry artifact cache directory"): + await artifact.materialize(cache._context_for(artifact.cache_key)) + + assert list(outside.iterdir()) == [] + @pytest.mark.anyio async def test_artifact_candidates_direct_squashfs_include_gzip_fallback( self, temp_cache_dir @@ -616,6 +689,39 @@ async def test_artifact_candidates_fall_back_to_gzip(self, temp_cache_dir): assert artifact.uri == "s3://bucket/path/site-packages.tar.gz" assert artifact.format == RegistryArtifactFormat.TAR_GZ + @pytest.mark.anyio + async def test_sidecar_lookup_failure_logs_only_redacted_uris( + self, temp_cache_dir: Path + ) -> None: + cache = RegistryArtifactCache(temp_cache_dir) + base_uri = "s3://affected-bucket/tenant/private/site-packages.tar.gz" + sidecar_uri = base_uri.removesuffix(".tar.gz") + ".squashfs" + + with ( + patch( + "tracecat.executor.registry_artifacts.blob.file_exists", + new_callable=AsyncMock, + side_effect=RuntimeError(f"failed for {sidecar_uri}"), + ), + patch("tracecat.executor.registry_artifacts.logger.warning") as warning, + ): + assert ( + await cache._sidecar_exists( + base_uri=base_uri, + sidecar_uri=sidecar_uri, + artifact_format=RegistryArtifactFormat.SQUASHFS, + ) + is False + ) + + warning.assert_called_once_with( + "Failed to check for registry artifact sidecar, falling back", + artifact_uri="s3://", + sidecar_uri="s3://", + artifact_format="squashfs", + error_type="RuntimeError", + ) + def test_can_try_squashfs_does_not_require_mount_binary(self, temp_cache_dir): """Prefer SquashFS whenever enabled; extraction may work without mounts.""" cache = RegistryArtifactCache(temp_cache_dir) @@ -693,7 +799,9 @@ async def mock_mount(self, ctx, image_path): ): result = await _materialize( cache, - "squashfs-key", + compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.tar.gz" + ), "s3://bucket/path/site-packages.tar.gz", ) @@ -912,7 +1020,8 @@ async def mock_download(self, ctx, path): path.write_bytes(_tarball_payload(size=1)) downloaded_paths.append(path) - def blocking_size_scan(path: Path) -> int: + def blocking_size_scan(path: Path, *, allocation_unit: int) -> int: + assert allocation_unit > 0 scan_started.set() scan_release.wait() input_present_at_finish.append(path.exists()) @@ -927,7 +1036,7 @@ def blocking_size_scan(path: Path) -> int: ), patch.object(TarballArtifact, "extract", new_callable=AsyncMock) as extract, ): - materializing = asyncio.create_task(cache._lease_artifact(artifact_uri)) + materializing = asyncio.create_task(_lease_and_release(cache, artifact_uri)) assert await asyncio.to_thread(scan_started.wait, 1) materializing.cancel() await asyncio.sleep(0) @@ -944,6 +1053,46 @@ def blocking_size_scan(path: Path) -> int: assert not downloaded_paths[0].exists() extract.assert_not_awaited() + @pytest.mark.anyio + async def test_failed_runtime_staging_cleanup_is_retried( + self, temp_cache_dir: Path + ) -> None: + """A failed extraction cleanup remains visible to later budget passes.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact = TarballArtifact( + uri="s3://bucket/path/failed-cleanup.tar.gz", + cache_key="failed-cleanup", + ) + ctx = cache._context_for(artifact.cache_key) + + async def mock_download(self, ctx, path): + del self, ctx + path.write_bytes(b"tarball") + + async def fail_extract(self, tarball_path, target_dir): + del self, tarball_path, target_dir + raise RuntimeError("extraction failed") + + with ( + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", fail_extract), + patch( + "tracecat.executor.registry_artifact_storage._delete_cache_path_off_loop", + new_callable=AsyncMock, + return_value=False, + ), + ): + with pytest.raises(RuntimeError, match="extraction failed"): + await artifact.materialize(ctx) + + assert len(cache._failed_startup_cleanup) == 1 + deferred_path = next(iter(cache._failed_startup_cleanup)) + assert deferred_path.is_dir() + + assert await cache._enforce_cache_budget() is True + assert cache._failed_startup_cleanup == set() + assert not deferred_path.exists() + @pytest.mark.anyio async def test_materialize_extracts_squashfs_when_mount_fails(self, temp_cache_dir): """Test that SquashFS mount failures fall back to unsquashfs extraction.""" @@ -982,7 +1131,9 @@ async def mock_extract(self, ctx, image_path): ): result = await _materialize( cache, - "fallback-key", + compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.tar.gz" + ), "s3://bucket/path/site-packages.tar.gz", ) @@ -1021,7 +1172,9 @@ async def mock_extract(self, ctx, image_path): ): result = await _materialize( cache, - "extract-key", + compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.tar.gz" + ), "s3://bucket/path/site-packages.tar.gz", ) @@ -1069,7 +1222,9 @@ async def mock_extract(self, ctx, image_path): ): result = await _materialize( cache, - "gzip-fallback-key", + compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.tar.gz" + ), "s3://bucket/path/site-packages.tar.gz", ) @@ -1093,7 +1248,7 @@ async def mock_download(self, ctx, path): with patch.object(TarballArtifact, "download", mock_download): result = await _materialize( cache, - "custom-key-test", + compute_registry_artifact_cache_key("s3://bucket/path/custom-key"), "s3://bucket/path/custom-key", ) @@ -1104,14 +1259,15 @@ async def mock_download(self, ctx, path): async def test_materialize_caches_result(self, temp_cache_dir): """Test that tarball extraction is cached.""" cache = RegistryArtifactCache(temp_cache_dir) - cache_key = "test-cache-key" + artifact_uri = "s3://bucket/test.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) target_dir = cache._paths_for(cache_key).tarball_target_dir target_dir.mkdir(parents=True) result = await _materialize( cache, cache_key, - "s3://bucket/test.tar.gz", + artifact_uri, ) assert result == [target_dir] @@ -1120,14 +1276,15 @@ async def test_materialize_caches_result(self, temp_cache_dir): async def test_materialize_concurrent_requests(self, temp_cache_dir): """Test that concurrent requests for same artifact do not race.""" cache = RegistryArtifactCache(temp_cache_dir) - cache_key = "concurrent-test" + artifact_uri = "s3://bucket/test.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) download_count = 0 async def mock_download(self, ctx, path): nonlocal download_count download_count += 1 await asyncio.sleep(0.1) - path.write_bytes(b"fake tarball content") + path.write_bytes(_tarball_payload(size=1)) async def mock_extract(self, tarball_path, target_dir): (target_dir / "extracted.txt").write_text("extracted") @@ -1137,9 +1294,9 @@ async def mock_extract(self, tarball_path, target_dir): patch.object(TarballArtifact, "extract", mock_extract), ): results = await asyncio.gather( - _materialize(cache, cache_key, "s3://bucket/test.tar.gz"), - _materialize(cache, cache_key, "s3://bucket/test.tar.gz"), - _materialize(cache, cache_key, "s3://bucket/test.tar.gz"), + _materialize(cache, cache_key, artifact_uri), + _materialize(cache, cache_key, artifact_uri), + _materialize(cache, cache_key, artifact_uri), ) assert all(r == results[0] for r in results) @@ -1252,6 +1409,7 @@ async def test_failed_first_admission_converges_deposited_image( artifact_uri = "s3://bucket/failed-first.squashfs" cache_key = compute_registry_artifact_cache_key(artifact_uri) artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) + empty_cache_bytes = cache._scan_cache_snapshot().total_bytes async def fail_after_deposit( artifact: SquashfsArtifact, @@ -1264,7 +1422,7 @@ async def fail_after_deposit( with ( patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 1), + patch(MAX_BYTES_CONFIG, empty_cache_bytes), patch.object( cache, "_artifact_candidates", @@ -1342,7 +1500,7 @@ async def blocked_unmount(requested_key: str) -> None: patch.object(cache, "_sidecar_exists", blocked_sidecar_lookup), patch.object(cache, "_unmount_idle_entry", blocked_unmount), ): - acquisition = asyncio.create_task(cache._lease_artifact(artifact_uri)) + acquisition = asyncio.create_task(_lease_and_release(cache, artifact_uri)) await lookup_started.wait() acquisition.cancel() await rollback_started.wait() @@ -1597,7 +1755,12 @@ async def hold_lease() -> None: with pytest.raises(asyncio.CancelledError): await holder - assert cleanup_calls == [*cache_keys, "converge"] + assert cleanup_calls == [ + cache_keys[1], + "converge", + cache_keys[0], + "converge", + ] assert all(cache._refcount(cache_key) == 0 for cache_key in cache_keys) @pytest.mark.anyio @@ -1702,7 +1865,7 @@ async def fail_download( del artifact, ctx, output_path raise RuntimeError("download failed") - tracked_lease_artifact = AsyncMock(wraps=cache._lease_artifact) + tracked_acquire_artifact = AsyncMock(wraps=cache._acquire_artifact) converge_cache_budget = AsyncMock() with ( @@ -1710,7 +1873,7 @@ async def fail_download( patch(SQUASHFS_ENABLED_CONFIG, False), patch.object(TarballArtifact, "download", fail_download), patch.object(cache, "_unmount", harness.unmount), - patch.object(cache, "_lease_artifact", tracked_lease_artifact), + patch.object(cache, "_acquire_artifact", tracked_acquire_artifact), patch.object( cache, "_converge_cache_budget", @@ -1722,7 +1885,8 @@ async def fail_download( pass requested_uris = [ - await_call.args[0] for await_call in tracked_lease_artifact.await_args_list + await_call.args[0].artifact_uri + for await_call in tracked_acquire_artifact.await_args_list ] assert requested_uris == [first_uri, failed_uri] assert cache._refcount(first_key) == 0 @@ -1732,7 +1896,7 @@ async def fail_download( assert cache._paths_for(first_key).squashfs_image_path.is_file() assert not cache._paths_for(failed_key).entry_dir.exists() assert untouched_path.is_dir() - converge_cache_budget.assert_awaited_once_with() + assert converge_cache_budget.await_count == 2 assert not cache.staging_dir.exists() or not any(cache.staging_dir.iterdir()) assert not cache.trash_dir.exists() or not any(cache.trash_dir.iterdir()) @@ -2001,7 +2165,7 @@ async def test_leased_entry_survives_eviction_of_idle_entry(self, temp_cache_dir ) async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") + path.write_bytes(_tarball_payload(size=1)) async def mock_extract(self, tarball_path, target_dir): (target_dir / "module.py").write_text("VALUE = 2") @@ -2025,11 +2189,19 @@ async def test_cold_download_reserves_space_before_writing( ) -> None: """Admission evicts idle bytes before a new download enters staging.""" cache = RegistryArtifactCache(temp_cache_dir) - idle = _write_image_entry(temp_cache_dir, "idle", size=80, mtime=100.0) + await cache.ensure_swept() + idle = _write_image_entry( + temp_cache_dir, + "idle", + size=128 * 1024, + mtime=100.0, + ) artifact_uri = "s3://bucket/new.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) payload = _tarball_payload(size=32) - max_bytes = len(payload) + 32 + allocation_unit = _filesystem_allocation_unit(temp_cache_dir) + structural_bytes = cache._scan_cache_snapshot().structural_bytes + max_bytes = structural_bytes + 7 * allocation_unit capacity_checked = False async def download_file_to_path( @@ -2039,10 +2211,14 @@ async def download_file_to_path( output_path: Path, max_bytes: int, ensure_capacity: Callable[[int], Awaitable[None]], + defer_cleanup: Callable[[Path], None] | None, + redact_log_identifiers: bool, ) -> int: del key, bucket nonlocal capacity_checked - assert max_bytes == len(payload) + 32 + assert max_bytes == structural_bytes + 7 * allocation_unit + assert defer_cleanup is not None + assert redact_log_identifiers is True await ensure_capacity(len(payload)) capacity_checked = True assert not idle.exists() @@ -2073,11 +2249,30 @@ async def test_impossible_tarball_reservation_preserves_warm_entry( ) -> None: """Impossible extraction cannot evict warm entries before rejection.""" cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() warm = _write_image_entry(temp_cache_dir, "warm", size=64, mtime=100.0) artifact_uri = "s3://bucket/compression-heavy.tar.gz" cache_key = compute_registry_artifact_cache_key(artifact_uri) payload = _tarball_payload(size=4096) - max_bytes = len(payload) + 256 + allocation_unit = _filesystem_allocation_unit(temp_cache_dir) + snapshot = cache._scan_cache_snapshot() + payload_reservation = allocated_size_bound( + len(payload), + allocation_unit=allocation_unit, + ) + # The download reservation includes one staging directory record. + max_bytes = snapshot.total_bytes + allocation_unit + payload_reservation + max_bytes += allocation_unit + extracted_size = ( + allocated_size_bound(4096, allocation_unit=allocation_unit) + + allocated_size_bound(0, allocation_unit=allocation_unit) + + allocated_size_bound( + 32 + len(os.fsencode("module.py")), + allocation_unit=allocation_unit, + ) + # Staged and canonical directory records remain allocated after rename. + + 2 * allocation_unit + ) async def download_file_to_path( *, @@ -2086,9 +2281,16 @@ async def download_file_to_path( output_path: Path, max_bytes: int, ensure_capacity: Callable[[int], Awaitable[None]], + defer_cleanup: Callable[[Path], None] | None, + redact_log_identifiers: bool, ) -> int: del key, bucket - assert max_bytes == len(payload) + 256 + assert ( + max_bytes + == snapshot.total_bytes + 2 * allocation_unit + payload_reservation + ) + assert defer_cleanup is not None + assert redact_log_identifiers is True await ensure_capacity(len(payload)) output_path.write_bytes(payload) return len(payload) @@ -2114,7 +2316,7 @@ async def download_file_to_path( async with cache.lease([artifact_uri]): pass - assert raised.value.additional_bytes == 4096 + assert raised.value.additional_bytes == extracted_size assert raised.value.max_bytes == max_bytes extract.assert_not_awaited() evict_entry.assert_not_awaited() @@ -2128,6 +2330,14 @@ async def test_squashfs_expansion_is_rejected_before_extraction( ) -> None: """SquashFS metadata is accounted before unsquashfs writes scratch.""" cache = RegistryArtifactCache(temp_cache_dir) + structural_bytes = cache._scan_cache_snapshot().structural_bytes + allocation_unit = _filesystem_allocation_unit(temp_cache_dir) + reserved_expansion = allocated_size_bound( + 101, + allocation_unit=allocation_unit, + ) + reserved_expansion += 2 * allocation_unit + max_bytes = structural_bytes + 2 * allocation_unit + reserved_expansion - 1 artifact_uri = "s3://bucket/oversized.squashfs" cache_key = compute_registry_artifact_cache_key(artifact_uri) @@ -2143,7 +2353,7 @@ async def download( with ( patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 100), + patch(MAX_BYTES_CONFIG, max_bytes), patch.object( RegistryArtifactMaterializationContext, "can_mount_squashfs", @@ -2166,7 +2376,7 @@ async def download( async with cache.lease([artifact_uri]): pass - assert raised.value.additional_bytes == 101 + assert raised.value.additional_bytes == reserved_expansion extract_image.assert_not_awaited() paths = cache._paths_for(cache_key) assert paths.squashfs_image_path.read_bytes() == b"image" @@ -2178,22 +2388,26 @@ async def test_successful_admission_enforces_actual_size_before_yield( ): """Post-publication enforcement sees the new entry's actual size.""" cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + allocation_unit = _filesystem_allocation_unit(temp_cache_dir) idle = _write_image_entry(temp_cache_dir, "idle", size=4096, mtime=100.0) + snapshot = cache._scan_cache_snapshot() + max_bytes = snapshot.total_bytes + 3 * allocation_unit new_uri = "s3://bucket/new.tar.gz" new_key = compute_registry_artifact_cache_key(new_uri) async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") + path.write_bytes(_tarball_payload(size=1)) async def mock_extract(self, tarball_path, target_dir): - (target_dir / "module.py").write_bytes(b"x" * 4096) + (target_dir / "module.py").write_bytes(b"x" * (2 * allocation_unit)) with ( patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 6000), + patch(MAX_BYTES_CONFIG, max_bytes), patch( "tracecat.executor.registry_artifacts._tarball_extracted_size", - return_value=4096, + return_value=allocation_unit, ), patch.object(TarballArtifact, "download", mock_download), patch.object(TarballArtifact, "extract", mock_extract), @@ -2218,7 +2432,7 @@ async def test_deletion_failure_does_not_block_materialization( cache_key = compute_registry_artifact_cache_key(artifact_uri) async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") + path.write_bytes(_tarball_payload(size=1)) async def mock_extract(self, tarball_path, target_dir): (target_dir / "module.py").write_text("VALUE = 2") @@ -2227,7 +2441,7 @@ async def mock_extract(self, tarball_path, target_dir): patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0), patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", return_value=False, ), patch.object(TarballArtifact, "download", mock_download), @@ -2269,6 +2483,8 @@ async def test_failed_physical_delete_retries_without_extra_eviction( size=16, mtime=300.0, ) + snapshot = cache._scan_cache_snapshot() + max_bytes = snapshot.structural_bytes + snapshot.entries["newest"].size_bytes real_delete = _delete_cache_path failed_once = False @@ -2281,9 +2497,9 @@ def fail_once(path: Path) -> bool: with ( patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 16), + patch(MAX_BYTES_CONFIG, max_bytes), patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", side_effect=fail_once, ), ): @@ -2311,6 +2527,8 @@ async def test_undeletable_trash_does_not_evict_warm_entries_for_admission( stale_trash = cache.trash_dir / "stale" stale_trash.mkdir(parents=True) (stale_trash / "image.squashfs").write_bytes(b"x" * 32) + snapshot = cache._scan_cache_snapshot() + allocation_unit = _filesystem_allocation_unit(temp_cache_dir) real_delete = _delete_cache_path def fail_stale_trash(path: Path) -> bool: @@ -2319,18 +2537,18 @@ def fail_stale_trash(path: Path) -> bool: return real_delete(path) with patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", side_effect=fail_stale_trash, ): with pytest.raises(RegistryArtifactCacheCapacityError) as raised: async with cache._admission_lock: await cache._ensure_cache_capacity( - additional_bytes=16, + additional_bytes=allocation_unit, protected_key="new", - max_bytes=64, + max_bytes=snapshot.total_bytes, ) - assert raised.value.current_bytes == 64 + assert raised.value.current_bytes == snapshot.total_bytes assert warm.exists() assert stale_trash.exists() @@ -2344,7 +2562,7 @@ async def test_rename_failure_does_not_block_materialization(self, temp_cache_di cache_key = compute_registry_artifact_cache_key(artifact_uri) async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") + path.write_bytes(_tarball_payload(size=1)) async def mock_extract(self, tarball_path, target_dir): (target_dir / "module.py").write_text("VALUE = 2") @@ -2353,7 +2571,7 @@ async def mock_extract(self, tarball_path, target_dir): patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0), patch( - "tracecat.executor.registry_artifacts._move_entry_to_trash", + "tracecat.executor.registry_artifact_storage._move_entry_to_trash", side_effect=OSError("rename failed"), ), patch.object(TarballArtifact, "download", mock_download), @@ -2376,7 +2594,7 @@ async def test_cache_scan_failure_does_not_block_materialization( cache_key = compute_registry_artifact_cache_key(artifact_uri) async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") + path.write_bytes(_tarball_payload(size=1)) async def mock_extract(self, tarball_path, target_dir): (target_dir / "module.py").write_text("VALUE = 2") @@ -2422,6 +2640,40 @@ async def test_releasing_a_lease_skips_the_scan_for_a_cache_hit( assert cache._budget_dirty is False + @pytest.mark.anyio + async def test_mutable_cache_hit_rescans_unknown_entry_growth( + self, temp_cache_dir: Path + ) -> None: + """Writable direct actions cannot grow a warm entry outside the cap.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/mutable-cached.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + target_dir = _write_tarball_entry(temp_cache_dir, cache_key) + entry_dir = cache._paths_for(cache_key).entry_dir + max_bytes = cache._scan_cache_snapshot().total_bytes + cache._budget_dirty = False + + with ( + patch(MAX_ENTRIES_CONFIG, 10), + patch(MAX_BYTES_CONFIG, max_bytes), + patch.object( + cache, + "_scan_cache_snapshot", + wraps=cache._scan_cache_snapshot, + ) as scan_cache_snapshot, + ): + async with cache.lease( + [artifact_uri], + paths_may_be_modified=True, + ) as registry_paths: + assert registry_paths == [target_dir] + (entry_dir / "action-output.bin").write_bytes(b"x" * 4096) + + assert scan_cache_snapshot.call_count == 1 + assert not entry_dir.exists() + assert cache._budget_dirty is False + @pytest.mark.anyio async def test_failed_cold_admission_keeps_warm_lru(self, temp_cache_dir): """A missing cold artifact must not evict an existing warm entry.""" @@ -2450,8 +2702,10 @@ async def mock_download(self, ctx, path): assert not cache._paths_for(missing_key).entry_dir.exists() @pytest.mark.anyio - async def test_failed_materialization_rearms_budget_dirty(self, temp_cache_dir): - """A failed materialization may leave a canonical image to evict.""" + async def test_failed_materialization_converges_deposited_image( + self, temp_cache_dir + ): + """A failed materialization still runs release-time budget enforcement.""" cache = RegistryArtifactCache(temp_cache_dir) await cache.ensure_swept() assert cache._budget_dirty is False @@ -2459,6 +2713,7 @@ async def test_failed_materialization_rearms_budget_dirty(self, temp_cache_dir): artifact_uri = "s3://bucket/path/site-packages.squashfs" cache_key = compute_registry_artifact_cache_key(artifact_uri) artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) + enforce_cache_budget = AsyncMock(return_value=True) async def mock_materialize(self, ctx): ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) @@ -2473,12 +2728,14 @@ async def mock_materialize(self, ctx): return_value=[artifact], ), patch.object(SquashfsArtifact, "materialize", mock_materialize), + patch.object(cache, "_enforce_cache_budget", enforce_cache_budget), ): with pytest.raises(RuntimeError, match="mount failed"): await _materialize(cache, cache_key, artifact_uri) assert cache._paths_for(cache_key).squashfs_image_path.is_file() - assert cache._budget_dirty is True + assert cache._budget_dirty is False + enforce_cache_budget.assert_awaited_once_with() @pytest.mark.anyio async def test_materialization_rearms_budget_dirty_consumed_mid_flight( @@ -2492,6 +2749,7 @@ async def test_materialization_rearms_budget_dirty_consumed_mid_flight( artifact_uri = "s3://bucket/path/site-packages.squashfs" cache_key = compute_registry_artifact_cache_key(artifact_uri) artifact = SquashfsArtifact(uri=artifact_uri, cache_key=cache_key) + enforce_cache_budget = AsyncMock(return_value=True) async def mock_materialize(self, ctx): # A concurrent lease release consumes the dirty signal and finishes @@ -2509,10 +2767,15 @@ async def mock_materialize(self, ctx): return_value=[artifact], ), patch.object(SquashfsArtifact, "materialize", mock_materialize), + patch.object(cache, "_enforce_cache_budget", enforce_cache_budget), ): await _materialize(cache, cache_key, artifact_uri) - assert cache._budget_dirty is True + assert cache._budget_dirty is False + assert enforce_cache_budget.await_args_list == [ + call(protected_key=cache_key), + call(), + ] @pytest.mark.anyio async def test_release_keeps_retrying_while_the_cache_stays_over_budget( @@ -2561,7 +2824,7 @@ async def mock_enforce_cache_budget( return True async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") + path.write_bytes(_tarball_payload(size=1)) async def mock_extract(self, tarball_path, target_dir): (target_dir / "module.py").write_text("VALUE = 1") @@ -2636,10 +2899,16 @@ async def test_enforce_budget_evicts_least_recently_used_until_under_max_bytes( os.utime(oldest_entry, (100.0, 100.0)) older = _write_image_entry(temp_cache_dir, "older", size=4096, mtime=200.0) newest = _write_image_entry(temp_cache_dir, "newest", size=4096, mtime=300.0) + snapshot = cache._scan_cache_snapshot() + max_bytes = ( + snapshot.structural_bytes + + snapshot.entries["older"].size_bytes + + snapshot.entries["newest"].size_bytes + ) with ( patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 9000), + patch(MAX_BYTES_CONFIG, max_bytes), ): within_budget = await cache._enforce_cache_budget(protected_key="pending") @@ -2704,9 +2973,9 @@ async def test_concurrent_budget_passes_only_evict_once(self, temp_cache_dir): scan_count_lock = threading.Lock() scan_count = 0 - def controlled_scan(): + def controlled_scan(*, allocation_unit: int | None = None): nonlocal scan_count - entries = original_scan() + entries = original_scan(allocation_unit=allocation_unit) with scan_count_lock: scan_index = scan_count scan_count += 1 @@ -2940,14 +3209,14 @@ def blocked_delete(path: Path) -> bool: return deleted async def mock_download(self, ctx, path): - path.write_bytes(b"fake tarball") + path.write_bytes(_tarball_payload(size=1)) async def mock_extract(self, tarball_path, target_dir): (target_dir / "module.py").write_text("VALUE = 2") with ( patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", side_effect=blocked_delete, ), patch( @@ -3046,7 +3315,7 @@ async def test_doomed_eviction_names_are_startup_scratch(self, temp_cache_dir): paths.tarball_target_dir.mkdir() with patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", return_value=True, ) as delete_cache_path: assert await cache._evict_entry(cache_key) == RegistryArtifactEviction( @@ -3100,20 +3369,17 @@ async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): assert not idle.exists() @pytest.mark.anyio - async def test_eviction_keeps_stable_runtime_state(self, temp_cache_dir): - """A key keeps one lock and zeroed lease state for the process lifetime.""" + async def test_eviction_retires_idle_runtime_state(self, temp_cache_dir): + """Eviction releases keyed lock metadata once every user is gone.""" cache = RegistryArtifactCache(temp_cache_dir) _write_tarball_entry(temp_cache_dir, "bookkeeping") cache._acquire_lease("bookkeeping") cache._release_lease("bookkeeping") - lock = cache._runtime_for("bookkeeping").lock assert await cache._evict_entry("bookkeeping") == RegistryArtifactEviction( retired=True, reclaimed=True ) - runtime = cache._runtime["bookkeeping"] - assert runtime.lock is lock - assert runtime.refcount == 0 + assert "bookkeeping" not in cache._runtime @pytest.mark.anyio async def test_eviction_skips_busy_key(self, temp_cache_dir): @@ -3168,6 +3434,23 @@ async def test_sweep_tolerates_missing_cache_dir(self, temp_cache_dir): assert cache.cache_dir == cache_dir assert not cache_dir.exists() + @pytest.mark.anyio + async def test_sweep_rejects_symlinked_cache_root( + self, temp_cache_dir: Path + ) -> None: + outside = temp_cache_dir / "outside" + orphaned = outside / "staging" / "orphaned.tmp" + orphaned.parent.mkdir(parents=True) + orphaned.write_bytes(b"keep") + symlinked_root = temp_cache_dir / "cache-link" + symlinked_root.symlink_to(outside, target_is_directory=True) + cache = RegistryArtifactCache(symlinked_root) + + with pytest.raises(OSError, match="Unsafe registry artifact cache directory"): + await cache.ensure_swept() + + assert orphaned.read_bytes() == b"keep" + @pytest.mark.anyio async def test_sweep_removes_orphaned_work(self, temp_cache_dir): """Interrupted work is reclaimed without touching unrelated paths.""" @@ -3190,6 +3473,30 @@ async def test_sweep_removes_orphaned_work(self, temp_cache_dir): assert unrelated_file.read_text() == "keep" assert entry_dir.is_dir() + @pytest.mark.anyio + async def test_sweep_reclaims_exact_legacy_top_level_layout( + self, temp_cache_dir: Path + ) -> None: + """Pre-entries cache paths cannot remain outside byte accounting.""" + cache_key = "0123456789abcdef" + legacy_image = temp_cache_dir / f"squashfs-{cache_key}.squashfs" + legacy_image.write_bytes(b"image") + legacy_extraction = temp_cache_dir / f"tarball-{cache_key}" + legacy_extraction.mkdir() + (legacy_extraction / "module.py").write_text("VALUE = 1") + legacy_staging = temp_cache_dir / f"{cache_key}.123.456.tmp" + legacy_staging.mkdir() + unrelated = temp_cache_dir / "squashfs-not-a-cache-key.squashfs" + unrelated.write_bytes(b"keep") + + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + + assert not legacy_image.exists() + assert not legacy_extraction.exists() + assert not legacy_staging.exists() + assert unrelated.read_bytes() == b"keep" + @pytest.mark.anyio async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): """A live mountpoint belongs to a running process and must survive.""" @@ -3360,7 +3667,7 @@ def fail_orphan(path: Path) -> bool: patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0), patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", side_effect=fail_orphan, ), ): @@ -3389,7 +3696,7 @@ def fail_once(path: Path) -> bool: return real_delete(path) with patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", side_effect=fail_once, ): await cache.ensure_swept() @@ -3425,7 +3732,7 @@ async def test_failed_startup_retirement_stays_dirty_and_retries( patch(MAX_ENTRIES_CONFIG, 1), patch(MAX_BYTES_CONFIG, 0), patch( - "tracecat.executor.registry_artifacts._move_entry_to_trash", + "tracecat.executor.registry_artifact_storage._move_entry_to_trash", side_effect=OSError("rename failed"), ), ): @@ -3469,6 +3776,8 @@ async def test_failed_startup_physical_delete_retries_exact_path( mtime=300.0, ) cache = RegistryArtifactCache(temp_cache_dir) + snapshot = cache._scan_cache_snapshot() + max_bytes = snapshot.structural_bytes + snapshot.entries["newest"].size_bytes real_delete = _delete_cache_path failed_once = False @@ -3481,9 +3790,9 @@ def fail_once(path: Path) -> bool: with ( patch(MAX_ENTRIES_CONFIG, 0), - patch(MAX_BYTES_CONFIG, 16), + patch(MAX_BYTES_CONFIG, max_bytes), patch( - "tracecat.executor.registry_artifacts._delete_cache_path", + "tracecat.executor.registry_artifact_storage._delete_cache_path", side_effect=fail_once, ), ): diff --git a/tests/unit/test_storage_blob.py b/tests/unit/test_storage_blob.py index c960a2cf5b..ce0afcef52 100644 --- a/tests/unit/test_storage_blob.py +++ b/tests/unit/test_storage_blob.py @@ -704,6 +704,70 @@ async def test_open_download_stream_yields_stream_and_length(self, mock_get_clie assert stream is mock_body assert length == 123 + @pytest.mark.anyio + @patch("tracecat.storage.blob.logger") + @patch("tracecat.storage.blob.get_storage_client") + async def test_open_download_stream_can_redact_transport_failure( + self, mock_get_client, mock_logger + ) -> None: + """Registry downloads suppress bucket, key, and provider messages.""" + sensitive_bucket = "affected-customer-bucket" + sensitive_key = "tenant/private/site-packages.squashfs" + mock_client = AsyncMock() + mock_get_client.return_value.__aenter__.return_value = mock_client + mock_client.get_object.side_effect = ClientError( + error_response={ + "Error": { + "Code": "AccessDenied", + "Message": f"denied {sensitive_bucket}/{sensitive_key}", + } + }, + operation_name="get_object", + ) + + with pytest.raises(blob_module.StorageDownloadError) as raised: + async with open_download_stream( + key=sensitive_key, + bucket=sensitive_bucket, + redact_log_identifiers=True, + ): + pass + + assert sensitive_bucket not in str(raised.value) + assert sensitive_key not in str(raised.value) + mock_logger.error.assert_called_once_with( + "Failed to open download stream", + key="", + bucket="", + error_code="AccessDenied", + error_type="ClientError", + ) + + @pytest.mark.anyio + @patch("tracecat.storage.blob.logger") + @patch("tracecat.storage.blob.get_storage_client") + async def test_redacted_download_rejects_provider_prose_as_error_code( + self, mock_get_client, mock_logger + ) -> None: + sensitive_message = "tenant/private/object is forbidden" + mock_client = AsyncMock() + mock_get_client.return_value.__aenter__.return_value = mock_client + mock_client.get_object.side_effect = ClientError( + error_response={"Error": {"Code": sensitive_message}}, + operation_name="get_object", + ) + + with pytest.raises(blob_module.StorageDownloadError) as raised: + async with open_download_stream( + key="tenant/private/object", + bucket="affected-bucket", + redact_log_identifiers=True, + ): + pass + + assert raised.value.error_code is None + assert sensitive_message not in str(mock_logger.mock_calls) + @pytest.mark.anyio async def test_download_file_to_path_writes_bytes( self, tmp_path: Path, monkeypatch @@ -722,7 +786,9 @@ async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 dummy_stream = DummyStream(chunks) @asynccontextmanager - async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 + async def _fake_open_download_stream( + *, key: str, bucket: str, redact_log_identifiers: bool = False + ): # noqa: ARG001 yield dummy_stream, sum(len(c) for c in chunks) monkeypatch.setattr( @@ -751,7 +817,9 @@ async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 yield b"should-not-write" @asynccontextmanager - async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 + async def _fake_open_download_stream( + *, key: str, bucket: str, redact_log_identifiers: bool = False + ): # noqa: ARG001 yield DummyStream(), 10 monkeypatch.setattr( @@ -782,7 +850,9 @@ async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 yield b"payload" @asynccontextmanager - async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 + async def _fake_open_download_stream( + *, key: str, bucket: str, redact_log_identifiers: bool = False + ): # noqa: ARG001 yield DummyStream(), 7 monkeypatch.setattr( @@ -819,7 +889,9 @@ async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 yield b"too-large" @asynccontextmanager - async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 + async def _fake_open_download_stream( + *, key: str, bucket: str, redact_log_identifiers: bool = False + ): # noqa: ARG001 yield DummyStream(), 5 monkeypatch.setattr( @@ -857,7 +929,9 @@ async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 yield b"hello" @asynccontextmanager - async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 + async def _fake_open_download_stream( + *, key: str, bucket: str, redact_log_identifiers: bool = False + ): # noqa: ARG001 yield DummyStream(), 5 monkeypatch.setattr( @@ -892,7 +966,9 @@ async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 await asyncio.Event().wait() @asynccontextmanager - async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 + async def _fake_open_download_stream( + *, key: str, bucket: str, redact_log_identifiers: bool = False + ): # noqa: ARG001 yield DummyStream(), None monkeypatch.setattr( @@ -918,6 +994,49 @@ async def _fake_open_download_stream(*, key: str, bucket: str): # noqa: ARG001 assert not out.exists() assert not (tmp_path / "out.bin.part").exists() + @pytest.mark.anyio + async def test_download_file_to_path_defers_failed_partial_cleanup( + self, tmp_path: Path, monkeypatch + ) -> None: + """A failed partial unlink remains discoverable for a later retry.""" + + class DummyStream: + async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 + yield b"partial" + raise RuntimeError("stream failed") + + @asynccontextmanager + async def _fake_open_download_stream( + *, key: str, bucket: str, redact_log_identifiers: bool + ): # noqa: ARG001 + assert redact_log_identifiers is True + yield DummyStream(), None + + monkeypatch.setattr( + "tracecat.storage.blob.open_download_stream", + _fake_open_download_stream, + ) + out = tmp_path / "out.bin" + partial_path = tmp_path / "out.bin.part" + deferred: list[Path] = [] + + with patch.object( + Path, + "unlink", + side_effect=PermissionError("cleanup denied"), + ): + with pytest.raises(RuntimeError, match="stream failed"): + await download_file_to_path( + key="tenant/private/object", + bucket="affected-bucket", + output_path=out, + defer_cleanup=deferred.append, + redact_log_identifiers=True, + ) + + assert deferred == [partial_path] + assert partial_path.is_file() + @pytest.mark.anyio @patch("tracecat.storage.blob.get_storage_client") async def test_ensure_bucket_exists_create_error_propagates(self, mock_get_client): diff --git a/tracecat/concurrency.py b/tracecat/concurrency.py index 3878f75c92..f518823c12 100644 --- a/tracecat/concurrency.py +++ b/tracecat/concurrency.py @@ -36,6 +36,35 @@ async def rejoin_future_on_cancel[T](future: asyncio.Future[T]) -> T: raise +async def rejoin_future_through_cancellation[T](future: asyncio.Future[T]) -> T: + """Rejoin a future through repeated cancellation without losing failures. + + A pending caller cancellation is propagated only after ``future`` finishes. + If cleanup also fails, cancellation remains the primary exception and the + cleanup failure is retained as its cause. + """ + pending_cancellation: asyncio.CancelledError | None = None + while not future.done(): + try: + await asyncio.shield(future) + except asyncio.CancelledError as e: + if future.cancelled(): + raise + pending_cancellation = e + except BaseException: + break + + try: + result = future.result() + except BaseException as future_error: + if pending_cancellation is not None: + raise pending_cancellation from future_error + raise + if pending_cancellation is not None: + raise pending_cancellation + return result + + async def run_blocking_rejoin_on_cancel[T](operation: Callable[[], T]) -> T: """Run blocking work without abandoning its worker thread.""" return await rejoin_future_on_cancel( diff --git a/tracecat/executor/action_runner.py b/tracecat/executor/action_runner.py index 63b9c3606d..1eea6a74bb 100644 --- a/tracecat/executor/action_runner.py +++ b/tracecat/executor/action_runner.py @@ -45,7 +45,10 @@ from tracecat.logger import logger from tracecat.sandbox.executor import ActionSandboxConfig, NsjailExecutor from tracecat.sandbox.types import ResourceLimits -from tracecat.sandbox.utils import communicate_process_group +from tracecat.sandbox.utils import ( + communicate_process_group, + terminate_supervised_process, +) from tracecat.secrets.common import apply_masks, apply_masks_object if TYPE_CHECKING: @@ -109,20 +112,23 @@ def _is_sandbox_available() -> bool: def _direct_subprocess_command(minimal_runner_path: Path) -> list[str]: - """Build the direct action subprocess command with new privileges disabled.""" + """Build a contained direct-action command with privileges disabled.""" runner_command = [sys.executable, str(minimal_runner_path)] - if sys.platform != "linux": - return runner_command - setpriv = shutil.which("setpriv") if setpriv is None: raise RuntimeError("setpriv is required for direct action subprocess isolation") + supervisor_path = Path(__file__).with_name("process_supervisor.py") return [ setpriv, "--no-new-privs", "--inh-caps=-all", "--ambient-caps=-all", + sys.executable, + # Keep registry-controlled PYTHONPATH out of the supervisor interpreter. + # Isolated mode leaves the environment intact for the nested action. + "-I", + str(supervisor_path), *runner_command, ] @@ -172,15 +178,19 @@ async def execute_action( """ timeout = timeout or config.TRACECAT__EXECUTOR_CLIENT_TIMEOUT + # Direct subprocesses receive host paths and can modify extracted + # artifacts. NsJail exposes the same paths through read-only bind mounts. + use_sandbox = force_sandbox or ( + config.TRACECAT__EXECUTOR_SANDBOX_ENABLED and _is_sandbox_available() + ) + # Materialize each registry artifact, collect paths in deterministic order. # The lease is held for the whole subprocess execution so cache eviction # cannot delete a directory the subprocess is still importing from. - async with self.registry_artifacts.lease(artifact_uris) as registry_paths: - # Check if sandbox execution is enabled and available - # force_sandbox=True overrides config (used by ephemeral backend) - use_sandbox = force_sandbox or ( - config.TRACECAT__EXECUTOR_SANDBOX_ENABLED and _is_sandbox_available() - ) + async with self.registry_artifacts.lease( + artifact_uris, + paths_may_be_modified=not use_sandbox, + ) as registry_paths: logger.debug( "Using sandbox execution", use_sandbox=use_sandbox, @@ -405,6 +415,7 @@ async def _execute_direct( if existing_pythonpath: pythonpath_parts.append(existing_pythonpath) env["PYTHONPATH"] = ":".join(pythonpath_parts) if pythonpath_parts else "" + env["PYTHONDONTWRITEBYTECODE"] = "1" # Get path to minimal_runner.py for subprocess execution from tracecat.executor import minimal_runner as minimal_runner_module @@ -443,6 +454,7 @@ async def _execute_direct( proc, input=input_json, timeout=timeout, + terminate=terminate_supervised_process, ) elapsed_ms = (time.monotonic() - start_time) * 1000 logger.info( diff --git a/tracecat/executor/backends/base.py b/tracecat/executor/backends/base.py index 082d4dfe7e..b9d35afe24 100644 --- a/tracecat/executor/backends/base.py +++ b/tracecat/executor/backends/base.py @@ -148,9 +148,14 @@ async def _execute_run_python( ) # The lease is held for the whole sandbox run so cache eviction cannot - # delete a directory the script is still importing from. + # delete a directory the script is still importing from. The sandbox + # service can select UnsafePidExecutor, which exposes host paths writable, + # so conservatively rescan their footprint after every run. registry_artifacts = get_action_runner().registry_artifacts - async with registry_artifacts.lease(artifact_uris) as registry_paths: + async with registry_artifacts.lease( + artifact_uris, + paths_may_be_modified=True, + ) as registry_paths: return await self._run_python_in_sandbox( script=script, args=args, diff --git a/tracecat/executor/backends/test.py b/tracecat/executor/backends/test.py index ce7b2025ae..6c7c4aa3f5 100644 --- a/tracecat/executor/backends/test.py +++ b/tracecat/executor/backends/test.py @@ -41,6 +41,9 @@ from tracecat.executor.action_runner import get_action_runner from tracecat.executor.backends.base import ExecutorBackend from tracecat.executor.backends.registry_helpers import get_registry_artifact_uris +from tracecat.executor.registry_artifacts import ( + compute_registry_artifact_cache_key, +) from tracecat.executor.schemas import ( ActionImplementation, ExecutorActionErrorInfo, @@ -301,13 +304,16 @@ async def _lease_registry_artifacts( for artifact_uri in artifact_uris: try: artifact_paths = await leases.enter_async_context( - registry_artifacts.lease([artifact_uri]) + registry_artifacts.lease( + [artifact_uri], + paths_may_be_modified=True, + ) ) except Exception as e: logger.warning( "Failed to materialize artifact for test execution", - artifact_uri=artifact_uri, - error=str(e), + cache_key=compute_registry_artifact_cache_key(artifact_uri), + error_type=type(e).__name__, ) continue extracted_paths.extend(str(path) for path in artifact_paths) diff --git a/tracecat/executor/process_supervisor.py b/tracecat/executor/process_supervisor.py new file mode 100644 index 0000000000..81cf618582 --- /dev/null +++ b/tracecat/executor/process_supervisor.py @@ -0,0 +1,276 @@ +"""Contain descendants of one Linux subprocess. + +The outer process remains the child observed by its caller. A detached +subreaper monitor owns the actual action process and all descendants orphaned by +it. Closing the control pipe asks that monitor to kill and reap its complete +child tree before the outer process exits. + +This module intentionally uses only the Python standard library so invoking it +does not import the Tracecat application in an untrusted action process. +""" + +from __future__ import annotations + +import ctypes +import os +import select +import signal +import sys +from collections.abc import Sequence +from contextlib import suppress +from pathlib import Path +from types import FrameType + +_PR_SET_CHILD_SUBREAPER = 36 +_PR_SET_PDEATHSIG = 1 +_PARENT_POLL_INTERVAL_MS = 10 + + +def _exit_code(wait_status: int) -> int: + """Convert a wait status into a shell-compatible non-negative exit code.""" + code = os.waitstatus_to_exitcode(wait_status) + return code if code >= 0 else 128 - code + + +def _set_child_subreaper() -> None: + """Make this process the reparenting boundary for orphaned descendants.""" + libc = ctypes.CDLL(None, use_errno=True) + prctl = libc.prctl + prctl.argtypes = [ + ctypes.c_int, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ] + prctl.restype = ctypes.c_int + if prctl(_PR_SET_CHILD_SUBREAPER, 1, 0, 0, 0) != 0: + error_number = ctypes.get_errno() + raise OSError(error_number, os.strerror(error_number)) + + +def _set_parent_death_signal(signal_number: int) -> None: + """Ask Linux to signal this process when its current parent exits.""" + libc = ctypes.CDLL(None, use_errno=True) + prctl = libc.prctl + prctl.argtypes = [ + ctypes.c_int, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ] + prctl.restype = ctypes.c_int + if prctl(_PR_SET_PDEATHSIG, signal_number, 0, 0, 0, 0) != 0: + error_number = ctypes.get_errno() + raise OSError(error_number, os.strerror(error_number)) + + +def _direct_child_pids_from_children_file() -> list[int]: + """Return direct child PIDs using procfs's optional children file.""" + children_path = Path(f"/proc/self/task/{os.getpid()}/children") + contents = children_path.read_text().strip() + return [int(pid) for pid in contents.split()] if contents else [] + + +def _proc_parent_pid(stat_path: Path) -> int: + """Read one process's parent PID from its procfs stat record.""" + contents = stat_path.read_text() + closing_paren = contents.rfind(")") + fields = contents[closing_paren + 1 :].split() + if closing_paren < 0 or len(fields) < 2: + raise RuntimeError("Malformed procfs stat record") + return int(fields[1]) + + +def _direct_child_pids_from_proc_stat() -> list[int]: + """Return direct child PIDs by scanning portable procfs stat records.""" + parent_pid = os.getpid() + _proc_parent_pid(Path("/proc/self/stat")) + child_pids: list[int] = [] + + with os.scandir("/proc") as entries: + for entry in entries: + if not entry.name.isdecimal(): + continue + try: + candidate_pid = int(entry.name) + candidate_parent_pid = _proc_parent_pid(Path(entry.path) / "stat") + except (FileNotFoundError, PermissionError, ProcessLookupError): + continue + if candidate_parent_pid == parent_pid: + child_pids.append(candidate_pid) + + return child_pids + + +def _direct_child_pids() -> list[int]: + """Return direct child PIDs from either available procfs interface.""" + try: + return _direct_child_pids_from_children_file() + except OSError: + return _direct_child_pids_from_proc_stat() + + +def _waitpid(pid: int, options: int = 0) -> tuple[int, int]: + """Wait for a child while tolerating signal interruptions.""" + while True: + try: + return os.waitpid(pid, options) + except InterruptedError: + continue + + +def _kill_and_reap_children() -> None: + """Kill every adopted child, including descendants orphaned while reaping.""" + while child_pids := _direct_child_pids(): + for child_pid in child_pids: + with suppress(ProcessLookupError): + os.kill(child_pid, signal.SIGKILL) + for child_pid in child_pids: + with suppress(ChildProcessError): + _waitpid(child_pid) + + +def _exec(command: Sequence[str], control_fd: int) -> None: + """Replace this child with the actual action command.""" + os.close(control_fd) + try: + os.execvpe(command[0], list(command), os.environ) + except OSError: + with suppress(OSError): + os.write(2, b"Failed to execute supervised process\n") + os._exit(127) + + +def _run_monitor( + control_fd: int, + command: Sequence[str], + *, + supervisor_pid: int, +) -> int: + """Run the action below a detached subreaper and contain its descendants.""" + try: + parent_cleanup_requested = False + + def request_parent_cleanup( + _signal: int, + _frame: FrameType | None, + ) -> None: + nonlocal parent_cleanup_requested + parent_cleanup_requested = True + + # The action can stop this same-UID monitor and kill its outer parent. + # SIGCONT both resumes the monitor and records that it must clean up. + signal.signal(signal.SIGCONT, request_parent_cleanup) + _set_parent_death_signal(signal.SIGCONT) + # Close the fork-to-prctl race before any untrusted command is started. + if os.getppid() != supervisor_pid or parent_cleanup_requested: + return 128 + signal.SIGTERM + + os.setsid() + _set_child_subreaper() + _direct_child_pids() + + action_pid = os.fork() + if action_pid == 0: + _exec(command, control_fd) + + poller = select.poll() + poller.register(control_fd, select.POLLIN | select.POLLHUP | select.POLLERR) + action_status: int | None = None + parent_closed = False + + while ( + action_status is None and not parent_closed and not parent_cleanup_requested + ): + waited_pid, wait_status = _waitpid(action_pid, os.WNOHANG) + if waited_pid == action_pid: + action_status = wait_status + break + if poller.poll(_PARENT_POLL_INTERVAL_MS): + parent_closed = os.read(control_fd, 1) == b"" + + _kill_and_reap_children() + if parent_closed or parent_cleanup_requested: + return 128 + signal.SIGTERM + if action_status is None: + raise RuntimeError("Supervised action exited without a wait status") + return _exit_code(action_status) + except BaseException: + with suppress(BaseException): + _kill_and_reap_children() + with suppress(OSError): + os.write(2, b"Process supervisor failed\n") + return 1 + finally: + with suppress(OSError): + os.close(control_fd) + + +def supervise(command: Sequence[str]) -> int: + """Run one command and return only after all of its descendants are reaped.""" + if sys.platform != "linux": + raise RuntimeError("The process supervisor requires Linux") + if not command: + raise ValueError("A supervised command is required") + + # The monitor is intentionally visible as the action's parent. Make this + # outer process a second subreaper so killing that monitor only reparents + # the action tree here, where it is still killed and reaped before return. + _set_child_subreaper() + _direct_child_pids() + + control_read_fd, control_write_fd = os.pipe() + writer_open = True + monitor_pid: int | None = None + + def close_control_pipe() -> None: + nonlocal writer_open + if writer_open: + with suppress(OSError): + os.close(control_write_fd) + writer_open = False + + def request_cleanup(_signal: int, _frame: FrameType | None) -> None: + close_control_pipe() + # The action shares the monitor's UID and can stop it. SIGKILL the + # monitor session so the outer subreaper adopts and reaps every member, + # including descendants that detached into their own sessions. + if monitor_pid is not None: + with suppress(ProcessLookupError): + os.killpg(monitor_pid, signal.SIGKILL) + + signal.signal(signal.SIGTERM, request_cleanup) + signal.signal(signal.SIGINT, request_cleanup) + + supervisor_pid = os.getpid() + monitor_pid = os.fork() + if monitor_pid == 0: + signal.signal(signal.SIGTERM, signal.SIG_DFL) + signal.signal(signal.SIGINT, signal.SIG_DFL) + os.close(control_write_fd) + os._exit( + _run_monitor( + control_read_fd, + command, + supervisor_pid=supervisor_pid, + ) + ) + + os.close(control_read_fd) + try: + _, monitor_status = _waitpid(monitor_pid) + return _exit_code(monitor_status) + finally: + close_control_pipe() + _kill_and_reap_children() + + +def main() -> int: + """CLI entry point.""" + return supervise(sys.argv[1:]) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tracecat/executor/registry_artifact_budget.py b/tracecat/executor/registry_artifact_budget.py new file mode 100644 index 0000000000..c86a72b531 --- /dev/null +++ b/tracecat/executor/registry_artifact_budget.py @@ -0,0 +1,85 @@ +"""Pure cache snapshots and eviction planning for registry artifacts.""" + +from __future__ import annotations + +from collections.abc import Mapping, Set +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactCacheEntry: + """Measured on-disk footprint and recency for one artifact key.""" + + cache_key: str + size_bytes: int + last_used: float + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactCacheSnapshot: + """One internally consistent measurement of cache-owned storage.""" + + entries: dict[str, RegistryArtifactCacheEntry] + structural_bytes: int + staging_bytes: int + trash_bytes: int + + @property + def total_bytes(self) -> int: + """Return all measured bytes owned by the cache.""" + return ( + self.structural_bytes + + self.staging_bytes + + self.trash_bytes + + sum(entry.size_bytes for entry in self.entries.values()) + ) + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactCacheBudget: + """Shared entry and byte limits for every enforcement path.""" + + max_entries: int + max_bytes: int + additional_bytes: int = 0 + + def fits(self, *, entry_count: int, total_bytes: int) -> bool: + """Return whether a measured cache satisfies this budget.""" + return (self.max_entries <= 0 or entry_count <= self.max_entries) and ( + self.max_bytes <= 0 or total_bytes + self.additional_bytes <= self.max_bytes + ) + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactEvictionPlan: + """Ordered eligible entries and whether deleting all could fit the budget.""" + + candidates: tuple[RegistryArtifactCacheEntry, ...] + can_fit: bool + + +def plan_registry_artifact_evictions( + entries: Mapping[str, RegistryArtifactCacheEntry], + *, + total_bytes: int, + budget: RegistryArtifactCacheBudget, + excluded: Set[str], + effective_last_used: Mapping[str, float] | None = None, +) -> RegistryArtifactEvictionPlan: + """Return one deterministic LRU plan without mutating cache state.""" + recency = effective_last_used or {} + candidates = tuple( + sorted( + (entry for entry in entries.values() if entry.cache_key not in excluded), + key=lambda entry: recency.get(entry.cache_key, entry.last_used), + ) + ) + projected_bytes = total_bytes - sum(entry.size_bytes for entry in candidates) + projected_entries = len(entries) - len(candidates) + return RegistryArtifactEvictionPlan( + candidates=candidates, + can_fit=budget.fits( + entry_count=projected_entries, + total_bytes=projected_bytes, + ), + ) diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py new file mode 100644 index 0000000000..945c1f7a60 --- /dev/null +++ b/tracecat/executor/registry_artifact_storage.py @@ -0,0 +1,1192 @@ +"""Filesystem ownership and budget enforcement for registry artifacts.""" + +from __future__ import annotations + +import asyncio +import contextlib +import functools +import os +import re +import shutil +import stat +import threading +import time +from collections.abc import AsyncGenerator, Awaitable, Callable, Iterable +from contextlib import asynccontextmanager +from dataclasses import dataclass, field +from pathlib import Path + +from tracecat import config +from tracecat.concurrency import ( + drain_future_through_cancellation, + rejoin_future_on_cancel, + run_blocking_rejoin_on_cancel, +) +from tracecat.executor.registry_artifact_budget import ( + RegistryArtifactCacheBudget, + RegistryArtifactCacheEntry, + RegistryArtifactCacheSnapshot, + plan_registry_artifact_evictions, +) +from tracecat.logger import logger + +__all__ = ( + "BASE_PYTHONPATH_DIR_NAME", + "CACHE_ENTRIES_DIR_NAME", + "CACHE_STAGING_DIR_NAME", + "CACHE_TRASH_DIR_NAME", + "RegistryArtifactAdmission", + "RegistryArtifactCacheCapacityError", + "RegistryArtifactCacheLoopError", + "RegistryArtifactCacheStorage", + "RegistryArtifactEviction", + "RegistryArtifactEvictionPass", + "RegistryArtifactMaterializationContext", + "RegistryArtifactPaths", + "RegistryArtifactRuntimeState", + "allocated_size_bound", + "communicate_rejoin_on_cancel", + "ensure_cache_entry_directory", + "ensure_real_directory", + "is_reusable_cache_directory", + "is_reusable_cache_file", + "remove_file_or_defer", + "remove_tree_rejoin_on_cancel", + "unique_work_path", + "validate_cache_entry_path", +) + +BASE_PYTHONPATH_DIR_NAME = "base" +"""Cache subdirectory used when no registry artifact is requested.""" + +CACHE_ENTRIES_DIR_NAME = "entries" +"""Directory containing one atomic subdirectory per cache key.""" + +CACHE_STAGING_DIR_NAME = "staging" +"""Directory containing in-progress materialization scratch.""" + +CACHE_TRASH_DIR_NAME = "trash" +"""Directory containing retired entries pending physical deletion.""" + +_LEGACY_CACHE_PATH_PATTERN = re.compile( + r"(?:squashfs|unsquashfs|tarball)-[0-9a-f]{16}(?:\.squashfs)?" + r"|[0-9a-f]{16}\.\d+\.\d+\.(?:squashfs|unsquashfs|tar\.gz|tmp)" +) +"""Exact pre-entries-layout cache names that are safe to reclaim.""" + + +class RegistryArtifactCacheLoopError(RuntimeError): + """A registry artifact cache was used outside its owning event loop.""" + + +class RegistryArtifactCacheCapacityError(RuntimeError): + """A cold artifact cannot fit within the configured cache byte budget.""" + + def __init__( + self, + *, + current_bytes: int, + additional_bytes: int, + max_bytes: int, + ) -> None: + super().__init__( + "Registry artifact admission exceeds the cache byte budget: " + f"current_bytes={current_bytes}, additional_bytes={additional_bytes}, " + f"max_bytes={max_bytes}" + ) + self.current_bytes = current_bytes + self.additional_bytes = additional_bytes + self.max_bytes = max_bytes + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactPaths: + """Executor-local cache paths for one registry artifact key.""" + + entry_dir: Path + squashfs_image_path: Path + squashfs_mount_dir: Path + squashfs_extract_dir: Path + tarball_target_dir: Path + + +@dataclass(slots=True) +class RegistryArtifactRuntimeState: + """Process-local synchronization and lease state for one cache key.""" + + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + refcount: int = 0 + last_used: float = 0.0 + lock_users: int = 0 + retire_when_idle: bool = False + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactEviction: + """Outcome of atomically retiring and physically deleting one entry.""" + + retired: bool + reclaimed: bool + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactEvictionPass: + """Result of applying one budget policy to evictable entries.""" + + total_bytes: int + fits: bool + exhausted_candidates: bool + + +@dataclass(frozen=True, slots=True) +class RegistryArtifactAdmission: + """Byte-bound admission hook shared by one cold materialization.""" + + max_bytes: int + allocation_unit: int + ensure_capacity: Callable[[int], Awaitable[None]] + + +@dataclass(slots=True) +class RegistryArtifactMaterializationContext: + """Shared runtime state for artifact materialization.""" + + cache_key: str + staging_dir: Path + paths: RegistryArtifactPaths + defer_cleanup: Callable[[Path], None] + admission: RegistryArtifactAdmission | None = None + + def can_mount_squashfs(self) -> bool: + return config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED and ( + shutil.which("mount") is not None + ) + + +async def _kill_and_reap_subprocess(process: asyncio.subprocess.Process) -> None: + """Kill a subprocess and wait until its child state is reaped.""" + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + + +async def communicate_rejoin_on_cancel( + process: asyncio.subprocess.Process, +) -> tuple[bytes, bytes]: + """Communicate without allowing cancellation to abandon child cleanup.""" + try: + stdout, stderr = await process.communicate() + except asyncio.CancelledError: + reaper = asyncio.ensure_future(_kill_and_reap_subprocess(process)) + await drain_future_through_cancellation(reaper) + raise + + if stdout is None or stderr is None: + raise RuntimeError("Captured subprocess output is required") + return stdout, stderr + + +def is_reusable_cache_directory(path: Path) -> bool: + """Return whether a cache path is a real directory, never a symlink.""" + try: + return stat.S_ISDIR(path.lstat().st_mode) + except FileNotFoundError: + return False + + +def is_reusable_cache_file(path: Path) -> bool: + """Return whether a cache path is a regular file, never a symlink.""" + try: + return stat.S_ISREG(path.lstat().st_mode) + except FileNotFoundError: + return False + + +def ensure_real_directory(path: Path) -> None: + """Create a cache directory without accepting a symlink redirect.""" + if os.path.lexists(path): + if not is_reusable_cache_directory(path): + raise OSError(f"Unsafe registry artifact cache directory: {path}") + return + path.mkdir(parents=True, exist_ok=True) + if not is_reusable_cache_directory(path): + raise OSError(f"Unsafe registry artifact cache directory: {path}") + + +def _validate_cache_root(cache_dir: Path) -> None: + """Reject a configured cache root redirected through its final component.""" + if os.path.lexists(cache_dir) and not is_reusable_cache_directory(cache_dir): + raise OSError(f"Unsafe registry artifact cache directory: {cache_dir}") + + +def _validate_cache_child_directory(path: Path) -> None: + """Reject a fixed cache child redirected outside its real cache root.""" + _validate_cache_root(path.parent) + if os.path.lexists(path) and not is_reusable_cache_directory(path): + raise OSError(f"Unsafe registry artifact cache directory: {path}") + + +def validate_cache_entry_path(paths: RegistryArtifactPaths) -> None: + """Reject an entry root redirected through cache-controlled symlinks.""" + entries_dir = paths.entry_dir.parent + _validate_cache_child_directory(entries_dir) + if os.path.lexists(paths.entry_dir) and not is_reusable_cache_directory( + paths.entry_dir + ): + raise OSError(f"Unsafe registry artifact cache directory: {paths.entry_dir}") + + +def ensure_cache_entry_directory(paths: RegistryArtifactPaths) -> None: + """Create a canonical entry only below real cache and entries roots.""" + entries_dir = paths.entry_dir.parent + cache_dir = entries_dir.parent + validate_cache_entry_path(paths) + ensure_real_directory(cache_dir) + ensure_real_directory(entries_dir) + ensure_real_directory(paths.entry_dir) + validate_cache_entry_path(paths) + + +def allocated_size_bound(size_bytes: int, *, allocation_unit: int) -> int: + """Round one filesystem object up to its minimum allocated footprint.""" + if size_bytes < 0: + raise ValueError("size_bytes must be non-negative") + if allocation_unit <= 0: + raise ValueError("allocation_unit must be positive") + return ( + max(1, (size_bytes + allocation_unit - 1) // allocation_unit) * allocation_unit + ) + + +def _filesystem_allocation_unit(path: Path) -> int: + """Return the allocation unit for a path or nearest existing parent.""" + candidate = path + while True: + try: + filesystem = os.statvfs(candidate) + return filesystem.f_frsize or filesystem.f_bsize or 1 + except FileNotFoundError: + parent = candidate.parent + if parent == candidate: + raise + candidate = parent + + +def _allocated_stat_size( + file_stat: os.stat_result, + *, + allocation_unit: int, +) -> int: + """Return allocated bytes while charging at least one unit per inode.""" + if allocation_unit <= 0: + raise ValueError("allocation_unit must be positive") + return max(allocation_unit, file_stat.st_blocks * 512) + + +def _directory_footprint( + directory: Path, + *, + allocation_unit: int | None = None, + pruned_directories: Iterable[Path] = (), + include_root: bool = True, +) -> int: + """Return allocated bytes for unique inodes without following symlinks.""" + + def raise_walk_error(error: OSError) -> None: + raise error + + if allocation_unit is None: + allocation_unit = _filesystem_allocation_unit(directory) + + total_bytes = 0 + seen_inodes: set[tuple[int, int]] = set() + pruned_paths = frozenset(pruned_directories) + + def allocated_inode_size(file_stat: os.stat_result) -> int: + inode_key = (file_stat.st_dev, file_stat.st_ino) + if inode_key in seen_inodes: + return 0 + seen_inodes.add(inode_key) + return _allocated_stat_size(file_stat, allocation_unit=allocation_unit) + + try: + root_stat = os.lstat(directory) + if not stat.S_ISDIR(root_stat.st_mode): + return allocated_inode_size(root_stat) + + for root, dirs, files in os.walk(directory, onerror=raise_walk_error): + root_path = Path(root) + if include_root or root_path != directory: + try: + total_bytes += allocated_inode_size(os.lstat(root)) + except FileNotFoundError: + continue + for file_name in files: + try: + total_bytes += allocated_inode_size( + os.lstat(os.path.join(root, file_name)) + ) + except FileNotFoundError: + continue + traversed_directories: list[str] = [] + for directory_name in dirs: + child_path = root_path / directory_name + try: + directory_stat = child_path.lstat() + except FileNotFoundError: + continue + if stat.S_ISLNK(directory_stat.st_mode) or child_path in pruned_paths: + total_bytes += allocated_inode_size(directory_stat) + continue + traversed_directories.append(directory_name) + dirs[:] = traversed_directories + except FileNotFoundError: + return 0 + return total_bytes + + +def _delete_cache_path(path: Path) -> bool: + """Best-effort delete one path without following a root symlink.""" + try: + path_stat = path.lstat() + if stat.S_ISDIR(path_stat.st_mode): + shutil.rmtree(path) + else: + path.unlink(missing_ok=True) + except FileNotFoundError: + return True + except OSError as e: + logger.warning( + "Failed to delete registry artifact cache path", + path=str(path), + error=str(e), + ) + return False + return True + + +async def _delete_cache_path_off_loop(path: Path) -> bool: + """Delete one path and rejoin its worker through repeated cancellation.""" + return await run_blocking_rejoin_on_cancel( + functools.partial(_delete_cache_path, path) + ) + + +async def remove_tree_rejoin_on_cancel( + path: Path, + *, + defer_cleanup: Callable[[Path], None], +) -> None: + """Remove a tree off-loop and retain failed paths for a later retry.""" + if not os.path.lexists(path): + return + try: + removed = await _delete_cache_path_off_loop(path) + except asyncio.CancelledError: + if os.path.lexists(path): + defer_cleanup(path) + raise + if not removed: + defer_cleanup(path) + + +def remove_file_or_defer( + path: Path, + *, + defer_cleanup: Callable[[Path], None], +) -> None: + """Remove one staging file without masking materialization outcomes.""" + try: + path.unlink(missing_ok=True) + except OSError as e: + defer_cleanup(path) + logger.warning( + "Deferred failed registry artifact file cleanup", + path=str(path), + error_type=type(e).__name__, + ) + + +def unique_work_path(root: Path, cache_key: str, *, suffix: str = "") -> Path: + """Return a unique path beneath a validated cache work directory.""" + _validate_cache_child_directory(root) + ensure_real_directory(root.parent) + ensure_real_directory(root) + _validate_cache_child_directory(root) + unique_id = time.time_ns() + while True: + path = root / f"{cache_key}.{os.getpid()}.{unique_id}{suffix}" + if not path.exists(): + return path + unique_id += 1 + + +def _move_entry_to_trash(entry_dir: Path, trash_dir: Path, cache_key: str) -> Path: + """Atomically retire one cache entry and return its trash path.""" + trash_path = unique_work_path(trash_dir, cache_key) + entry_dir.rename(trash_path) + return trash_path + + +class RegistryArtifactCacheStorage: + """Own cache state, filesystem lifecycle, and one shared budget policy.""" + + def __init__(self, cache_dir: Path): + self.cache_dir = cache_dir + self.entries_dir = cache_dir / CACHE_ENTRIES_DIR_NAME + self.staging_dir = cache_dir / CACHE_STAGING_DIR_NAME + self.trash_dir = cache_dir / CACHE_TRASH_DIR_NAME + self._runtime: dict[str, RegistryArtifactRuntimeState] = {} + self._owner_binding_lock = threading.Lock() + self._owner_loop: asyncio.AbstractEventLoop | None = None + self._owner_thread_id: int | None = None + self._admission_lock = asyncio.Lock() + self._swept = False + self._sweep_task: asyncio.Task[None] | None = None + self._sweep_lock = asyncio.Lock() + self._failed_startup_cleanup: set[Path] = set() + self._budget_dirty = True + + async def ensure_swept(self) -> None: + """Run the cancellation-safe startup sweep exactly once.""" + self._assert_owner_loop() + if self._swept: + return + async with self._sweep_lock: + if self._swept: + return + sweep_task = self._sweep_task + if sweep_task is None or ( + sweep_task.done() + and (sweep_task.cancelled() or sweep_task.exception() is not None) + ): + sweep_task = asyncio.ensure_future( + asyncio.to_thread(self._sweep_startup_state) + ) + self._sweep_task = sweep_task + try: + await asyncio.shield(sweep_task) + except asyncio.CancelledError: + if sweep_task.cancelled() and self._sweep_task is sweep_task: + self._sweep_task = None + raise + except Exception: + if self._sweep_task is sweep_task: + self._sweep_task = None + raise + self._swept = True + + def _assert_owner_loop(self) -> None: + """Bind to the current loop or reject another loop/thread.""" + current_loop = asyncio.get_running_loop() + current_thread_id = threading.get_ident() + with self._owner_binding_lock: + if self._owner_loop is None: + self._owner_loop = current_loop + self._owner_thread_id = current_thread_id + return + if ( + self._owner_loop is current_loop + and self._owner_thread_id == current_thread_id + ): + return + raise RegistryArtifactCacheLoopError( + "RegistryArtifactCache is bound to one event loop and thread " + f"(owner_loop={id(self._owner_loop)}, " + f"owner_thread={self._owner_thread_id}, " + f"current_loop={id(current_loop)}, " + f"current_thread={current_thread_id})" + ) + + def _runtime_for(self, cache_key: str) -> RegistryArtifactRuntimeState: + """Return stable process-local state for one cache key.""" + if runtime := self._runtime.get(cache_key): + return runtime + runtime = RegistryArtifactRuntimeState() + self._runtime[cache_key] = runtime + return runtime + + @asynccontextmanager + async def _runtime_lock( + self, + cache_key: str, + *, + wait: bool = True, + ) -> AsyncGenerator[RegistryArtifactRuntimeState | None]: + """Track holders and waiters while serializing one cache key. + + Tracking begins before waiting on the lock, so an eviction cannot + retire the state and let a later caller create a second lock while a + waiter still references the first one. + """ + runtime = self._runtime_for(cache_key) + runtime.lock_users += 1 + try: + if not wait and runtime.lock.locked(): + yield None + return + async with runtime.lock: + yield runtime + finally: + runtime.lock_users -= 1 + self._retire_runtime_if_idle(cache_key, runtime) + + def _request_runtime_retirement( + self, + cache_key: str, + runtime: RegistryArtifactRuntimeState | None = None, + ) -> None: + """Retire evicted key state after every holder and waiter is gone.""" + if runtime is None: + runtime = self._runtime.get(cache_key) + if runtime is None: + return + runtime.retire_when_idle = True + self._retire_runtime_if_idle(cache_key, runtime) + + def _retire_runtime_if_idle( + self, + cache_key: str, + runtime: RegistryArtifactRuntimeState, + ) -> None: + """Drop one retired state only when no task can still use its lock.""" + if ( + runtime.retire_when_idle + and runtime.refcount == 0 + and runtime.lock_users == 0 + and self._runtime.get(cache_key) is runtime + ): + del self._runtime[cache_key] + + def _request_runtime_retirement_if_entry_missing(self, cache_key: str) -> None: + """Retire failed-admission state once its empty entry shell is gone.""" + if not os.path.lexists(self._paths_for(cache_key).entry_dir): + self._request_runtime_retirement(cache_key) + + def _context_for( + self, + cache_key: str, + *, + admission: RegistryArtifactAdmission | None = None, + ) -> RegistryArtifactMaterializationContext: + """Return materialization state for one cache key.""" + return RegistryArtifactMaterializationContext( + cache_key=cache_key, + staging_dir=self.staging_dir, + paths=self._paths_for(cache_key), + defer_cleanup=self._failed_startup_cleanup.add, + admission=admission, + ) + + def _admission_for(self, cache_key: str) -> RegistryArtifactAdmission | None: + """Return byte-bound admission controls for one cold cache key.""" + max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES + if max_bytes <= 0: + return None + allocation_unit = _filesystem_allocation_unit(self.cache_dir) + + async def ensure_capacity(additional_bytes: int) -> None: + await self._ensure_cache_capacity( + additional_bytes=allocated_size_bound( + additional_bytes, + allocation_unit=allocation_unit, + ), + protected_key=cache_key, + max_bytes=max_bytes, + ) + + return RegistryArtifactAdmission( + max_bytes=max_bytes, + allocation_unit=allocation_unit, + ensure_capacity=ensure_capacity, + ) + + def _base_pythonpath_dir(self) -> Path: + """Return the base PYTHONPATH directory for an artifact-free action.""" + base_dir = self.cache_dir / BASE_PYTHONPATH_DIR_NAME + _validate_cache_root(self.cache_dir) + ensure_real_directory(self.cache_dir) + _validate_cache_child_directory(base_dir) + ensure_real_directory(base_dir) + return base_dir + + def _acquire_lease(self, cache_key: str) -> None: + """Pin a cache entry against eviction and mark it recently used.""" + runtime = self._runtime_for(cache_key) + runtime.refcount += 1 + runtime.last_used = time.time() + self._touch_entry(cache_key) + + def _release_lease(self, cache_key: str) -> bool: + """Release one pin and return whether the entry became idle.""" + runtime = self._runtime.get(cache_key) + if runtime is None or runtime.refcount == 0: + return False + runtime.refcount -= 1 + runtime.last_used = time.time() + return runtime.refcount == 0 + + async def _unmount_idle_entry(self, cache_key: str) -> None: + """Best-effort unmount an idle entry while retaining its image.""" + async with self._runtime_lock(cache_key, wait=False) as runtime: + if runtime is None: + return + if self._refcount(cache_key) > 0: + return + mount_dir = self._paths_for(cache_key).squashfs_mount_dir + try: + mounted = mount_dir.is_mount() + except OSError as e: + logger.warning( + "Failed to inspect registry artifact mount state", + cache_key=cache_key, + mount_dir=str(mount_dir), + error=str(e), + ) + return + if not mounted: + return + if not await self._unmount(mount_dir): + logger.warning( + "Failed to unmount registry artifact for loop-device reclamation", + cache_key=cache_key, + mount_dir=str(mount_dir), + ) + return + logger.info( + "Unmounted idle registry artifact", + cache_key=cache_key, + mount_dir=str(mount_dir), + ) + + def _refcount(self, cache_key: str) -> int: + """Return the number of live leases on an entry.""" + runtime = self._runtime.get(cache_key) + return 0 if runtime is None else runtime.refcount + + def _touch_entry(self, cache_key: str) -> None: + """Best-effort refresh of the entry-root mtime for restart-safe LRU.""" + try: + os.utime(self._paths_for(cache_key).entry_dir) + except OSError: + logger.debug( + "Could not refresh registry artifact entry mtime", + cache_key=cache_key, + ) + + def _paths_for(self, cache_key: str) -> RegistryArtifactPaths: + """Return local paths for a registry artifact key.""" + entry_dir = self.entries_dir / cache_key + return RegistryArtifactPaths( + entry_dir=entry_dir, + squashfs_image_path=entry_dir / "image.squashfs", + squashfs_mount_dir=entry_dir / "mount", + squashfs_extract_dir=entry_dir / "extracted", + tarball_target_dir=entry_dir / "tarball", + ) + + async def _converge_cache_budget(self) -> None: + """Bring an idle cache back under budget after a lease release.""" + while self._budget_dirty: + self._budget_dirty = False + try: + within_budget = await self._enforce_cache_budget() + except OSError as e: + logger.warning( + "Failed to converge registry artifact cache to budget", + cache_dir=str(self.cache_dir), + error=str(e), + ) + self._budget_dirty = True + break + except BaseException: + self._budget_dirty = True + raise + if not within_budget: + self._budget_dirty = True + break + + async def _enforce_cache_budget( + self, + *, + protected_key: str | None = None, + ) -> bool: + """Evict idle LRU entries until the measured cache fits.""" + async with self._admission_lock: + return await self._enforce_cache_budget_locked(protected_key=protected_key) + + async def _reclaim_pending_work(self) -> tuple[bool, bool]: + """Retry pending trash and deferred cleanup off the event loop.""" + return await rejoin_future_on_cancel( + asyncio.gather( + asyncio.to_thread(self._clear_work_dir, self.trash_dir), + asyncio.to_thread(self._retry_failed_startup_cleanup), + ) + ) + + async def _enforce_cache_budget_locked( + self, + *, + protected_key: str | None, + ) -> bool: + """Enforce entry and byte limits while admission is serialized.""" + trash_clean, startup_clean = await self._reclaim_pending_work() + if not (trash_clean and startup_clean): + return False + + budget = RegistryArtifactCacheBudget( + max_entries=config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES, + max_bytes=config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES, + ) + if budget.max_entries <= 0 and budget.max_bytes <= 0: + return True + + snapshot = await asyncio.to_thread(self._scan_cache_snapshot) + eviction_pass = await self._evict_until_fits( + snapshot.entries, + total_bytes=snapshot.total_bytes, + excluded=set() if protected_key is None else {protected_key}, + budget=budget, + ) + if not eviction_pass.fits and eviction_pass.exhausted_candidates: + logger.warning( + "Registry artifact cache is over budget but every entry is in use", + cache_dir=str(self.cache_dir), + entries=len(snapshot.entries), + max_entries=budget.max_entries, + total_bytes=eviction_pass.total_bytes, + max_bytes=budget.max_bytes, + ) + return eviction_pass.fits + + async def _ensure_cache_capacity( + self, + *, + additional_bytes: int, + protected_key: str, + max_bytes: int, + ) -> None: + """Reserve peak bytes for one serialized cold writer.""" + if additional_bytes < 0: + raise ValueError("additional_bytes must be non-negative") + + def capacity_error(current_bytes: int) -> RegistryArtifactCacheCapacityError: + return RegistryArtifactCacheCapacityError( + current_bytes=current_bytes, + additional_bytes=additional_bytes, + max_bytes=max_bytes, + ) + + trash_clean, startup_clean = await self._reclaim_pending_work() + snapshot = await asyncio.to_thread(self._scan_cache_snapshot) + entries = snapshot.entries + total_bytes = snapshot.total_bytes + non_evictable_bytes = ( + snapshot.structural_bytes + + snapshot.staging_bytes + + snapshot.trash_bytes + + sum( + entry.size_bytes + for entry in entries.values() + if entry.cache_key == protected_key + or self._refcount(entry.cache_key) > 0 + or ( + (runtime := self._runtime.get(entry.cache_key)) is not None + and runtime.lock.locked() + ) + ) + ) + if non_evictable_bytes + additional_bytes > max_bytes: + raise capacity_error(non_evictable_bytes) + if ( + not (trash_clean and startup_clean) + and total_bytes + additional_bytes > max_bytes + ): + raise capacity_error(total_bytes) + eviction_pass = await self._evict_until_fits( + entries, + total_bytes=total_bytes, + excluded={protected_key}, + budget=RegistryArtifactCacheBudget( + max_entries=0, + max_bytes=max_bytes, + additional_bytes=additional_bytes, + ), + ) + if not eviction_pass.fits: + raise capacity_error(eviction_pass.total_bytes) + + async def _evict_until_fits( + self, + entries: dict[str, RegistryArtifactCacheEntry], + *, + total_bytes: int, + excluded: set[str], + budget: RegistryArtifactCacheBudget, + ) -> RegistryArtifactEvictionPass: + """Apply the shared LRU policy until a measured cache fits.""" + skipped = set(excluded) + while not budget.fits(entry_count=len(entries), total_bytes=total_bytes): + plan = plan_registry_artifact_evictions( + entries, + total_bytes=total_bytes, + budget=budget, + excluded=skipped + | { + entry.cache_key + for entry in entries.values() + if self._refcount(entry.cache_key) > 0 + }, + effective_last_used={ + entry.cache_key: max( + entry.last_used, + runtime.last_used, + ) + for entry in entries.values() + if (runtime := self._runtime.get(entry.cache_key)) is not None + }, + ) + if not plan.candidates: + return RegistryArtifactEvictionPass( + total_bytes=total_bytes, + fits=False, + exhausted_candidates=True, + ) + candidate = plan.candidates[0] + eviction = await self._evict_entry(candidate.cache_key) + if not eviction.retired: + skipped.add(candidate.cache_key) + continue + del entries[candidate.cache_key] + if not eviction.reclaimed: + return RegistryArtifactEvictionPass( + total_bytes=total_bytes, + fits=False, + exhausted_candidates=False, + ) + total_bytes -= candidate.size_bytes + + return RegistryArtifactEvictionPass( + total_bytes=total_bytes, + fits=True, + exhausted_candidates=False, + ) + + async def _evict_entry(self, cache_key: str) -> RegistryArtifactEviction: + """Atomically retire and physically delete one idle cache entry.""" + async with self._runtime_lock(cache_key, wait=False) as runtime: + if runtime is None: + logger.debug( + "Skipping eviction of busy registry artifact", + cache_key=cache_key, + ) + return RegistryArtifactEviction(retired=False, reclaimed=False) + + if self._refcount(cache_key) > 0: + return RegistryArtifactEviction(retired=False, reclaimed=False) + paths = self._paths_for(cache_key) + validate_cache_entry_path(paths) + if not paths.entry_dir.exists(): + self._request_runtime_retirement(cache_key, runtime) + return RegistryArtifactEviction(retired=True, reclaimed=True) + if paths.squashfs_mount_dir.is_mount() and not await self._unmount( + paths.squashfs_mount_dir + ): + logger.warning( + "Failed to unmount registry artifact, skipping eviction", + cache_key=cache_key, + mount_dir=str(paths.squashfs_mount_dir), + ) + return RegistryArtifactEviction(retired=False, reclaimed=False) + try: + trash_path = _move_entry_to_trash( + paths.entry_dir, + self.trash_dir, + cache_key, + ) + except OSError as e: + self._budget_dirty = True + logger.warning( + "Failed to retire registry artifact cache entry", + cache_key=cache_key, + entry_dir=str(paths.entry_dir), + error=str(e), + ) + return RegistryArtifactEviction(retired=False, reclaimed=False) + self._request_runtime_retirement(cache_key, runtime) + + try: + deleted = await _delete_cache_path_off_loop(trash_path) + except BaseException: + self._budget_dirty = True + raise + if not deleted: + self._budget_dirty = True + logger.warning( + "Registry artifact eviction remains pending physical deletion", + cache_key=cache_key, + trash_path=str(trash_path), + ) + return RegistryArtifactEviction(retired=True, reclaimed=False) + logger.info("Evicted registry artifact from cache", cache_key=cache_key) + return RegistryArtifactEviction(retired=True, reclaimed=True) + + async def _unmount(self, mount_dir: Path) -> bool: + """Unmount a SquashFS artifact directory, releasing its loop device.""" + umount = shutil.which("umount") + if umount is None: + return False + proc = await asyncio.create_subprocess_exec( + umount, + str(mount_dir), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await communicate_rejoin_on_cancel(proc) + if proc.returncode == 0 or not mount_dir.is_mount(): + return True + logger.warning( + "umount command failed", + mount_dir=str(mount_dir), + output=(stderr or stdout).decode(errors="replace").strip(), + ) + return False + + def _cache_structural_footprint(self, *, allocation_unit: int) -> int: + """Measure cache roots and non-entry data exactly once.""" + return _directory_footprint( + self.cache_dir, + allocation_unit=allocation_unit, + pruned_directories=( + self.entries_dir, + self.staging_dir, + self.trash_dir, + ), + ) + + def _prepare_budget_roots(self) -> None: + """Create fixed roots before measuring or retiring an entry.""" + ensure_real_directory(self.cache_dir) + for root in (self.entries_dir, self.staging_dir, self.trash_dir): + _validate_cache_child_directory(root) + ensure_real_directory(root) + _validate_cache_child_directory(root) + + def _scan_cache_snapshot(self) -> RegistryArtifactCacheSnapshot: + """Measure entries, work directories, and structure together.""" + self._prepare_budget_roots() + allocation_unit = _filesystem_allocation_unit(self.cache_dir) + return RegistryArtifactCacheSnapshot( + entries=self._scan_cache_entries(allocation_unit=allocation_unit), + structural_bytes=self._cache_structural_footprint( + allocation_unit=allocation_unit + ), + staging_bytes=_directory_footprint( + self.staging_dir, + allocation_unit=allocation_unit, + include_root=False, + ), + trash_bytes=_directory_footprint( + self.trash_dir, + allocation_unit=allocation_unit, + include_root=False, + ), + ) + + def _scan_cache_entries( + self, + *, + allocation_unit: int | None = None, + ) -> dict[str, RegistryArtifactCacheEntry]: + """Measure every registry artifact entry currently on disk.""" + if allocation_unit is None: + allocation_unit = _filesystem_allocation_unit(self.cache_dir) + return { + cache_key: self._measure_entry( + cache_key, + allocation_unit=allocation_unit, + ) + for cache_key in self._discover_cache_keys() + } + + def _discover_cache_keys(self) -> set[str]: + """Return cache keys represented by real atomic entry directories.""" + _validate_cache_child_directory(self.entries_dir) + try: + entries = list(os.scandir(self.entries_dir)) + except FileNotFoundError: + return set() + return { + entry.name + for entry in entries + if entry.name and entry.is_dir(follow_symlinks=False) + } + + def _measure_entry( + self, + cache_key: str, + *, + allocation_unit: int | None = None, + ) -> RegistryArtifactCacheEntry: + """Measure all entry-owned inodes, pruning active mount contents.""" + paths = self._paths_for(cache_key) + validate_cache_entry_path(paths) + if allocation_unit is None: + allocation_unit = _filesystem_allocation_unit(self.cache_dir) + try: + mount_is_active = paths.squashfs_mount_dir.is_mount() + except FileNotFoundError: + mount_is_active = False + pruned_directories = (paths.squashfs_mount_dir,) if mount_is_active else () + size_bytes = _directory_footprint( + paths.entry_dir, + allocation_unit=allocation_unit, + pruned_directories=pruned_directories, + ) + try: + last_used = paths.entry_dir.stat().st_mtime + except FileNotFoundError: + last_used = 0.0 + return RegistryArtifactCacheEntry( + cache_key=cache_key, + size_bytes=size_bytes, + last_used=last_used, + ) + + def _sweep_startup_state(self) -> None: + """Reclaim orphaned work, legacy layout, and over-budget entries.""" + _validate_cache_root(self.cache_dir) + if not self.cache_dir.is_dir(): + self._budget_dirty = False + return + try: + staging_clean = self._clear_work_dir( + self.staging_dir, + remember_failures=True, + ) + trash_clean = self._clear_work_dir(self.trash_dir) + legacy_clean = self._clear_legacy_cache_layout() + cleanup_complete = staging_clean and trash_clean and legacy_clean + within_budget = cleanup_complete and self._trim_startup_cache() + self._budget_dirty = not (cleanup_complete and within_budget) + except OSError as e: + logger.warning( + "Failed to sweep registry artifact cache", + cache_dir=str(self.cache_dir), + error=str(e), + ) + raise + + def _clear_legacy_cache_layout(self) -> bool: + """Reclaim exact top-level artifacts from the pre-entries layout.""" + _validate_cache_root(self.cache_dir) + try: + paths = list(self.cache_dir.iterdir()) + except FileNotFoundError: + return True + deleted = True + for path in paths: + if _LEGACY_CACHE_PATH_PATTERN.fullmatch(path.name) is None: + continue + try: + mounted = is_reusable_cache_directory(path) and path.is_mount() + except OSError: + mounted = True + if mounted or not _delete_cache_path(path): + deleted = False + self._failed_startup_cleanup.add(path) + continue + self._failed_startup_cleanup.discard(path) + logger.info("Removed legacy registry artifact cache path", path=str(path)) + return deleted + + def _clear_work_dir( + self, + work_dir: Path, + *, + remember_failures: bool = False, + ) -> bool: + """Best-effort remove every child of a staging or trash directory.""" + _validate_cache_child_directory(work_dir) + try: + paths = list(work_dir.iterdir()) + except FileNotFoundError: + return True + except OSError as e: + logger.warning( + "Failed to inspect registry artifact work directory", + path=str(work_dir), + error=str(e), + ) + raise + deleted = True + for path in paths: + if _delete_cache_path(path): + if remember_failures: + self._failed_startup_cleanup.discard(path) + logger.info("Removed registry artifact work path", path=str(path)) + else: + deleted = False + if remember_failures: + self._failed_startup_cleanup.add(path) + return deleted + + def _retry_failed_startup_cleanup(self) -> bool: + """Retry exact deferred paths without sweeping live staging work.""" + for path in tuple(self._failed_startup_cleanup): + if _delete_cache_path(path): + self._failed_startup_cleanup.discard(path) + return not self._failed_startup_cleanup + + def _trim_startup_cache(self) -> bool: + """Apply the shared measured budget during startup.""" + budget = RegistryArtifactCacheBudget( + max_entries=config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES, + max_bytes=config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES, + ) + if budget.max_entries <= 0 and budget.max_bytes <= 0: + return True + + snapshot = self._scan_cache_snapshot() + entries = snapshot.entries + total_bytes = snapshot.total_bytes + mounted_keys = { + entry.cache_key + for entry in entries.values() + if self._paths_for(entry.cache_key).squashfs_mount_dir.is_mount() + } + plan = plan_registry_artifact_evictions( + entries, + total_bytes=total_bytes, + budget=budget, + excluded=mounted_keys, + ) + for entry in plan.candidates: + if budget.fits(entry_count=len(entries), total_bytes=total_bytes): + break + paths = self._paths_for(entry.cache_key) + try: + trash_path = _move_entry_to_trash( + paths.entry_dir, + self.trash_dir, + entry.cache_key, + ) + except OSError as e: + logger.warning( + "Failed to retire stale registry artifact during startup sweep", + cache_key=entry.cache_key, + entry_dir=str(paths.entry_dir), + error=str(e), + ) + return False + del entries[entry.cache_key] + if _delete_cache_path(trash_path): + total_bytes -= entry.size_bytes + else: + return False + logger.info( + "Evicted stale registry artifact during startup sweep", + cache_key=entry.cache_key, + size_bytes=entry.size_bytes, + ) + return budget.fits(entry_count=len(entries), total_bytes=total_bytes) diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index c43eeb2013..f060f1507e 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -4,35 +4,72 @@ import asyncio import contextlib -import functools import hashlib import os import shutil import sysconfig import tarfile -import threading import time from abc import ABC, abstractmethod -from collections.abc import AsyncIterator, Awaitable, Callable, Iterable +from collections.abc import AsyncGenerator, Awaitable, Callable from contextlib import asynccontextmanager -from dataclasses import dataclass, field +from dataclasses import dataclass from enum import StrEnum -from pathlib import Path +from pathlib import Path, PurePosixPath +from types import TracebackType +from urllib.parse import urlsplit, urlunsplit import httpx import tracecat_registry from tracecat import config from tracecat.concurrency import ( - drain_future_through_cancellation, - rejoin_future_on_cancel, + rejoin_future_through_cancellation, run_blocking_rejoin_on_cancel, ) +from tracecat.executor.registry_artifact_storage import ( + RegistryArtifactAdmission, + RegistryArtifactCacheCapacityError, + RegistryArtifactCacheLoopError, + RegistryArtifactCacheStorage, + RegistryArtifactEviction, + RegistryArtifactMaterializationContext, + allocated_size_bound, + communicate_rejoin_on_cancel, + ensure_cache_entry_directory, + ensure_real_directory, + is_reusable_cache_directory, + is_reusable_cache_file, + remove_file_or_defer, + remove_tree_rejoin_on_cancel, + unique_work_path, + validate_cache_entry_path, +) from tracecat.logger import logger from tracecat.registry.artifact_keys import parse_s3_uri from tracecat.registry.constants import DEFAULT_REGISTRY_ORIGIN from tracecat.storage import blob +__all__ = ( + "BUNDLED_BUILTIN_REGISTRY_URI_PREFIX", + "SQUASHFS_MOUNT_OPTIONS", + "BuiltinArtifact", + "RegistryArtifact", + "RegistryArtifactAdmission", + "RegistryArtifactCache", + "RegistryArtifactCacheCapacityError", + "RegistryArtifactCacheLoopError", + "RegistryArtifactEviction", + "RegistryArtifactFormat", + "RegistryArtifactMaterializationContext", + "RegistryArtifactUriError", + "SquashfsArtifact", + "SquashfsMountCommandError", + "TarballArtifact", + "bundled_builtin_registry_uri", + "compute_registry_artifact_cache_key", +) + class RegistryArtifactFormat(StrEnum): """Executor-supported registry artifact encodings.""" @@ -53,18 +90,6 @@ class RegistryArtifactFormat(StrEnum): BUNDLED_BUILTIN_REGISTRY_URI_PREFIX = f"tracecat-builtin://{DEFAULT_REGISTRY_ORIGIN}/" """Pseudo-URI for the builtin registry already installed in the executor image.""" -BASE_PYTHONPATH_DIR_NAME = "base" -"""Cache subdirectory used as the PYTHONPATH entry when no artifact is requested.""" - -CACHE_ENTRIES_DIR_NAME = "entries" -"""Directory containing one atomic subdirectory per cache key.""" - -CACHE_STAGING_DIR_NAME = "staging" -"""Directory containing in-progress materialization scratch.""" - -CACHE_TRASH_DIR_NAME = "trash" -"""Directory containing atomically retired entries pending physical deletion.""" - class SquashfsMountCommandError(RuntimeError): """The ``mount`` command itself failed for a SquashFS registry artifact. @@ -75,88 +100,8 @@ class SquashfsMountCommandError(RuntimeError): """ -class RegistryArtifactCacheLoopError(RuntimeError): - """A registry artifact cache was used outside its owning event loop.""" - - -class RegistryArtifactCacheCapacityError(RuntimeError): - """A cold artifact cannot fit within the configured cache byte budget.""" - - def __init__( - self, - *, - current_bytes: int, - additional_bytes: int, - max_bytes: int, - ) -> None: - super().__init__( - "Registry artifact admission exceeds the cache byte budget: " - f"current_bytes={current_bytes}, additional_bytes={additional_bytes}, " - f"max_bytes={max_bytes}" - ) - self.current_bytes = current_bytes - self.additional_bytes = additional_bytes - self.max_bytes = max_bytes - - -@dataclass(frozen=True, slots=True) -class RegistryArtifactPaths: - """Executor-local cache paths for one registry artifact key.""" - - entry_dir: Path - squashfs_image_path: Path - squashfs_mount_dir: Path - squashfs_extract_dir: Path - tarball_target_dir: Path - - -@dataclass(slots=True) -class RegistryArtifactRuntimeState: - """Process-local synchronization and lease state for one cache key.""" - - lock: asyncio.Lock = field(default_factory=asyncio.Lock) - refcount: int = 0 - last_used: float = 0.0 - - -@dataclass(frozen=True, slots=True) -class RegistryArtifactEviction: - """Outcome of atomically retiring and physically deleting one entry.""" - - retired: bool - reclaimed: bool - - -@dataclass(frozen=True, slots=True) -class RegistryArtifactCacheEntry: - """Measured on-disk footprint and recency for one registry artifact key.""" - - cache_key: str - size_bytes: int - last_used: float - - -@dataclass(frozen=True, slots=True) -class RegistryArtifactAdmission: - """Byte-bound admission hook shared by one cold materialization.""" - - max_bytes: int - ensure_capacity: Callable[[int], Awaitable[None]] - - -@dataclass(slots=True) -class RegistryArtifactMaterializationContext: - """Shared runtime state for artifact materialization.""" - - cache_key: str - staging_dir: Path - paths: RegistryArtifactPaths - admission: RegistryArtifactAdmission | None = None - - def can_mount_squashfs(self) -> bool: - return config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED and ( - shutil.which("mount") is not None - ) +class RegistryArtifactUriError(ValueError): + """A registry artifact URI is malformed, with identifiers suppressed.""" @dataclass(frozen=True, slots=True) @@ -188,30 +133,7 @@ def _temp_path( ctx: RegistryArtifactMaterializationContext, suffix: str, ) -> Path: - return _unique_work_path(ctx.staging_dir, self.cache_key, suffix=suffix) - - -async def _kill_and_reap_subprocess(process: asyncio.subprocess.Process) -> None: - """Kill a subprocess and wait until its child state is reaped.""" - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - - -async def _communicate_rejoin_on_cancel( - process: asyncio.subprocess.Process, -) -> tuple[bytes, bytes]: - """Communicate without allowing cancellation to abandon child cleanup.""" - try: - stdout, stderr = await process.communicate() - except asyncio.CancelledError: - reaper = asyncio.ensure_future(_kill_and_reap_subprocess(process)) - await drain_future_through_cancellation(reaper) - raise - - if stdout is None or stderr is None: - raise RuntimeError("Captured subprocess output is required") - return stdout, stderr + return unique_work_path(ctx.staging_dir, self.cache_key, suffix=suffix) @dataclass(frozen=True, slots=True) @@ -253,13 +175,23 @@ def format(self) -> RegistryArtifactFormat: def cached_path( self, ctx: RegistryArtifactMaterializationContext ) -> list[Path] | None: - if ctx.paths.squashfs_mount_dir.is_mount(): + validate_cache_entry_path(ctx.paths) + mount_dir = ctx.paths.squashfs_mount_dir + if os.path.lexists(mount_dir) and not is_reusable_cache_directory(mount_dir): + remove_file_or_defer( + mount_dir, + defer_cleanup=ctx.defer_cleanup, + ) + if is_reusable_cache_directory(mount_dir) and mount_dir.is_mount(): logger.debug( "Using cached SquashFS registry mount", cache_key=ctx.cache_key, ) - return [ctx.paths.squashfs_mount_dir] - if ctx.paths.squashfs_extract_dir.exists(): + return [mount_dir] + if _is_reusable_extraction_dir( + ctx.paths.squashfs_extract_dir, + defer_cleanup=ctx.defer_cleanup, + ): logger.debug( "Using cached SquashFS registry extraction", cache_key=ctx.cache_key, @@ -278,7 +210,7 @@ async def materialize( logger.warning( "Failed to mount SquashFS registry artifact, trying extraction", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, error=str(e), ) @@ -291,10 +223,14 @@ async def download( image_path: Path, ) -> float: """Ensure the SquashFS image exists locally and return download time.""" - if image_path.exists(): + validate_cache_entry_path(ctx.paths) + if await _reuse_or_reclaim_cache_file( + image_path, + defer_cleanup=ctx.defer_cleanup, + ): return 0.0 - image_path.parent.mkdir(parents=True, exist_ok=True) + ensure_cache_entry_directory(ctx.paths) temp_image = self._temp_path(ctx, ".squashfs") try: download_start = time.monotonic() @@ -302,6 +238,8 @@ async def download( self.uri, temp_image, admission=ctx.admission, + defer_cleanup=ctx.defer_cleanup, + published_path=image_path, ) try: temp_image.rename(image_path) @@ -310,7 +248,10 @@ async def download( raise return (time.monotonic() - download_start) * 1000 finally: - temp_image.unlink(missing_ok=True) + remove_file_or_defer( + temp_image, + defer_cleanup=ctx.defer_cleanup, + ) async def mount( self, @@ -330,17 +271,25 @@ async def mount( SquashfsMountCommandError: The ``mount`` command failed. Exception: The image could not be downloaded or prepared. """ + validate_cache_entry_path(ctx.paths) target_dir = ctx.paths.squashfs_mount_dir - if target_dir.is_mount(): + if os.path.lexists(target_dir) and not is_reusable_cache_directory(target_dir): + remove_file_or_defer( + target_dir, + defer_cleanup=ctx.defer_cleanup, + ) + if os.path.lexists(target_dir) and not is_reusable_cache_directory(target_dir): + raise OSError("Failed to reclaim malformed SquashFS mount target") + if is_reusable_cache_directory(target_dir) and target_dir.is_mount(): return target_dir - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) - target_dir.mkdir(parents=True, exist_ok=True) + ensure_cache_entry_directory(ctx.paths) + ensure_real_directory(target_dir) logger.info( "Materializing SquashFS registry artifact", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, ) start_time = time.monotonic() @@ -354,7 +303,7 @@ async def mount( logger.info( "SquashFS registry artifact mounted", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, download_ms=f"{download_elapsed:.1f}", mount_ms=f"{mount_elapsed:.1f}", @@ -367,16 +316,20 @@ async def extract( ctx: RegistryArtifactMaterializationContext, image_path: Path, ) -> Path: + validate_cache_entry_path(ctx.paths) target_dir = ctx.paths.squashfs_extract_dir - if target_dir.exists(): + if _is_reusable_extraction_dir( + target_dir, + defer_cleanup=ctx.defer_cleanup, + ): return target_dir - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + ensure_cache_entry_directory(ctx.paths) logger.info( "Extracting SquashFS registry artifact", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, ) start_time = time.monotonic() @@ -385,7 +338,14 @@ async def extract( temp_dir = self._temp_path(ctx, ".unsquashfs") try: if ctx.admission is not None: - extracted_size = await self._squashfs_extracted_size(image_path) + extracted_size = await self._squashfs_extracted_size( + image_path, + allocation_unit=ctx.admission.allocation_unit, + ) + extracted_size += _directory_records_size_bound( + (temp_dir.name, target_dir.name), + allocation_unit=ctx.admission.allocation_unit, + ) await ctx.admission.ensure_capacity(extracted_size) extract_start = time.monotonic() temp_dir.mkdir(parents=True, exist_ok=True) @@ -398,25 +358,30 @@ async def extract( logger.info( "SquashFS registry artifact extracted", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, download_ms=f"{download_elapsed:.1f}", extract_ms=f"{extract_elapsed:.1f}", total_ms=f"{total_elapsed:.1f}", ) except OSError: - if target_dir.exists(): + if _is_reusable_extraction_dir( + target_dir, + defer_cleanup=ctx.defer_cleanup, + ): logger.debug( "SquashFS already extracted by another process", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, ) else: raise finally: - if temp_dir.exists(): - shutil.rmtree(temp_dir, ignore_errors=True) + await remove_tree_rejoin_on_cancel( + temp_dir, + defer_cleanup=ctx.defer_cleanup, + ) return target_dir @@ -449,7 +414,7 @@ async def _mount_image(self, image_path: Path, target_dir: Path) -> None: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - stdout, stderr = await _communicate_rejoin_on_cancel(proc) + stdout, stderr = await communicate_rejoin_on_cancel(proc) if proc.returncode == 0 or target_dir.is_mount(): return @@ -477,7 +442,7 @@ async def _extract_image(self, image_path: Path, target_dir: Path) -> None: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - stdout, stderr = await _communicate_rejoin_on_cancel(proc) + stdout, stderr = await communicate_rejoin_on_cancel(proc) if proc.returncode == 0: return @@ -485,8 +450,13 @@ async def _extract_image(self, image_path: Path, target_dir: Path) -> None: output = (stderr or stdout).decode(errors="replace").strip() raise RuntimeError(output or "unsquashfs command failed") - async def _squashfs_extracted_size(self, image_path: Path) -> int: - """Return a conservative logical size for an extracted SquashFS image.""" + async def _squashfs_extracted_size( + self, + image_path: Path, + *, + allocation_unit: int, + ) -> int: + """Return an allocated-size bound for an extracted SquashFS image.""" unsquashfs = shutil.which("unsquashfs") if unsquashfs is None: raise RuntimeError("unsquashfs command is not installed") @@ -498,12 +468,12 @@ async def _squashfs_extracted_size(self, image_path: Path) -> int: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - stdout, stderr = await _communicate_rejoin_on_cancel(proc) + stdout, stderr = await communicate_rejoin_on_cancel(proc) if proc.returncode != 0: output = (stderr or stdout).decode(errors="replace").strip() raise RuntimeError(output or "unsquashfs listing failed") - return _squashfs_listing_size(stdout) + return _squashfs_listing_size(stdout, allocation_unit=allocation_unit) @dataclass(frozen=True, slots=True) @@ -517,7 +487,11 @@ def format(self) -> RegistryArtifactFormat: def cached_path( self, ctx: RegistryArtifactMaterializationContext ) -> list[Path] | None: - if ctx.paths.tarball_target_dir.exists(): + validate_cache_entry_path(ctx.paths) + if _is_reusable_extraction_dir( + ctx.paths.tarball_target_dir, + defer_cleanup=ctx.defer_cleanup, + ): logger.debug( "Using cached tarball extraction", cache_key=ctx.cache_key, @@ -528,11 +502,12 @@ def cached_path( async def materialize( self, ctx: RegistryArtifactMaterializationContext ) -> list[Path]: + validate_cache_entry_path(ctx.paths) target_dir = ctx.paths.tarball_target_dir logger.info( "Materializing tarball registry artifact", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, ) start_time = time.monotonic() @@ -541,17 +516,24 @@ async def materialize( temp_dir = self._temp_path(ctx, ".tmp") try: - ctx.paths.entry_dir.mkdir(parents=True, exist_ok=True) + ensure_cache_entry_directory(ctx.paths) download_start = time.monotonic() await self.download(ctx, temp_tarball) download_elapsed = (time.monotonic() - download_start) * 1000 - if ctx.admission is not None: + if (admission := ctx.admission) is not None: extracted_size = await run_blocking_rejoin_on_cancel( - lambda: _tarball_extracted_size(temp_tarball) + lambda: _tarball_extracted_size( + temp_tarball, + allocation_unit=admission.allocation_unit, + ) ) - await ctx.admission.ensure_capacity(extracted_size) + extracted_size += _directory_records_size_bound( + (temp_dir.name, target_dir.name), + allocation_unit=admission.allocation_unit, + ) + await admission.ensure_capacity(extracted_size) extract_start = time.monotonic() temp_dir.mkdir(parents=True, exist_ok=True) @@ -564,27 +546,36 @@ async def materialize( logger.info( "Tarball extracted and cached", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, download_ms=f"{download_elapsed:.1f}", extract_ms=f"{extract_elapsed:.1f}", total_ms=f"{total_elapsed:.1f}", ) except OSError: - if target_dir.exists(): + if _is_reusable_extraction_dir( + target_dir, + defer_cleanup=ctx.defer_cleanup, + ): logger.debug( "Tarball already extracted by another process", cache_key=ctx.cache_key, - artifact_uri=self.uri, + artifact_uri=_artifact_uri_for_logging(self.uri), artifact_format=self.format.value, ) else: raise finally: - if temp_dir.exists(): - shutil.rmtree(temp_dir, ignore_errors=True) - if temp_tarball.exists(): - temp_tarball.unlink(missing_ok=True) + try: + await remove_tree_rejoin_on_cancel( + temp_dir, + defer_cleanup=ctx.defer_cleanup, + ) + finally: + remove_file_or_defer( + temp_tarball, + defer_cleanup=ctx.defer_cleanup, + ) return [target_dir] @@ -597,6 +588,7 @@ async def download( self.uri, output_path, admission=ctx.admission, + defer_cleanup=ctx.defer_cleanup, ) async def extract(self, tarball_path: Path, target_dir: Path) -> None: @@ -629,22 +621,48 @@ async def _download_s3_artifact( output_path: Path, *, admission: RegistryArtifactAdmission | None = None, + defer_cleanup: Callable[[Path], None] | None = None, + published_path: Path | None = None, ) -> None: """Download an S3 registry artifact to a local path.""" - bucket, key = parse_s3_uri(artifact_uri) + try: + bucket, key = parse_s3_uri(artifact_uri) + except ValueError: + raise RegistryArtifactUriError("Invalid registry artifact URI") from None + + ensure_capacity: Callable[[int], Awaitable[None]] | None = None + if admission is not None: + metadata_bytes = _directory_entry_size_bound( + f"{output_path.name}.part", + allocation_unit=admission.allocation_unit, + ) + if published_path is not None: + metadata_bytes += _directory_entry_size_bound( + published_path.name, + allocation_unit=admission.allocation_unit, + ) + + async def ensure_download_capacity(content_bytes: int) -> None: + await admission.ensure_capacity(content_bytes + metadata_bytes) + + ensure_capacity = ensure_download_capacity + try: await blob.download_file_to_path( key=key, bucket=bucket, output_path=output_path, max_bytes=None if admission is None else admission.max_bytes, - ensure_capacity=None if admission is None else admission.ensure_capacity, + ensure_capacity=ensure_capacity, + defer_cleanup=defer_cleanup, + redact_log_identifiers=True, ) except FileNotFoundError as e: - request = httpx.Request("GET", artifact_uri) + safe_uri = _artifact_uri_for_logging(artifact_uri) + request = httpx.Request("GET", safe_uri) response = httpx.Response(status_code=404, request=request) raise httpx.HTTPStatusError( - f"Registry artifact not found: {artifact_uri}", + f"Registry artifact not found: {safe_uri}", request=request, response=response, ) from e @@ -738,204 +756,268 @@ def _is_cache_entry_uri(artifact_uri: str) -> bool: return _bundled_builtin_registry_version(artifact_uri) is None -def _tarball_extracted_size(tarball_path: Path) -> int: - """Return a conservative logical size for all tarball members.""" +def _artifact_uri_for_logging(artifact_uri: str) -> str: + """Retain only a non-sensitive artifact URI scheme for diagnostics.""" + try: + parsed = urlsplit(artifact_uri) + except ValueError: + return "" + if not parsed.scheme: + return "" + return urlunsplit((parsed.scheme, "", "", "", "")) + + +def _is_reusable_extraction_dir( + path: Path, + *, + defer_cleanup: Callable[[Path], None], +) -> bool: + """Accept canonical directories and reclaim file or symlink targets.""" + if is_reusable_cache_directory(path): + return True + if os.path.lexists(path): + remove_file_or_defer(path, defer_cleanup=defer_cleanup) + return False + + +async def _reuse_or_reclaim_cache_file( + path: Path, + *, + defer_cleanup: Callable[[Path], None], +) -> bool: + """Reuse a regular file or reclaim a malformed canonical target.""" + if is_reusable_cache_file(path): + return True + if not os.path.lexists(path): + return False + if is_reusable_cache_directory(path): + await remove_tree_rejoin_on_cancel(path, defer_cleanup=defer_cleanup) + else: + remove_file_or_defer(path, defer_cleanup=defer_cleanup) + if os.path.lexists(path): + raise OSError("Failed to reclaim malformed registry artifact cache target") + return False + + +_DIRECTORY_ENTRY_OVERHEAD_BYTES = 32 +"""Conservative per-child filesystem directory-record overhead.""" + + +def _directory_entry_size_bound( + child_name: str, + *, + allocation_unit: int, +) -> int: + """Reserve directory storage for one unique child name.""" + return allocated_size_bound( + _DIRECTORY_ENTRY_OVERHEAD_BYTES + len(os.fsencode(child_name)), + allocation_unit=allocation_unit, + ) + + +def _directory_records_size_bound( + child_names: tuple[str, ...], + *, + allocation_unit: int, +) -> int: + """Reserve parent-directory storage retained across staged publication.""" + return sum( + _directory_entry_size_bound(name, allocation_unit=allocation_unit) + for name in child_names + ) + + +def _tarball_extracted_size( + tarball_path: Path, + *, + allocation_unit: int = 1, +) -> int: + """Return a conservative allocated-size bound for a tarball extraction.""" total_bytes = 0 + root_path = PurePosixPath(".") + required_parent_dirs: set[PurePosixPath] = set() + explicit_dirs: set[PurePosixPath] = set() + directory_children: dict[PurePosixPath, set[str]] = {} + + def record_directory_child(path: PurePosixPath) -> None: + if path != root_path: + directory_children.setdefault(path.parent, set()).add(path.name) + with tarfile.open(tarball_path, "r:gz") as tar: for member in tar: if member.size < 0: raise ValueError( f"Registry tarball member has a negative size: {member.name}" ) - total_bytes += member.size + total_bytes += allocated_size_bound( + member.size, + allocation_unit=allocation_unit, + ) + member_path = PurePosixPath(member.name) + record_directory_child(member_path) + if member.isdir(): + explicit_dirs.add(member_path) + for parent in member_path.parents: + if parent == root_path: + break + required_parent_dirs.add(parent) + record_directory_child(parent) + + # Extraction creates a target root even when the archive omits it. + total_bytes += allocated_size_bound(0, allocation_unit=allocation_unit) + total_bytes += len(required_parent_dirs - explicit_dirs) * allocation_unit + total_bytes += sum( + _directory_entry_size_bound(child_name, allocation_unit=allocation_unit) + for child_names in directory_children.values() + for child_name in child_names + ) return total_bytes -def _squashfs_listing_size(output: bytes) -> int: - """Sum file sizes from ``unsquashfs -lln`` output, failing closed.""" +def _squashfs_listing_size(output: bytes, *, allocation_unit: int = 1) -> int: + """Bound allocated bytes from ``unsquashfs -lln`` output, failing closed.""" total_bytes = 0 - for raw_line in output.decode(errors="strict").splitlines(): + parsed_entries = 0 + root_path = PurePosixPath(".") + listing_root: str | None = None + required_dirs: set[PurePosixPath] = {root_path} + listed_dirs: set[PurePosixPath] = set() + directory_children: dict[PurePosixPath, set[str]] = {} + + def record_directory_child(path: PurePosixPath) -> None: + if path != root_path: + directory_children.setdefault(path.parent, set()).add(path.name) + + for raw_line in output.decode(errors="replace").splitlines(): line = raw_line.strip() if not line: continue - fields = line.split(maxsplit=4) + fields = line.split(maxsplit=5) mode = fields[0] - if len(mode) != 10 or mode[0] not in "-l": + if len(mode) != 10 or mode[0] not in "bcdlps-": continue - if len(fields) < 5 or "/" not in fields[1] or not fields[2].isdigit(): - raise ValueError(f"Could not parse SquashFS listing line: {line}") - total_bytes += int(fields[2]) - return total_bytes - - -def _directory_footprint(directory: Path) -> int: - """Return the total file size of a cache directory. - - Args: - directory: Cache directory to measure. - - Returns: - Total byte size of contained files, or zero when the directory is - missing. - """ - - def raise_walk_error(error: OSError) -> None: - raise error - - total_bytes = 0 - try: - walker = os.walk(directory, onerror=raise_walk_error) - for root, _dirs, files in walker: - for file_name in files: - try: - total_bytes += os.lstat(os.path.join(root, file_name)).st_size - except FileNotFoundError: - continue - except FileNotFoundError: - return 0 - return total_bytes - - -def _delete_cache_path(path: Path) -> bool: - """Best-effort delete one cache path while reporting filesystem failures.""" - try: - if path.is_dir(): - shutil.rmtree(path) - else: - path.unlink(missing_ok=True) - except OSError as e: - logger.warning( - "Failed to delete registry artifact cache path", - path=str(path), - error=str(e), + if len(fields) < 6 or "/" not in fields[1] or not fields[2].isdigit(): + raise ValueError("Could not parse SquashFS listing line") + + listed_path_text = fields[5] + if mode[0] == "l": + listed_path_text = listed_path_text.split(" -> ", maxsplit=1)[0] + listed_path = PurePosixPath(listed_path_text) + if listed_path.is_absolute() or not listed_path.parts: + raise ValueError("Could not parse SquashFS listing path") + if listing_root is None: + listing_root = listed_path.parts[0] + elif listed_path.parts[0] != listing_root: + raise ValueError("Inconsistent SquashFS listing root") + + relative_parts = listed_path.parts[1:] + entry_path = PurePosixPath(*relative_parts) if relative_parts else root_path + if entry_path == root_path and mode[0] != "d": + raise ValueError("Could not parse SquashFS listing root") + + parsed_entries += 1 + total_bytes += allocated_size_bound( + int(fields[2]), + allocation_unit=allocation_unit, ) - return False - return True - - -async def _delete_cache_path_off_loop(path: Path) -> bool: - """Delete one path without abandoning its worker thread on cancellation.""" - # A worker thread cannot be killed. Rejoin it so no live deletion can race a - # later trash-directory scan, even through repeated caller cancellation. - return await run_blocking_rejoin_on_cancel( - functools.partial(_delete_cache_path, path) + record_directory_child(entry_path) + if mode[0] == "d": + listed_dirs.add(entry_path) + for parent in entry_path.parents: + required_dirs.add(parent) + if parent == root_path: + break + record_directory_child(parent) + + if parsed_entries == 0: + raise ValueError("Could not parse any SquashFS listing entries") + total_bytes += len(required_dirs - listed_dirs) * allocation_unit + total_bytes += sum( + _directory_entry_size_bound(child_name, allocation_unit=allocation_unit) + for child_names in directory_children.values() + for child_name in child_names ) + return total_bytes -def _unique_work_path(root: Path, cache_key: str, *, suffix: str = "") -> Path: - """Return a unique path beneath a cache work directory.""" - root.mkdir(parents=True, exist_ok=True) - unique_id = time.time_ns() - while True: - path = root / f"{cache_key}.{os.getpid()}.{unique_id}{suffix}" - if not path.exists(): - return path - unique_id += 1 - - -def _move_entry_to_trash(entry_dir: Path, trash_dir: Path, cache_key: str) -> Path: - """Atomically retire one cache entry and return its trash path.""" - trash_path = _unique_work_path(trash_dir, cache_key) - entry_dir.rename(trash_path) - return trash_path - +@dataclass(slots=True) +class _RegistryArtifactLease: + """Own exactly one artifact pin and its cancellation-safe release.""" + + cache: RegistryArtifactCache + artifact_uri: str + cache_key: str | None + paths: list[Path] | None = None + paths_may_be_modified: bool = False + _acquired: bool = False + _closed: bool = False + + def mark_acquired(self) -> None: + """Record the point after which this handle must release a pin.""" + if self.cache_key is None or self._acquired: + raise RuntimeError("Invalid registry artifact lease acquisition") + self._acquired = True + + async def __aenter__(self) -> list[Path]: + try: + self.paths = await self.cache._acquire_artifact(self) + except BaseException as operation_error: + try: + await self.aclose() + except BaseException as cleanup_error: + raise operation_error from cleanup_error + raise + return self.paths -class RegistryArtifactCache: - """Materializes registry artifacts into executor-local Python paths.""" + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + del exc_type + try: + await self.aclose() + except BaseException as cleanup_error: + if exc_value is not None: + raise exc_value.with_traceback(traceback) from cleanup_error + raise - def __init__(self, cache_dir: Path): - self.cache_dir = cache_dir - self.entries_dir = cache_dir / CACHE_ENTRIES_DIR_NAME - self.staging_dir = cache_dir / CACHE_STAGING_DIR_NAME - self.trash_dir = cache_dir / CACHE_TRASH_DIR_NAME - # Runtime states live for the process lifetime so every operation for a - # key always serializes on the same lock. - self._runtime: dict[str, RegistryArtifactRuntimeState] = {} - # The cache contains asyncio locks, tasks, and multi-step lease state. - # Bind the public API to one loop/thread so a future synchronous - # Temporal activity fails immediately instead of corrupting that state - # through a thread-local event loop. - self._owner_binding_lock = threading.Lock() - self._owner_loop: asyncio.AbstractEventLoop | None = None - self._owner_thread_id: int | None = None - # Cold materializations and budget passes share this outer lock. It - # keeps byte reservations stable while downloads and extraction write. - self._admission_lock = asyncio.Lock() - # Guard the off-loop startup sweep independently from cache operations. - self._swept: bool = False - self._sweep_task: asyncio.Task[None] | None = None - self._sweep_lock = asyncio.Lock() - # Startup is the only time the whole staging directory is swept. Exact - # paths that could not be removed are safe to retry later. - self._failed_startup_cleanup: set[Path] = set() - # Whether the on-disk cache may exceed its budget. Set when a new entry - # is materialized and cleared once enforcement measures a cache that - # fits, so steady-state cache hits never pay for a disk scan. - self._budget_dirty = True - - async def ensure_swept(self) -> None: - """Run the startup sweep exactly once successfully, off the event loop. - - Idempotent and cancellation-safe under concurrency: every caller joins - one stored sweep task, and cancelling a waiter never abandons or - restarts its live sweep. The lock is deliberately held while awaiting - that shared task so queued callers observe its result before proceeding. - Failures clear the task so the next caller retries. The sweep runs - before the first lease or materialization, so it never observes - in-flight cache entries. - """ - self._assert_owner_loop() - if self._swept: + async def aclose(self) -> None: + """Release the pin exactly once and finish all resulting maintenance.""" + if self._closed: return - async with self._sweep_lock: - if self._swept: - return - sweep_task = self._sweep_task - if sweep_task is None or ( - sweep_task.done() - and (sweep_task.cancelled() or sweep_task.exception() is not None) - ): - sweep_task = asyncio.ensure_future( - asyncio.to_thread(self._sweep_startup_state) + self._closed = True + cache_key = self.cache_key + if cache_key is None or not self._acquired: + return + self._acquired = False + + if self.paths_may_be_modified: + self.cache._budget_dirty = True + idle = self.cache._release_lease(cache_key) + try: + if idle or self.cache._budget_dirty: + cleanup_task = asyncio.ensure_future( + self.cache._finish_lease_cleanup( + [cache_key] if idle else [], + ) ) - self._sweep_task = sweep_task - try: - await asyncio.shield(sweep_task) - except asyncio.CancelledError: - if sweep_task.cancelled() and self._sweep_task is sweep_task: - self._sweep_task = None - raise - except Exception: - if self._sweep_task is sweep_task: - self._sweep_task = None - raise - self._swept = True - - def _assert_owner_loop(self) -> None: - """Bind to the current loop or reject use from another loop/thread.""" - current_loop = asyncio.get_running_loop() - current_thread_id = threading.get_ident() - with self._owner_binding_lock: - if self._owner_loop is None: - self._owner_loop = current_loop - self._owner_thread_id = current_thread_id - return - if ( - self._owner_loop is current_loop - and self._owner_thread_id == current_thread_id - ): - return + await rejoin_future_through_cancellation(cleanup_task) + finally: + self.cache._request_runtime_retirement_if_entry_missing(cache_key) - raise RegistryArtifactCacheLoopError( - "RegistryArtifactCache is bound to one event loop and thread " - f"(owner_loop={id(self._owner_loop)}, " - f"owner_thread={self._owner_thread_id}, " - f"current_loop={id(current_loop)}, " - f"current_thread={current_thread_id})" - ) + +class RegistryArtifactCache(RegistryArtifactCacheStorage): + """Materializes registry artifacts into executor-local Python paths.""" @asynccontextmanager - async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Path]]: + async def lease( + self, + artifact_uris: list[str] | None, + *, + paths_may_be_modified: bool = False, + ) -> AsyncGenerator[list[Path]]: """Materialize registry artifacts and pin them for the life of the context. Leased cache entries are never evicted, so callers may keep importing @@ -944,6 +1026,8 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat Args: artifact_uris: Registry artifact URIs in deterministic PYTHONPATH order, or None to use the base PYTHONPATH directory. + paths_may_be_modified: Whether the consumer can write to returned + paths. Mutable leases re-arm budget convergence after use. Yields: Importable Python paths for the requested artifacts. @@ -955,39 +1039,30 @@ async def lease(self, artifact_uris: list[str] | None) -> AsyncIterator[list[Pat yield [self._base_pythonpath_dir()] return - leased_keys: list[str] = [] - try: + async with contextlib.AsyncExitStack() as leases: + handles: list[_RegistryArtifactLease] = [] registry_paths: list[Path] = [] for artifact_uri in artifact_uris: - cache_key, artifact_paths = await self._lease_artifact(artifact_uri) - if cache_key is not None: - leased_keys.append(cache_key) - registry_paths.extend(artifact_paths) + cache_key = ( + compute_registry_artifact_cache_key(artifact_uri) + if _is_cache_entry_uri(artifact_uri) + else None + ) + handle = _RegistryArtifactLease( + cache=self, + artifact_uri=artifact_uri, + cache_key=cache_key, + ) + registry_paths.extend(await leases.enter_async_context(handle)) + handles.append(handle) logger.info( "Using registry artifact environments", count=len(registry_paths), ) + if paths_may_be_modified: + for handle in handles: + handle.paths_may_be_modified = True yield registry_paths - finally: - idle_keys = [ - cache_key for cache_key in leased_keys if self._release_lease(cache_key) - ] - if idle_keys or self._budget_dirty: - cleanup_task = asyncio.ensure_future( - self._finish_lease_cleanup(idle_keys) - ) - pending_cancellation: asyncio.CancelledError | None = None - while True: - try: - await asyncio.shield(cleanup_task) - break - except asyncio.CancelledError as e: - if cleanup_task.cancelled(): - raise - pending_cancellation = e - - if pending_cancellation is not None: - raise pending_cancellation async def _finish_lease_cleanup(self, idle_keys: list[str]) -> None: """Unmount every newly idle entry and converge the cache budget.""" @@ -995,57 +1070,55 @@ async def _finish_lease_cleanup(self, idle_keys: list[str]) -> None: await self._unmount_idle_entry(cache_key) await self._converge_cache_budget() - async def _lease_artifact(self, artifact_uri: str) -> tuple[str | None, list[Path]]: - """Pin and materialize one artifact, returning its releasable cache key.""" - cache_key = compute_registry_artifact_cache_key(artifact_uri) - ctx = self._context_for(cache_key) - if not _is_cache_entry_uri(artifact_uri): + async def _acquire_artifact( + self, + lease: _RegistryArtifactLease, + ) -> list[Path]: + """Materialize the artifact and transfer its pin to ``lease``.""" + artifact_uri = lease.artifact_uri + cache_key = lease.cache_key + if cache_key is None: + builtin_key = compute_registry_artifact_cache_key(artifact_uri) + ctx = self._context_for(builtin_key) candidates = await self._artifact_candidates(ctx, artifact_uri) - return None, await self._materialize_candidates(ctx, candidates) + return await self._materialize_candidates(ctx, candidates) - lock = self._runtime_for(cache_key).lock - lease_acquired = False - try: - async with lock: - self._acquire_lease(cache_key) - lease_acquired = True + ctx = self._context_for(cache_key) + async with self._runtime_lock(cache_key): + self._acquire_lease(cache_key) + lease.mark_acquired() + if cached_paths := self._locally_cached_path(ctx, artifact_uri): + return cached_paths + + async with self._admission_lock: + async with self._runtime_lock(cache_key): if cached_paths := self._locally_cached_path(ctx, artifact_uri): - return cache_key, cached_paths - - async with self._admission_lock: - async with lock: - if cached_paths := self._locally_cached_path(ctx, artifact_uri): - return cache_key, cached_paths - ctx = self._context_for( - cache_key, - admission=self._admission_for(cache_key), - ) - candidates = await self._artifact_candidates(ctx, artifact_uri) - # Recheck after acquiring the lock. - if cached_paths := self._first_cached_path(candidates, ctx): - return cache_key, cached_paths - paths = await self._materialize_candidates(ctx, candidates) - self._touch_entry(cache_key) - - # Enforce entry count and final measured size after publication, - # outside the per-key lock. Peak-byte reservations may already - # have reclaimed idle LRU entries when that was required to keep - # staging and extraction within the hard byte cap. - try: - await self._enforce_cache_budget(protected_key=cache_key) - except OSError as e: - # Cache maintenance must never block artifact admission. - logger.warning( - "Failed to enforce registry artifact cache budget", - cache_dir=str(self.cache_dir), - error=str(e), + return cached_paths + ctx = self._context_for( + cache_key, + admission=self._admission_for(cache_key), ) - return cache_key, paths - except BaseException: - if lease_acquired and self._release_lease(cache_key): - rollback = asyncio.ensure_future(self._unmount_idle_entry(cache_key)) - await drain_future_through_cancellation(rollback) - raise + candidates = await self._artifact_candidates(ctx, artifact_uri) + # Recheck after acquiring the lock. + if cached_paths := self._first_cached_path(candidates, ctx): + return cached_paths + paths = await self._materialize_candidates(ctx, candidates) + self._touch_entry(cache_key) + + # Enforce entry count and final measured size after publication, + # outside the per-key lock. Peak-byte reservations may already + # have reclaimed idle LRU entries when that was required to keep + # staging and extraction within the hard byte cap. + try: + await self._enforce_cache_budget(protected_key=cache_key) + except OSError as e: + # Cache maintenance must never block artifact admission. + logger.warning( + "Failed to enforce registry artifact cache budget", + cache_dir=str(self.cache_dir), + error=str(e), + ) + return paths async def _materialize_candidates( self, @@ -1062,7 +1135,7 @@ async def _materialize_candidates( logger.info( "Trying registry artifact candidate", cache_key=cache_key, - artifact_uri=artifact.uri, + artifact_uri=_artifact_uri_for_logging(artifact.uri), artifact_format=artifact.format.value, candidate=index + 1, candidates=len(candidates), @@ -1085,154 +1158,13 @@ async def _materialize_candidates( logger.warning( "Failed to materialize registry artifact candidate, trying fallback", cache_key=cache_key, - artifact_uri=artifact.uri, + artifact_uri=_artifact_uri_for_logging(artifact.uri), artifact_format=artifact.format.value, - error=str(e), + error_type=type(e).__name__, ) raise RuntimeError(f"No registry artifact candidates for {ctx.cache_key}") - def _runtime_for(self, cache_key: str) -> RegistryArtifactRuntimeState: - """Return the process-local state for one cache key.""" - if runtime := self._runtime.get(cache_key): - return runtime - runtime = RegistryArtifactRuntimeState() - self._runtime[cache_key] = runtime - return runtime - - def _context_for( - self, - cache_key: str, - *, - admission: RegistryArtifactAdmission | None = None, - ) -> RegistryArtifactMaterializationContext: - """Return a materialization context for a registry artifact key.""" - return RegistryArtifactMaterializationContext( - cache_key=cache_key, - staging_dir=self.staging_dir, - paths=self._paths_for(cache_key), - admission=admission, - ) - - def _admission_for(self, cache_key: str) -> RegistryArtifactAdmission | None: - """Return byte-bound admission controls for one cold cache key.""" - max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES - if max_bytes <= 0: - return None - - async def ensure_capacity(additional_bytes: int) -> None: - await self._ensure_cache_capacity( - additional_bytes=additional_bytes, - protected_key=cache_key, - max_bytes=max_bytes, - ) - - return RegistryArtifactAdmission( - max_bytes=max_bytes, - ensure_capacity=ensure_capacity, - ) - - def _base_pythonpath_dir(self) -> Path: - """Return the base PYTHONPATH directory used when no artifact is requested.""" - base_dir = self.cache_dir / BASE_PYTHONPATH_DIR_NAME - base_dir.mkdir(parents=True, exist_ok=True) - return base_dir - - def _acquire_lease(self, cache_key: str) -> None: - """Pin a cache entry against eviction and mark it as recently used. - - Callers must hold the per-key lock so the increment is ordered against - in-flight eviction of the same key. - """ - runtime = self._runtime_for(cache_key) - runtime.refcount += 1 - runtime.last_used = time.time() - self._touch_entry(cache_key) - - def _release_lease(self, cache_key: str) -> bool: - """Release one pin and return whether the entry became idle.""" - runtime = self._runtime.get(cache_key) - if runtime is None or runtime.refcount == 0: - return False - runtime.refcount -= 1 - runtime.last_used = time.time() - return runtime.refcount == 0 - - async def _unmount_idle_entry(self, cache_key: str) -> None: - """Best-effort unmount one idle entry while retaining its reusable image. - - Loop-device reclamation is independent from disk-budget eviction. The - per-key lock and lease recheck prevent an entry from being unmounted - while an action is importing from it. The image and empty mount - directory remain cached so a later admission can remount without - downloading the artifact again. - - Args: - cache_key: Cache key whose mounted artifact should be released. - """ - try: - lock = self._runtime_for(cache_key).lock - if lock.locked(): - logger.debug( - "Skipping unmount of busy registry artifact", - cache_key=cache_key, - ) - return - - async with lock: - if self._refcount(cache_key) > 0: - return - - mount_dir = self._paths_for(cache_key).squashfs_mount_dir - if not mount_dir.is_mount(): - return - if not await self._unmount(mount_dir): - logger.warning( - "Failed to unmount registry artifact for loop-device reclamation", - cache_key=cache_key, - mount_dir=str(mount_dir), - ) - return - - logger.info( - "Unmounted idle registry artifact", - cache_key=cache_key, - mount_dir=str(mount_dir), - ) - except OSError as e: - logger.warning( - "Failed to release idle registry artifact mount", - cache_key=cache_key, - error=str(e), - ) - - def _refcount(self, cache_key: str) -> int: - """Return the number of live leases on a cache entry.""" - runtime = self._runtime.get(cache_key) - return 0 if runtime is None else runtime.refcount - - def _touch_entry(self, cache_key: str) -> None: - """Best-effort refresh of the entry-root mtime for restart-safe LRU.""" - entry_dir = self._paths_for(cache_key).entry_dir - try: - os.utime(entry_dir) - except OSError: - logger.debug( - "Could not refresh registry artifact entry mtime", - cache_key=cache_key, - ) - - def _paths_for(self, cache_key: str) -> RegistryArtifactPaths: - """Return local cache paths for a registry artifact key.""" - entry_dir = self.entries_dir / cache_key - return RegistryArtifactPaths( - entry_dir=entry_dir, - squashfs_image_path=entry_dir / "image.squashfs", - squashfs_mount_dir=entry_dir / "mount", - squashfs_extract_dir=entry_dir / "extracted", - tarball_target_dir=entry_dir / "tarball", - ) - def _first_cached_path( self, candidates: list[RegistryArtifact], @@ -1324,6 +1256,7 @@ def _remove_unpublished_entry( directories, so canonical artifacts and unknown contents are preserved. """ paths = ctx.paths + validate_cache_entry_path(paths) try: if paths.squashfs_mount_dir.is_mount(): return @@ -1333,6 +1266,7 @@ def _remove_unpublished_entry( for directory in (paths.squashfs_mount_dir, paths.entry_dir): with contextlib.suppress(OSError): directory.rmdir() + self._request_runtime_retirement_if_entry_missing(ctx.cache_key) async def _artifact_candidates( self, @@ -1377,23 +1311,23 @@ async def _sidecar_exists( artifact_format: RegistryArtifactFormat, ) -> bool: """Return whether a registry sidecar exists, logging lookup failures.""" - bucket, key = parse_s3_uri(sidecar_uri) try: + bucket, key = parse_s3_uri(sidecar_uri) if await blob.file_exists(key=key, bucket=bucket): logger.debug( "Using registry artifact sidecar", - artifact_uri=base_uri, - sidecar_uri=sidecar_uri, + artifact_uri=_artifact_uri_for_logging(base_uri), + sidecar_uri=_artifact_uri_for_logging(sidecar_uri), artifact_format=artifact_format.value, ) return True except Exception as e: logger.warning( "Failed to check for registry artifact sidecar, falling back", - artifact_uri=base_uri, - sidecar_uri=sidecar_uri, + artifact_uri=_artifact_uri_for_logging(base_uri), + sidecar_uri=_artifact_uri_for_logging(sidecar_uri), artifact_format=artifact_format.value, - error=str(e), + error_type=type(e).__name__, ) return False @@ -1401,501 +1335,3 @@ async def _sidecar_exists( def _can_try_squashfs(self) -> bool: """Return whether this process should prefer SquashFS artifacts.""" return config.TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED - - async def _converge_cache_budget(self) -> None: - """Bring an idle cache back under budget after a lease is released. - - Successful materialization enforces the budget after publication while - protecting the new entry. The cache can still sit over budget while - entries are leased. This runs on release, when every newly idle entry is - evictable. The scan is skipped entirely unless a materialization attempt - has occurred since the last successful enforcement. - - Each successful pass consumes the dirty signal before its awaited scan. - A follow-up pass therefore occurs only when a concurrent materialization - sets the flag again. Without new materializations the loop terminates, - while an over-budget or failed scan restores the flag and breaks so it - cannot spin while entries remain leased. Cancellation also restores the - consumed flag before propagating. - """ - while self._budget_dirty: - self._budget_dirty = False - try: - within_budget = await self._enforce_cache_budget() - except OSError as e: - logger.warning( - "Failed to converge registry artifact cache to budget", - cache_dir=str(self.cache_dir), - error=str(e), - ) - self._budget_dirty = True - break - except BaseException: - # Cancellation must re-arm the consumed dirty signal before propagating. - self._budget_dirty = True - raise - - if not within_budget: - self._budget_dirty = True - break - - async def _enforce_cache_budget(self, *, protected_key: str | None = None) -> bool: - """Evict least-recently-used idle entries until the cache fits its budget. - - The admission lock excludes cold writers during the scan/select/evict - pass. Callers invoke enforcement without holding a per-key lock. Cold - writers already hold the admission lock and use ``_ensure_cache_capacity`` - for their staged reservations instead. - - Args: - protected_key: Newly materialized cache key. It is counted against - the budget when present but never evicted. None when enforcing - against idle entries after leases are released. - - Returns: - Whether the cache is within budget once eviction has finished. - """ - async with self._admission_lock: - return await self._enforce_cache_budget_locked(protected_key=protected_key) - - async def _reclaim_pending_work(self) -> tuple[bool, bool]: - """Retry pending trash and startup cleanup off the event loop.""" - return await rejoin_future_on_cancel( - asyncio.gather( - asyncio.to_thread(self._clear_work_dir, self.trash_dir), - asyncio.to_thread(self._retry_failed_startup_cleanup), - ) - ) - - async def _enforce_cache_budget_locked( - self, - *, - protected_key: str | None, - ) -> bool: - """Enforce entry and byte limits while the admission lock is held.""" - trash_clean, startup_clean = await self._reclaim_pending_work() - cleanup_complete = trash_clean and startup_clean - if not cleanup_complete: - return False - - max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES - max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES - if max_entries <= 0 and max_bytes <= 0: - return True - - entries = await asyncio.to_thread(self._scan_cache_entries) - total_bytes = sum(entry.size_bytes for entry in entries.values()) - protected = set() if protected_key is None else {protected_key} - skipped: set[str] = set() - - while (max_entries > 0 and len(entries) > max_entries) or ( - max_bytes > 0 and total_bytes > max_bytes - ): - candidate = self._least_recently_used( - entries.values(), - excluded=skipped | protected, - ) - if candidate is None: - logger.warning( - "Registry artifact cache is over budget but every entry is in use", - cache_dir=str(self.cache_dir), - entries=len(entries), - max_entries=max_entries, - total_bytes=total_bytes, - max_bytes=max_bytes, - ) - return False - - eviction = await self._evict_entry(candidate.cache_key) - if eviction.retired: - del entries[candidate.cache_key] - if not eviction.reclaimed: - return False - total_bytes -= candidate.size_bytes - else: - skipped.add(candidate.cache_key) - - return True - - async def _ensure_cache_capacity( - self, - *, - additional_bytes: int, - protected_key: str, - max_bytes: int, - ) -> None: - """Reserve peak bytes for a cold writer without exceeding the cap. - - The caller must hold the admission lock and its key lock so reservations - stay serialized with normal budget passes and other cold writers. - """ - if additional_bytes < 0: - raise ValueError("additional_bytes must be non-negative") - - def capacity_error(current_bytes: int) -> RegistryArtifactCacheCapacityError: - return RegistryArtifactCacheCapacityError( - current_bytes=current_bytes, - additional_bytes=additional_bytes, - max_bytes=max_bytes, - ) - - trash_clean, startup_clean = await self._reclaim_pending_work() - entries = await asyncio.to_thread(self._scan_cache_entries) - staging_bytes, trash_bytes = await asyncio.gather( - asyncio.to_thread(_directory_footprint, self.staging_dir), - asyncio.to_thread(_directory_footprint, self.trash_dir), - ) - total_bytes = ( - sum(entry.size_bytes for entry in entries.values()) - + staging_bytes - + trash_bytes - ) - non_evictable_bytes = ( - staging_bytes - + trash_bytes - + sum( - entry.size_bytes - for entry in entries.values() - if entry.cache_key == protected_key - or self._refcount(entry.cache_key) > 0 - or ( - (runtime := self._runtime.get(entry.cache_key)) is not None - and runtime.lock.locked() - ) - ) - ) - if non_evictable_bytes + additional_bytes > max_bytes: - raise capacity_error(non_evictable_bytes) - if ( - not (trash_clean and startup_clean) - and total_bytes + additional_bytes > max_bytes - ): - raise capacity_error(total_bytes) - skipped = {protected_key} - - while total_bytes + additional_bytes > max_bytes: - candidate = self._least_recently_used( - entries.values(), - excluded=skipped, - ) - if candidate is None: - raise capacity_error(total_bytes) - - eviction = await self._evict_entry(candidate.cache_key) - if eviction.retired: - del entries[candidate.cache_key] - if not eviction.reclaimed: - raise capacity_error(total_bytes) - total_bytes -= candidate.size_bytes - else: - skipped.add(candidate.cache_key) - - def _least_recently_used( - self, - entries: Iterable[RegistryArtifactCacheEntry], - *, - excluded: set[str], - ) -> RegistryArtifactCacheEntry | None: - """Return the least recently used idle entry eligible for eviction.""" - eligible = [ - entry - for entry in entries - if entry.cache_key not in excluded and self._refcount(entry.cache_key) == 0 - ] - if not eligible: - return None - return min(eligible, key=self._recency) - - def _recency(self, entry: RegistryArtifactCacheEntry) -> float: - """Return the most recent known use time for a cache entry.""" - runtime = self._runtime.get(entry.cache_key) - if runtime is None: - return entry.last_used - return max(entry.last_used, runtime.last_used) - - async def _evict_entry(self, cache_key: str) -> RegistryArtifactEviction: - """Remove one cache entry from disk, unmounting it first. - - The entry is skipped rather than forced when it is leased, busy, or - cannot be unmounted: deleting the image file behind a live mount would - leave an open-file zombie holding the loop device. - - After unmounting, the entry root is atomically renamed into ``trash`` - under the per-key lock. The lock is then released before physical - deletion runs in a worker thread. - - Args: - cache_key: Cache key to evict. - - Returns: - Whether the entry was retired and its bytes were reclaimed. - """ - lock = self._runtime_for(cache_key).lock - if lock.locked(): - logger.debug( - "Skipping eviction of busy registry artifact", - cache_key=cache_key, - ) - return RegistryArtifactEviction(retired=False, reclaimed=False) - - async with lock: - if self._refcount(cache_key) > 0: - return RegistryArtifactEviction(retired=False, reclaimed=False) - - paths = self._paths_for(cache_key) - if not paths.entry_dir.exists(): - return RegistryArtifactEviction(retired=True, reclaimed=True) - if paths.squashfs_mount_dir.is_mount() and not await self._unmount( - paths.squashfs_mount_dir - ): - logger.warning( - "Failed to unmount registry artifact, skipping eviction", - cache_key=cache_key, - mount_dir=str(paths.squashfs_mount_dir), - ) - return RegistryArtifactEviction(retired=False, reclaimed=False) - - try: - trash_path = _move_entry_to_trash( - paths.entry_dir, - self.trash_dir, - cache_key, - ) - except OSError as e: - self._budget_dirty = True - logger.warning( - "Failed to retire registry artifact cache entry", - cache_key=cache_key, - entry_dir=str(paths.entry_dir), - error=str(e), - ) - return RegistryArtifactEviction(retired=False, reclaimed=False) - - try: - deleted = await _delete_cache_path_off_loop(trash_path) - except BaseException: - self._budget_dirty = True - raise - if not deleted: - self._budget_dirty = True - logger.warning( - "Registry artifact eviction remains pending physical deletion", - cache_key=cache_key, - trash_path=str(trash_path), - ) - return RegistryArtifactEviction(retired=True, reclaimed=False) - - logger.info("Evicted registry artifact from cache", cache_key=cache_key) - return RegistryArtifactEviction(retired=True, reclaimed=True) - - async def _unmount(self, mount_dir: Path) -> bool: - """Unmount a SquashFS artifact directory, releasing its loop device. - - Cancellation kills and reaps the umount subprocess before propagating, - so the caller's per-key lock covers the complete unmount lifecycle. If - umount never took effect, the mounted entry stays consistent and can be - reused; if it already took effect, the missing extraction directory - makes the entry a plain cache miss on the next admission. - """ - umount = shutil.which("umount") - if umount is None: - return False - - proc = await asyncio.create_subprocess_exec( - umount, - str(mount_dir), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout, stderr = await _communicate_rejoin_on_cancel(proc) - if proc.returncode == 0 or not mount_dir.is_mount(): - return True - - logger.warning( - "umount command failed", - mount_dir=str(mount_dir), - output=(stderr or stdout).decode(errors="replace").strip(), - ) - return False - - def _scan_cache_entries(self) -> dict[str, RegistryArtifactCacheEntry]: - """Measure every registry artifact entry currently on disk.""" - return { - cache_key: self._measure_entry(cache_key) - for cache_key in self._discover_cache_keys() - } - - def _discover_cache_keys(self) -> set[str]: - """Return cache keys represented by atomic entry directories.""" - try: - entries = list(os.scandir(self.entries_dir)) - except FileNotFoundError: - return set() - - return { - entry.name - for entry in entries - if entry.name and entry.is_dir(follow_symlinks=False) - } - - def _measure_entry(self, cache_key: str) -> RegistryArtifactCacheEntry: - """Measure the on-disk footprint and recency of one cache entry. - - The mount directory is excluded because a mounted view only costs the - image file that backs it. The image is measured with a single ``stat`` - so a concurrent eviction deleting it cannot fail the scan. - """ - paths = self._paths_for(cache_key) - size_bytes = 0 - - try: - image_stat = paths.squashfs_image_path.stat() - except FileNotFoundError: - pass - else: - size_bytes += image_stat.st_size - - for directory in (paths.squashfs_extract_dir, paths.tarball_target_dir): - size_bytes += _directory_footprint(directory) - - try: - last_used = paths.entry_dir.stat().st_mtime - except FileNotFoundError: - last_used = 0.0 - - return RegistryArtifactCacheEntry( - cache_key=cache_key, - size_bytes=size_bytes, - last_used=last_used, - ) - - def _sweep_startup_state(self) -> None: - """Reclaim orphaned cache state left behind by a previous process. - - Scratch and trash paths from interrupted work are removed, and active - entries are trimmed to budget using entry-root mtimes as LRU order. - - The worker warms this sweep before activities can run; lazy first-use - sweeping remains a safe fallback. - """ - if not self.cache_dir.is_dir(): - self._budget_dirty = False - return - - try: - staging_clean = self._clear_work_dir( - self.staging_dir, - remember_failures=True, - ) - trash_clean = self._clear_work_dir(self.trash_dir) - cleanup_complete = staging_clean and trash_clean - within_budget = cleanup_complete and self._trim_startup_cache() - self._budget_dirty = not (cleanup_complete and within_budget) - except OSError as e: - logger.warning( - "Failed to sweep registry artifact cache", - cache_dir=str(self.cache_dir), - error=str(e), - ) - raise - - def _clear_work_dir( - self, - work_dir: Path, - *, - remember_failures: bool = False, - ) -> bool: - """Best-effort remove every child of a staging or trash directory.""" - try: - paths = list(work_dir.iterdir()) - except FileNotFoundError: - return True - except OSError as e: - logger.warning( - "Failed to inspect registry artifact work directory", - path=str(work_dir), - error=str(e), - ) - raise - - deleted = True - for path in paths: - if _delete_cache_path(path): - if remember_failures: - self._failed_startup_cleanup.discard(path) - logger.info( - "Removed registry artifact work path", - path=str(path), - ) - else: - deleted = False - if remember_failures: - self._failed_startup_cleanup.add(path) - return deleted - - def _retry_failed_startup_cleanup(self) -> bool: - """Retry exact startup paths without sweeping live staging work.""" - for path in tuple(self._failed_startup_cleanup): - if _delete_cache_path(path): - self._failed_startup_cleanup.discard(path) - return not self._failed_startup_cleanup - - def _trim_startup_cache(self) -> bool: - """Trim the cache to budget before any artifact is leased. - - Returns whether active entries and pending physical deletion fit within - the configured budget. - """ - max_entries = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES - max_bytes = config.TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES - if max_entries <= 0 and max_bytes <= 0: - return True - - entries = self._scan_cache_entries() - total_bytes = sum(entry.size_bytes for entry in entries.values()) - # Mounted entries belong to a live process sharing this cache directory. - candidates = sorted( - ( - entry - for entry in entries.values() - if not self._paths_for(entry.cache_key).squashfs_mount_dir.is_mount() - ), - key=lambda entry: entry.last_used, - ) - - def within_budget() -> bool: - return (max_entries <= 0 or len(entries) <= max_entries) and ( - max_bytes <= 0 or total_bytes <= max_bytes - ) - - for entry in candidates: - if within_budget(): - break - paths = self._paths_for(entry.cache_key) - try: - trash_path = _move_entry_to_trash( - paths.entry_dir, - self.trash_dir, - entry.cache_key, - ) - except OSError as e: - logger.warning( - "Failed to retire stale registry artifact during startup sweep", - cache_key=entry.cache_key, - entry_dir=str(paths.entry_dir), - error=str(e), - ) - return False - - del entries[entry.cache_key] - if _delete_cache_path(trash_path): - total_bytes -= entry.size_bytes - else: - return False - logger.info( - "Evicted stale registry artifact during startup sweep", - cache_key=entry.cache_key, - size_bytes=entry.size_bytes, - ) - - return within_budget() diff --git a/tracecat/sandbox/utils.py b/tracecat/sandbox/utils.py index 0904235a48..a8272c12c9 100644 --- a/tracecat/sandbox/utils.py +++ b/tracecat/sandbox/utils.py @@ -11,9 +11,11 @@ import shutil import signal import subprocess +from collections.abc import Awaitable, Callable from contextlib import suppress from pathlib import Path +from tracecat.concurrency import rejoin_future_through_cancellation from tracecat.config import ( TRACECAT__DISABLE_NSJAIL, TRACECAT__SANDBOX_NSJAIL_PATH, @@ -41,14 +43,31 @@ async def terminate_process_group(process: asyncio.subprocess.Process) -> None: await process.wait() +async def terminate_supervised_process(process: asyncio.subprocess.Process) -> None: + """Request descendant cleanup from a Linux process supervisor.""" + if process.returncode is None: + with suppress(ProcessLookupError): + os.kill(process.pid, signal.SIGTERM) + # An untrusted action runs under the supervisor's UID and can stop the + # outer supervisor. SIGTERM remains pending for a stopped process, so + # resume it before waiting for its cleanup handler to run. + with suppress(ProcessLookupError): + os.kill(process.pid, signal.SIGCONT) + # The supervisor exits only after its detached subreaper has killed and + # reaped the action's complete descendant tree. Do not impose a shorter + # wait that could release registry leases while cleanup is still active. + await process.wait() + + async def _finish_process_group_cleanup( process: asyncio.subprocess.Process, communicate_task: asyncio.Task[tuple[bytes | None, bytes | None]], - termination_task: asyncio.Task[None] | None, + termination_task: asyncio.Future[None] | None, + terminate: Callable[[asyncio.subprocess.Process], Awaitable[None]], ) -> None: """Finish process termination and consume the communication task.""" if termination_task is None: - termination_task = asyncio.create_task(terminate_process_group(process)) + termination_task = asyncio.ensure_future(terminate(process)) try: await termination_task finally: @@ -58,34 +77,12 @@ async def _finish_process_group_cleanup( await communicate_task -async def _rejoin_cleanup_through_cancellation( - cleanup_task: asyncio.Task[None], -) -> None: - """Wait for cleanup despite repeated caller cancellation.""" - pending_cancellation: asyncio.CancelledError | None = None - while not cleanup_task.done(): - try: - await asyncio.shield(cleanup_task) - except asyncio.CancelledError as e: - if cleanup_task.cancelled(): - raise - pending_cancellation = e - - try: - cleanup_task.result() - except BaseException as cleanup_error: - if pending_cancellation is not None: - raise pending_cancellation from cleanup_error - raise - if pending_cancellation is not None: - raise pending_cancellation - - async def communicate_process_group( process: asyncio.subprocess.Process, *, input: bytes | None = None, # noqa: A002 timeout: float | None = None, + terminate: Callable[[asyncio.subprocess.Process], Awaitable[None]] | None = None, ) -> tuple[bytes, bytes]: """Communicate with a process while containing its process group. @@ -93,16 +90,18 @@ async def communicate_process_group( ``communicate()`` and asyncio's ``wait()`` to wait after the leader exits. Polling ``returncode`` observes the leader exit independently, letting us terminate the group immediately and close those pipes. Cancellation also - terminates the group before it propagates. + terminates the group before it propagates. Callers with a stronger process + supervisor can supply its cleanup function as ``terminate``. """ + terminator = terminate or terminate_process_group communicate_task = asyncio.create_task(process.communicate(input=input)) - termination_task: asyncio.Task[None] | None = None + termination_task: asyncio.Future[None] | None = None operation_error: BaseException | None = None try: async with asyncio.timeout(timeout): while process.returncode is None: await asyncio.sleep(_PROCESS_EXIT_POLL_INTERVAL_SECONDS) - termination_task = asyncio.create_task(terminate_process_group(process)) + termination_task = asyncio.ensure_future(terminator(process)) await asyncio.shield(termination_task) stdout, stderr = await communicate_task except BaseException as e: @@ -114,10 +113,11 @@ async def communicate_process_group( process, communicate_task, termination_task, + terminator, ) ) try: - await _rejoin_cleanup_through_cancellation(cleanup_task) + await rejoin_future_through_cancellation(cleanup_task) except BaseException as cleanup_error: if operation_error is not None: raise operation_error from cleanup_error diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index 1fccec7362..1e82cc707c 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -7,7 +7,7 @@ import os import threading import weakref -from collections.abc import AsyncIterator, Awaitable, Callable +from collections.abc import AsyncGenerator, Awaitable, Callable from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass, field from pathlib import Path @@ -17,7 +17,7 @@ import aiofiles from aiobotocore.config import AioConfig from boto3.s3.transfer import TransferConfig -from botocore.exceptions import ClientError +from botocore.exceptions import BotoCoreError, ClientError from tracecat import config from tracecat.logger import logger @@ -35,6 +35,39 @@ DEFAULT_UPLOAD_CHUNK_SIZE_BYTES = 8 * 1024 * 1024 # 8MB DEFAULT_UPLOAD_MAX_CONCURRENCY = 4 DEFAULT_UPLOAD_MAX_IO_QUEUE_SIZE = 2 +_REDACTED_STORAGE_IDENTIFIER = "" + + +class StorageDownloadError(RuntimeError): + """A storage download failed without exposing object identifiers.""" + + def __init__(self, *, error_code: str | None) -> None: + super().__init__("Storage download failed") + self.error_code = error_code + + +def _download_log_identifiers( + key: str, + bucket: str, + *, + redact: bool, +) -> tuple[str, str]: + """Return storage identifiers that are safe for logs and errors.""" + if redact: + return _REDACTED_STORAGE_IDENTIFIER, _REDACTED_STORAGE_IDENTIFIER + return key, bucket + + +def _safe_storage_error_code(value: object) -> str | None: + """Return a bounded machine code, never provider-controlled prose.""" + if not isinstance(value, str) or not value or len(value) > 64: + return None + if not value.isascii() or any( + not (character.isalnum() or character in "._-") for character in value + ): + return None + return value + # Shared S3/MinIO client config: explicit standard-mode retries so transient # failures (throttling, 5xx, connection resets) are retried with backoff instead @@ -205,7 +238,7 @@ async def close_storage_client_cache() -> None: @asynccontextmanager -async def get_storage_client() -> AsyncIterator[S3Client]: +async def get_storage_client() -> AsyncGenerator[S3Client]: """Get a configured S3 client for either AWS S3. Yields: @@ -703,7 +736,9 @@ async def download_file_range( async def open_download_stream( key: str, bucket: str, -) -> AsyncIterator[tuple[StreamingBody, int | None]]: + *, + redact_log_identifiers: bool = False, +) -> AsyncGenerator[tuple[StreamingBody, int | None]]: """Open a streaming download for an S3/MinIO object. This is safer for very large objects because it allows callers to @@ -718,14 +753,21 @@ async def open_download_stream( Args: key: The S3 object key. bucket: Bucket name (required). + redact_log_identifiers: Hide the key and bucket in logs and errors. Yields: Tuple of (streaming body, content_length). Raises: ClientError: If the download fails. + StorageDownloadError: If a redacted download fails. FileNotFoundError: If the file doesn't exist. """ + log_key, log_bucket = _download_log_identifiers( + key, + bucket, + redact=redact_log_identifiers, + ) try: async with get_storage_client() as s3_client: response = await s3_client.get_object(Bucket=bucket, Key=key) @@ -734,13 +776,25 @@ async def open_download_stream( async with body: yield body, content_length except ClientError as e: - if e.response.get("Error", {}).get("Code") == "NoSuchKey": + error_code = _safe_storage_error_code(e.response.get("Error", {}).get("Code")) + if error_code == "NoSuchKey": logger.warning( "File not found in storage", - key=key, - bucket=bucket, + key=log_key, + bucket=log_bucket, ) + if redact_log_identifiers: + raise FileNotFoundError from None raise FileNotFoundError from e + if redact_log_identifiers: + logger.error( + "Failed to open download stream", + key=log_key, + bucket=log_bucket, + error_code=error_code, + error_type=type(e).__name__, + ) + raise StorageDownloadError(error_code=error_code) from None logger.error( "Failed to open download stream", key=key, @@ -748,6 +802,17 @@ async def open_download_stream( error=str(e), ) raise + except BotoCoreError as e: + if not redact_log_identifiers: + raise + logger.error( + "Failed to open download stream", + key=log_key, + bucket=log_bucket, + error_code=None, + error_type=type(e).__name__, + ) + raise StorageDownloadError(error_code=None) from None async def download_file_to_path( @@ -759,6 +824,8 @@ async def download_file_to_path( max_bytes: int | None = None, expected_sha256: str | None = None, ensure_capacity: Callable[[int], Awaitable[None]] | None = None, + defer_cleanup: Callable[[Path], None] | None = None, + redact_log_identifiers: bool = False, ) -> int: """Stream an S3/MinIO object to a local file. @@ -775,6 +842,8 @@ async def download_file_to_path( ensure_capacity: Optional callback invoked before the first disk write with the maximum number of bytes the download may occupy. When the server omits ContentLength, max_bytes is required to provide that bound. + defer_cleanup: Optional callback retaining failed partial-file cleanup. + redact_log_identifiers: Hide the key and bucket in logs and errors. Returns: Total bytes written. @@ -784,19 +853,25 @@ async def download_file_to_path( hasher = hashlib.sha256() if expected_sha256 is not None else None bytes_written = 0 + log_key, log_bucket = _download_log_identifiers( + key, + bucket, + redact=redact_log_identifiers, + ) try: - async with open_download_stream(key=key, bucket=bucket) as ( - stream, - content_length, - ): + async with open_download_stream( + key=key, + bucket=bucket, + redact_log_identifiers=redact_log_identifiers, + ) as (stream, content_length): if ( max_bytes is not None and content_length is not None and content_length > max_bytes ): raise ValueError( - f"Refusing to download {bucket}/{key} to disk: " + f"Refusing to download {log_bucket}/{log_key} to disk: " f"ContentLength={content_length} exceeds max_bytes={max_bytes}" ) @@ -807,7 +882,7 @@ async def download_file_to_path( if max_bytes is None: raise ValueError( "Cannot reserve disk capacity for a download without " - f"ContentLength or max_bytes: {bucket}/{key}" + f"ContentLength or max_bytes: {log_bucket}/{log_key}" ) reserved_bytes = max_bytes await ensure_capacity(reserved_bytes) @@ -820,7 +895,7 @@ async def download_file_to_path( bytes_written += len(chunk) if download_limit is not None and bytes_written > download_limit: raise ValueError( - f"Refusing to download {bucket}/{key} to disk: " + f"Refusing to download {log_bucket}/{log_key} to disk: " f"bytes_written={bytes_written} exceeds " f"max_bytes={download_limit}" ) @@ -832,7 +907,7 @@ async def download_file_to_path( actual_sha256 = hasher.hexdigest() if actual_sha256 != expected_sha256: raise ValueError( - f"Integrity check failed for {bucket}/{key}: " + f"Integrity check failed for {log_bucket}/{log_key}: " f"expected {expected_sha256}, got {actual_sha256}" ) @@ -840,17 +915,20 @@ async def download_file_to_path( except BaseException: try: temp_path.unlink(missing_ok=True) - except Exception: + except Exception as cleanup_error: + if defer_cleanup is not None: + defer_cleanup(temp_path) logger.warning( "Failed to cleanup partial download", temp_path=str(temp_path), + error_type=type(cleanup_error).__name__, ) raise logger.debug( "File streamed to disk successfully", - key=key, - bucket=bucket, + key=log_key, + bucket=log_bucket, output_path=str(output_path), size=bytes_written, ) From 2f7a4d5742fb01958fb541e669e1275bc7a32655 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:45:36 -0700 Subject: [PATCH 152/161] fix(executor): close registry cache correctness gaps --- tests/unit/test_action_runner.py | 5 +- tests/unit/test_registry_artifacts.py | 205 ++++++++++++++++-- tests/unit/test_storage_blob.py | 45 ++++ tracecat/executor/registry_artifact_mounts.py | 24 ++ .../executor/registry_artifact_storage.py | 21 +- tracecat/executor/registry_artifacts.py | 74 +++++-- tracecat/storage/blob.py | 23 +- 7 files changed, 336 insertions(+), 61 deletions(-) create mode 100644 tracecat/executor/registry_artifact_mounts.py diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index eafa81af77..d3845ead79 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -808,7 +808,10 @@ async def release_mount(mount_dir: Path) -> bool: return True with ( - patch.object(Path, "is_mount", lambda path: path in mounted), + patch( + "tracecat.executor.registry_artifact_mounts.is_mount", + lambda path: path in mounted, + ), patch.object( action_runner, "_direct_subprocess_command", diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 6d05782f68..bc672f34dd 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -31,6 +31,7 @@ RegistryArtifactCacheCapacityError, RegistryArtifactCacheLoopError, RegistryArtifactEviction, + RegistryArtifactExtractionError, RegistryArtifactFormat, RegistryArtifactUriError, SquashfsArtifact, @@ -54,6 +55,7 @@ "tracecat.executor.registry_artifacts.config" ".TRACECAT__EXECUTOR_REGISTRY_SQUASHFS_ENABLED" ) +MOUNT_CHECK = "tracecat.executor.registry_artifact_mounts.is_mount" def _write_tarball_entry(cache_dir: Path, cache_key: str) -> Path: @@ -282,6 +284,15 @@ def test_compute_registry_artifact_cache_key_case_sensitive(self): assert key1 != key2 + def test_compute_registry_artifact_cache_key_whitespace_sensitive(self): + """Cache identity preserves whitespace that can belong to an S3 key.""" + key = compute_registry_artifact_cache_key("s3://bucket/path/file.tar.gz") + whitespace_key = compute_registry_artifact_cache_key( + "s3://bucket/path/file.tar.gz " + ) + + assert key != whitespace_key + def test_compute_registry_artifact_cache_key_empty(self): """Test that empty URI returns the base cache key.""" assert compute_registry_artifact_cache_key("") == "base" @@ -1003,6 +1014,81 @@ def blocking_extractall(*args: object, **kwargs: object) -> None: assert first_cancellation_propagated_early is False assert second_cancellation_propagated_early is False + @pytest.mark.parametrize("operation", ["extract", "size-command", "size-parse"]) + @pytest.mark.anyio + async def test_squashfs_failures_sanitize_subprocess_output( + self, + temp_cache_dir: Path, + operation: str, + ) -> None: + sensitive_output = b"synthetic-customer/repository/member.py" + artifact = SquashfsArtifact( + uri="s3://bucket/path/custom.squashfs", + cache_key="malformed-squashfs", + ) + image_path = temp_cache_dir / "image.squashfs" + image_path.write_bytes(b"squashfs") + target_dir = temp_cache_dir / "target" + target_dir.mkdir() + process = AsyncMock() + process.returncode = 0 if operation == "size-parse" else 1 + stdout = sensitive_output if operation == "size-parse" else b"" + stderr = sensitive_output if operation != "size-parse" else b"" + + with ( + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/usr/bin/unsquashfs", + ), + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ), + patch( + "tracecat.executor.registry_artifacts.communicate_rejoin_on_cancel", + new_callable=AsyncMock, + return_value=(stdout, stderr), + ), + pytest.raises(RegistryArtifactExtractionError) as raised, + ): + if operation == "extract": + await artifact._extract_image(image_path, target_dir) + else: + await artifact._squashfs_extracted_size( + image_path, + allocation_unit=1, + ) + + assert str(raised.value) == "Registry artifact extraction failed" + assert sensitive_output.decode() not in str(raised.value) + assert raised.value.__cause__ is None + + @pytest.mark.anyio + async def test_tarball_extract_sanitizes_member_failures( + self, + temp_cache_dir: Path, + ) -> None: + sensitive_member = "../../synthetic-customer/repository/module.py" + artifact = TarballArtifact( + uri="s3://bucket/path/site-packages.tar.gz", + cache_key="malformed-tarball", + ) + tarball_path = temp_cache_dir / "artifact.tar.gz" + target_dir = temp_cache_dir / "target" + target_dir.mkdir() + with tarfile.open(tarball_path, "w:gz") as tar: + member = tarfile.TarInfo(sensitive_member) + member.size = 1 + tar.addfile(member, io.BytesIO(b"x")) + + with pytest.raises(RegistryArtifactExtractionError) as raised: + await artifact.extract(tarball_path, target_dir) + + assert str(raised.value) == "Registry artifact extraction failed" + assert sensitive_member not in str(raised.value) + assert raised.value.__cause__ is None + @pytest.mark.anyio async def test_repeatedly_cancelled_tarball_size_scan_rejoins_thread( self, temp_cache_dir: Path @@ -1553,10 +1639,17 @@ async def test_lease_without_uris_returns_base_pythonpath_dir(self, temp_cache_d """No artifact URIs still yields the base PYTHONPATH directory.""" cache = RegistryArtifactCache(temp_cache_dir) - async with cache.lease(None) as registry_paths: - assert registry_paths == [temp_cache_dir / "base"] - assert registry_paths[0].is_dir() + with patch.object( + cache, + "ensure_swept", + new_callable=AsyncMock, + side_effect=AssertionError("cache-free leases must not sweep"), + ) as ensure_swept: + async with cache.lease(None) as registry_paths: + assert registry_paths == [temp_cache_dir / "base"] + assert registry_paths[0].is_dir() + ensure_swept.assert_not_awaited() assert cache._runtime == {} @pytest.mark.anyio @@ -1605,7 +1698,7 @@ async def hold_lease(index: int) -> None: await releases[index].wait() with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(MOUNT_CHECK, lambda path: path in harness.mounted), patch(SQUASHFS_ENABLED_CONFIG, True), patch( "tracecat.executor.registry_artifacts.shutil.which", @@ -1667,7 +1760,7 @@ async def hold_lease(index: int) -> None: await releases[index].wait() with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(MOUNT_CHECK, lambda path: path in harness.mounted), patch(SQUASHFS_ENABLED_CONFIG, True), patch( "tracecat.executor.registry_artifacts.shutil.which", @@ -1763,6 +1856,35 @@ async def hold_lease() -> None: ] assert all(cache._refcount(cache_key) == 0 for cache_key in cache_keys) + @pytest.mark.anyio + async def test_cleanup_failure_preserves_successful_lease_outcome( + self, temp_cache_dir: Path + ) -> None: + """Post-lease maintenance cannot replace a successful caller result.""" + cache = RegistryArtifactCache(temp_cache_dir) + artifact_uri = "s3://bucket/cleanup-failure.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + target_dir = _write_tarball_entry(temp_cache_dir, cache_key) + + async def fail_cleanup(idle_keys: list[str]) -> None: + del idle_keys + raise RuntimeError("cleanup failed with sensitive details") + + with ( + patch.object(cache, "_finish_lease_cleanup", fail_cleanup), + patch("tracecat.executor.registry_artifacts.logger.error") as log_error, + ): + async with cache.lease([artifact_uri]) as registry_paths: + result = registry_paths + + assert result == [target_dir] + assert cache._refcount(cache_key) == 0 + log_error.assert_called_once_with( + "Registry artifact lease cleanup failed; preserving caller outcome", + cache_dir=str(temp_cache_dir), + error_type="RuntimeError", + ) + @pytest.mark.anyio async def test_new_lease_racing_final_release_prevents_stale_unmount( self, temp_cache_dir: Path @@ -1806,7 +1928,7 @@ async def new_holder() -> None: await release_newcomer.wait() with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(MOUNT_CHECK, lambda path: path in harness.mounted), patch.object(cache, "_unmount", harness.unmount), patch.object( cache, @@ -1869,7 +1991,7 @@ async def fail_download( converge_cache_budget = AsyncMock() with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(MOUNT_CHECK, lambda path: path in harness.mounted), patch(SQUASHFS_ENABLED_CONFIG, False), patch.object(TarballArtifact, "download", fail_download), patch.object(cache, "_unmount", harness.unmount), @@ -2000,7 +2122,7 @@ async def test_duplicate_uri_balances_each_acquisition_and_release( harness = _SquashfsMountHarness(cache) with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(MOUNT_CHECK, lambda path: path in harness.mounted), patch(SQUASHFS_ENABLED_CONFIG, True), patch( "tracecat.executor.registry_artifacts.shutil.which", @@ -2065,7 +2187,7 @@ async def take_lease() -> None: leased_path_exists.append(registry_paths[0].is_dir()) with ( - patch.object(Path, "is_mount", lambda self: self in mounted), + patch(MOUNT_CHECK, lambda path: path in mounted), patch(SQUASHFS_ENABLED_CONFIG, True), patch( "tracecat.executor.registry_artifacts.shutil.which", @@ -2113,15 +2235,24 @@ async def test_builtin_artifact_is_exempt_from_cache_accounting( cache = RegistryArtifactCache(temp_cache_dir) - with patch.object( - cache, - "_enforce_cache_budget", - new_callable=AsyncMock, - ) as enforce_cache_budget: + with ( + patch.object( + cache, + "ensure_swept", + new_callable=AsyncMock, + side_effect=AssertionError("builtin leases must not sweep"), + ) as ensure_swept, + patch.object( + cache, + "_enforce_cache_budget", + new_callable=AsyncMock, + ) as enforce_cache_budget, + ): async with cache.lease([bundled_builtin_registry_uri(version)]) as paths: assert paths == [site_packages.resolve()] assert cache._runtime == {} + ensure_swept.assert_not_awaited() enforce_cache_budget.assert_not_awaited() @@ -2939,7 +3070,7 @@ async def mock_umount(*args, **kwargs): return process with ( - patch.object(Path, "is_mount", lambda self: self in mounted), + patch(MOUNT_CHECK, lambda path: path in mounted), patch( "tracecat.executor.registry_artifacts.shutil.which", return_value="/sbin/umount", @@ -3081,6 +3212,31 @@ async def test_enforce_budget_proceeds_over_budget_when_everything_is_leased( assert within_budget is False assert leased.exists() + @pytest.mark.anyio + async def test_eviction_fails_closed_when_mount_inspection_fails( + self, temp_cache_dir: Path + ) -> None: + """Unknown mount state cannot be mistaken for an unmounted entry.""" + cache = RegistryArtifactCache(temp_cache_dir) + paths = cache._paths_for("unknown-mount-state") + paths.entry_dir.mkdir(parents=True) + paths.squashfs_image_path.write_bytes(b"squashfs") + paths.squashfs_mount_dir.mkdir() + real_lstat = Path.lstat + + def fail_mount_lstat(path: Path): + if path == paths.squashfs_mount_dir: + raise PermissionError("mount inspection denied") + return real_lstat(path) + + with patch.object(Path, "lstat", fail_mount_lstat): + with pytest.raises(PermissionError, match="mount inspection denied"): + await cache._evict_entry("unknown-mount-state") + + assert paths.entry_dir.is_dir() + assert paths.squashfs_image_path.is_file() + assert paths.squashfs_mount_dir.is_dir() + @pytest.mark.anyio async def test_eviction_unmounts_before_deleting_the_image(self, temp_cache_dir): """Unlinking a mounted image would strand an open-file zombie.""" @@ -3102,7 +3258,7 @@ async def mock_umount(*args, **kwargs): return process with ( - patch.object(Path, "is_mount", lambda self: self in mounted), + patch(MOUNT_CHECK, lambda path: path in mounted), patch( "tracecat.executor.registry_artifacts.shutil.which", return_value="/sbin/umount", @@ -3153,7 +3309,7 @@ async def mock_umount(*args, **kwargs): return released_process with ( - patch.object(Path, "is_mount", lambda self: self in mounted), + patch(MOUNT_CHECK, lambda path: path in mounted), patch( "tracecat.executor.registry_artifacts.shutil.which", return_value="/sbin/umount", @@ -3349,7 +3505,7 @@ async def test_eviction_skips_entry_when_unmount_fails(self, temp_cache_dir): process.returncode = 32 with ( - patch.object(Path, "is_mount", lambda self: self in mounted), + patch(MOUNT_CHECK, lambda path: path in mounted), patch( "tracecat.executor.registry_artifacts.shutil.which", return_value="/sbin/umount", @@ -3506,7 +3662,7 @@ async def test_sweep_keeps_mounted_dirs(self, temp_cache_dir): mount_dir = paths.squashfs_mount_dir mount_dir.mkdir() - with patch.object(Path, "is_mount", lambda self: self == mount_dir): + with patch(MOUNT_CHECK, lambda path: path == mount_dir): await cache.ensure_swept() assert mount_dir.is_dir() @@ -3576,12 +3732,17 @@ def blocking_sweep() -> None: @pytest.mark.anyio async def test_lease_triggers_startup_sweep(self, temp_cache_dir): - """Lease admission reclaims startup scratch before yielding paths.""" + """Cache-backed lease admission reclaims scratch before yielding paths.""" cache = RegistryArtifactCache(temp_cache_dir) orphaned_dir = cache.staging_dir / "abc123.999999.4321" orphaned_dir.mkdir(parents=True) + artifact_uri = "s3://bucket/cached.tar.gz" + _write_tarball_entry( + temp_cache_dir, + compute_registry_artifact_cache_key(artifact_uri), + ) - async with cache.lease(None): + async with cache.lease([artifact_uri]): assert not orphaned_dir.exists() @pytest.mark.anyio @@ -3838,7 +3999,7 @@ async def test_loop_device_exhaustion_isolated_sticky_extraction_fallback( ) with ( - patch.object(Path, "is_mount", lambda path: path in harness.mounted), + patch(MOUNT_CHECK, lambda path: path in harness.mounted), patch(SQUASHFS_ENABLED_CONFIG, True), patch( "tracecat.executor.registry_artifacts.shutil.which", diff --git a/tests/unit/test_storage_blob.py b/tests/unit/test_storage_blob.py index ce0afcef52..3e8f7dec01 100644 --- a/tests/unit/test_storage_blob.py +++ b/tests/unit/test_storage_blob.py @@ -878,6 +878,51 @@ async def ensure_capacity(size_bytes: int) -> None: assert reservations == [7] assert out.read_bytes() == b"payload" + @pytest.mark.anyio + async def test_download_file_to_path_grows_unknown_length_reservation( + self, tmp_path: Path, monkeypatch + ): + """Unknown-length downloads reserve each chunk above current usage.""" + + class DummyStream: + async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 + yield b"abc" + yield b"defg" + + @asynccontextmanager + async def _fake_open_download_stream( + *, key: str, bucket: str, redact_log_identifiers: bool = False + ): # noqa: ARG001 + yield DummyStream(), None + + monkeypatch.setattr( + "tracecat.storage.blob.open_download_stream", + _fake_open_download_stream, + ) + + out = tmp_path / "out.bin" + temp_path = tmp_path / "out.bin.part" + reservations: list[int] = [] + partial_sizes: list[int] = [] + + async def ensure_capacity(size_bytes: int) -> None: + partial_size = temp_path.stat().st_size + assert 2 + partial_size + size_bytes <= 10 + reservations.append(size_bytes) + partial_sizes.append(partial_size) + + await download_file_to_path( + key="k", + bucket="b", + output_path=out, + max_bytes=10, + ensure_capacity=ensure_capacity, + ) + + assert reservations == [3, 4] + assert partial_sizes == [0, 3] + assert out.read_bytes() == b"abcdefg" + @pytest.mark.anyio async def test_download_file_to_path_never_exceeds_reserved_length( self, tmp_path: Path, monkeypatch diff --git a/tracecat/executor/registry_artifact_mounts.py b/tracecat/executor/registry_artifact_mounts.py new file mode 100644 index 0000000000..5829211b1c --- /dev/null +++ b/tracecat/executor/registry_artifact_mounts.py @@ -0,0 +1,24 @@ +"""Fail-closed mount-state inspection for registry artifact caches.""" + +from __future__ import annotations + +import os +import stat +from pathlib import Path + + +def is_mount(path: Path) -> bool: + """Return whether a path is mounted without hiding inspection failures.""" + try: + path_stat = path.lstat() + except FileNotFoundError: + return False + + if stat.S_ISLNK(path_stat.st_mode): + return False + + parent = Path(os.path.realpath(path / "..", strict=True)) + parent_stat = parent.lstat() + return ( + path_stat.st_dev != parent_stat.st_dev or path_stat.st_ino == parent_stat.st_ino + ) diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 945c1f7a60..a0a8f42ecc 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -22,6 +22,7 @@ rejoin_future_on_cancel, run_blocking_rejoin_on_cancel, ) +from tracecat.executor import registry_artifact_mounts from tracecat.executor.registry_artifact_budget import ( RegistryArtifactCacheBudget, RegistryArtifactCacheEntry, @@ -635,7 +636,7 @@ async def _unmount_idle_entry(self, cache_key: str) -> None: return mount_dir = self._paths_for(cache_key).squashfs_mount_dir try: - mounted = mount_dir.is_mount() + mounted = registry_artifact_mounts.is_mount(mount_dir) except OSError as e: logger.warning( "Failed to inspect registry artifact mount state", @@ -889,9 +890,9 @@ async def _evict_entry(self, cache_key: str) -> RegistryArtifactEviction: if not paths.entry_dir.exists(): self._request_runtime_retirement(cache_key, runtime) return RegistryArtifactEviction(retired=True, reclaimed=True) - if paths.squashfs_mount_dir.is_mount() and not await self._unmount( + if registry_artifact_mounts.is_mount( paths.squashfs_mount_dir - ): + ) and not await self._unmount(paths.squashfs_mount_dir): logger.warning( "Failed to unmount registry artifact, skipping eviction", cache_key=cache_key, @@ -943,7 +944,7 @@ async def _unmount(self, mount_dir: Path) -> bool: stderr=asyncio.subprocess.PIPE, ) stdout, stderr = await communicate_rejoin_on_cancel(proc) - if proc.returncode == 0 or not mount_dir.is_mount(): + if proc.returncode == 0 or not registry_artifact_mounts.is_mount(mount_dir): return True logger.warning( "umount command failed", @@ -1034,7 +1035,9 @@ def _measure_entry( if allocation_unit is None: allocation_unit = _filesystem_allocation_unit(self.cache_dir) try: - mount_is_active = paths.squashfs_mount_dir.is_mount() + mount_is_active = registry_artifact_mounts.is_mount( + paths.squashfs_mount_dir + ) except FileNotFoundError: mount_is_active = False pruned_directories = (paths.squashfs_mount_dir,) if mount_is_active else () @@ -1089,7 +1092,9 @@ def _clear_legacy_cache_layout(self) -> bool: if _LEGACY_CACHE_PATH_PATTERN.fullmatch(path.name) is None: continue try: - mounted = is_reusable_cache_directory(path) and path.is_mount() + mounted = is_reusable_cache_directory( + path + ) and registry_artifact_mounts.is_mount(path) except OSError: mounted = True if mounted or not _delete_cache_path(path): @@ -1153,7 +1158,9 @@ def _trim_startup_cache(self) -> bool: mounted_keys = { entry.cache_key for entry in entries.values() - if self._paths_for(entry.cache_key).squashfs_mount_dir.is_mount() + if registry_artifact_mounts.is_mount( + self._paths_for(entry.cache_key).squashfs_mount_dir + ) } plan = plan_registry_artifact_evictions( entries, diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index f060f1507e..2f4ed450e0 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -27,6 +27,7 @@ rejoin_future_through_cancellation, run_blocking_rejoin_on_cancel, ) +from tracecat.executor import registry_artifact_mounts from tracecat.executor.registry_artifact_storage import ( RegistryArtifactAdmission, RegistryArtifactCacheCapacityError, @@ -58,6 +59,7 @@ "RegistryArtifactAdmission", "RegistryArtifactCache", "RegistryArtifactCacheCapacityError", + "RegistryArtifactExtractionError", "RegistryArtifactCacheLoopError", "RegistryArtifactEviction", "RegistryArtifactFormat", @@ -104,6 +106,13 @@ class RegistryArtifactUriError(ValueError): """A registry artifact URI is malformed, with identifiers suppressed.""" +class RegistryArtifactExtractionError(RuntimeError): + """A registry archive could not be inspected or extracted safely.""" + + def __init__(self) -> None: + super().__init__("Registry artifact extraction failed") + + @dataclass(frozen=True, slots=True) class RegistryArtifact(ABC): """An executor-local materializable registry artifact.""" @@ -182,7 +191,9 @@ def cached_path( mount_dir, defer_cleanup=ctx.defer_cleanup, ) - if is_reusable_cache_directory(mount_dir) and mount_dir.is_mount(): + if is_reusable_cache_directory(mount_dir) and registry_artifact_mounts.is_mount( + mount_dir + ): logger.debug( "Using cached SquashFS registry mount", cache_key=ctx.cache_key, @@ -280,7 +291,9 @@ async def mount( ) if os.path.lexists(target_dir) and not is_reusable_cache_directory(target_dir): raise OSError("Failed to reclaim malformed SquashFS mount target") - if is_reusable_cache_directory(target_dir) and target_dir.is_mount(): + if is_reusable_cache_directory( + target_dir + ) and registry_artifact_mounts.is_mount(target_dir): return target_dir ensure_cache_entry_directory(ctx.paths) @@ -400,7 +413,7 @@ async def _mount_image(self, image_path: Path, target_dir: Path) -> None: Raises: SquashfsMountCommandError: The ``mount`` command failed. """ - if target_dir.is_mount(): + if registry_artifact_mounts.is_mount(target_dir): return proc = await asyncio.create_subprocess_exec( @@ -416,7 +429,7 @@ async def _mount_image(self, image_path: Path, target_dir: Path) -> None: ) stdout, stderr = await communicate_rejoin_on_cancel(proc) - if proc.returncode == 0 or target_dir.is_mount(): + if proc.returncode == 0 or registry_artifact_mounts.is_mount(target_dir): return output = (stderr or stdout).decode(errors="replace").strip() @@ -442,13 +455,12 @@ async def _extract_image(self, image_path: Path, target_dir: Path) -> None: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - stdout, stderr = await communicate_rejoin_on_cancel(proc) + await communicate_rejoin_on_cancel(proc) if proc.returncode == 0: return - output = (stderr or stdout).decode(errors="replace").strip() - raise RuntimeError(output or "unsquashfs command failed") + raise RegistryArtifactExtractionError() async def _squashfs_extracted_size( self, @@ -468,12 +480,14 @@ async def _squashfs_extracted_size( stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - stdout, stderr = await communicate_rejoin_on_cancel(proc) + stdout, _ = await communicate_rejoin_on_cancel(proc) if proc.returncode != 0: - output = (stderr or stdout).decode(errors="replace").strip() - raise RuntimeError(output or "unsquashfs listing failed") - return _squashfs_listing_size(stdout, allocation_unit=allocation_unit) + raise RegistryArtifactExtractionError() + try: + return _squashfs_listing_size(stdout, allocation_unit=allocation_unit) + except Exception: + raise RegistryArtifactExtractionError() from None @dataclass(frozen=True, slots=True) @@ -523,12 +537,15 @@ async def materialize( download_elapsed = (time.monotonic() - download_start) * 1000 if (admission := ctx.admission) is not None: - extracted_size = await run_blocking_rejoin_on_cancel( - lambda: _tarball_extracted_size( - temp_tarball, - allocation_unit=admission.allocation_unit, + try: + extracted_size = await run_blocking_rejoin_on_cancel( + lambda: _tarball_extracted_size( + temp_tarball, + allocation_unit=admission.allocation_unit, + ) ) - ) + except Exception: + raise RegistryArtifactExtractionError() from None extracted_size += _directory_records_size_bound( (temp_dir.name, target_dir.name), allocation_unit=admission.allocation_unit, @@ -607,7 +624,10 @@ def _do_extract() -> None: raise ValueError(f"Unsupported tarball format: {tarball_path}") - await run_blocking_rejoin_on_cancel(_do_extract) + try: + await run_blocking_rejoin_on_cancel(_do_extract) + except Exception: + raise RegistryArtifactExtractionError() from None logger.debug( "Tarball extracted", @@ -672,9 +692,8 @@ def compute_registry_artifact_cache_key(artifact_uri: str) -> str: """Compute the local cache key for a registry artifact URI.""" if not artifact_uri: return "base" - # S3 keys are case-sensitive, so preserve URI case when hashing. - content = artifact_uri.strip() - return hashlib.sha256(content.encode()).hexdigest()[:16] + # S3 keys are byte-sensitive, so hash the exact URI used for retrieval. + return hashlib.sha256(artifact_uri.encode()).hexdigest()[:16] def bundled_builtin_registry_uri(version: str) -> str: @@ -981,7 +1000,13 @@ async def __aexit__( except BaseException as cleanup_error: if exc_value is not None: raise exc_value.with_traceback(traceback) from cleanup_error - raise + if not isinstance(cleanup_error, Exception): + raise + logger.error( + "Registry artifact lease cleanup failed; preserving caller outcome", + cache_dir=str(self.cache.cache_dir), + error_type=type(cleanup_error).__name__, + ) async def aclose(self) -> None: """Release the pin exactly once and finish all resulting maintenance.""" @@ -1032,13 +1057,14 @@ async def lease( Yields: Importable Python paths for the requested artifacts. """ - await self.ensure_swept() - if not artifact_uris: logger.info("No registry artifact URIs provided, using base PYTHONPATH") yield [self._base_pythonpath_dir()] return + if any(_is_cache_entry_uri(uri) for uri in artifact_uris): + await self.ensure_swept() + async with contextlib.AsyncExitStack() as leases: handles: list[_RegistryArtifactLease] = [] registry_paths: list[Path] = [] @@ -1258,7 +1284,7 @@ def _remove_unpublished_entry( paths = ctx.paths validate_cache_entry_path(paths) try: - if paths.squashfs_mount_dir.is_mount(): + if registry_artifact_mounts.is_mount(paths.squashfs_mount_dir): return except OSError: return diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index 1e82cc707c..065e1a2aa0 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -839,9 +839,9 @@ async def download_file_to_path( chunk_size: Chunk size for streaming reads (default: 8MB). max_bytes: Optional guardrail; raise if the object exceeds this size. expected_sha256: Optional integrity check; raise if computed SHA-256 differs. - ensure_capacity: Optional callback invoked before the first disk write with - the maximum number of bytes the download may occupy. When the server - omits ContentLength, max_bytes is required to provide that bound. + ensure_capacity: Optional callback invoked before disk writes. When the + server omits ContentLength, max_bytes is required and capacity is + checked incrementally before each chunk is written. defer_cleanup: Optional callback retaining failed partial-file cleanup. redact_log_identifiers: Hide the key and bucket in logs and errors. @@ -876,6 +876,7 @@ async def download_file_to_path( ) download_limit = max_bytes + grow_reservation_by_chunk = False if ensure_capacity is not None: reserved_bytes = content_length if reserved_bytes is None: @@ -884,11 +885,17 @@ async def download_file_to_path( "Cannot reserve disk capacity for a download without " f"ContentLength or max_bytes: {log_bucket}/{log_key}" ) - reserved_bytes = max_bytes - await ensure_capacity(reserved_bytes) - download_limit = reserved_bytes + grow_reservation_by_chunk = True + else: + await ensure_capacity(reserved_bytes) + download_limit = ( + reserved_bytes + if download_limit is None + else min(download_limit, reserved_bytes) + ) - async with aiofiles.open(temp_path, "wb") as f: + # Unbuffered writes keep prior chunks visible to capacity scans. + async with aiofiles.open(temp_path, "wb", buffering=0) as f: async for chunk in stream.iter_chunks(chunk_size=chunk_size): if not chunk: continue @@ -899,6 +906,8 @@ async def download_file_to_path( f"bytes_written={bytes_written} exceeds " f"max_bytes={download_limit}" ) + if grow_reservation_by_chunk and ensure_capacity is not None: + await ensure_capacity(len(chunk)) if hasher is not None: hasher.update(chunk) await f.write(chunk) From b929c7d141bf5d7e2d324d39a569cb3ed99562fd Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:54:25 -0700 Subject: [PATCH 153/161] fix(executor): harden cache accounting and writes --- tests/unit/test_registry_artifacts.py | 20 +++++ tests/unit/test_storage_blob.py | 76 +++++++++++++++++++ .../executor/registry_artifact_storage.py | 32 +++++++- tracecat/storage/blob.py | 12 ++- 4 files changed, 135 insertions(+), 5 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index bc672f34dd..23abf2ffe8 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -2259,6 +2259,26 @@ async def test_builtin_artifact_is_exempt_from_cache_accounting( class TestRegistryArtifactCacheEviction: """Tests for bounded eviction of registry artifact cache entries.""" + def test_snapshot_accounts_for_invalid_entry_children( + self, temp_cache_dir: Path + ) -> None: + """Files and symlinks directly under entries remain budget-visible.""" + cache = RegistryArtifactCache(temp_cache_dir) + before = cache._scan_cache_snapshot() + invalid_file = cache.entries_dir / "invalid-file" + invalid_file.write_bytes(b"x" * 65536) + invalid_link = cache.entries_dir / "invalid-link" + invalid_link.symlink_to(invalid_file) + allocation_unit = _filesystem_allocation_unit(temp_cache_dir) + invalid_bytes = _allocated_stat_size( + invalid_file.lstat(), allocation_unit=allocation_unit + ) + _allocated_stat_size(invalid_link.lstat(), allocation_unit=allocation_unit) + + after = cache._scan_cache_snapshot() + + assert after.entries == {} + assert after.structural_bytes >= before.structural_bytes + invalid_bytes + def test_delete_cache_path_reports_directory_failure(self, temp_cache_dir): """Physical deletion failures are observable instead of suppressed.""" entry_dir = temp_cache_dir / "entry" diff --git a/tests/unit/test_storage_blob.py b/tests/unit/test_storage_blob.py index 3e8f7dec01..d28a654ca8 100644 --- a/tests/unit/test_storage_blob.py +++ b/tests/unit/test_storage_blob.py @@ -2,6 +2,7 @@ import asyncio import hashlib +import threading from contextlib import asynccontextmanager from pathlib import Path from unittest.mock import AsyncMock, patch @@ -1039,6 +1040,81 @@ async def _fake_open_download_stream( assert not out.exists() assert not (tmp_path / "out.bin.part").exists() + @pytest.mark.anyio + async def test_download_file_to_path_rejoins_cancelled_write_before_cleanup( + self, tmp_path: Path, monkeypatch + ) -> None: + """Cancellation cannot unlink a partial file while its write is running.""" + write_started = threading.Event() + finish_write = threading.Event() + file_closed = threading.Event() + + class DummyStream: + async def iter_chunks(self, *, chunk_size: int): # noqa: ARG002 + yield b"partial" + + @asynccontextmanager + async def _fake_open_download_stream( + *, key: str, bucket: str, redact_log_identifiers: bool = False + ): # noqa: ARG001 + yield DummyStream(), 7 + + out = tmp_path / "out.bin" + partial_path = tmp_path / "out.bin.part" + original_open = Path.open + original_unlink = Path.unlink + + class BlockingFile: + def __init__(self, path: Path, *args, **kwargs) -> None: + self._file = original_open(path, *args, **kwargs) + + def __enter__(self): + return self + + def __exit__(self, *args) -> None: + self._file.close() + file_closed.set() + + def write(self, chunk: bytes) -> int: + write_started.set() + assert finish_write.wait(timeout=5) + return self._file.write(chunk) + + def blocking_open(path: Path, *args, **kwargs): + if path == partial_path: + return BlockingFile(path, *args, **kwargs) + return original_open(path, *args, **kwargs) + + def checked_unlink(path: Path, *args, **kwargs): + if path == partial_path: + assert file_closed.is_set() + return original_unlink(path, *args, **kwargs) + + monkeypatch.setattr( + "tracecat.storage.blob.open_download_stream", + _fake_open_download_stream, + ) + monkeypatch.setattr(Path, "open", blocking_open) + monkeypatch.setattr(Path, "unlink", checked_unlink) + + download = asyncio.create_task( + download_file_to_path(key="k", bucket="b", output_path=out) + ) + assert await asyncio.to_thread(write_started.wait, 2) + download.cancel() + await asyncio.sleep(0) + download.cancel() + await asyncio.sleep(0) + assert not download.done() + + finish_write.set() + with pytest.raises(asyncio.CancelledError): + await download + + assert file_closed.is_set() + assert not out.exists() + assert not partial_path.exists() + @pytest.mark.anyio async def test_download_file_to_path_defers_failed_partial_cleanup( self, tmp_path: Path, monkeypatch diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index a0a8f42ecc..b8f5aacc2c 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -955,7 +955,7 @@ async def _unmount(self, mount_dir: Path) -> bool: def _cache_structural_footprint(self, *, allocation_unit: int) -> int: """Measure cache roots and non-entry data exactly once.""" - return _directory_footprint( + cache_structure = _directory_footprint( self.cache_dir, allocation_unit=allocation_unit, pruned_directories=( @@ -964,6 +964,36 @@ def _cache_structural_footprint(self, *, allocation_unit: int) -> int: self.trash_dir, ), ) + return cache_structure + self._invalid_entry_footprint( + allocation_unit=allocation_unit + ) + + def _invalid_entry_footprint(self, *, allocation_unit: int) -> int: + """Measure non-directory children that cannot be cache entries.""" + _validate_cache_child_directory(self.entries_dir) + try: + entries = list(os.scandir(self.entries_dir)) + except FileNotFoundError: + return 0 + + total_bytes = 0 + seen_inodes: set[tuple[int, int]] = set() + for entry in entries: + try: + entry_stat = entry.stat(follow_symlinks=False) + except FileNotFoundError: + continue + if stat.S_ISDIR(entry_stat.st_mode): + continue + inode_key = (entry_stat.st_dev, entry_stat.st_ino) + if inode_key in seen_inodes: + continue + seen_inodes.add(inode_key) + total_bytes += _allocated_stat_size( + entry_stat, + allocation_unit=allocation_unit, + ) + return total_bytes def _prepare_budget_roots(self) -> None: """Create fixed roots before measuring or retiring an entry.""" diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index 065e1a2aa0..fb8c896f42 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import functools import hashlib import os import threading @@ -14,12 +15,12 @@ from typing import TYPE_CHECKING import aioboto3 -import aiofiles from aiobotocore.config import AioConfig from boto3.s3.transfer import TransferConfig from botocore.exceptions import BotoCoreError, ClientError from tracecat import config +from tracecat.concurrency import run_blocking_rejoin_on_cancel from tracecat.logger import logger if TYPE_CHECKING: @@ -894,8 +895,9 @@ async def download_file_to_path( else min(download_limit, reserved_bytes) ) - # Unbuffered writes keep prior chunks visible to capacity scans. - async with aiofiles.open(temp_path, "wb", buffering=0) as f: + # Unbuffered writes keep prior chunks visible to capacity scans. Each + # executor-backed write is rejoined before the file is closed or unlinked. + with temp_path.open("wb", buffering=0) as output_file: async for chunk in stream.iter_chunks(chunk_size=chunk_size): if not chunk: continue @@ -910,7 +912,9 @@ async def download_file_to_path( await ensure_capacity(len(chunk)) if hasher is not None: hasher.update(chunk) - await f.write(chunk) + await run_blocking_rejoin_on_cancel( + functools.partial(output_file.write, chunk) + ) if hasher is not None: actual_sha256 = hasher.hexdigest() From a7abbe2f0bb00f22002dfb14f71084cd29e333e5 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:26:38 -0700 Subject: [PATCH 154/161] fix(executor): contain artifact utility process groups --- tests/unit/test_registry_artifacts.py | 86 +++++++++++++++---- .../executor/registry_artifact_storage.py | 30 +------ tracecat/executor/registry_artifacts.py | 11 ++- 3 files changed, 81 insertions(+), 46 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 23abf2ffe8..0cbbb3b014 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -5,6 +5,7 @@ import asyncio import io import os +import signal import tarfile import tempfile import threading @@ -170,6 +171,7 @@ class _BlockingSubprocess: """Fake subprocess that blocks in communicate until it is cancelled.""" def __init__(self, *, block_wait: bool = False) -> None: + self.pid = 999_999_999 self.communicate_started = asyncio.Event() self.wait_started = asyncio.Event() self.release_wait = asyncio.Event() @@ -177,8 +179,12 @@ def __init__(self, *, block_wait: bool = False) -> None: self.returncode: int | None = None self._block_wait = block_wait - async def communicate(self) -> tuple[bytes, bytes]: + async def communicate( + self, + input: bytes | None = None, # noqa: A002 + ) -> tuple[bytes, bytes]: """Block until the task awaiting subprocess completion is cancelled.""" + del input self.communicate_started.set() await asyncio.Event().wait() return b"", b"" @@ -210,9 +216,17 @@ def returncode(self) -> int | None: """Return the wrapped subprocess exit status.""" return self.process.returncode - async def communicate(self) -> tuple[bytes, bytes]: + @property + def pid(self) -> int: + """Return the wrapped subprocess process-group identifier.""" + return self.process.pid + + async def communicate( + self, + input: bytes | None = None, # noqa: A002 + ) -> tuple[bytes, bytes]: """Wait for the wrapped subprocess and collect its output.""" - return await self.process.communicate() + return await self.process.communicate(input=input) def kill(self) -> None: """Kill the wrapped subprocess and record the signal.""" @@ -839,14 +853,21 @@ async def test_mount_squashfs_uses_hardened_read_only_options( image_path.write_bytes(b"squashfs") target_dir.mkdir() process = AsyncMock() + process.pid = 1234 process.communicate.return_value = (b"", b"") process.returncode = 0 - with patch( - "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", - new_callable=AsyncMock, - return_value=process, - ) as create_subprocess_exec: + with ( + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ) as create_subprocess_exec, + patch( + "tracecat.sandbox.utils.terminate_process_group", + new_callable=AsyncMock, + ) as terminate_process_group, + ): await artifact.mount(ctx, image_path) create_subprocess_exec.assert_awaited_once_with( @@ -859,7 +880,9 @@ async def test_mount_squashfs_uses_hardened_read_only_options( str(target_dir), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + start_new_session=True, ) + terminate_process_group.assert_awaited_once_with(process) @pytest.mark.anyio async def test_repeatedly_cancelled_mount_reaps_subprocess(self, temp_cache_dir): @@ -878,10 +901,13 @@ async def test_repeatedly_cancelled_mount_reaps_subprocess(self, temp_cache_dir) target_dir.mkdir() process = _BlockingSubprocess(block_wait=True) - with patch( - "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", - new_callable=AsyncMock, - return_value=process, + with ( + patch( + "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ) as create_subprocess_exec, + patch("tracecat.sandbox.utils.os.killpg") as kill_group, ): mounting = asyncio.create_task( artifact._mount_image(image_path, target_dir) @@ -898,6 +924,10 @@ async def test_repeatedly_cancelled_mount_reaps_subprocess(self, temp_cache_dir) await mounting assert process.cleanup_calls == ["kill", "wait"] + await_args = create_subprocess_exec.await_args + assert await_args is not None + assert await_args.kwargs["start_new_session"] is True + kill_group.assert_called_once_with(process.pid, signal.SIGKILL) assert target_dir.is_dir() assert not target_dir.is_mount() @@ -925,12 +955,14 @@ async def test_cancelled_squashfs_extract_kills_and_reaps_subprocess( async def create_sleep_subprocess( *args: object, **kwargs: object ) -> _CapturedSubprocess: - del args, kwargs + del args + assert kwargs["start_new_session"] is True process = await real_create_subprocess_exec( "/bin/sleep", "30", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + start_new_session=True, ) captured = _CapturedSubprocess(process) captured_processes.append(captured) @@ -1044,9 +1076,9 @@ async def test_squashfs_failures_sanitize_subprocess_output( "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process, - ), + ) as create_subprocess_exec, patch( - "tracecat.executor.registry_artifacts.communicate_rejoin_on_cancel", + "tracecat.executor.registry_artifacts.communicate_process_group", new_callable=AsyncMock, return_value=(stdout, stderr), ), @@ -1061,6 +1093,9 @@ async def test_squashfs_failures_sanitize_subprocess_output( ) assert str(raised.value) == "Registry artifact extraction failed" + await_args = create_subprocess_exec.await_args + assert await_args is not None + assert await_args.kwargs["start_new_session"] is True assert sensitive_output.decode() not in str(raised.value) assert raised.value.__cause__ is None @@ -2165,6 +2200,7 @@ async def test_lease_is_never_admitted_across_an_in_flight_eviction( umount_process.returncode = 0 async def mock_umount(*args, **kwargs): + assert kwargs["start_new_session"] is True umount_started.set() await finish_umount.wait() mounted.discard(paths.squashfs_mount_dir) @@ -2197,6 +2233,10 @@ async def take_lease() -> None: "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", side_effect=mock_umount, ), + patch( + "tracecat.sandbox.utils.terminate_process_group", + new_callable=AsyncMock, + ), patch.object(SquashfsArtifact, "mount", mock_mount), ): eviction = asyncio.create_task(cache._evict_entry(cache_key)) @@ -3086,6 +3126,7 @@ async def test_final_lease_release_unmounts_and_retains_image(self, temp_cache_d process.returncode = 0 async def mock_umount(*args, **kwargs): + assert kwargs["start_new_session"] is True mounted.discard(paths.squashfs_mount_dir) return process @@ -3100,6 +3141,10 @@ async def mock_umount(*args, **kwargs): "create_subprocess_exec", side_effect=mock_umount, ), + patch( + "tracecat.sandbox.utils.terminate_process_group", + new_callable=AsyncMock, + ), ): async with cache.lease([artifact_uri]) as registry_paths: assert registry_paths == [paths.squashfs_mount_dir] @@ -3273,6 +3318,7 @@ async def test_eviction_unmounts_before_deleting_the_image(self, temp_cache_dir) process.returncode = 0 async def mock_umount(*args, **kwargs): + assert kwargs["start_new_session"] is True image_present_at_umount.append(paths.squashfs_image_path.exists()) mounted.discard(paths.squashfs_mount_dir) return process @@ -3287,6 +3333,10 @@ async def mock_umount(*args, **kwargs): "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", side_effect=mock_umount, ) as create_subprocess_exec, + patch( + "tracecat.sandbox.utils.terminate_process_group", + new_callable=AsyncMock, + ), ): evicted = await cache._evict_entry("mounted") @@ -3322,6 +3372,7 @@ async def test_repeatedly_cancelled_unmount_reaps_before_releasing_key_lock( async def mock_umount(*args, **kwargs): nonlocal unmount_attempts + assert kwargs["start_new_session"] is True unmount_attempts += 1 if unmount_attempts == 1: return blocked_process @@ -3338,6 +3389,7 @@ async def mock_umount(*args, **kwargs): "tracecat.executor.registry_artifacts.asyncio.create_subprocess_exec", side_effect=mock_umount, ), + patch("tracecat.sandbox.utils.os.killpg") as kill_group, ): eviction = asyncio.create_task(cache._evict_entry(cache_key)) await blocked_process.communicate_started.wait() @@ -3352,6 +3404,10 @@ async def mock_umount(*args, **kwargs): await eviction assert blocked_process.cleanup_calls == ["kill", "wait"] + kill_group.assert_called_once_with( + blocked_process.pid, + signal.SIGKILL, + ) async with cache.lease([artifact_uri]) as registry_paths: assert registry_paths == [paths.squashfs_mount_dir] assert registry_paths[0].is_dir() diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index b8f5aacc2c..d39a431065 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio -import contextlib import functools import os import re @@ -18,7 +17,6 @@ from tracecat import config from tracecat.concurrency import ( - drain_future_through_cancellation, rejoin_future_on_cancel, run_blocking_rejoin_on_cancel, ) @@ -30,6 +28,7 @@ plan_registry_artifact_evictions, ) from tracecat.logger import logger +from tracecat.sandbox.utils import communicate_process_group __all__ = ( "BASE_PYTHONPATH_DIR_NAME", @@ -46,7 +45,6 @@ "RegistryArtifactPaths", "RegistryArtifactRuntimeState", "allocated_size_bound", - "communicate_rejoin_on_cancel", "ensure_cache_entry_directory", "ensure_real_directory", "is_reusable_cache_directory", @@ -164,29 +162,6 @@ def can_mount_squashfs(self) -> bool: ) -async def _kill_and_reap_subprocess(process: asyncio.subprocess.Process) -> None: - """Kill a subprocess and wait until its child state is reaped.""" - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - - -async def communicate_rejoin_on_cancel( - process: asyncio.subprocess.Process, -) -> tuple[bytes, bytes]: - """Communicate without allowing cancellation to abandon child cleanup.""" - try: - stdout, stderr = await process.communicate() - except asyncio.CancelledError: - reaper = asyncio.ensure_future(_kill_and_reap_subprocess(process)) - await drain_future_through_cancellation(reaper) - raise - - if stdout is None or stderr is None: - raise RuntimeError("Captured subprocess output is required") - return stdout, stderr - - def is_reusable_cache_directory(path: Path) -> bool: """Return whether a cache path is a real directory, never a symlink.""" try: @@ -942,8 +917,9 @@ async def _unmount(self, mount_dir: Path) -> bool: str(mount_dir), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + start_new_session=True, ) - stdout, stderr = await communicate_rejoin_on_cancel(proc) + stdout, stderr = await communicate_process_group(proc) if proc.returncode == 0 or not registry_artifact_mounts.is_mount(mount_dir): return True logger.warning( diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 2f4ed450e0..2befc07ef7 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -36,7 +36,6 @@ RegistryArtifactEviction, RegistryArtifactMaterializationContext, allocated_size_bound, - communicate_rejoin_on_cancel, ensure_cache_entry_directory, ensure_real_directory, is_reusable_cache_directory, @@ -49,6 +48,7 @@ from tracecat.logger import logger from tracecat.registry.artifact_keys import parse_s3_uri from tracecat.registry.constants import DEFAULT_REGISTRY_ORIGIN +from tracecat.sandbox.utils import communicate_process_group from tracecat.storage import blob __all__ = ( @@ -426,8 +426,9 @@ async def _mount_image(self, image_path: Path, target_dir: Path) -> None: str(target_dir), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + start_new_session=True, ) - stdout, stderr = await communicate_rejoin_on_cancel(proc) + stdout, stderr = await communicate_process_group(proc) if proc.returncode == 0 or registry_artifact_mounts.is_mount(target_dir): return @@ -454,8 +455,9 @@ async def _extract_image(self, image_path: Path, target_dir: Path) -> None: str(image_path), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + start_new_session=True, ) - await communicate_rejoin_on_cancel(proc) + await communicate_process_group(proc) if proc.returncode == 0: return @@ -479,8 +481,9 @@ async def _squashfs_extracted_size( str(image_path), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + start_new_session=True, ) - stdout, _ = await communicate_rejoin_on_cancel(proc) + stdout, _ = await communicate_process_group(proc) if proc.returncode != 0: raise RegistryArtifactExtractionError() From 001dbea4dc430fc74ca690a24b0e7895a0e6e3ce Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:29:57 -0700 Subject: [PATCH 155/161] fix(executor): avoid mutable empty cache leases --- tests/unit/test_action_runner.py | 12 ++++-- tests/unit/test_registry_artifacts.py | 38 +++++++++++++------ .../executor/registry_artifact_storage.py | 13 ------- tracecat/executor/registry_artifacts.py | 6 +-- 4 files changed, 37 insertions(+), 32 deletions(-) diff --git a/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index d3845ead79..6bd029b846 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -148,15 +148,19 @@ async def communicate( return communication @pytest.mark.anyio - async def test_lease_without_artifacts_yields_base_pythonpath(self, temp_cache_dir): - """Test that the base PYTHONPATH directory is used without artifacts.""" + async def test_lease_without_artifacts_yields_no_registry_paths( + self, temp_cache_dir + ): + """An artifact-free action receives no extra registry import paths.""" runner = ActionRunner(cache_dir=temp_cache_dir) async with runner.registry_artifacts.lease(None) as registry_paths: - assert registry_paths == [temp_cache_dir / "base"] + assert registry_paths == [] async with runner.registry_artifacts.lease([]) as registry_paths: - assert registry_paths == [temp_cache_dir / "base"] + assert registry_paths == [] + + assert not (temp_cache_dir / "base").exists() @pytest.mark.anyio async def test_execute_action_timeout( diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 0cbbb3b014..2637c883a5 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1467,7 +1467,7 @@ def use_cache_from_another_thread() -> None: # A rejected caller must not poison the owning loop's cache. async with cache.lease(None) as registry_paths: - assert registry_paths == [temp_cache_dir / "base"] + assert registry_paths == [] def test_touch_entry_refreshes_tarball_root_mtime(self, temp_cache_dir): """Touching a tarball-only entry persists its restart-safe recency.""" @@ -1670,22 +1670,36 @@ async def take_lease() -> None: cache._release_lease(cache_key) @pytest.mark.anyio - async def test_lease_without_uris_returns_base_pythonpath_dir(self, temp_cache_dir): - """No artifact URIs still yields the base PYTHONPATH directory.""" + async def test_mutable_lease_without_uris_returns_no_paths(self, temp_cache_dir): + """An empty mutable lease exposes no unaccounted cache directory.""" cache = RegistryArtifactCache(temp_cache_dir) + cache._budget_dirty = False - with patch.object( - cache, - "ensure_swept", - new_callable=AsyncMock, - side_effect=AssertionError("cache-free leases must not sweep"), - ) as ensure_swept: - async with cache.lease(None) as registry_paths: - assert registry_paths == [temp_cache_dir / "base"] - assert registry_paths[0].is_dir() + with ( + patch.object( + cache, + "ensure_swept", + new_callable=AsyncMock, + side_effect=AssertionError("cache-free leases must not sweep"), + ) as ensure_swept, + patch.object( + cache, + "_converge_cache_budget", + new_callable=AsyncMock, + side_effect=AssertionError("empty leases must not converge"), + ) as converge_cache_budget, + ): + async with cache.lease( + None, + paths_may_be_modified=True, + ) as registry_paths: + assert registry_paths == [] ensure_swept.assert_not_awaited() + converge_cache_budget.assert_not_awaited() + assert cache._budget_dirty is False assert cache._runtime == {} + assert not (temp_cache_dir / "base").exists() @pytest.mark.anyio async def test_lease_preserves_uri_order(self, temp_cache_dir): diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index d39a431065..d600ff6094 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -31,7 +31,6 @@ from tracecat.sandbox.utils import communicate_process_group __all__ = ( - "BASE_PYTHONPATH_DIR_NAME", "CACHE_ENTRIES_DIR_NAME", "CACHE_STAGING_DIR_NAME", "CACHE_TRASH_DIR_NAME", @@ -55,9 +54,6 @@ "validate_cache_entry_path", ) -BASE_PYTHONPATH_DIR_NAME = "base" -"""Cache subdirectory used when no registry artifact is requested.""" - CACHE_ENTRIES_DIR_NAME = "entries" """Directory containing one atomic subdirectory per cache key.""" @@ -577,15 +573,6 @@ async def ensure_capacity(additional_bytes: int) -> None: ensure_capacity=ensure_capacity, ) - def _base_pythonpath_dir(self) -> Path: - """Return the base PYTHONPATH directory for an artifact-free action.""" - base_dir = self.cache_dir / BASE_PYTHONPATH_DIR_NAME - _validate_cache_root(self.cache_dir) - ensure_real_directory(self.cache_dir) - _validate_cache_child_directory(base_dir) - ensure_real_directory(base_dir) - return base_dir - def _acquire_lease(self, cache_key: str) -> None: """Pin a cache entry against eviction and mark it recently used.""" runtime = self._runtime_for(cache_key) diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 2befc07ef7..496607b030 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -1053,7 +1053,7 @@ async def lease( Args: artifact_uris: Registry artifact URIs in deterministic PYTHONPATH - order, or None to use the base PYTHONPATH directory. + order. Empty input requests no additional import paths. paths_may_be_modified: Whether the consumer can write to returned paths. Mutable leases re-arm budget convergence after use. @@ -1061,8 +1061,8 @@ async def lease( Importable Python paths for the requested artifacts. """ if not artifact_uris: - logger.info("No registry artifact URIs provided, using base PYTHONPATH") - yield [self._base_pythonpath_dir()] + logger.info("No registry artifact URIs provided") + yield [] return if any(_is_cache_entry_uri(uri) for uri in artifact_uris): From 4f4c2d29434252c81a91b172e8a92c9765630060 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:51:30 -0700 Subject: [PATCH 156/161] fix(executor): expose registry cache byte limit --- deployments/fargate/main.tf | 1 + deployments/fargate/modules/ecs/locals.tf | 12 +++++++----- deployments/fargate/modules/ecs/variables.tf | 11 +++++++++++ deployments/fargate/variables.tf | 11 +++++++++++ docker-compose.dev.yml | 2 ++ docker-compose.local.yml | 2 ++ docker-compose.yml | 2 ++ 7 files changed, 36 insertions(+), 5 deletions(-) diff --git a/deployments/fargate/main.tf b/deployments/fargate/main.tf index 420b1dcffa..35ae0255d0 100644 --- a/deployments/fargate/main.tf +++ b/deployments/fargate/main.tf @@ -155,6 +155,7 @@ module "ecs" { executor_desired_count = var.executor_desired_count executor_client_timeout = var.executor_client_timeout executor_queue = var.executor_queue + executor_registry_cache_max_bytes = var.executor_registry_cache_max_bytes executor_max_concurrent_activities = var.executor_max_concurrent_activities executor_threadpool_max_workers = var.executor_threadpool_max_workers executor_for_each_max_concurrency = var.executor_for_each_max_concurrency diff --git a/deployments/fargate/modules/ecs/locals.tf b/deployments/fargate/modules/ecs/locals.tf index 13f4c45405..a0b2cbd46e 100644 --- a/deployments/fargate/modules/ecs/locals.tf +++ b/deployments/fargate/modules/ecs/locals.tf @@ -164,11 +164,12 @@ locals { local.tracecat_db_configs, local.tracecat_db_configs_executor, { - TRACECAT__API_URL = local.internal_api_url - TRACECAT__DB_ENDPOINT = local.core_db_hostname - TRACECAT__SERVICE_NAME = "executor" - TRACECAT__EXECUTOR_BACKEND = "direct" - TRACECAT__EXECUTOR_QUEUE = var.executor_queue + TRACECAT__API_URL = local.internal_api_url + TRACECAT__DB_ENDPOINT = local.core_db_hostname + TRACECAT__SERVICE_NAME = "executor" + TRACECAT__EXECUTOR_BACKEND = "direct" + TRACECAT__EXECUTOR_QUEUE = var.executor_queue + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES = var.executor_registry_cache_max_bytes # Executor concurrency tuning (see tracecat/executor/worker.py and tracecat/executor/service.py) TRACECAT__EXECUTOR_MAX_CONCURRENT_ACTIVITIES = var.executor_max_concurrent_activities TRACECAT__EXECUTOR_THREADPOOL_MAX_WORKERS = var.executor_threadpool_max_workers @@ -199,6 +200,7 @@ locals { TRACECAT__AGENT_QUEUE = var.agent_queue TRACECAT__AGENT_EXECUTOR_QUEUE = var.agent_executor_queue TRACECAT__EXECUTOR_QUEUE = var.executor_queue + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES = var.executor_registry_cache_max_bytes TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES = var.agent_executor_max_concurrent_activities TRACECAT__EXECUTOR_CLIENT_TIMEOUT = var.executor_client_timeout TRACECAT__LLM_PROXY_READ_TIMEOUT = var.llm_proxy_read_timeout diff --git a/deployments/fargate/modules/ecs/variables.tf b/deployments/fargate/modules/ecs/variables.tf index b2633a470f..727f5bab02 100644 --- a/deployments/fargate/modules/ecs/variables.tf +++ b/deployments/fargate/modules/ecs/variables.tf @@ -622,6 +622,17 @@ variable "executor_queue" { default = "shared-action-queue" } +variable "executor_registry_cache_max_bytes" { + type = number + description = "Maximum executor-local registry artifact cache size in bytes (TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES). Set to 0 to disable size-based limits." + default = 10737418240 + + validation { + condition = var.executor_registry_cache_max_bytes >= 0 && floor(var.executor_registry_cache_max_bytes) == var.executor_registry_cache_max_bytes + error_message = "executor_registry_cache_max_bytes must be a non-negative integer." + } +} + variable "executor_max_concurrent_activities" { type = number description = "Max concurrent activities per executor task (TRACECAT__EXECUTOR_MAX_CONCURRENT_ACTIVITIES)." diff --git a/deployments/fargate/variables.tf b/deployments/fargate/variables.tf index f8295e9832..07c47cf6bb 100644 --- a/deployments/fargate/variables.tf +++ b/deployments/fargate/variables.tf @@ -590,6 +590,17 @@ variable "executor_queue" { default = "shared-action-queue" } +variable "executor_registry_cache_max_bytes" { + type = number + description = "Maximum executor-local registry artifact cache size in bytes (TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES). Set to 0 to disable size-based limits." + default = 10737418240 + + validation { + condition = var.executor_registry_cache_max_bytes >= 0 && floor(var.executor_registry_cache_max_bytes) == var.executor_registry_cache_max_bytes + error_message = "executor_registry_cache_max_bytes must be a non-negative integer." + } +} + variable "executor_max_concurrent_activities" { type = number description = "Max concurrent activities per executor task (TRACECAT__EXECUTOR_MAX_CONCURRENT_ACTIVITIES)." diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index 25cd8c94da..2151cfb410 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -208,6 +208,7 @@ services: TRACECAT__COLLECTION_MANIFESTS_ENABLED: ${TRACECAT__COLLECTION_MANIFESTS_ENABLED:-true} TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES: ${TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES} # Registry + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES:-10737418240} TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false} # Local registry TRACECAT__LOCAL_REPOSITORY_PATH: ${TRACECAT__LOCAL_REPOSITORY_PATH} @@ -395,6 +396,7 @@ services: TRACECAT__COLLECTION_MANIFESTS_ENABLED: ${TRACECAT__COLLECTION_MANIFESTS_ENABLED:-true} TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES: ${TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES} # Registry + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES:-10737418240} TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false} # Local registry TRACECAT__LOCAL_REPOSITORY_PATH: ${TRACECAT__LOCAL_REPOSITORY_PATH} diff --git a/docker-compose.local.yml b/docker-compose.local.yml index 3fbc99311d..8519002935 100644 --- a/docker-compose.local.yml +++ b/docker-compose.local.yml @@ -224,6 +224,7 @@ services: TRACECAT__COLLECTION_MANIFESTS_ENABLED: ${TRACECAT__COLLECTION_MANIFESTS_ENABLED:-true} TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES: ${TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES} # Registry + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES:-10737418240} TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false} # Local registry TRACECAT__LOCAL_REPOSITORY_PATH: ${TRACECAT__LOCAL_REPOSITORY_PATH} @@ -417,6 +418,7 @@ services: TRACECAT__COLLECTION_MANIFESTS_ENABLED: ${TRACECAT__COLLECTION_MANIFESTS_ENABLED:-true} TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES: ${TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES} # Registry + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES:-10737418240} TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false} # Local registry TRACECAT__LOCAL_REPOSITORY_PATH: ${TRACECAT__LOCAL_REPOSITORY_PATH} diff --git a/docker-compose.yml b/docker-compose.yml index b5dd69b9dc..1026fd5e4f 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -216,6 +216,7 @@ services: TRACECAT__COLLECTION_MANIFESTS_ENABLED: ${TRACECAT__COLLECTION_MANIFESTS_ENABLED:-true} TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES: ${TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES} # Registry + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES:-10737418240} TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false} # Local registry TRACECAT__LOCAL_REPOSITORY_PATH: ${TRACECAT__LOCAL_REPOSITORY_PATH} @@ -408,6 +409,7 @@ services: TRACECAT__COLLECTION_MANIFESTS_ENABLED: ${TRACECAT__COLLECTION_MANIFESTS_ENABLED:-true} TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES: ${TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES} # Registry + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES:-10737418240} TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false} # Local registry TRACECAT__LOCAL_REPOSITORY_PATH: ${TRACECAT__LOCAL_REPOSITORY_PATH} From 88c22ea95e5ef77e6387a3e10e4fb1273df6165f Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 4 Aug 2026 10:48:06 -0700 Subject: [PATCH 157/161] fix(executor): expose registry cache entry limit --- deployments/fargate/main.tf | 1 + deployments/fargate/modules/ecs/locals.tf | 14 ++++++++------ deployments/fargate/modules/ecs/variables.tf | 11 +++++++++++ deployments/fargate/variables.tf | 11 +++++++++++ docker-compose.dev.yml | 2 ++ docker-compose.local.yml | 2 ++ docker-compose.yml | 2 ++ 7 files changed, 37 insertions(+), 6 deletions(-) diff --git a/deployments/fargate/main.tf b/deployments/fargate/main.tf index 35ae0255d0..ce111a2a61 100644 --- a/deployments/fargate/main.tf +++ b/deployments/fargate/main.tf @@ -155,6 +155,7 @@ module "ecs" { executor_desired_count = var.executor_desired_count executor_client_timeout = var.executor_client_timeout executor_queue = var.executor_queue + executor_registry_cache_max_entries = var.executor_registry_cache_max_entries executor_registry_cache_max_bytes = var.executor_registry_cache_max_bytes executor_max_concurrent_activities = var.executor_max_concurrent_activities executor_threadpool_max_workers = var.executor_threadpool_max_workers diff --git a/deployments/fargate/modules/ecs/locals.tf b/deployments/fargate/modules/ecs/locals.tf index a0b2cbd46e..9496793c3a 100644 --- a/deployments/fargate/modules/ecs/locals.tf +++ b/deployments/fargate/modules/ecs/locals.tf @@ -164,12 +164,13 @@ locals { local.tracecat_db_configs, local.tracecat_db_configs_executor, { - TRACECAT__API_URL = local.internal_api_url - TRACECAT__DB_ENDPOINT = local.core_db_hostname - TRACECAT__SERVICE_NAME = "executor" - TRACECAT__EXECUTOR_BACKEND = "direct" - TRACECAT__EXECUTOR_QUEUE = var.executor_queue - TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES = var.executor_registry_cache_max_bytes + TRACECAT__API_URL = local.internal_api_url + TRACECAT__DB_ENDPOINT = local.core_db_hostname + TRACECAT__SERVICE_NAME = "executor" + TRACECAT__EXECUTOR_BACKEND = "direct" + TRACECAT__EXECUTOR_QUEUE = var.executor_queue + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES = var.executor_registry_cache_max_entries + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES = var.executor_registry_cache_max_bytes # Executor concurrency tuning (see tracecat/executor/worker.py and tracecat/executor/service.py) TRACECAT__EXECUTOR_MAX_CONCURRENT_ACTIVITIES = var.executor_max_concurrent_activities TRACECAT__EXECUTOR_THREADPOOL_MAX_WORKERS = var.executor_threadpool_max_workers @@ -200,6 +201,7 @@ locals { TRACECAT__AGENT_QUEUE = var.agent_queue TRACECAT__AGENT_EXECUTOR_QUEUE = var.agent_executor_queue TRACECAT__EXECUTOR_QUEUE = var.executor_queue + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES = var.executor_registry_cache_max_entries TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES = var.executor_registry_cache_max_bytes TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES = var.agent_executor_max_concurrent_activities TRACECAT__EXECUTOR_CLIENT_TIMEOUT = var.executor_client_timeout diff --git a/deployments/fargate/modules/ecs/variables.tf b/deployments/fargate/modules/ecs/variables.tf index 727f5bab02..6254d08a65 100644 --- a/deployments/fargate/modules/ecs/variables.tf +++ b/deployments/fargate/modules/ecs/variables.tf @@ -622,6 +622,17 @@ variable "executor_queue" { default = "shared-action-queue" } +variable "executor_registry_cache_max_entries" { + type = number + description = "Maximum number of entries in the executor-local registry artifact cache (TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES). Set to 0 to disable entry-count eviction." + default = 64 + + validation { + condition = var.executor_registry_cache_max_entries >= 0 && floor(var.executor_registry_cache_max_entries) == var.executor_registry_cache_max_entries + error_message = "executor_registry_cache_max_entries must be a non-negative integer." + } +} + variable "executor_registry_cache_max_bytes" { type = number description = "Maximum executor-local registry artifact cache size in bytes (TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES). Set to 0 to disable size-based limits." diff --git a/deployments/fargate/variables.tf b/deployments/fargate/variables.tf index 07c47cf6bb..275418bd20 100644 --- a/deployments/fargate/variables.tf +++ b/deployments/fargate/variables.tf @@ -590,6 +590,17 @@ variable "executor_queue" { default = "shared-action-queue" } +variable "executor_registry_cache_max_entries" { + type = number + description = "Maximum number of entries in the executor-local registry artifact cache (TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES). Set to 0 to disable entry-count eviction." + default = 64 + + validation { + condition = var.executor_registry_cache_max_entries >= 0 && floor(var.executor_registry_cache_max_entries) == var.executor_registry_cache_max_entries + error_message = "executor_registry_cache_max_entries must be a non-negative integer." + } +} + variable "executor_registry_cache_max_bytes" { type = number description = "Maximum executor-local registry artifact cache size in bytes (TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES). Set to 0 to disable size-based limits." diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index 2151cfb410..fb47bc0ae0 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -208,6 +208,7 @@ services: TRACECAT__COLLECTION_MANIFESTS_ENABLED: ${TRACECAT__COLLECTION_MANIFESTS_ENABLED:-true} TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES: ${TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES} # Registry + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES:-64} TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES:-10737418240} TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false} # Local registry @@ -396,6 +397,7 @@ services: TRACECAT__COLLECTION_MANIFESTS_ENABLED: ${TRACECAT__COLLECTION_MANIFESTS_ENABLED:-true} TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES: ${TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES} # Registry + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES:-64} TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES:-10737418240} TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false} # Local registry diff --git a/docker-compose.local.yml b/docker-compose.local.yml index 8519002935..b025cf9bfd 100644 --- a/docker-compose.local.yml +++ b/docker-compose.local.yml @@ -224,6 +224,7 @@ services: TRACECAT__COLLECTION_MANIFESTS_ENABLED: ${TRACECAT__COLLECTION_MANIFESTS_ENABLED:-true} TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES: ${TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES} # Registry + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES:-64} TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES:-10737418240} TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false} # Local registry @@ -418,6 +419,7 @@ services: TRACECAT__COLLECTION_MANIFESTS_ENABLED: ${TRACECAT__COLLECTION_MANIFESTS_ENABLED:-true} TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES: ${TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES} # Registry + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES:-64} TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES:-10737418240} TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false} # Local registry diff --git a/docker-compose.yml b/docker-compose.yml index 1026fd5e4f..75355e9d5a 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -216,6 +216,7 @@ services: TRACECAT__COLLECTION_MANIFESTS_ENABLED: ${TRACECAT__COLLECTION_MANIFESTS_ENABLED:-true} TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES: ${TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES} # Registry + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES:-64} TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES:-10737418240} TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false} # Local registry @@ -409,6 +410,7 @@ services: TRACECAT__COLLECTION_MANIFESTS_ENABLED: ${TRACECAT__COLLECTION_MANIFESTS_ENABLED:-true} TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES: ${TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES} # Registry + TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_ENTRIES:-64} TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES: ${TRACECAT__EXECUTOR_REGISTRY_CACHE_MAX_BYTES:-10737418240} TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false} # Local registry From 01185be301e0eca89dd0047dc6e28953983b5a5b Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 4 Aug 2026 10:51:13 -0700 Subject: [PATCH 158/161] fix(sandbox): supervise unsafe pid fallback --- tests/unit/test_unsafe_pid_executor.py | 131 ++++++++++++++++++++++-- tracecat/sandbox/unsafe_pid_executor.py | 39 ++++++- 2 files changed, 155 insertions(+), 15 deletions(-) diff --git a/tests/unit/test_unsafe_pid_executor.py b/tests/unit/test_unsafe_pid_executor.py index 00d5f46834..1a962297a1 100644 --- a/tests/unit/test_unsafe_pid_executor.py +++ b/tests/unit/test_unsafe_pid_executor.py @@ -5,6 +5,7 @@ import logging import os import signal +import sys from pathlib import Path from unittest.mock import AsyncMock, patch @@ -70,9 +71,54 @@ def main(pid_file, wait): """ +def _detached_background_process_script() -> str: + return """ +import subprocess +import sys +from pathlib import Path + +def main(pid_file): + child = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(30)"], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + Path(pid_file).write_text(str(child.pid)) + return child.pid +""" + + class TestUnsafePidExecutor: @pytest.fixture - def executor(self, tmp_path) -> UnsafePidExecutor: + def executor( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> UnsafePidExecutor: + executor = UnsafePidExecutor(cache_dir=str(tmp_path)) + if sys.platform != "linux": + # Production execution is Linux-only. Keep wrapper behavior covered + # on developer machines without adding a non-Linux runtime fallback. + async def build_test_execution_cmd( + python_path: str, + wrapper_path: Path, + ) -> unsafe_pid_executor._ExecutionCommand: + return unsafe_pid_executor._ExecutionCommand( + argv=[python_path, str(wrapper_path)], + supervised=False, + ) + + monkeypatch.setattr( + executor, + "_build_execution_cmd", + build_test_execution_cmd, + ) + return executor + + @pytest.fixture + def command_executor(self, tmp_path: Path) -> UnsafePidExecutor: return UnsafePidExecutor(cache_dir=str(tmp_path)) @pytest.mark.parametrize( @@ -137,7 +183,9 @@ def process_disappeared( @pytest.mark.anyio async def test_build_execution_cmd_with_pid_namespace( - self, executor: UnsafePidExecutor, monkeypatch: pytest.MonkeyPatch + self, + command_executor: UnsafePidExecutor, + monkeypatch: pytest.MonkeyPatch, ) -> None: async def pid_namespace_available() -> bool: return True @@ -146,15 +194,42 @@ async def pid_namespace_available() -> bool: "tracecat.sandbox.unsafe_pid_executor.pid_namespace_available", pid_namespace_available, ) - cmd = await executor._build_execution_cmd( - "python3", executor.cache_dir / "wrapper.py" + command = await command_executor._build_execution_cmd( + "python3", command_executor.cache_dir / "wrapper.py" + ) + assert command.argv[:4] == ["unshare", "--pid", "--fork", "--kill-child"] + assert command.supervised is False + + @pytest.mark.anyio + async def test_build_execution_cmd_without_pid_namespace_uses_supervisor( + self, + command_executor: UnsafePidExecutor, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + async def pid_namespace_unavailable() -> bool: + return False + + monkeypatch.setattr( + unsafe_pid_executor, + "pid_namespace_available", + pid_namespace_unavailable, + ) + + wrapper_path = command_executor.cache_dir / "wrapper.py" + command = await command_executor._build_execution_cmd( + "python3", + wrapper_path, ) - assert cmd[:4] == ["unshare", "--pid", "--fork", "--kill-child"] + + assert command.argv[:2] == [sys.executable, "-I"] + assert Path(command.argv[2]).name == "process_supervisor.py" + assert command.argv[-2:] == ["python3", str(wrapper_path)] + assert command.supervised is True @pytest.mark.anyio async def test_pid_isolation_warning_logged_once( self, - executor: UnsafePidExecutor, + command_executor: UnsafePidExecutor, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, ) -> None: @@ -178,11 +253,11 @@ def warning(self, message: str, **kwargs: object) -> None: monkeypatch.setattr("tracecat.sandbox.unsafe_pid_executor.logger", FakeLogger()) caplog.set_level(logging.WARNING, logger="tracecat.sandbox.unsafe_pid_executor") - await executor._build_execution_cmd( - "python3", executor.cache_dir / "wrapper.py" + await command_executor._build_execution_cmd( + "python3", command_executor.cache_dir / "wrapper.py" ) - await executor._build_execution_cmd( - "python3", executor.cache_dir / "wrapper.py" + await command_executor._build_execution_cmd( + "python3", command_executor.cache_dir / "wrapper.py" ) warnings = [ @@ -303,6 +378,42 @@ async def test_execute_kills_background_descendants_holding_output_pipes( with contextlib.suppress(ProcessLookupError): os.kill(child_pid, signal.SIGKILL) + @pytest.mark.skipif( + sys.platform != "linux", + reason="Detached fallback containment uses the Linux subreaper supervisor", + ) + @pytest.mark.anyio + async def test_execute_without_pid_namespace_reaps_detached_descendant( + self, + executor: UnsafePidExecutor, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + async def pid_namespace_unavailable() -> bool: + return False + + monkeypatch.setattr( + unsafe_pid_executor, + "pid_namespace_available", + pid_namespace_unavailable, + ) + pid_file = tmp_path / "detached-child.pid" + result = await asyncio.wait_for( + executor.execute( + script=_detached_background_process_script(), + inputs={"pid_file": str(pid_file)}, + ), + timeout=5, + ) + + assert result.success + child_pid = int(pid_file.read_text()) + try: + await _wait_for_process_exit(child_pid) + finally: + with contextlib.suppress(ProcessLookupError): + os.kill(child_pid, signal.SIGKILL) + @pytest.mark.anyio async def test_execute_kills_descendants_on_timeout( self, executor: UnsafePidExecutor, tmp_path: Path diff --git a/tracecat/sandbox/unsafe_pid_executor.py b/tracecat/sandbox/unsafe_pid_executor.py index 2be02133cd..eb81c85d31 100644 --- a/tracecat/sandbox/unsafe_pid_executor.py +++ b/tracecat/sandbox/unsafe_pid_executor.py @@ -11,8 +11,10 @@ import logging import os import shutil +import sys import tempfile import time +from dataclasses import dataclass from pathlib import Path from typing import Any @@ -33,6 +35,7 @@ communicate_process_group, pid_namespace_available, pid_namespace_probe_error, + terminate_supervised_process, ) module_logger = logging.getLogger(__name__) @@ -240,6 +243,14 @@ def main(): ''' +@dataclass(frozen=True, slots=True) +class _ExecutionCommand: + """Command metadata for one unsafe executor subprocess.""" + + argv: list[str] + supervised: bool + + class UnsafePidExecutor: """Executor for Python scripts without nsjail, using subprocess isolation.""" @@ -307,17 +318,27 @@ def _with_python_paths( async def _build_execution_cmd( self, python_path: str, wrapper_path: Path - ) -> list[str]: + ) -> _ExecutionCommand: base_cmd = [python_path, str(wrapper_path)] if await pid_namespace_available(): - return ["unshare", "--pid", "--fork", "--kill-child", *base_cmd] + return _ExecutionCommand( + argv=["unshare", "--pid", "--fork", "--kill-child", *base_cmd], + supervised=False, + ) if not self._pid_isolation_warning_emitted: message = "PID namespace isolation unavailable; running script without PID isolation" logger.warning(message, reason=pid_namespace_probe_error()) module_logger.warning(message) self._pid_isolation_warning_emitted = True - return base_cmd + + supervisor_path = ( + Path(__file__).resolve().parents[1] / "executor" / "process_supervisor.py" + ) + return _ExecutionCommand( + argv=[sys.executable, "-I", str(supervisor_path), *base_cmd], + supervised=True, + ) async def _create_venv(self, venv_path: Path) -> None: create_cmd = ["uv", "venv", str(venv_path), "--python", "3.12"] @@ -464,9 +485,12 @@ async def execute( if execution_env_vars: exec_env.update(execution_env_vars) - cmd = await self._build_execution_cmd(python_path, wrapper_path) + execution_command = await self._build_execution_cmd( + python_path, + wrapper_path, + ) process = await asyncio.create_subprocess_exec( - *cmd, + *execution_command.argv, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=str(work_dir), @@ -477,6 +501,11 @@ async def execute( stdout_bytes, stderr_bytes = await communicate_process_group( process, timeout=timeout_seconds, + terminate=( + terminate_supervised_process + if execution_command.supervised + else None + ), ) except TimeoutError as e: raise SandboxTimeoutError( From ebb221b531f272dc46af704718a4b7a26a956282 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Wed, 5 Aug 2026 12:45:52 -0700 Subject: [PATCH 159/161] fix(executor): make deferred cleanup generation-safe --- tests/unit/test_registry_artifacts.py | 41 ++++++++++- .../executor/registry_artifact_storage.py | 70 ++++++++++++++++--- 2 files changed, 98 insertions(+), 13 deletions(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 2637c883a5..25fb98c2a5 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1211,9 +1211,44 @@ async def fail_extract(self, tarball_path, target_dir): assert deferred_path.is_dir() assert await cache._enforce_cache_budget() is True - assert cache._failed_startup_cleanup == set() + assert cache._failed_startup_cleanup == {} assert not deferred_path.exists() + @pytest.mark.anyio + async def test_deferred_cleanup_does_not_delete_replacement( + self, temp_cache_dir: Path + ) -> None: + """A stale cleanup record cannot delete a newly published artifact.""" + cache = RegistryArtifactCache(temp_cache_dir) + await cache.ensure_swept() + artifact_uri = "s3://bucket/path/replacement.tar.gz" + cache_key = compute_registry_artifact_cache_key(artifact_uri) + target_dir = cache._paths_for(cache_key).tarball_target_dir + target_dir.parent.mkdir(parents=True) + target_dir.write_text("malformed") + cache._context_for(cache_key).defer_cleanup(target_dir) + target_dir.unlink() + + async def mock_download(self, ctx, path): + del self, ctx + path.write_bytes(_tarball_payload(size=1)) + + async def mock_extract(self, tarball_path, output_dir): + del self, tarball_path + (output_dir / "module.py").write_text("VALUE = 1") + + with ( + patch(SQUASHFS_ENABLED_CONFIG, False), + patch.object(TarballArtifact, "download", mock_download), + patch.object(TarballArtifact, "extract", mock_extract), + ): + async with cache.lease([artifact_uri]) as registry_paths: + assert registry_paths == [target_dir] + assert (target_dir / "module.py").is_file() + + assert target_dir.is_dir() + assert cache._failed_startup_cleanup == {} + @pytest.mark.anyio async def test_materialize_extracts_squashfs_when_mount_fails(self, temp_cache_dir): """Test that SquashFS mount failures fall back to unsquashfs extraction.""" @@ -3952,13 +3987,13 @@ def fail_once(path: Path) -> bool: ): await cache.ensure_swept() assert orphaned.is_file() - assert cache._failed_startup_cleanup == {orphaned} + assert set(cache._failed_startup_cleanup) == {orphaned} assert cache._budget_dirty is True assert await cache._enforce_cache_budget() is True assert not orphaned.exists() - assert cache._failed_startup_cleanup == set() + assert cache._failed_startup_cleanup == {} @pytest.mark.anyio async def test_failed_startup_retirement_stays_dirty_and_retries( diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index d600ff6094..63ba77695e 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -105,6 +105,15 @@ class RegistryArtifactPaths: tarball_target_dir: Path +@dataclass(frozen=True, slots=True) +class _RegistryArtifactCleanupIdentity: + """Filesystem identity for one deferred cleanup target generation.""" + + device: int + inode: int + file_type: int + + @dataclass(slots=True) class RegistryArtifactRuntimeState: """Process-local synchronization and lease state for one cache key.""" @@ -174,6 +183,19 @@ def is_reusable_cache_file(path: Path) -> bool: return False +def _cleanup_identity(path: Path) -> _RegistryArtifactCleanupIdentity | None: + """Return a replacement-sensitive identity for one cleanup target.""" + try: + path_stat = path.lstat() + except FileNotFoundError: + return None + return _RegistryArtifactCleanupIdentity( + device=path_stat.st_dev, + inode=path_stat.st_ino, + file_type=stat.S_IFMT(path_stat.st_mode), + ) + + def ensure_real_directory(path: Path) -> None: """Create a cache directory without accepting a symlink redirect.""" if os.path.lexists(path): @@ -416,7 +438,7 @@ def __init__(self, cache_dir: Path): self._swept = False self._sweep_task: asyncio.Task[None] | None = None self._sweep_lock = asyncio.Lock() - self._failed_startup_cleanup: set[Path] = set() + self._failed_startup_cleanup: dict[Path, _RegistryArtifactCleanupIdentity] = {} self._budget_dirty = True async def ensure_swept(self) -> None: @@ -546,7 +568,7 @@ def _context_for( cache_key=cache_key, staging_dir=self.staging_dir, paths=self._paths_for(cache_key), - defer_cleanup=self._failed_startup_cleanup.add, + defer_cleanup=self._defer_cleanup, admission=admission, ) @@ -1092,9 +1114,9 @@ def _clear_legacy_cache_layout(self) -> bool: mounted = True if mounted or not _delete_cache_path(path): deleted = False - self._failed_startup_cleanup.add(path) + self._defer_cleanup(path) continue - self._failed_startup_cleanup.discard(path) + self._failed_startup_cleanup.pop(path, None) logger.info("Removed legacy registry artifact cache path", path=str(path)) return deleted @@ -1121,19 +1143,47 @@ def _clear_work_dir( for path in paths: if _delete_cache_path(path): if remember_failures: - self._failed_startup_cleanup.discard(path) + self._failed_startup_cleanup.pop(path, None) logger.info("Removed registry artifact work path", path=str(path)) else: deleted = False if remember_failures: - self._failed_startup_cleanup.add(path) + self._defer_cleanup(path) return deleted + def _defer_cleanup(self, path: Path) -> None: + """Remember only the current filesystem generation for later deletion.""" + try: + identity = _cleanup_identity(path) + except OSError as e: + self._failed_startup_cleanup.pop(path, None) + logger.warning( + "Could not identify deferred registry artifact cleanup path", + path=str(path), + error_type=type(e).__name__, + ) + return + if identity is None: + self._failed_startup_cleanup.pop(path, None) + return + self._failed_startup_cleanup[path] = identity + def _retry_failed_startup_cleanup(self) -> bool: - """Retry exact deferred paths without sweeping live staging work.""" - for path in tuple(self._failed_startup_cleanup): - if _delete_cache_path(path): - self._failed_startup_cleanup.discard(path) + """Retry deferred objects without deleting replacement generations.""" + for path, expected_identity in tuple(self._failed_startup_cleanup.items()): + try: + current_identity = _cleanup_identity(path) + except OSError: + continue + if current_identity != expected_identity: + if self._failed_startup_cleanup.get(path) == expected_identity: + self._failed_startup_cleanup.pop(path, None) + continue + if ( + _delete_cache_path(path) + and self._failed_startup_cleanup.get(path) == expected_identity + ): + self._failed_startup_cleanup.pop(path, None) return not self._failed_startup_cleanup def _trim_startup_cache(self) -> bool: From 6b0784c016143fc5032a55c90b72b32eb3a70745 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Wed, 5 Aug 2026 12:47:24 -0700 Subject: [PATCH 160/161] fix(executor): persist final lease release recency --- tests/unit/test_registry_artifacts.py | 18 ++++++++++++++++++ tracecat/executor/registry_artifact_storage.py | 5 ++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_registry_artifacts.py b/tests/unit/test_registry_artifacts.py index 25fb98c2a5..ed48f87dc7 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -1516,6 +1516,24 @@ def test_touch_entry_refreshes_tarball_root_mtime(self, temp_cache_dir): assert entry_dir.stat().st_mtime > 100.0 + def test_final_lease_release_persists_restart_safe_recency( + self, temp_cache_dir: Path + ) -> None: + """Only the final release persists the entry's latest use on disk.""" + cache = RegistryArtifactCache(temp_cache_dir) + cache_key = "nested-leases" + _write_tarball_entry(temp_cache_dir, cache_key) + entry_dir = cache._paths_for(cache_key).entry_dir + cache._acquire_lease(cache_key) + cache._acquire_lease(cache_key) + os.utime(entry_dir, (100.0, 100.0)) + + assert cache._release_lease(cache_key) is False + assert entry_dir.stat().st_mtime == 100.0 + + assert cache._release_lease(cache_key) is True + assert entry_dir.stat().st_mtime > 100.0 + @pytest.mark.anyio async def test_lease_refcounts_and_touches_image_mtime(self, temp_cache_dir): """A lease pins its entry and refreshes the restart-safe LRU timestamp.""" diff --git a/tracecat/executor/registry_artifact_storage.py b/tracecat/executor/registry_artifact_storage.py index 63ba77695e..97d98e9686 100644 --- a/tracecat/executor/registry_artifact_storage.py +++ b/tracecat/executor/registry_artifact_storage.py @@ -609,7 +609,10 @@ def _release_lease(self, cache_key: str) -> bool: return False runtime.refcount -= 1 runtime.last_used = time.time() - return runtime.refcount == 0 + became_idle = runtime.refcount == 0 + if became_idle: + self._touch_entry(cache_key) + return became_idle async def _unmount_idle_entry(self, cache_key: str) -> None: """Best-effort unmount an idle entry while retaining its image.""" From f51ad27073b3e1f859205998c80337968a32fa6c Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Wed, 5 Aug 2026 13:09:28 -0700 Subject: [PATCH 161/161] refactor(executor): clarify squashfs sidecar flag --- tracecat/executor/registry_artifacts.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tracecat/executor/registry_artifacts.py b/tracecat/executor/registry_artifacts.py index 496607b030..393b8ba58c 100644 --- a/tracecat/executor/registry_artifacts.py +++ b/tracecat/executor/registry_artifacts.py @@ -1211,7 +1211,7 @@ def _locally_cached_path( artifact_uri: str, ) -> list[Path] | None: """Return a reusable local candidate without probing remote sidecars.""" - include_unverified_sidecar = ( + include_squashfs_sidecar = ( _bundled_builtin_registry_version(artifact_uri) is None and _artifact_format(artifact_uri) == RegistryArtifactFormat.TAR_GZ and self._can_try_squashfs() @@ -1219,7 +1219,7 @@ def _locally_cached_path( candidates = self._candidate_artifacts( ctx, artifact_uri, - include_unverified_sidecar=include_unverified_sidecar, + include_squashfs_sidecar=include_squashfs_sidecar, ) return self._first_cached_path(candidates, ctx) @@ -1228,7 +1228,7 @@ def _candidate_artifacts( ctx: RegistryArtifactMaterializationContext, artifact_uri: str, *, - include_unverified_sidecar: bool, + include_squashfs_sidecar: bool, ) -> list[RegistryArtifact]: """Build artifact candidates in executor preference order.""" if version := _bundled_builtin_registry_version(artifact_uri): @@ -1258,7 +1258,7 @@ def _candidate_artifacts( return candidates candidates = [] - if include_unverified_sidecar and ( + if include_squashfs_sidecar and ( squashfs_uri := _squashfs_sidecar_uri(artifact_uri) ): candidates.append( @@ -1307,17 +1307,17 @@ async def _artifact_candidates( return self._candidate_artifacts( ctx, artifact_uri, - include_unverified_sidecar=False, + include_squashfs_sidecar=False, ) artifact_format = _artifact_format(artifact_uri) - include_unverified_sidecar = False + include_squashfs_sidecar = False if ( artifact_format == RegistryArtifactFormat.TAR_GZ and self._can_try_squashfs() and (squashfs_uri := _squashfs_sidecar_uri(artifact_uri)) ): - include_unverified_sidecar = ( + include_squashfs_sidecar = ( ctx.paths.squashfs_image_path.exists() or await self._sidecar_exists( base_uri=artifact_uri, @@ -1329,7 +1329,7 @@ async def _artifact_candidates( return self._candidate_artifacts( ctx, artifact_uri, - include_unverified_sidecar=include_unverified_sidecar, + include_squashfs_sidecar=include_squashfs_sidecar, ) async def _sidecar_exists(