feat(environments): allocate virtual machines by group or session

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,
-        )

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)

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(

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}")

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:

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)

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")