diff --git a/deployments/fargate/main.tf b/deployments/fargate/main.tf index 420b1dcffa..ce111a2a61 100644 --- a/deployments/fargate/main.tf +++ b/deployments/fargate/main.tf @@ -155,6 +155,8 @@ 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 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..9496793c3a 100644 --- a/deployments/fargate/modules/ecs/locals.tf +++ b/deployments/fargate/modules/ecs/locals.tf @@ -164,11 +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__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 @@ -199,6 +201,8 @@ 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 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..6254d08a65 100644 --- a/deployments/fargate/modules/ecs/variables.tf +++ b/deployments/fargate/modules/ecs/variables.tf @@ -622,6 +622,28 @@ 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." + 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..275418bd20 100644 --- a/deployments/fargate/variables.tf +++ b/deployments/fargate/variables.tf @@ -590,6 +590,28 @@ 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." + 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..fb47bc0ae0 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -208,6 +208,8 @@ 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 TRACECAT__LOCAL_REPOSITORY_PATH: ${TRACECAT__LOCAL_REPOSITORY_PATH} @@ -395,6 +397,8 @@ 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 TRACECAT__LOCAL_REPOSITORY_PATH: ${TRACECAT__LOCAL_REPOSITORY_PATH} diff --git a/docker-compose.local.yml b/docker-compose.local.yml index 3fbc99311d..b025cf9bfd 100644 --- a/docker-compose.local.yml +++ b/docker-compose.local.yml @@ -224,6 +224,8 @@ 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 TRACECAT__LOCAL_REPOSITORY_PATH: ${TRACECAT__LOCAL_REPOSITORY_PATH} @@ -417,6 +419,8 @@ 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 TRACECAT__LOCAL_REPOSITORY_PATH: ${TRACECAT__LOCAL_REPOSITORY_PATH} diff --git a/docker-compose.yml b/docker-compose.yml index b5dd69b9dc..75355e9d5a 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -216,6 +216,8 @@ 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 TRACECAT__LOCAL_REPOSITORY_PATH: ${TRACECAT__LOCAL_REPOSITORY_PATH} @@ -408,6 +410,8 @@ 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 TRACECAT__LOCAL_REPOSITORY_PATH: ${TRACECAT__LOCAL_REPOSITORY_PATH} diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 1f5c7b92c8..9e1ce7edc3 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -180,8 +180,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_registry_artifact_cache_mount_lifecycle.py b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py new file mode 100644 index 0000000000..6982282480 --- /dev/null +++ b/tests/integration/test_registry_artifact_cache_mount_lifecycle.py @@ -0,0 +1,448 @@ +"""Dockerized lifecycle test for executor registry artifact SquashFS mounts. + +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. +""" + +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: + """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"): + pytest.skip(f"SquashFS mounts unsupported in this container: {skipped}") + + # 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 + + # 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 + + # 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 + + # 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"] == 2 + + # 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.parent.mkdir(parents=True, exist_ok=True) + 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) 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: + 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) + device = mounts[str(mount_dir)] + lease_mounts_observed += 1 + + 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() + ) + + 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 + ) + + # (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_from_retained_image"] = registry_paths == [ + retained_paths.squashfs_mount_dir + ] + payload["remount_released"] = not retained_paths.squashfs_mount_dir.is_mount() + + # (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 + payload["evicted_paths_removed"] = not ( + evicted_paths.squashfs_image_path.exists() + or evicted_paths.squashfs_mount_dir.exists() + ) + + # (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) + _build_squashfs_image( + root / "source-3", + fourth_paths.squashfs_image_path, + "module_3.py", + ) + async with cache.lease([fourth_uri]): + mounts_before_converge = _squashfs_mounts(cache_dir) + converged_device = mounts_before_converge[ + str(fourth_paths.squashfs_mount_dir) + ] + 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. + 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): + 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(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() + evicted_paths = sweep_cache._paths_for(sweep_keys[0]) + retained_paths = sweep_cache._paths_for(sweep_keys[1]) + payload["startup_sweep_trimmed"] = ( + 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()): + 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_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/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_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 caf22119cf..a096ffbbe8 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,42 @@ 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.paths_may_be_modified = False + self.leased = False - async def resolve_registry_paths( - self, artifact_uris: list[str] | None = None - ) -> list[Path]: + @asynccontextmanager + async def lease( + self, + artifact_uris: list[str] | None = None, + *, + paths_may_be_modified: bool = False, + ) -> AsyncIterator[list[Path]]: self.artifact_uris = artifact_uris - return self.paths + self.paths_may_be_modified = paths_may_be_modified + 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 +1197,49 @@ 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 + assert fake_runner.registry_artifacts.paths_may_be_modified is True + + @pytest.mark.anyio async def test_run_python_backend_fails_without_registry_artifacts( monkeypatch: pytest.MonkeyPatch, @@ -1221,14 +1292,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..1e346061bc 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,14 @@ from __future__ import annotations +import asyncio +import sys +import threading 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 @@ -23,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 @@ -242,6 +252,188 @@ 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, + *, + 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 + 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_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, + *, + 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] + 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/tests/unit/test_action_runner.py b/tests/unit/test_action_runner.py index 8b06ca69dc..6bd029b846 100644 --- a/tests/unit/test_action_runner.py +++ b/tests/unit/test_action_runner.py @@ -6,8 +6,11 @@ from __future__ import annotations 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 @@ -22,6 +25,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, @@ -30,6 +34,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 @@ -87,18 +92,45 @@ 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 ) -> 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, + 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), timeout=timeout, @@ -116,15 +148,19 @@ 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_no_registry_paths( + self, temp_cache_dir + ): + """An artifact-free action receives no extra registry import paths.""" 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 == [] - result = await runner.ensure_registry_environment("") - assert result == [] + async with runner.registry_artifacts.lease([]) as registry_paths: + assert registry_paths == [] + + assert not (temp_cache_dir / "base").exists() @pytest.mark.anyio async def test_execute_action_timeout( @@ -174,6 +210,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 @@ -466,9 +555,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() @@ -486,7 +576,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", @@ -512,8 +601,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( @@ -573,3 +671,204 @@ 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 = 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 + ) + + 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 + + @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 + real_terminate = sandbox_utils.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] = [] + + async def capture_subprocess(*args, **kwargs): + nonlocal process + process = await real_create_subprocess_exec(*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(requested_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( + "tracecat.executor.registry_artifact_mounts.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( + action_runner, + "terminate_supervised_process", + side_effect=controlled_termination, + ), + 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 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() + + 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_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 22e91f96e7..f59a3ba1c1 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 @@ -46,6 +47,8 @@ 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 +from tracecat.sandbox.types import SandboxConfig _DOCKER_CHILD_ENV = "TRACECAT__EXECUTOR_ACTION_SMOKE_DOCKER_CHILD" _SKIP_SENTINEL = "TRACE_CAT_EXECUTOR_ACTION_SMOKE_SKIP:" @@ -76,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() @@ -102,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(): @@ -519,9 +538,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( *, @@ -595,7 +615,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() @@ -604,7 +624,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: @@ -616,6 +637,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) @@ -648,16 +670,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() @@ -671,6 +700,104 @@ async def _run_current_builtin_smoke_case( } +@pytest.mark.parametrize("operation", list(CancelledNsjailOperation)) +@pytest.mark.anyio +async def test_cancelled_nsjail_operation_kills_process_group( + tmp_path: Path, + operation: CancelledNsjailOperation, +) -> None: + """Cancellation propagates only after nsjail and its child are gone.""" + 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 + 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( + 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() + return process + + with patch( + "tracecat.sandbox.executor.asyncio.create_subprocess_exec", + side_effect=capture_subprocess, + ): + 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() + 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): + await execution + + 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( "smoke_case", [ @@ -700,7 +827,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 @@ -722,7 +851,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( diff --git a/tests/unit/test_multitenant_registry.py b/tests/unit/test_multitenant_registry.py index 34ff62757f..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 @@ -32,6 +33,20 @@ 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 + + +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 # ============================================================================= @@ -62,7 +77,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 @@ -70,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") @@ -81,10 +95,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 @@ -117,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") @@ -128,8 +142,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,11 +168,11 @@ 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") + _write_empty_tarball(path) async def mock_extract(self, tarball_path: Path, target_dir: Path): raise RuntimeError("Extraction failed - corrupt tarball") @@ -169,14 +182,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,14 +207,13 @@ 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] 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") @@ -209,11 +223,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 6954ef57ae..ed48f87dc7 100644 --- a/tests/unit/test_registry_artifacts.py +++ b/tests/unit/test_registry_artifacts.py @@ -3,26 +3,242 @@ from __future__ import annotations import asyncio +import io +import os +import signal 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 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, + RegistryArtifactCacheCapacityError, + RegistryArtifactCacheLoopError, + RegistryArtifactEviction, + RegistryArtifactExtractionError, RegistryArtifactFormat, + RegistryArtifactUriError, SquashfsArtifact, + SquashfsMountCommandError, TarballArtifact, + _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" +) +MOUNT_CHECK = "tracecat.executor.registry_artifact_mounts.is_mount" + + +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]: + """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 + + +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: + """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 + @pytest.fixture def temp_cache_dir(): @@ -82,6 +298,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" @@ -104,7 +329,15 @@ 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, + 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( @@ -122,7 +355,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.""" @@ -141,12 +374,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.""" @@ -169,12 +401,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.""" @@ -182,7 +413,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( @@ -208,6 +440,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.""" @@ -241,56 +493,171 @@ 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() - 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 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, - "_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 cache.materialize( - 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) + + 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 result == [cached_path] - assert artifact_candidates.await_count == 2 - assert seen_candidates == [ - [RegistryArtifactFormat.TAR_GZ], - [RegistryArtifactFormat.SQUASHFS, RegistryArtifactFormat.TAR_GZ], + 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", + ] + ) + + # 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 @@ -347,6 +714,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) @@ -422,8 +822,11 @@ async def mock_mount(self, ctx, image_path): new_callable=AsyncMock, ) as tarball_materialize, ): - result = await cache.materialize( - "squashfs-key", + result = await _materialize( + cache, + compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.tar.gz" + ), "s3://bucket/path/site-packages.tar.gz", ) @@ -446,17 +849,25 @@ 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() + 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( @@ -469,7 +880,374 @@ 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): + """Repeated cancellation cannot abandon a killed mount process.""" + 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(block_wait=True) + + 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) + ) + 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 + + 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() + + @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 + 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) + 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.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, + ) as create_subprocess_exec, + patch( + "tracecat.executor.registry_artifacts.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, + allocation_unit=1, + ) + + 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 + + @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 + ) -> 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, *, allocation_unit: int) -> int: + assert allocation_unit > 0 + 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(_lease_and_release(cache, 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_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 == {} + 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): @@ -477,7 +1255,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 @@ -507,14 +1285,17 @@ async def mock_extract(self, ctx, image_path): new_callable=AsyncMock, ) as tarball_materialize, ): - result = await cache.materialize( - "fallback-key", + result = await _materialize( + cache, + compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.tar.gz" + ), "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 @@ -545,14 +1326,17 @@ async def mock_extract(self, ctx, image_path): ), patch.object(SquashfsArtifact, "extract", mock_extract), ): - result = await cache.materialize( - "extract-key", + result = await _materialize( + cache, + compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.tar.gz" + ), "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( @@ -569,7 +1353,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") @@ -592,14 +1376,17 @@ 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( - "gzip-fallback-key", + result = await _materialize( + cache, + compute_registry_artifact_cache_key( + "s3://bucket/path/site-packages.tar.gz" + ), "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): @@ -615,8 +1402,9 @@ 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( - "custom-key-test", + result = await _materialize( + cache, + compute_registry_artifact_cache_key("s3://bucket/path/custom-key"), "s3://bucket/path/custom-key", ) @@ -627,11 +1415,16 @@ 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" - target_dir = temp_cache_dir / f"tarball-{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 cache.materialize(cache_key, "s3://bucket/test.tar.gz") + result = await _materialize( + cache, + cache_key, + artifact_uri, + ) assert result == [target_dir] @@ -639,14 +1432,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") @@ -656,30 +1450,2785 @@ 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, artifact_uri), + _materialize(cache, cache_key, artifact_uri), + _materialize(cache, cache_key, artifact_uri), ) 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 == [] + + 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 + + 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.""" + 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) + empty_cache_bytes = cache._scan_cache_snapshot().total_bytes + + 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, empty_cache_bytes), + 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 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(_lease_and_release(cache, 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_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, + 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): + """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(MOUNT_CHECK, 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(MOUNT_CHECK, 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[1], + "converge", + cache_keys[0], + "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]) -> 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 + ) -> 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(MOUNT_CHECK, 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_acquire_artifact = AsyncMock(wraps=cache._acquire_artifact) + converge_cache_budget = AsyncMock() + + with ( + 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), + patch.object(cache, "_acquire_artifact", tracked_acquire_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].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 + 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() + 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()) + + @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_lock_for_same_key(self, temp_cache_dir): - """Test that same cache key returns same lock.""" + 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) - lock1 = await cache._lock_for("key1") - lock2 = await cache._lock_for("key1") + with ( + patch(MOUNT_CHECK, 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 lock1 is lock2 + 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_lock_for_different_keys(self, temp_cache_dir): - """Test that different cache keys return different locks.""" + 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): + assert kwargs["start_new_session"] is True + 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(MOUNT_CHECK, lambda path: path 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( + "tracecat.sandbox.utils.terminate_process_group", + new_callable=AsyncMock, + ), + 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) - lock1 = await cache._lock_for("key1") - lock2 = await cache._lock_for("key2") + 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, + "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() + + +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" + 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(_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 _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) + 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) + 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( + *, + key: str, + bucket: str, + 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 == 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() + 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) + 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) + 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( + *, + key: str, + bucket: str, + 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 + == 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) + + 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 == extracted_size + 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_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) + 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) + + 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, max_bytes), + 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 == reserved_expansion + 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) + 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(_tarball_payload(size=1)) + + async def mock_extract(self, tarball_path, target_dir): + (target_dir / "module.py").write_bytes(b"x" * (2 * allocation_unit)) + + with ( + patch(MAX_ENTRIES_CONFIG, 0), + patch(MAX_BYTES_CONFIG, max_bytes), + patch( + "tracecat.executor.registry_artifacts._tarball_extracted_size", + return_value=allocation_unit, + ), + 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_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, + ) + snapshot = cache._scan_cache_snapshot() + max_bytes = snapshot.structural_bytes + snapshot.entries["newest"].size_bytes + 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, max_bytes), + 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: 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) + 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: + 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: + async with cache._admission_lock: + await cache._ensure_cache_capacity( + additional_bytes=allocation_unit, + protected_key="new", + max_bytes=snapshot.total_bytes, + ) + + assert raised.value.current_bytes == snapshot.total_bytes + 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 _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(_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_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_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.""" + 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_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 + + 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) + 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), + 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 False + enforce_cache_budget.assert_awaited_once_with() + + @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) + enforce_cache_budget = AsyncMock(return_value=True) + + 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), + patch.object(cache, "_enforce_cache_budget", enforce_cache_budget), + ): + await _materialize(cache, cache_key, artifact_uri) + + 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( + 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 _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) + 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, max_bytes), + ): + 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): + assert kwargs["start_new_session"] is True + mounted.discard(paths.squashfs_mount_dir) + return process + + with ( + patch(MOUNT_CHECK, lambda path: path in mounted), + patch( + "tracecat.executor.registry_artifacts.shutil.which", + return_value="/sbin/umount", + ), + patch.object( + asyncio, + "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] + 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(*, allocation_unit: int | None = None): + nonlocal scan_count + entries = original_scan(allocation_unit=allocation_unit) + 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_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.""" + 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): + 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(MOUNT_CHECK, lambda path: path 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, + patch( + "tracecat.sandbox.utils.terminate_process_group", + new_callable=AsyncMock, + ), + ): + 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_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(MOUNT_CHECK, lambda path: path 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, + ), + 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_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.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: + 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), + 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._admission_lock.locked() + finally: + cleanup_release.set() + + with pytest.raises(asyncio.CancelledError): + await running + + assert cleanup_finished.is_set() + 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.""" + 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(MOUNT_CHECK, lambda path: path 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_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") + + 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() + + +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_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.""" + 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_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.""" + 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(MOUNT_CHECK, lambda path: path == 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): + """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([artifact_uri]): + 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_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 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 == {} + + @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) + snapshot = cache._scan_cache_snapshot() + max_bytes = snapshot.structural_bytes + snapshot.entries["newest"].size_bytes + 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, max_bytes), + 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(MOUNT_CHECK, 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 lock1 is not lock2 + assert mount_attempts == ["first", "second"] diff --git a/tests/unit/test_storage_blob.py b/tests/unit/test_storage_blob.py index c8df650a66..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 @@ -704,6 +705,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 +787,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 +818,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( @@ -771,6 +840,130 @@ 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, redact_log_identifiers: bool = False + ): # 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_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 + ): + """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, redact_log_identifiers: bool = False + ): # 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 @@ -782,7 +975,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( @@ -803,6 +998,166 @@ 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, 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" + 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 + 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 + ) -> 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/tests/unit/test_unsafe_pid_executor.py b/tests/unit/test_unsafe_pid_executor.py index 6667d749b0..1a962297a1 100644 --- a/tests/unit/test_unsafe_pid_executor.py +++ b/tests/unit/test_unsafe_pid_executor.py @@ -5,7 +5,9 @@ import logging import os import signal +import sys from pathlib import Path +from unittest.mock import AsyncMock, patch import pytest @@ -69,11 +71,99 @@ 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( + ("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: @@ -93,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 @@ -102,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: @@ -134,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 = [ @@ -259,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/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/concurrency.py b/tracecat/concurrency.py index e0dbd5d8a5..f518823c12 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,67 @@ 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 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( + 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/config.py b/tracecat/config.py index 2251c87c63..465914c34c 100644 --- a/tracecat/config.py +++ b/tracecat/config.py @@ -200,6 +200,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 64 +) +"""Maximum number of registry artifacts kept in the executor-local cache. + +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. + +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/action_runner.py b/tracecat/executor/action_runner.py index 2ba55885ab..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, ] @@ -142,38 +148,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. - - 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.""" - 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 - - 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] - async def execute_action( self, input: RunActionInput, @@ -204,39 +178,43 @@ 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) + # 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() ) - 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, + # 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: + logger.debug( + "Using sandbox execution", + use_sandbox=use_sandbox, + force_sandbox=force_sandbox, ) - 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, - ) - 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, @@ -401,7 +379,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 @@ -433,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 @@ -471,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 cf5a614573..b9d35afe24 100644 --- a/tracecat/executor/backends/base.py +++ b/tracecat/executor/backends/base.py @@ -14,7 +14,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 @@ -131,10 +131,61 @@ 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. 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, + paths_may_be_modified=True, + ) 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( @@ -164,36 +215,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..6c7c4aa3f5 100644 --- a/tracecat/executor/backends/test.py +++ b/tracecat/executor/backends/test.py @@ -18,9 +18,10 @@ from __future__ import annotations import asyncio +import functools 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 @@ -28,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, @@ -39,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, @@ -54,7 +59,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 @@ -121,20 +126,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", @@ -224,13 +233,27 @@ 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. + """ + 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.""" if not action_impl.module or not action_impl.name: @@ -247,10 +270,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,21 +298,25 @@ 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], + 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) logger.debug( "Materialized registry artifacts for test execution", 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_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 new file mode 100644 index 0000000000..97d98e9686 --- /dev/null +++ b/tracecat/executor/registry_artifact_storage.py @@ -0,0 +1,1245 @@ +"""Filesystem ownership and budget enforcement for registry artifacts.""" + +from __future__ import annotations + +import asyncio +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 ( + 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, + RegistryArtifactCacheSnapshot, + plan_registry_artifact_evictions, +) +from tracecat.logger import logger +from tracecat.sandbox.utils import communicate_process_group + +__all__ = ( + "CACHE_ENTRIES_DIR_NAME", + "CACHE_STAGING_DIR_NAME", + "CACHE_TRASH_DIR_NAME", + "RegistryArtifactAdmission", + "RegistryArtifactCacheCapacityError", + "RegistryArtifactCacheLoopError", + "RegistryArtifactCacheStorage", + "RegistryArtifactEviction", + "RegistryArtifactEvictionPass", + "RegistryArtifactMaterializationContext", + "RegistryArtifactPaths", + "RegistryArtifactRuntimeState", + "allocated_size_bound", + "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", +) + +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(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.""" + + 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 + ) + + +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 _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): + 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: dict[Path, _RegistryArtifactCleanupIdentity] = {} + 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._defer_cleanup, + 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 _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() + 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.""" + 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 = registry_artifact_mounts.is_mount(mount_dir) + 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 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) + 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, + 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 _cache_structural_footprint(self, *, allocation_unit: int) -> int: + """Measure cache roots and non-entry data exactly once.""" + cache_structure = _directory_footprint( + self.cache_dir, + allocation_unit=allocation_unit, + pruned_directories=( + self.entries_dir, + self.staging_dir, + 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.""" + 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 = 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 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 registry_artifact_mounts.is_mount(path) + except OSError: + mounted = True + if mounted or not _delete_cache_path(path): + deleted = False + self._defer_cleanup(path) + continue + self._failed_startup_cleanup.pop(path, None) + 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.pop(path, None) + logger.info("Removed registry artifact work path", path=str(path)) + else: + deleted = False + if remember_failures: + 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 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: + """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 registry_artifact_mounts.is_mount( + self._paths_for(entry.cache_key).squashfs_mount_dir + ) + } + 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 2a82a28490..393b8ba58c 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 shutil @@ -10,19 +11,67 @@ import tarfile import time from abc import ABC, abstractmethod +from collections.abc import AsyncGenerator, Awaitable, Callable +from contextlib import asynccontextmanager 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 ( + 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, + RegistryArtifactCacheLoopError, + RegistryArtifactCacheStorage, + RegistryArtifactEviction, + RegistryArtifactMaterializationContext, + allocated_size_bound, + 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.sandbox.utils import communicate_process_group from tracecat.storage import blob +__all__ = ( + "BUNDLED_BUILTIN_REGISTRY_URI_PREFIX", + "SQUASHFS_MOUNT_OPTIONS", + "BuiltinArtifact", + "RegistryArtifact", + "RegistryArtifactAdmission", + "RegistryArtifactCache", + "RegistryArtifactCacheCapacityError", + "RegistryArtifactExtractionError", + "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.""" @@ -44,41 +93,24 @@ class RegistryArtifactFormat(StrEnum): """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.""" +class SquashfsMountCommandError(RuntimeError): + """The ``mount`` command itself failed for a SquashFS registry artifact. - squashfs_image_path: Path - squashfs_mount_dir: Path - squashfs_extract_dir: Path - tarball_target_dir: Path + 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(slots=True) -class SquashfsMountState: - """Shared process-local SquashFS mount state.""" +class RegistryArtifactUriError(ValueError): + """A registry artifact URI is malformed, with identifiers suppressed.""" - disabled: bool = False +class RegistryArtifactExtractionError(RuntimeError): + """A registry archive could not be inspected or extracted safely.""" -@dataclass(slots=True) -class RegistryArtifactMaterializationContext: - """Shared runtime state for artifact materialization.""" - - cache_key: str - cache_dir: Path - paths: RegistryArtifactPaths - squashfs_mount_state: SquashfsMountState - - 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) - ) - - def disable_squashfs_mount(self) -> None: - self.squashfs_mount_state.disabled = True + def __init__(self) -> None: + super().__init__("Registry artifact extraction failed") @dataclass(frozen=True, slots=True) @@ -110,8 +142,7 @@ def _temp_path( ctx: RegistryArtifactMaterializationContext, suffix: str, ) -> Path: - unique_id = id(asyncio.current_task()) - return ctx.cache_dir / f"{self.cache_key}.{os.getpid()}.{unique_id}{suffix}" + return unique_work_path(ctx.staging_dir, self.cache_key, suffix=suffix) @dataclass(frozen=True, slots=True) @@ -153,13 +184,25 @@ 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 registry_artifact_mounts.is_mount( + mount_dir + ): 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, @@ -174,12 +217,11 @@ async def materialize( if ctx.can_mount_squashfs(): try: return [await self.mount(ctx, image_path)] - except Exception as e: - ctx.disable_squashfs_mount() + except SquashfsMountCommandError as e: 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), ) @@ -192,13 +234,24 @@ 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 + ensure_cache_entry_directory(ctx.paths) 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, + defer_cleanup=ctx.defer_cleanup, + published_path=image_path, + ) try: temp_image.rename(image_path) except OSError: @@ -206,24 +259,50 @@ 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, 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. + """ + 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 registry_artifact_mounts.is_mount(target_dir): return target_dir - ctx.cache_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() @@ -237,7 +316,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}", @@ -250,16 +329,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.cache_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() @@ -267,6 +350,16 @@ 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, + 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) await self._extract_image(image_path, temp_dir) @@ -278,31 +371,49 @@ 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 async def _mount_image(self, image_path: Path, target_dir: Path) -> None: - """Mount a SquashFS image read-only at target_dir.""" - if target_dir.is_mount(): + """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( @@ -315,17 +426,23 @@ 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 proc.communicate() + 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() - 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.""" + """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") @@ -338,13 +455,42 @@ 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 proc.communicate() + 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, + 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") + + 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) @@ -358,7 +504,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, @@ -369,11 +519,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() @@ -382,12 +533,28 @@ async def materialize( temp_dir = self._temp_path(ctx, ".tmp") try: - ctx.cache_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 (admission := ctx.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 + 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) await self.extract(temp_tarball, temp_dir) @@ -399,27 +566,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] @@ -428,10 +604,20 @@ 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, + defer_cleanup=ctx.defer_cleanup, + ) 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"): @@ -441,7 +627,11 @@ def _do_extract() -> None: raise ValueError(f"Unsupported tarball format: {tarball_path}") - await asyncio.to_thread(_do_extract) + try: + await run_blocking_rejoin_on_cancel(_do_extract) + except Exception: + raise RegistryArtifactExtractionError() from None + logger.debug( "Tarball extracted", target=str(target_dir), @@ -449,20 +639,53 @@ 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, + 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=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 @@ -472,9 +695,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: @@ -547,84 +769,430 @@ def _artifact_format(artifact_uri: str) -> RegistryArtifactFormat: return RegistryArtifactFormat.TAR_GZ -class RegistryArtifactCache: - """Materializes registry artifacts into executor-local Python paths.""" +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 _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 += 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, *, 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: + 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 + + +@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 + + 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 + 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.""" + if self._closed: + return + 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 [], + ) + ) + await rejoin_future_through_cancellation(cleanup_task) + finally: + self.cache._request_runtime_retirement_if_entry_missing(cache_key) - def __init__(self, cache_dir: Path): - self.cache_dir = cache_dir - self._locks: dict[str, asyncio.Lock] = {} - self._locks_lock = asyncio.Lock() - self._squashfs_mount_state = SquashfsMountState() - - async def ensure_environment(self, artifact_uri: str | None) -> list[Path]: - """Materialize an optional registry artifact and return PYTHONPATH entries.""" - if not artifact_uri: - return [] - cache_key = compute_registry_artifact_cache_key(artifact_uri) - return await self.materialize(cache_key, artifact_uri) - - async def materialize(self, cache_key: str, artifact_uri: str) -> list[Path]: - """Materialize a registry artifact as local importable directories.""" - ctx = self._context_for(cache_key) - candidates = await self._artifact_candidates(ctx, artifact_uri) - if cached_paths := self._first_cached_path(candidates, ctx): - return cached_paths +class RegistryArtifactCache(RegistryArtifactCacheStorage): + """Materializes registry artifacts into executor-local Python paths.""" - lock = await self._lock_for(cache_key) - async with lock: + @asynccontextmanager + 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 + from the returned paths until the context exits. + + Args: + artifact_uris: Registry artifact URIs in deterministic PYTHONPATH + 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. + + Yields: + Importable Python paths for the requested artifacts. + """ + if not artifact_uris: + logger.info("No registry artifact URIs provided") + yield [] + 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] = [] + for artifact_uri in artifact_uris: + 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 + + 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 _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) - if cached_paths := self._first_cached_path(candidates, ctx): + return await self._materialize_candidates(ctx, candidates) + + 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 - 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), - ) - return await artifact.materialize(ctx) - 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), - ) + async with self._admission_lock: + async with self._runtime_lock(cache_key): + if cached_paths := self._locally_cached_path(ctx, artifact_uri): + return 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 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 - 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 _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, - paths=self._paths_for(cache_key), - squashfs_mount_state=self._squashfs_mount_state, - ) + async def _materialize_candidates( + self, + ctx: RegistryArtifactMaterializationContext, + candidates: list[RegistryArtifact], + ) -> list[Path]: + """Materialize the first viable artifact candidate. - def _paths_for(self, cache_key: str) -> RegistryArtifactPaths: - """Return local cache paths for a registry artifact 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}", - ) + Callers hold the cache key's lock for evictable entries. + """ + 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_for_logging(artifact.uri), + artifact_format=artifact.format.value, + 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: + raise + logger.warning( + "Failed to materialize registry artifact candidate, trying fallback", + cache_key=cache_key, + artifact_uri=_artifact_uri_for_logging(artifact.uri), + artifact_format=artifact.format.value, + error_type=type(e).__name__, + ) + + raise RuntimeError(f"No registry artifact candidates for {ctx.cache_key}") def _first_cached_path( self, @@ -637,12 +1205,32 @@ def _first_cached_path( return cached_paths return None - async def _artifact_candidates( + def _locally_cached_path( + self, + ctx: RegistryArtifactMaterializationContext, + artifact_uri: str, + ) -> list[Path] | None: + """Return a reusable local candidate without probing remote sidecars.""" + include_squashfs_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_squashfs_sidecar=include_squashfs_sidecar, + ) + return self._first_cached_path(candidates, ctx) + + def _candidate_artifacts( self, ctx: RegistryArtifactMaterializationContext, artifact_uri: str, + *, + include_squashfs_sidecar: bool, ) -> list[RegistryArtifact]: - """Return artifact candidates in executor preference order.""" + """Build artifact candidates in executor preference order.""" if version := _bundled_builtin_registry_version(artifact_uri): return [ BuiltinArtifact( @@ -654,7 +1242,7 @@ async def _artifact_candidates( artifact_format = _artifact_format(artifact_uri) if artifact_format == RegistryArtifactFormat.SQUASHFS: - candidates = [ + candidates: list[RegistryArtifact] = [ SquashfsArtifact( uri=artifact_uri, cache_key=ctx.cache_key, @@ -669,29 +1257,16 @@ async def _artifact_candidates( ) 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( - base_uri=artifact_uri, - sidecar_uri=squashfs_uri, - artifact_format=RegistryArtifactFormat.SQUASHFS, - ): - candidates.append( - SquashfsArtifact( - uri=squashfs_uri, - cache_key=ctx.cache_key, - ) - ) - + candidates = [] + if include_squashfs_sidecar 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, @@ -700,6 +1275,63 @@ async def _artifact_candidates( ) return candidates + 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 + validate_cache_entry_path(paths) + try: + if registry_artifact_mounts.is_mount(paths.squashfs_mount_dir): + return + except OSError: + return + + 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, + ctx: RegistryArtifactMaterializationContext, + artifact_uri: str, + ) -> list[RegistryArtifact]: + """Return artifact candidates in executor preference order.""" + if _bundled_builtin_registry_version(artifact_uri) is not None: + return self._candidate_artifacts( + ctx, + artifact_uri, + include_squashfs_sidecar=False, + ) + + artifact_format = _artifact_format(artifact_uri) + include_squashfs_sidecar = False + if ( + artifact_format == RegistryArtifactFormat.TAR_GZ + and self._can_try_squashfs() + and (squashfs_uri := _squashfs_sidecar_uri(artifact_uri)) + ): + include_squashfs_sidecar = ( + ctx.paths.squashfs_image_path.exists() + or await self._sidecar_exists( + base_uri=artifact_uri, + sidecar_uri=squashfs_uri, + artifact_format=RegistryArtifactFormat.SQUASHFS, + ) + ) + + return self._candidate_artifacts( + ctx, + artifact_uri, + include_squashfs_sidecar=include_squashfs_sidecar, + ) + async def _sidecar_exists( self, *, @@ -708,23 +1340,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 diff --git a/tracecat/executor/worker.py b/tracecat/executor/worker.py index c553916e0b..be9873a10d 100644 --- a/tracecat/executor/worker.py +++ b/tracecat/executor/worker.py @@ -54,6 +54,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, @@ -130,6 +131,18 @@ async def main(shutdown_event: asyncio.Event | None = None) -> None: # socket path is available in their immutable process environment. await action_gateway.start() + # Warm the registry artifact cache sweep before the backend spawns + # workers or activities run; cache construction itself is cheap. + 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() diff --git a/tracecat/sandbox/executor.py b/tracecat/sandbox/executor.py index 5ad8b8c4d1..faf72df5ec 100644 --- a/tracecat/sandbox/executor.py +++ b/tracecat/sandbox/executor.py @@ -30,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.""" @@ -466,20 +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 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 @@ -612,18 +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 TimeoutError as e: - process.kill() - await process.wait() raise SandboxTimeoutError( f"Package installation timed out after {timeout_seconds}s" ) from e @@ -879,19 +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 TimeoutError as e: - process.kill() - await process.wait() raise SandboxTimeoutError( f"Action execution timed out after {config.timeout_seconds}s" ) from e diff --git a/tracecat/sandbox/unsafe_pid_executor.py b/tracecat/sandbox/unsafe_pid_executor.py index ccd5cda937..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"] @@ -330,12 +351,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) + _, 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( @@ -371,16 +391,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 TimeoutError as e: - process.kill() - await process.wait() raise PackageInstallError( f"Package installation timed out after {timeout_seconds}s" ) from e @@ -466,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), @@ -479,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( diff --git a/tracecat/sandbox/utils.py b/tracecat/sandbox/utils.py index 360b1349e3..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,11 +43,46 @@ 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 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. @@ -53,24 +90,38 @@ 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)) - group_terminated = False + 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) - await terminate_process_group(process) - group_terminated = True + termination_task = asyncio.ensure_future(terminator(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, + terminator, + ) + ) + try: + await rejoin_future_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") diff --git a/tracecat/storage/blob.py b/tracecat/storage/blob.py index a66c801cfc..fb8c896f42 100644 --- a/tracecat/storage/blob.py +++ b/tracecat/storage/blob.py @@ -3,23 +3,24 @@ from __future__ import annotations import asyncio +import functools import hashlib import os import threading import weakref -from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator, Awaitable, Callable from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING import aioboto3 -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.concurrency import run_blocking_rejoin_on_cancel from tracecat.logger import logger if TYPE_CHECKING: @@ -35,6 +36,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 +239,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 +737,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 +754,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 +777,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 +803,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( @@ -758,6 +824,9 @@ 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, + defer_cleanup: Callable[[Path], None] | None = None, + redact_log_identifiers: bool = False, ) -> int: """Stream an S3/MinIO object to a local file. @@ -771,6 +840,11 @@ 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 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. Returns: Total bytes written. @@ -780,59 +854,94 @@ 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}" ) - async with aiofiles.open(temp_path, "wb") as f: + 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}" + ) + 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) + ) + + # 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 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"Refusing to download {log_bucket}/{log_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 f.write(chunk) + await run_blocking_rejoin_on_cancel( + functools.partial(output_file.write, chunk) + ) if hasher is not 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}" ) os.replace(temp_path, output_path) - except Exception: + 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, )