diff --git a/opencollab/adapters/_docker_snapshot.py b/opencollab/adapters/_docker_snapshot.py index e22f07f8..e08bef6f 100644 --- a/opencollab/adapters/_docker_snapshot.py +++ b/opencollab/adapters/_docker_snapshot.py @@ -6,12 +6,14 @@ import shlex import shutil import tempfile +from contextlib import asynccontextmanager from pathlib import Path from typing import Any from opencollab.adapters._candidate_physical import _initialized_submodule_paths from opencollab.adapters._env_docker import DockerEnvironment from opencollab.adapters._env_local import LocalEnvironment +from opencollab.adapters._env_vm import VMEnvironment from opencollab.adapters._git_patch_bytes import _complete, _patch_temporary from opencollab.adapters._team_candidate_environment import _TeamCandidateEnvironment from opencollab.application.async_timeout import await_owned_operation @@ -26,8 +28,11 @@ async def _copy_to_host(environment: Any, source: str, target: Path, *, director operation = asyncio.to_thread(shutil.copyfile, source, target) await await_owned_operation(operation, propagate_cancellation=True) return + if isinstance(delegate, VMEnvironment): + await delegate.copy_to_host(source, target, directory=directory) + return if not isinstance(delegate, DockerEnvironment): - raise TypeError("Docker isolation requires a local or Docker source environment") + raise TypeError("Workspace snapshots require a local, Docker, or VM source environment") await delegate.setup() source_path = source.rstrip("/") + "/." if directory else source result = await delegate._docker("cp", "--", f"{delegate._container_id}:{source_path}", str(target)) @@ -43,7 +48,8 @@ def _remove_marker(path: Path) -> None: marker.unlink() -async def _import_snapshot(source: Any, target: DockerEnvironment) -> str: +@asynccontextmanager +async def _snapshot_contents(source: Any): probe = await source.exec_cmd(f"git -C {shlex.quote(source.workspace)} rev-parse --show-toplevel") _complete(probe, "Docker source repository probe", allowed=(0, 128)) repository = probe.returncode == 0 @@ -68,7 +74,7 @@ async def _import_snapshot(source: Any, target: DockerEnvironment) -> str: bundle = directory / "repository.bundle" async with _patch_temporary(source, root, "", prefix=".docker-import-", suffix=".bundle") as remote: _complete(await source.exec_cmd( - f"git -C {shlex.quote(source_root)} bundle create {shlex.quote(remote)} --tags HEAD", + f"git -C {shlex.quote(source_root)} bundle create - --tags HEAD > {shlex.quote(remote)}", timeout=900, ), "Docker source Git bundle") await _copy_to_host(source, remote, bundle) @@ -113,6 +119,18 @@ async def _import_snapshot(source: Any, target: DockerEnvironment) -> str: ), "Docker directory snapshot") finally: await local.cleanup() + yield contents, prefix + finally: + await await_owned_operation( + asyncio.to_thread(shutil.rmtree, directory), propagate_cancellation=True, + ) + + +async def _import_snapshot(source: Any, target: DockerEnvironment | VMEnvironment) -> str: + async with _snapshot_contents(source) as (contents, prefix): + if isinstance(target, VMEnvironment): + await target.copy_from_host(contents, target.workspace, directory=True) + return posixpath.normpath(posixpath.join(target.workspace, prefix)) container = await target.setup() _complete(await target.exec_cmd(f"mkdir -p -- {shlex.quote(target.workspace)}"), "Docker workspace directory") copied = await target._docker("cp", "--", str(contents) + "/.", f"{container}:{target.workspace}", timeout=900) @@ -130,7 +148,3 @@ async def _import_snapshot(source: Any, target: DockerEnvironment) -> str: if ownership.returncode or ownership.stdout_dropped_bytes or ownership.stderr_dropped_bytes: raise OSError("Docker isolated workspace ownership could not be assigned") return posixpath.normpath(posixpath.join(target.workspace, prefix)) - finally: - await await_owned_operation( - asyncio.to_thread(shutil.rmtree, directory), propagate_cancellation=True, - ) diff --git a/opencollab/adapters/_env_vm.py b/opencollab/adapters/_env_vm.py index 21b4c9fb..c83a5da3 100644 --- a/opencollab/adapters/_env_vm.py +++ b/opencollab/adapters/_env_vm.py @@ -75,6 +75,7 @@ def __init__(self, connection: VMConnection, workspace: str, *, state_directory: self._uncertain: set[str] = set() self._edit_locks: set[str] = set() self._temporary_files: dict[str, list[int]] = {} + self._mutable_temporaries: set[str] = set() self._downloads: set[asyncio.Task[None]] = set() self._transfers: dict[ProcessRegistry, tuple[ asyncio.Task[asyncio.subprocess.Process], asyncio.Task[tuple[bytes, int]] | None, @@ -309,6 +310,24 @@ async def remove_file(self, path: str) -> None: await self._operate("remove", {"path": target, "identity": self._temporary_files[target]}) self._temporary_files.pop(target) + async def _remove_mutable_temporary(self, operation_id: str) -> None: + if operation_id in self._operations: + await self._cancel_operation(operation_id) + await self._operate("remove_mutable_temporary", {"operation_id": operation_id}, require_active=False) + self._mutable_temporaries.discard(operation_id) + + @contextlib.asynccontextmanager + async def _mutable_temporary(self, content: str, *, prefix: str, suffix: str) -> AsyncIterator[str]: + if not re.fullmatch(r"[A-Za-z0-9._-]+", prefix) or not re.fullmatch(r"[A-Za-z0-9._-]+", suffix): + raise ValueError("temporary file prefix and suffix must be simple path components") + operation_id = uuid.uuid4().hex + self._mutable_temporaries.add(operation_id) + try: + path = await self._operate("mutable_temporary", {"content": content, "name": prefix + suffix}, operation_id) + yield path + finally: + await await_owned_operation(self._remove_mutable_temporary(operation_id), propagate_cancellation=True) + @contextlib.asynccontextmanager async def file_write_lock(self, path: str) -> AsyncIterator[None]: self._ensure_active() @@ -498,6 +517,8 @@ async def cleanup(self) -> None: except BaseException: self._uncertain.add(operation_id) raise + for operation_id in tuple(self._mutable_temporaries): + await self._remove_mutable_temporary(operation_id) for path in tuple(self._temporary_files): await self._operate("remove", {"path": path, "identity": self._temporary_files[path]}, require_active=False) diff --git a/opencollab/adapters/_git_patch_bytes.py b/opencollab/adapters/_git_patch_bytes.py index 13f5c202..2b35ea8d 100644 --- a/opencollab/adapters/_git_patch_bytes.py +++ b/opencollab/adapters/_git_patch_bytes.py @@ -10,6 +10,7 @@ from typing import Any from opencollab.adapters._env_local import LocalEnvironment +from opencollab.adapters._env_vm import VMEnvironment from opencollab.application.async_timeout import await_owned_operation from opencollab.patches import git_diff_endpoint, git_header_tokens, patch_entries, split_patch_blocks @@ -82,7 +83,11 @@ async def _patch_temporary( cleanup_best_effort: bool = False, repository: bool = True, ) -> AsyncIterator[str]: - if isinstance(environment, LocalEnvironment): + delegate = getattr(environment, "_delegate", environment) + if isinstance(delegate, VMEnvironment): + async with delegate._mutable_temporary(content, prefix=prefix, suffix=suffix) as path: + yield path + elif isinstance(environment, LocalEnvironment): git_directory = None if repository: git_directory = _complete( diff --git a/opencollab/adapters/_vm_leases.py b/opencollab/adapters/_vm_leases.py new file mode 100644 index 00000000..07823adb --- /dev/null +++ b/opencollab/adapters/_vm_leases.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import asyncio +import math +import sqlite3 +from collections.abc import Iterator, Sequence +from contextlib import closing, contextmanager +from pathlib import Path + +from opencollab.adapters.vm_config import VMResource + + +class VMCapacityError(RuntimeError): + pass + + +class _VMLeases: + def __init__(self, database: str, resources: Sequence[VMResource], wait_timeout: float | None) -> None: + if not isinstance(database, str) or not database.strip() or database == ":memory:": + raise ValueError("VM lease_database must name a shared SQLite file") + if wait_timeout is not None and ( + isinstance(wait_timeout, bool) or not isinstance(wait_timeout, (int, float)) + or not math.isfinite(wait_timeout) or wait_timeout < 0 + ): + raise ValueError("VM wait_timeout must be finite and nonnegative or None") + self.database = str(Path(database).expanduser().absolute()) + self.resources = tuple(resources) + self.wait_timeout = wait_timeout + Path(self.database).parent.mkdir(parents=True, exist_ok=True) + with self._connect() as connection: + connection.execute( + "CREATE TABLE IF NOT EXISTS vm_leases " + "(resource_id TEXT PRIMARY KEY, owner TEXT NOT NULL, workspace TEXT NOT NULL)" + ) + + @contextmanager + def _connect(self) -> Iterator[sqlite3.Connection]: + with closing(sqlite3.connect(self.database, timeout=0.1)) as connection, connection: + yield connection + + def _try_acquire(self, owner: str) -> VMResource | None: + with self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + for resource in self.resources: + claimed = connection.execute( + "INSERT OR IGNORE INTO vm_leases (resource_id, owner, workspace) VALUES (?, ?, ?)", + (resource.resource_id, owner, f"{resource.workspace_root}/{owner}"), + ) + if claimed.rowcount: + return resource + return None + + async def acquire(self, owner: str, closing: asyncio.Event | None = None) -> VMResource: + deadline = None if self.wait_timeout is None else asyncio.get_running_loop().time() + self.wait_timeout + while True: + if closing is not None and closing.is_set(): + raise RuntimeError("VM isolation pool closed while waiting for capacity") + # The transaction is short and synchronous, so cancellation cannot + # abandon a successful claim in a worker thread. + try: + resource = self._try_acquire(owner) + except sqlite3.OperationalError as exc: + if "locked" not in str(exc): + raise + resource = None + if resource is not None: + return resource + remaining = None if deadline is None else deadline - asyncio.get_running_loop().time() + if remaining is not None and remaining <= 0: + raise VMCapacityError("VM capacity exhausted; release a lease or configure cancellable waiting") + await asyncio.sleep(0.05 if remaining is None else min(0.05, remaining)) + + def release(self, resource: VMResource, owner: str) -> None: + with self._connect() as connection: + removed = connection.execute( + "DELETE FROM vm_leases WHERE resource_id = ? AND owner = ?", (resource.resource_id, owner), + ) + if removed.rowcount != 1: + raise RuntimeError(f"VM lease ownership changed for {resource.resource_id}") diff --git a/opencollab/adapters/_vm_remote.py b/opencollab/adapters/_vm_remote.py index 32f59af5..e15b1d11 100644 --- a/opencollab/adapters/_vm_remote.py +++ b/opencollab/adapters/_vm_remote.py @@ -12,6 +12,7 @@ import re import shutil import signal +import stat import struct import subprocess import sys @@ -208,6 +209,54 @@ async def _perform(request: dict[str, Any], directory: Path) -> Any: return asdict(await env.read_text_range(**args)) if action == "write": return await env.write_file_with_change(args["path"], args["content"]) + if action == "mutable_temporary": + if not re.fullmatch(r"[A-Za-z0-9._-]+", args["name"]): + raise ValueError("invalid mutable temporary name") + temporary = directory / "temporary" + temporary.mkdir(mode=0o700) + fd = open_directory_anchor(temporary) + try: + info = os.fstat(fd) + _write(directory, "temporary.json", [info.st_dev, info.st_ino]) + create_regular_bytes_atomic_at(fd, temporary, args["name"], args["content"].encode(), + max_bytes=ENV_FILE_WRITE_LIMIT_BYTES) + finally: + os.close(fd) + return str(temporary / args["name"]) + if action == "remove_mutable_temporary": + operation_id = args["operation_id"] + if not isinstance(operation_id, str) or not _ID.fullmatch(operation_id): + raise ValueError("invalid mutable temporary operation") + owner = directory.parent / operation_id + parent = open_directory_anchor(owner) + try: + try: + expected = _read(owner, "temporary.json") + fd = open_directory_anchor(owner / "temporary") + except FileNotFoundError: + if os.path.lexists(owner / "temporary"): + raise OSError("mutable temporary ownership is incomplete") from None + return None + try: + info = os.fstat(fd) + if [info.st_dev, info.st_ino] != expected: + raise OSError("mutable temporary directory was replaced") + for name in os.listdir(fd): + mode = os.stat(name, dir_fd=fd, follow_symlinks=False).st_mode + if stat.S_ISDIR(mode): + shutil.rmtree(f"/proc/self/fd/{fd}/{name}") + else: + os.unlink(name, dir_fd=fd) + current = os.stat("temporary", dir_fd=parent, follow_symlinks=False) + if [current.st_dev, current.st_ino] != expected: + raise OSError("mutable temporary directory changed during cleanup") + os.rmdir("temporary", dir_fd=parent) + os.fsync(parent) + finally: + os.close(fd) + finally: + os.close(parent) + return None if action in {"temporary", "remove"}: fd = open_directory_anchor(env.workspace) try: diff --git a/opencollab/adapters/candidate_workspace.py b/opencollab/adapters/candidate_workspace.py index cc0dbb70..ed8e0e3c 100644 --- a/opencollab/adapters/candidate_workspace.py +++ b/opencollab/adapters/candidate_workspace.py @@ -31,6 +31,7 @@ from opencollab.adapters._env_docker import DockerEnvironment from opencollab.adapters._env_local import LocalEnvironment from opencollab.adapters._env_process import ProcessCleanupError +from opencollab.adapters._env_vm import VMEnvironment from opencollab.adapters._env_worktree_submodules import _initialize_source_available_submodules from opencollab.adapters._git_patch_bytes import _complete, _complete_patch from opencollab.adapters._git_patch_bytes import _patch_temporary as _candidate_temporary @@ -505,7 +506,11 @@ async def _candidate_environment(self, path: str) -> Any: command_prefix=self._environment._command_prefix, timeout_returncode=self._environment._timeout_returncode, ) - raise RuntimeError("candidate worktrees require a local or Docker environment") + if isinstance(self._environment, VMEnvironment): + return VMEnvironment( + self._environment.connection, path, state_directory=self._environment.state_directory, + ) + raise RuntimeError("candidate worktrees require a local, Docker, or VM environment") async def acquire(self, label: str) -> _CandidateLease: return await self._acquire(label) diff --git a/opencollab/adapters/vm_config.py b/opencollab/adapters/vm_config.py index 2c80da97..b3e0dfca 100644 --- a/opencollab/adapters/vm_config.py +++ b/opencollab/adapters/vm_config.py @@ -42,3 +42,29 @@ def validate_vm_path(value: str) -> str: ): raise ValueError("VM paths must be normalized absolute non-root paths") return value + + +@dataclass(frozen=True) +class VMResource: + resource_id: str + connection: VMConnection + workspace_root: str = "/workspace/opencollab" + state_directory: str = "/tmp/opencollab-vm" + reset_command: str | None = None + + def __post_init__(self) -> None: + if ( + not isinstance(self.resource_id, str) or not self.resource_id.strip() + or any(ord(character) < 32 for character in self.resource_id) + ): + raise ValueError("VM resource_id must be a stable nonempty identifier without control characters") + if not isinstance(self.connection, VMConnection): + raise TypeError("VM resource connection must be a VMConnection") + validate_vm_path(self.workspace_root) + validate_vm_path(self.state_directory) + if self.state_directory == self.workspace_root or self.state_directory.startswith(self.workspace_root + "/"): + raise ValueError("VM operation state must be outside its workspace root") + if self.reset_command is not None and ( + not isinstance(self.reset_command, str) or not self.reset_command.strip() or "\0" in self.reset_command + ): + raise ValueError("VM reset_command must be a nonempty shell command") diff --git a/opencollab/adapters/vm_isolation.py b/opencollab/adapters/vm_isolation.py new file mode 100644 index 00000000..7ba48cbd --- /dev/null +++ b/opencollab/adapters/vm_isolation.py @@ -0,0 +1,308 @@ +from __future__ import annotations + +import asyncio +import os +import posixpath +import shlex +import shutil +import tempfile +import uuid +from collections.abc import Sequence +from contextlib import nullcontext +from dataclasses import dataclass, field +from pathlib import Path + +from opencollab.adapters._docker_snapshot import _import_snapshot, _snapshot_contents +from opencollab.adapters._env_vm import VMEnvironment, VMOperationError +from opencollab.adapters._git_patch_bytes import _complete +from opencollab.adapters._team_candidate_environment import _TeamCandidateEnvironment +from opencollab.adapters._vm_leases import VMCapacityError as VMCapacityError +from opencollab.adapters._vm_leases import _VMLeases +from opencollab.adapters.env import Environment +from opencollab.adapters.vm_config import VMResource +from opencollab.adapters.worktree_pool import WorktreePool +from opencollab.application.async_timeout import await_owned_operation +from opencollab.application.exception_notes import add_exception_note +from opencollab.domain.identity import role_storage_slug, validate_role_identity + + +@dataclass +class _VMOwner: + resource: VMResource + token: str + control: VMEnvironment + root: str + source_pools: list[WorktreePool] = field(default_factory=list) + views: list[VMEnvironment] = field(default_factory=list) + pool: WorktreePool | None = None + reset: VMEnvironment | None = None + reset_step: int = 0 + reset_operation: str | None = None + cleaned: bool = False + + +class VMIsolationPool: + def __init__( + self, workspace: str, *, resources: Sequence[VMResource], lease_database: str, + base_environment: Environment | None = None, scope: str = "group", wait_timeout: float | None = 0, + owned_save_dir: str | None = None, enable_adoption: bool = False, + owned_output_paths: Sequence[str] = (), + ) -> None: + if scope not in {"group", "session"}: + raise ValueError("VM isolation scope must be group or session") + resources = tuple(resources) + if not resources or any(not isinstance(resource, VMResource) for resource in resources): + raise ValueError("VM isolation requires at least one VMResource") + if len({resource.resource_id for resource in resources}) != len(resources): + raise ValueError("VM resources must have unique stable resource IDs") + if any(resource.reset_command is None for resource in resources): + raise ValueError("VM isolation requires an explicit reset_command for every resource before reuse") + self._registry = _VMLeases(lease_database, resources, wait_timeout) + self._workspace = workspace + self._base_environment = base_environment + self._scope = scope + self._owned_save_dir = owned_save_dir + self._owned_output_paths = tuple(owned_output_paths) + self.artifact_directory = str(Path(owned_save_dir).absolute()) if owned_save_dir else tempfile.mkdtemp( + prefix="opencollab-vm-results-", + ) + self._owners: list[_VMOwner] = [] + self._leases: dict[Environment, _VMOwner] = {} + self._exports: dict[Environment, Path] = {} + self._exported: set[Environment] = set() + self._group: _VMOwner | None = None + self._group_lock = asyncio.Lock() + self._lifecycle = asyncio.Condition() + self._inflight = 0 + self._closed = False + self._closing = asyncio.Event() + self._release_lock = asyncio.Lock() + self._delivery = WorktreePool( + workspace, use_worktrees=False, base_environment=base_environment, enable_adoption=enable_adoption, + ) + + @property + def retained_resources(self) -> tuple[str, ...]: + return tuple(owner.resource.resource_id for owner in self._owners) + + async def _create_owner(self, role: str, source: Environment | None) -> _VMOwner: + token = uuid.uuid4().hex + resource = await self._registry.acquire(token, self._closing) + control = VMEnvironment(resource.connection, resource.workspace_root, state_directory=resource.state_directory) + owner = _VMOwner(resource, token, control, f"{resource.workspace_root}/{token}") + self._owners.append(owner) + try: + await self._prepare_owner(owner, role, source) + except BaseException as failure: + try: + await await_owned_operation(self._cleanup_owner(owner)) + except BaseException as cleanup_error: + add_exception_note(failure, f"VM lease retained for cleanup retry: {cleanup_error}") + raise + return owner + + async def _prepare_owner(self, owner: _VMOwner, role: str, source: Environment | None) -> None: + control = owner.control + await control.setup() + source_pool = WorktreePool( + self._workspace, use_worktrees=True, base_environment=self._base_environment, + enable_adoption=True, owned_save_dir=self._owned_save_dir, + owned_output_paths=self._owned_output_paths, + ) + owner.source_pools.append(source_pool) + candidate = await source_pool.acquire(role, source=source) + view = await self._import(owner, candidate, "source") + owner.pool = WorktreePool(view.workspace, use_worktrees=True, base_environment=view, enable_adoption=True) + await source_pool.release() + owner.source_pools.remove(source_pool) + + async def _import(self, owner: _VMOwner, source: Environment, component: str) -> VMEnvironment: + path = posixpath.join(owner.root, component) + _complete(await owner.control.exec_cmd(f"mkdir -p -- {shlex.quote(path)}"), "VM workspace directory") + target = VMEnvironment(owner.resource.connection, path, state_directory=owner.resource.state_directory) + owner.views.append(target) + workspace = await _import_snapshot(source, target) + if workspace == target.workspace: + return target + view = VMEnvironment(owner.resource.connection, workspace, state_directory=owner.resource.state_directory) + owner.views.append(view) + await view.setup() + return view + + async def acquire(self, role: str, *, source: Environment | None = None) -> Environment: + role = validate_role_identity(role) + async with self._lifecycle: + if self._closed: + raise RuntimeError("VM isolation pool is closed") + self._inflight += 1 + owner = None + created = False + try: + # Only group creation is serialized. A session waiting for capacity + # must leave sibling release operations free to return a VM. + async with (self._group_lock if self._scope == "group" else nullcontext()): + owner = self._group if self._scope == "group" else None + if owner is None: + owner = await self._create_owner(role, source) + created = True + assert owner.pool is not None + review_source = None + if source is not None and not created: + if self._leases.get(source) is owner: + review_source = source + else: + review_source = await self._import(owner, source, f"import-{uuid.uuid4().hex}") + environment = await owner.pool.acquire(role, source=review_source) + self._leases[environment] = owner + component = f"{role_storage_slug(role)}-{uuid.uuid4().hex[:8]}" + self._exports[environment] = ( + Path(self.artifact_directory) / "vm" / owner.token / component + ) + environment.resource_id = owner.resource.resource_id + environment.export_directory = str(self._exports[environment]) + if self._scope == "group": + self._group = owner + if self._closed: + raise RuntimeError("VM isolation pool closed during acquisition") + return environment + except BaseException as failure: + if created and owner is not None: + try: + await await_owned_operation(self._cleanup_owner(owner)) + except BaseException as cleanup_error: + add_exception_note(failure, f"VM lease retained for cleanup retry: {cleanup_error}") + raise + finally: + async with self._lifecycle: + self._inflight -= 1 + self._lifecycle.notify_all() + + async def _export(self, env: Environment, owner: _VMOwner) -> None: + if not isinstance(env, _TeamCandidateEnvironment): + raise TypeError("VM pool expected a candidate workspace") + for view in (owner.control, *owner.views, env._lease.base_environment): + for operation_id in view.pending_operations: + try: + await view.recover_operation(operation_id) + except BaseException: + if operation_id in view.pending_operations: + raise + if env in self._exported: + return + env.retain_changes() + await env._delegate.cleanup() + await env.capture_changes() + view = VMEnvironment(owner.resource.connection, env.workspace, state_directory=owner.resource.state_directory) + owner.views.append(view) + destination = self._exports[env] + destination.parent.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=".vm-export-", dir=destination.parent)) + try: + async with _snapshot_contents(view) as (contents, _prefix): + await await_owned_operation(asyncio.to_thread( + shutil.copytree, contents, staging, symlinks=True, dirs_exist_ok=True, + ), propagate_cancellation=True) + os.replace(staging, destination) + self._exported.add(env) + except BaseException: + env.retain_changes() + raise + finally: + if staging.exists(): + await await_owned_operation(asyncio.to_thread(shutil.rmtree, staging), propagate_cancellation=True) + + async def _cleanup_owner(self, owner: _VMOwner) -> None: + if not owner.cleaned: + for env, value in tuple(self._leases.items()): + if value is owner: + await self._export(env, owner) + if owner.pool is not None: + await owner.pool.release() + for source_pool in tuple(owner.source_pools): + await source_pool.release() + owner.source_pools.remove(source_pool) + for view in owner.views: + await view.cleanup() + await owner.control.cleanup() + owner.cleaned = True + if owner.reset is None: + owner.reset = VMEnvironment( + owner.resource.connection, owner.resource.workspace_root, + state_directory=owner.resource.state_directory, + ) + assert owner.resource.reset_command is not None + commands = (owner.resource.reset_command, f"rm -rf -- {shlex.quote(owner.root)}") + while owner.reset_step < len(commands): + try: + if owner.reset_operation is not None: + result = await owner.reset.recover_operation(owner.reset_operation) + else: + owner.reset_operation = uuid.uuid4().hex + result = await owner.reset.exec_operation( + owner.reset_operation, commands[owner.reset_step], timeout=900, + ) + except VMOperationError: + if owner.reset_operation not in owner.reset.pending_operations: + owner.reset_operation = None + raise + if result.returncode: + owner.reset_operation = None + _complete(result, "VM guest reset" if owner.reset_step == 0 else "VM workspace reset") + owner.reset_step += 1 + owner.reset_operation = None + await owner.reset.cleanup() + self._registry.release(owner.resource, owner.token) + self._owners.remove(owner) + for env, value in tuple(self._leases.items()): + if value is owner: + self._leases.pop(env) + self._exports.pop(env) + self._exported.discard(env) + if self._group is owner: + self._group = None + + async def adopt_captured(self, patch: str, physical_paths: Sequence[str]) -> None: + await self._delivery.adopt_captured(patch, physical_paths) + + async def adopt(self, patch: str) -> None: + await self._delivery.adopt(patch) + + async def release_env(self, env: Environment) -> None: + async with self._release_lock: + owner = self._leases.get(env) + if owner is None: + return + await await_owned_operation(self._export(env, owner), propagate_cancellation=True) + assert owner.pool is not None + await await_owned_operation(owner.pool.release_env(env), propagate_cancellation=True) + self._leases.pop(env) + self._exports.pop(env) + self._exported.discard(env) + if owner is not self._group and owner not in self._leases.values(): + await await_owned_operation(self._cleanup_owner(owner), propagate_cancellation=True) + + async def release(self) -> None: + self._closed = True + self._closing.set() + await await_owned_operation(self._release_owned(), propagate_cancellation=True) + + async def _release_owned(self) -> None: + async with self._lifecycle: + while self._inflight: + await self._lifecycle.wait() + async with self._release_lock: + errors: list[BaseException] = [] + for owner in reversed(tuple(self._owners)): + try: + await self._cleanup_owner(owner) + except BaseException as exc: + errors.append(exc) + if errors: + raise OSError("VM cleanup or reset failed; leases and workspaces retained for retry") from errors[0] + await self._delivery.release() + + async def cleanup(self) -> None: + await self.release() + + +__all__ = ["VMCapacityError", "VMIsolationPool"] diff --git a/opencollab/adapters/worktree_pool.py b/opencollab/adapters/worktree_pool.py index cef15aa2..24586d36 100644 --- a/opencollab/adapters/worktree_pool.py +++ b/opencollab/adapters/worktree_pool.py @@ -13,6 +13,7 @@ from contextvars import ContextVar from opencollab.adapters._candidate_physical import _physical_paths +from opencollab.adapters._env_vm import VMEnvironment from opencollab.adapters._git_patch_bytes import _complete from opencollab.adapters._team_candidate_environment import _TeamCandidateEnvironment from opencollab.adapters.candidate_workspace import EnvCandidateWorkspace, _apply_increment @@ -83,6 +84,7 @@ async def acquire(self, role: str, *, source: Environment | None = None) -> Envi """Lend an environment, optionally snapshotting a completed teammate's files.""" role = validate_role_identity(role) base = self._base_environment if source is None else source + base_delegate = base._delegate if isinstance(base, _TeamCandidateEnvironment) else base if not self._use_worktrees: # Without isolation every agent works where the run works, which is # the supplied environment when there is one and the host workspace @@ -100,7 +102,7 @@ async def acquire(self, role: str, *, source: Environment | None = None) -> Envi source_owner = None try: async with (self._delivery_lock if self._enable_adoption else nullcontext()): - if self._enable_adoption or source is not None: + if self._enable_adoption or source is not None or isinstance(base_delegate, VMEnvironment): candidate_source = self._delivery_environment() if source is None else source if source is not None: if getattr(source, "revoked", False): @@ -122,6 +124,14 @@ async def acquire(self, role: str, *, source: Environment | None = None) -> Envi timeout_returncode=candidate_source._timeout_returncode, ) candidate_source = source_owner + if isinstance(candidate_source, _TeamCandidateEnvironment): + candidate_source = candidate_source._delegate + if isinstance(candidate_source, VMEnvironment) and source_owner is None: + source_owner = VMEnvironment( + candidate_source.connection, candidate_source.workspace, + state_directory=candidate_source.state_directory, + ) + candidate_source = source_owner probe = await candidate_source.exec_cmd( f"git -C {shlex.quote(candidate_source.workspace)} rev-parse --show-toplevel", ) diff --git a/opencollab/bootstrap/vm.py b/opencollab/bootstrap/vm.py new file mode 100644 index 00000000..b5c99811 --- /dev/null +++ b/opencollab/bootstrap/vm.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence +from functools import partial + +from opencollab.adapters._env_vm import VMEnvironment +from opencollab.adapters.vm_config import VMConnection, VMResource +from opencollab.adapters.vm_isolation import VMIsolationPool +from opencollab.application.ports import WorktreePoolPort + + +def vm_isolation_provider( + resources: Sequence[VMResource], *, lease_database: str, + scope: str = "group", wait_timeout: float | None = 0, +) -> Callable[..., WorktreePoolPort]: + if scope not in {"group", "session"}: + raise ValueError("VM isolation scope must be group or session") + return partial(VMIsolationPool, resources=tuple(resources), lease_database=lease_database, + scope=scope, wait_timeout=wait_timeout) + + +__all__ = ["VMConnection", "VMEnvironment", "VMResource", "vm_isolation_provider"] diff --git a/tests/runtime/test_vm_git_temporary.py b/tests/runtime/test_vm_git_temporary.py new file mode 100644 index 00000000..9b4efb69 --- /dev/null +++ b/tests/runtime/test_vm_git_temporary.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +from opencollab.adapters import _vm_remote + + +@pytest.fixture +def temporary_workspace(tmp_path, monkeypatch): + workspace = tmp_path / "workspace" + workspace.mkdir() + state = tmp_path / "state" + state.mkdir() + identity = {"device": workspace.stat().st_dev, "inode": workspace.stat().st_ino} + monkeypatch.setattr(_vm_remote, "_identity", lambda path: identity) + + async def perform(operation, action, arguments): + directory = state / operation + directory.mkdir(exist_ok=True) + return await _vm_remote._perform({"workspace": str(workspace), "identity": identity, + "action": action, "args": arguments}, directory) + + return state, perform + + +async def test_mutable_temporary_pins_directory_and_cleans_replaced_files_without_following_links( + tmp_path, temporary_workspace, +): + state, perform = temporary_workspace + outside = tmp_path / "outside" + outside.write_text("keep") + path = Path(await perform("create", "mutable_temporary", {"name": "git.index", "content": "initial"})) + assert path.parent.stat().st_mode & 0o777 == 0o700 + inode = path.stat().st_ino + replacement = path.with_name("replacement") + replacement.write_text("Git rewrites its index") + os.replace(replacement, path) + assert path.stat().st_ino != inode + (path.parent / "external-link").symlink_to(outside) + await perform("remove", "remove_mutable_temporary", {"operation_id": "create"}) + assert not (state / "create" / "temporary").exists() + assert outside.read_text() == "keep" + + +async def test_mutable_temporary_retains_replaced_directory_until_original_is_restored(temporary_workspace): + state, perform = temporary_workspace + path = Path(await perform("create", "mutable_temporary", {"name": "git.index", "content": "initial"})) + saved = path.parent.with_name("saved") + path.parent.rename(saved) + path.parent.mkdir() + (path.parent / "replacement").write_text("retain") + with pytest.raises(OSError, match="directory was replaced"): + await perform("remove", "remove_mutable_temporary", {"operation_id": "create"}) + assert (path.parent / "replacement").read_text() == "retain" + (path.parent / "replacement").unlink() + path.parent.rmdir() + saved.rename(path.parent) + await perform("retry", "remove_mutable_temporary", {"operation_id": "create"}) + assert not (state / "create" / "temporary").exists() diff --git a/tests/runtime/test_vm_isolation.py b/tests/runtime/test_vm_isolation.py new file mode 100644 index 00000000..488e3658 --- /dev/null +++ b/tests/runtime/test_vm_isolation.py @@ -0,0 +1,400 @@ +from __future__ import annotations + +import asyncio +import errno +import os +import shutil +import sqlite3 +import subprocess +import sys +import tempfile +from contextlib import asynccontextmanager +from pathlib import Path + +import pytest + +from opencollab.adapters import candidate_workspace, vm_isolation, worktree_pool +from opencollab.adapters._env_vm import VMEnvironment, VMOperationError, VMTransportError +from opencollab.adapters._vm_leases import VMCapacityError, _VMLeases +from opencollab.adapters.env import LocalEnvironment +from opencollab.adapters.session_environments import SessionEnvironmentPool +from opencollab.adapters.vm_config import VMConnection, VMResource +from opencollab.adapters.vm_isolation import VMIsolationPool +from opencollab.bootstrap.vm import vm_isolation_provider + + +def _repository(path: Path) -> None: + path.mkdir() + subprocess.run(["git", "init", "--quiet", str(path)], check=True) + (path / "answer.txt").write_text("initial\n") + for arguments in (("add", "."), ("-c", "user.name=Tester", "-c", "user.email=test@example.com", + "commit", "--quiet", "-m", "initial")): + subprocess.run(["git", "-C", str(path), *arguments], check=True) + + +def _resource(path: Path, resource_id: str = "vm-one", *, reset_command: str = "true") -> VMResource: + path.mkdir(exist_ok=True) + return VMResource(resource_id, VMConnection("localhost", "test"), str(path), str(path.parent / "state"), + reset_command) + + +class _LocalVM(VMEnvironment): + copy_failure = False + + def __init__(self, connection, workspace, *, state_directory, run_id=None): + super().__init__(connection, workspace, state_directory=state_directory, run_id=run_id) + self.local = LocalEnvironment(workspace) + + async def setup(self, mount_dir=None): + self._ensure_active() + return self.workspace + + async def exec_cmd(self, cmd, timeout=120, *, input_bytes=None): + self._ensure_active() + return await self.local.exec_cmd(cmd, timeout, input_bytes=input_bytes) + + async def exec_operation(self, operation_id, cmd, timeout=120, *, input_bytes=None): + return await self.exec_cmd(cmd, timeout, input_bytes=input_bytes) + + async def write_temp_file(self, *args, **kwargs): + return await self.local.write_temp_file(*args, **kwargs) + + async def remove_file(self, path): + await self.local.remove_file(path) + + @asynccontextmanager + async def _mutable_temporary(self, content, *, prefix, suffix): + with tempfile.TemporaryDirectory(prefix="opencollab-vm-unit-") as directory: + path = Path(directory) / (prefix + suffix) + path.write_text(content) + yield str(path) + + async def copy_to_host(self, source, destination, *, directory=False): + if self.copy_failure: + raise OSError("transfer interrupted") + if directory: + shutil.copytree(source, destination, symlinks=True, dirs_exist_ok=True) + else: + shutil.copyfile(source, destination) + + async def copy_from_host(self, source, destination, *, directory=False): + if directory: + shutil.copytree(source, destination, symlinks=True, dirs_exist_ok=True) + else: + shutil.copyfile(source, destination) + + async def cleanup(self): + if self._uncertain: + raise VMOperationError("operation status unknown", operation_id="unknown", state="unknown") + self.revoke() + await self.local.cleanup() + + +@pytest.fixture +def local_vm(monkeypatch): + monkeypatch.setattr(vm_isolation, "VMEnvironment", _LocalVM) + monkeypatch.setattr(candidate_workspace, "VMEnvironment", _LocalVM) + monkeypatch.setattr(worktree_pool, "VMEnvironment", _LocalVM) + monkeypatch.setattr(_LocalVM, "copy_failure", False) + + +async def test_database_claim_is_exclusive_across_controller_processes_and_connection_aliases(tmp_path): + resource = _resource(tmp_path / "guest") + database = str(tmp_path / "leases.sqlite") + leases = _VMLeases(database, [resource], 0) + await leases.acquire("parent") + code = """ +import asyncio, sys +from opencollab.adapters._vm_leases import _VMLeases, VMCapacityError +from opencollab.adapters.vm_config import VMConnection, VMResource +resource = VMResource('vm-one', VMConnection('different-alias', 'test'), sys.argv[2], reset_command='true') +async def main(): + try: + await _VMLeases(sys.argv[1], [resource], 0).acquire('child') + except VMCapacityError: + print('exhausted') + else: + print('claimed') +asyncio.run(main()) +""" + def child(): + return subprocess.check_output( + [sys.executable, "-c", code, database, resource.workspace_root], text=True, + ).strip() + assert await asyncio.to_thread(child) == "exhausted" + leases.release(resource, "parent") + assert await asyncio.to_thread(child) == "claimed" + with sqlite3.connect(database) as connection: + assert connection.execute("SELECT resource_id, owner FROM vm_leases").fetchall() == [("vm-one", "child")] + + +async def test_resource_waiting_is_cancellable_without_claiming_the_occupied_vm(tmp_path): + resource = _resource(tmp_path / "guest") + database = str(tmp_path / "leases.sqlite") + immediate = _VMLeases(database, [resource], 0) + await immediate.acquire("parent") + waiting = _VMLeases(database, [resource], None) + task = asyncio.create_task(waiting.acquire("cancelled")) + await asyncio.sleep(0.08) + task.cancel("stop nested run") + with pytest.raises(asyncio.CancelledError) as cancelled: + await task + # asyncio forwards the cancellation message to awaiters from Python 3.11. + if sys.version_info >= (3, 11): + assert str(cancelled.value) == "stop nested run" + with sqlite3.connect(database) as connection: + assert connection.execute("SELECT owner FROM vm_leases").fetchall() == [("parent",)] + immediate.release(resource, "parent") + assert await waiting.acquire("next") == resource + + +async def test_group_sessions_are_separate_and_exports_keep_binary_symlinks_and_history( + tmp_path, local_vm, monkeypatch, +): + source = tmp_path / "source" + _repository(source) + (source / "payload.bin").write_bytes(b"\0\xff\x80") + (source / "shortcut").symlink_to("answer.txt") + (source / "executable").write_text("#!/bin/sh\nexit 0\n") + (source / "executable").chmod(0o755) + pool = VMIsolationPool(str(source), resources=[_resource(tmp_path / "guest")], + lease_database=str(tmp_path / "leases.sqlite"), owned_save_dir=str(tmp_path / "results")) + first, second = await pool.acquire("first"), await pool.acquire("second") + assert first.resource_id == second.resource_id == "vm-one" + assert first.workspace != second.workspace + await first.exec_cmd("printf 'changed\\n' > answer.txt && printf 'log\\n' > run.log") + assert (await second.exec_cmd("cat answer.txt")).stdout == "initial\n" + destination = Path(first.export_directory) + replace = os.replace + + def require_same_filesystem_staging(source, target): + if Path(source).parent != Path(target).parent: + raise OSError(errno.EXDEV, "cross-device link") + replace(source, target) + + monkeypatch.setattr(vm_isolation.os, "replace", require_same_filesystem_staging) + await pool.release() + assert not pool.retained_resources + assert (destination / "payload.bin").read_bytes() == b"\0\xff\x80" + assert (destination / "shortcut").is_symlink() + assert os.readlink(destination / "shortcut") == "answer.txt" + assert (destination / "executable").stat().st_mode & 0o777 == 0o755 + assert (destination / "run.log").read_text() == "log\n" + assert subprocess.check_output(["git", "-C", str(destination), "show", "HEAD:answer.txt"], text=True) == "initial\n" + assert (source / "answer.txt").read_text() == "initial\n" + + +async def test_session_scopes_claim_different_vms_and_nested_exhaustion_is_clear(tmp_path, local_vm): + source = tmp_path / "source" + _repository(source) + resources = [_resource(tmp_path / "guest-one", "one"), _resource(tmp_path / "guest-two", "two")] + provider = vm_isolation_provider(resources, lease_database=str(tmp_path / "leases.sqlite"), scope="session") + sessions = SessionEnvironmentPool(str(source), isolation="vm-session", providers={"vm-session": provider}) + first, second = await sessions.acquire("first"), await sessions.acquire("second") + assert first.resource_id != second.resource_id + child = sessions.child_pool(first, role="child") + with pytest.raises(VMCapacityError, match="capacity exhausted"): + await child.acquire("nested") + await child.release() + await sessions.release() + + +async def test_export_failure_retains_workspace_and_lease_until_retry(tmp_path, local_vm, monkeypatch): + source = tmp_path / "source" + _repository(source) + resource = _resource(tmp_path / "guest") + database = str(tmp_path / "leases.sqlite") + pool = VMIsolationPool(str(source), resources=[resource], lease_database=database) + environment = await pool.acquire("worker") + await environment.exec_cmd("printf 'candidate\\n' > answer.txt") + monkeypatch.setattr(_LocalVM, "copy_failure", True) + with pytest.raises(OSError, match="retained"): + await pool.release() + assert Path(environment.workspace).is_dir() + assert pool.retained_resources == (resource.resource_id,) + with pytest.raises(VMCapacityError): + await _VMLeases(database, [resource], 0).acquire("other") + monkeypatch.setattr(_LocalVM, "copy_failure", False) + await pool.release() + assert (Path(environment.export_directory) / "answer.txt").read_text() == "candidate\n" + assert not Path(environment.workspace).exists() + + +@pytest.mark.parametrize("stage", ["capture", "cleanup"]) +async def test_lost_candidate_reply_recovers_before_export_or_cleanup_retry(tmp_path, local_vm, monkeypatch, stage): + source = tmp_path / "source" + _repository(source) + resource = _resource(tmp_path / "guest") + database = str(tmp_path / "leases.sqlite") + pool = VMIsolationPool(str(source), resources=[resource], lease_database=database, scope="session") + environment = await pool.acquire("worker") + await environment.exec_cmd("printf candidate > answer.txt") + control = environment._lease.base_environment + execute = control.exec_cmd + completed = {} + recovered = [] + available = False + + async def lose_reply(cmd, timeout=120, *, input_bytes=None): + result = await execute(cmd, timeout, input_bytes=input_bytes) + selected = environment._delegate.revoked if stage == "capture" else "worktree list --porcelain" in cmd + if selected and not completed: + completed["candidate-reply"] = result + control._operations.add("candidate-reply") + control._uncertain.add("candidate-reply") + raise VMTransportError("candidate reply lost", operation_id="candidate-reply") + return result + + async def recover(operation_id): + if not available: + raise VMOperationError("operation status unknown", operation_id=operation_id, state="unknown") + recovered.append(operation_id) + control._operations.remove(operation_id) + control._uncertain.remove(operation_id) + return completed[operation_id] + + monkeypatch.setattr(control, "exec_cmd", lose_reply) + monkeypatch.setattr(control, "recover_operation", recover) + destination = Path(environment.export_directory) + with pytest.raises(VMTransportError): + await pool.release_env(environment) + assert destination.exists() is (stage == "cleanup") + with pytest.raises(VMOperationError): + await pool.release_env(environment) + with pytest.raises(VMCapacityError): + await _VMLeases(database, [resource], 0).acquire("other") + assert Path(environment.workspace).exists() + available = True + await pool.release_env(environment) + assert recovered == ["candidate-reply"] + assert (destination / "answer.txt").read_text() == "candidate" + assert not pool.retained_resources + assert await _VMLeases(database, [resource], 0).acquire("next") == resource + await pool.release() + + +async def test_reset_failure_retains_lease_and_restores_guest_before_next_claim(tmp_path, local_vm): + source = tmp_path / "source" + _repository(source) + marker, ready = tmp_path / "global-state", tmp_path / "reset-ready" + marker.write_text("dirty") + resource = _resource(tmp_path / "guest", reset_command=f"test -f '{ready}' && rm -f '{marker}'") + database = str(tmp_path / "leases.sqlite") + pool = VMIsolationPool(str(source), resources=[resource], lease_database=database) + await pool.acquire("worker") + with pytest.raises(OSError, match="retained"): + await pool.release() + assert marker.exists() + with pytest.raises(VMCapacityError): + await _VMLeases(database, [resource], 0).acquire("other") + ready.touch() + await pool.release() + assert not marker.exists() + assert await _VMLeases(database, [resource], 0).acquire("next") == resource + + +async def test_release_interrupts_configured_capacity_wait(tmp_path, local_vm): + resource = _resource(tmp_path / "guest") + database = str(tmp_path / "leases.sqlite") + registry = _VMLeases(database, [resource], 0) + await registry.acquire("outer") + pool = VMIsolationPool(str(tmp_path), resources=[resource], lease_database=database, wait_timeout=None) + acquisition = asyncio.create_task(pool.acquire("nested")) + await asyncio.sleep(0.08) + await asyncio.wait_for(pool.release(), 1) + with pytest.raises(RuntimeError, match="closed"): + await acquisition + registry.release(resource, "outer") + + +async def test_sessions_prepare_independently_and_release_capacity_to_waiters(tmp_path, local_vm, monkeypatch): + source = tmp_path / "source" + _repository(source) + resources = [_resource(tmp_path / "guest-one", "one"), _resource(tmp_path / "guest-two", "two")] + pool = VMIsolationPool(str(source), resources=resources, lease_database=str(tmp_path / "leases.sqlite"), + scope="session", wait_timeout=None) + entered, proceed = asyncio.Event(), asyncio.Event() + original = pool._prepare_owner + preparing = 0 + + async def prepare(owner, role, parent): + nonlocal preparing + preparing += 1 + if preparing == 2: + entered.set() + await proceed.wait() + await original(owner, role, parent) + + monkeypatch.setattr(pool, "_prepare_owner", prepare) + tasks = [asyncio.create_task(pool.acquire(role)) for role in ("first", "second")] + try: + await asyncio.wait_for(entered.wait(), 1) + finally: + proceed.set() + first, second = await asyncio.gather(*tasks) + waiter = asyncio.create_task(pool.acquire("third")) + await asyncio.sleep(0.08) + assert not waiter.done() + await pool.release_env(first) + third = await asyncio.wait_for(waiter, 10) + assert third.resource_id == first.resource_id + assert second.resource_id != third.resource_id + await pool.release() + + +async def test_unknown_execution_retains_the_vm_until_recovery(tmp_path, local_vm): + source = tmp_path / "source" + _repository(source) + resource = _resource(tmp_path / "guest") + database = str(tmp_path / "leases.sqlite") + pool = VMIsolationPool(str(source), resources=[resource], lease_database=database) + environment = await pool.acquire("worker") + environment._delegate._uncertain.add("unknown") + with pytest.raises(OSError, match="retained"): + await pool.release() + assert Path(environment.workspace).exists() + with pytest.raises(VMCapacityError): + await _VMLeases(database, [resource], 0).acquire("other") + environment._delegate._uncertain.clear() + await pool.release() + + +async def test_lost_reset_reply_recovers_the_same_operation_without_repeating_reset(tmp_path, local_vm, monkeypatch): + source = tmp_path / "source" + _repository(source) + counter = tmp_path / "reset-count" + command = f"printf x >> '{counter}'" + resource = _resource(tmp_path / "guest", reset_command=command) + pool = VMIsolationPool(str(source), resources=[resource], lease_database=str(tmp_path / "leases.sqlite")) + await pool.acquire("worker") + original = _LocalVM.exec_operation + completed = {} + + async def lose_reply(self, operation_id, cmd, timeout=120, *, input_bytes=None): + result = await original(self, operation_id, cmd, timeout, input_bytes=input_bytes) + if cmd == command: + completed[operation_id] = result + self._operations.add(operation_id) + self._uncertain.add(operation_id) + raise VMTransportError("reset response lost", operation_id=operation_id) + return result + + async def recover(self, operation_id): + self._operations.remove(operation_id) + self._uncertain.remove(operation_id) + return completed[operation_id] + + monkeypatch.setattr(_LocalVM, "exec_operation", lose_reply) + monkeypatch.setattr(_LocalVM, "recover_operation", recover) + with pytest.raises(OSError, match="retained"): + await pool.release() + assert counter.read_text() == "x" + await pool.release() + assert counter.read_text() == "x" + assert not pool.retained_resources + + +def test_reusable_pool_requires_explicit_guest_reset(tmp_path): + resource = VMResource("one", VMConnection("localhost", "test")) + with pytest.raises(ValueError, match="reset_command"): + VMIsolationPool(str(tmp_path), resources=[resource], lease_database=str(tmp_path / "leases.sqlite")) diff --git a/tests/runtime/test_vm_pool_ssh.py b/tests/runtime/test_vm_pool_ssh.py new file mode 100644 index 00000000..19d5ba66 --- /dev/null +++ b/tests/runtime/test_vm_pool_ssh.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +import asyncio +import json +import os +import shlex +import subprocess +import sys +import uuid +from dataclasses import asdict +from pathlib import Path + +import pytest +import pytest_asyncio + +from opencollab.adapters._env_vm import VMEnvironment, VMTransportError +from opencollab.adapters._git_patch_bytes import _patch_temporary +from opencollab.adapters._vm_leases import VMCapacityError, _VMLeases +from opencollab.adapters.vm_config import VMConnection, VMResource +from opencollab.adapters.vm_isolation import VMIsolationPool +from opencollab.adapters.worktree_pool import WorktreePool + + +def _repository(path): + path.mkdir() + subprocess.run(["git", "init", "--quiet", str(path)], check=True) + (path / "answer.txt").write_text("initial\n") + subprocess.run(["git", "-C", str(path), "add", "."], check=True) + subprocess.run(["git", "-C", str(path), "-c", "user.name=Tester", "-c", "user.email=test@example.com", + "commit", "--quiet", "-m", "initial"], check=True) + subprocess.run(["git", "-C", str(path), "tag", "v1.0"], check=True) + + +@pytest_asyncio.fixture +async def resources(): + configuration = os.environ.get("OPENCOLLAB_TEST_VM_POOL_CONFIG") + if configuration is None: + pytest.skip("OPENCOLLAB_TEST_VM_POOL_CONFIG must configure two borrowed Linux guests") + entries = json.loads(Path(configuration).read_text())["resources"] + selected = [] + parents = [] + for entry in entries: + connection = VMConnection(**entry["connection"]) + token = uuid.uuid4().hex + parent = VMEnvironment(connection, entry["workspace_root"], state_directory=entry["state_directory"]) + root = f"{parent.workspace}/oc-pr2-{token}" + assert (await parent.exec_cmd(f"mkdir -m 700 -- {shlex.quote(root)}")).returncode == 0 + parents.append((parent, root)) + selected.append(VMResource(entry["resource_id"], connection, root, entry["state_directory"], + f"rm -f -- {shlex.quote(root + '.global')}")) + assert len(selected) >= 2 + try: + yield selected + finally: + for parent, root in parents: + await parent.exec_cmd(f"rm -rf -- {shlex.quote(root)} {shlex.quote(root + '.global')}") + await parent.cleanup() + + +async def test_ssh_group_worktrees_and_self_contained_candidate_exports(tmp_path, resources): + source = tmp_path / "source" + _repository(source) + (source / "payload.bin").write_bytes(bytes(range(256)) * 4097) + (source / "shortcut").symlink_to("answer.txt") + executable = source / "executable" + executable.write_text("#!/bin/sh\nprintf executable\n") + executable.chmod(0o755) + pool = VMIsolationPool(str(source), resources=resources[-1:], lease_database=str(tmp_path / "leases.sqlite"), + owned_save_dir=str(tmp_path / "results")) + first, second = await pool.acquire("first"), await pool.acquire("second") + assert first.resource_id == second.resource_id + assert first.workspace != second.workspace + assert (await first.exec_cmd("printf changed > answer.txt; printf log > run.log")).returncode == 0 + assert (await second.exec_cmd("cat answer.txt")).stdout == "initial\n" + child_pool = WorktreePool(first.workspace, use_worktrees=True, base_environment=first) + child = await child_pool.acquire("nested") + assert child.workspace != first.workspace + assert (await child.exec_cmd("cat answer.txt")).stdout == "changed" + assert (await child.exec_cmd("printf child > answer.txt")).returncode == 0 + assert (await first.exec_cmd("cat answer.txt")).stdout == "changed" + await child_pool.release() + before = await first.capture_changes() + assert (await first.exec_cmd("printf second > answer.txt; " + "python3 -c \"from pathlib import Path; p=Path('payload.bin'); " + "data=bytearray(p.read_bytes()); data[3]=77; p.write_bytes(data)\"")).returncode == 0 + after = await first.capture_changes() + increment = await first.diff_captured(before, after) + assert "+second" in increment + assert "GIT binary patch" in increment + destination = Path(first.export_directory) + await pool.release() + expected = bytearray((source / "payload.bin").read_bytes()) + expected[3] = 77 + assert (destination / "payload.bin").read_bytes() == expected + assert (destination / "shortcut").is_symlink() + assert (destination / "executable").stat().st_mode & 0o777 == 0o755 + assert (destination / "run.log").read_text() == "log" + assert subprocess.check_output(["git", "-C", str(destination), "show", "HEAD:answer.txt"], text=True) == "initial\n" + assert subprocess.check_output(["git", "-C", str(destination), "describe", "--tags"], text=True).strip() == "v1.0" + + +async def test_ssh_two_guests_and_independent_controller_compete_for_stable_resource(tmp_path, resources): + source = tmp_path / "source" + _repository(source) + database = str(tmp_path / "leases.sqlite") + pool = VMIsolationPool(str(source), resources=resources, lease_database=database, scope="session") + first, second = await asyncio.gather(pool.acquire("first"), pool.acquire("second")) + assert first.resource_id != second.resource_id + identities = await asyncio.gather(first.exec_cmd("cat /proc/sys/kernel/random/boot_id"), + second.exec_cmd("cat /proc/sys/kernel/random/boot_id")) + assert identities[0].stdout.strip() != identities[1].stdout.strip() + encoded = json.dumps({"resources": [asdict(resource) for resource in resources], "database": database, + "workspace": str(source)}) + code = """ +import asyncio, json, sys +from opencollab.bootstrap.vm import VMConnection, VMResource +from opencollab.adapters.vm_isolation import VMIsolationPool, VMCapacityError +config = json.loads(sys.argv[1]) +resources = [VMResource(**{**item, 'connection': VMConnection(**item['connection'])}) for item in config['resources']] +async def main(): + pool = VMIsolationPool(config['workspace'], resources=resources, lease_database=config['database']) + try: + await pool.acquire('independent-controller') + except VMCapacityError: + print('capacity exhausted') + finally: + await pool.release() +asyncio.run(main()) +""" + result = await asyncio.to_thread(subprocess.check_output, [sys.executable, "-c", code, encoded], text=True) + assert result.strip() == "capacity exhausted" + with pytest.raises(VMCapacityError): + await pool.acquire("nested-needs-another-guest", source=first) + marker = shlex.quote(resources[0].workspace_root + ".global") + assert (await first.exec_cmd("printf dirty > " + marker)).returncode == 0 + await pool.release_env(first) + next_environment = await pool.acquire("next") + assert next_environment.resource_id == first.resource_id + assert (await next_environment.exec_cmd("test ! -e " + marker)).returncode == 0 + await pool.release() + + +async def test_ssh_unknown_operation_export_failure_and_reset_failure_keep_lease(tmp_path, resources, monkeypatch): + source = tmp_path / "source" + _repository(source) + original = resources[-1] + ready = original.workspace_root + "/reset-ready" + marker = original.workspace_root + ".global" + resource = VMResource(original.resource_id, original.connection, original.workspace_root, original.state_directory, + f"test -f {shlex.quote(ready)} && rm -f -- {shlex.quote(marker)}") + database = str(tmp_path / "leases.sqlite") + pool = VMIsolationPool(str(source), resources=[resource], lease_database=database) + environment = await pool.acquire("worker") + delegate = environment._delegate + operation_id = uuid.uuid4().hex + assert (await environment.exec_cmd("printf dirty > " + shlex.quote(marker))).returncode == 0 + await delegate.setup() + await delegate._rpc({"method": "start", "operation_id": operation_id, "request": { + "workspace": delegate.workspace, "identity": delegate._identity, "action": "exec", + "args": {"cmd": "printf candidate > answer.txt; sleep 30", "timeout": 60, "input": None}}}) + delegate._operations.add(operation_id) + delegate._uncertain.add(operation_id) + rpc = delegate._rpc + + async def disconnected(payload): + raise VMTransportError("connection interrupted", operation_id=operation_id) + + monkeypatch.setattr(delegate, "_rpc", disconnected) + with pytest.raises(OSError, match="retained"): + await pool.release() + with pytest.raises(VMCapacityError): + await _VMLeases(database, [resource], 0).acquire("other-controller") + monkeypatch.setattr(delegate, "_rpc", rpc) + transfer = VMEnvironment.copy_to_host + + async def fail_export(self, source, target, *, directory=False): + raise OSError("export destination unavailable") + + monkeypatch.setattr(VMEnvironment, "copy_to_host", fail_export) + with pytest.raises(OSError, match="retained"): + await pool.release() + observer = VMEnvironment(resource.connection, resource.workspace_root, state_directory=resource.state_directory) + assert (await observer.exec_cmd("test -d " + shlex.quote(environment.workspace))).returncode == 0 + monkeypatch.setattr(VMEnvironment, "copy_to_host", transfer) + with pytest.raises(OSError, match="retained"): + await pool.release() + assert (await observer.exec_cmd("test -f " + shlex.quote(marker))).returncode == 0 + with pytest.raises(VMCapacityError): + await _VMLeases(database, [resource], 0).acquire("other-controller") + assert (await observer.exec_cmd("touch " + shlex.quote(ready))).returncode == 0 + await pool.release() + assert (Path(environment.export_directory) / "answer.txt").read_text() == "candidate" + assert (await observer.exec_cmd("test ! -e " + shlex.quote(marker))).returncode == 0 + assert not pool.retained_resources + await observer.cleanup() + + +async def test_ssh_git_temporary_allows_atomic_files_and_rejects_replaced_directory(resources): + resource = resources[-1] + environment = VMEnvironment(resource.connection, resource.workspace_root, state_directory=resource.state_directory) + outside = resource.workspace_root + "/outside" + assert (await environment.exec_cmd("printf keep > " + shlex.quote(outside))).returncode == 0 + async with _patch_temporary(environment, environment.workspace, "before", prefix="git-", suffix=".index", + repository=False) as path: + directory = str(Path(path).parent) + old_inode = (await environment.exec_cmd("stat -c %i " + shlex.quote(path))).stdout + assert (await environment.exec_cmd("stat -c %a " + shlex.quote(directory))).stdout.strip() == "700" + command = (f"printf replacement > {shlex.quote(path + '.new')} && " + f"mv {shlex.quote(path + '.new')} {shlex.quote(path)} && " + f"mkdir {shlex.quote(path + '.parts')} && " + f"ln -s {shlex.quote(outside)} {shlex.quote(path + '.parts/link')}") + assert (await environment.exec_cmd(command)).returncode == 0 + assert (await environment.exec_cmd("stat -c %i " + shlex.quote(path))).stdout != old_inode + assert (await environment.exec_cmd("test ! -e " + shlex.quote(directory))).returncode == 0 + assert (await environment.exec_cmd("cat " + shlex.quote(outside))).stdout == "keep" + context = environment._mutable_temporary("before", prefix="git-", suffix=".index") + path = await context.__aenter__() + directory = str(Path(path).parent) + command = (f"mv {shlex.quote(directory)} {shlex.quote(directory + '.saved')} && " + f"mkdir {shlex.quote(directory)} && printf replacement > {shlex.quote(directory + '/other')}") + assert (await environment.exec_cmd(command)).returncode == 0 + with pytest.raises(OSError, match="directory was replaced"): + await context.__aexit__(None, None, None) + assert (await environment.exec_cmd("cat " + shlex.quote(directory + "/other"))).stdout == "replacement" + assert environment._mutable_temporaries + command = f"rm -r {shlex.quote(directory)} && mv {shlex.quote(directory + '.saved')} {shlex.quote(directory)}" + assert (await environment.exec_cmd(command)).returncode == 0 + await environment.cleanup() + assert not environment._mutable_temporaries