cs249r_book/binder/cli/core/workspace.py

"""Disposable git worktrees that keep concurrent binder builds apart.

A build writes at the Quarto project root (``_quarto.yml``, ``index.qmd``,
LaTeX intermediates, the ``.quarto`` cache), so two builds in one checkout
collide. A workspace is a detached ``git worktree`` of a snapshot of the
invoking checkout, created under the system temporary directory and removed
when the run ends. The snapshot carries staged and unstaged edits, and
untracked, non-ignored files are copied in, so a workspace builds what is on
disk rather than only what is committed.
"""

from __future__ import annotations

import os
import re
import shutil
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Tuple

#: Parent directory of every workspace binder creates. Workspaces are deleted
#: recursively, so :func:`create_workspace` and :func:`remove_workspace`
#: refuse paths outside it.
WORKSPACE_ROOT = Path(tempfile.gettempdir()) / "binder-workspaces"

#: Empty hooks directory, shared by a run's workspaces, that ``git worktree add``
#: is pointed at so repository hooks do not run in a throwaway checkout.
_NO_HOOKS = ".no-hooks"


def workspace_path(run_id: str, name: str) -> Path:
    """Return the directory for workspace *name* of run *run_id* under :data:`WORKSPACE_ROOT`."""
    return WORKSPACE_ROOT / run_id / name


def _git(repo_root: Path, *args: str, check: bool = True) -> str:
    """Run ``git -C repo_root <args>`` and return its standard output.

    Raises:
        RuntimeError: If *check* is set and git exits non-zero.
    """
    result = subprocess.run(["git", "-C", str(repo_root), *args], capture_output=True, text=True)
    if check and result.returncode != 0:
        detail = (result.stderr or result.stdout).strip()
        raise RuntimeError(f"git {' '.join(args)} failed: {detail}")
    return result.stdout


@dataclass(frozen=True)
class Snapshot:
    """Content that every workspace of one run is created from.

    Attributes:
        repo_root: Checkout the snapshot was taken from.
        commit: Commit holding the tracked content, uncommitted edits included.
        untracked: Repository-relative paths of untracked, non-ignored files.
    """

    repo_root: Path
    commit: str
    untracked: Tuple[str, ...] = ()


def take_snapshot(repo_root: Path) -> Snapshot:
    """Capture *repo_root*'s working tree for workspace creation.

    ``git stash create`` records staged and unstaged edits as an unreferenced
    commit without touching the shared stash list; a clean tree falls back to
    ``HEAD``. Untracked files are listed, not committed, and copied into each
    workspace by :func:`create_workspace`.
    """
    repo_root = Path(repo_root).resolve()
    commit = ""
    for attempt in range(5):
        try:
            commit = _git(repo_root, "stash", "create").strip()
            break
        except RuntimeError as err:
            if ("could not write index" in str(err) or "index.lock" in str(err)) and attempt < 4:
                import time
                time.sleep(0.3)
                continue
            raise
    if not commit:
        commit = _git(repo_root, "rev-parse", "HEAD").strip()
    listed = _git(repo_root, "ls-files", "--others", "--exclude-standard", "-z")
    return Snapshot(repo_root, commit, tuple(path for path in listed.split("\0") if path))


@dataclass
class Workspace:
    """A worktree created by :func:`create_workspace`."""

    path: Path
    snapshot: Snapshot


def create_workspace(snapshot: Snapshot, path: Path) -> Workspace:
    """Check out *snapshot* as a detached worktree at *path*.

    Repository hooks are disabled for the checkout; they set up developer
    tooling that a throwaway build tree does not need. Untracked files from the
    snapshot are copied in afterwards.

    Raises:
        ValueError: If *path* is not inside :data:`WORKSPACE_ROOT`.
        RuntimeError: If ``git worktree add`` fails.
    """
    path = _checked_path(path)
    no_hooks = path.parent / _NO_HOOKS
    no_hooks.mkdir(parents=True, exist_ok=True)
    try:
        _git(snapshot.repo_root, "-c", f"core.hooksPath={no_hooks}",
             "worktree", "add", "--detach", "--force", str(path), snapshot.commit)
    except RuntimeError:
        _remove_empty_run_dir(path.parent)
        raise
    for relative in snapshot.untracked:
        source = snapshot.repo_root / relative
        if source.is_file():
            target = path / relative
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(source, target)
    return Workspace(path, snapshot)


def remove_workspace(workspace: Workspace) -> None:
    """Delete a workspace directory and its worktree registration.

    Anything left inside the workspace, including build output, is discarded;
    callers collect what they need first. The run directory holding the
    workspace goes too once its last workspace is removed.

    Raises:
        ValueError: If the workspace path is not inside :data:`WORKSPACE_ROOT`.
    """
    path = _checked_path(workspace.path)
    repo_root = workspace.snapshot.repo_root
    _git(repo_root, "worktree", "remove", "--force", str(path), check=False)
    if path.exists():
        shutil.rmtree(path, ignore_errors=True)
    _git(repo_root, "worktree", "prune", check=False)
    _remove_empty_run_dir(path.parent)


#: Run directories are named ``<YYYYMMDD>-<HHMMSS>-<pid>`` by
#: ``cli.core.parallel.new_run_id``; the pid says whether the run is still alive.
_RUN_DIR_NAME = re.compile(r"^\d{8}-\d{6}-(\d+)$")

#: Marker left in a run directory whose workspaces were kept on purpose
#: (``--keep-workspaces``); the automatic sweep leaves such runs alone.
KEEP_MARKER = ".keep"


def mark_kept(workspace: Workspace) -> None:
    """Record that *workspace*'s run was kept for inspection."""
    (workspace.path.parent / KEEP_MARKER).touch()


def _pid_alive(pid: int) -> bool:
    """Return True if a process with *pid* exists (or might: permission denied)."""
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        return True
    return True


def stale_runs(root: Optional[Path] = None, include_kept: bool = False) -> List[Path]:
    """Return run directories left behind by binder processes that have exited.

    A run is stale when its directory name ends in a pid that no longer exists.
    Runs of the current process, runs whose process is alive, and directories
    that do not look like run directories are never returned. A recycled pid
    makes a dead run look alive, which only delays its removal. Runs kept
    with ``--keep-workspaces`` are skipped unless *include_kept* is set.
    """
    root = Path(root or WORKSPACE_ROOT).resolve()
    if not root.is_dir():
        return []
    stale = []
    for entry in sorted(root.iterdir()):
        match = _RUN_DIR_NAME.match(entry.name)
        if not match or not entry.is_dir() or entry.is_symlink():
            continue
        pid = int(match.group(1))
        if not include_kept and (entry / KEEP_MARKER).exists():
            continue
        if pid != os.getpid() and not _pid_alive(pid):
            stale.append(entry)
    return stale


def remove_stale_runs(repo_root: Optional[Path] = None, include_kept: bool = False) -> List[Path]:
    """Delete run directories of exited binder processes and prune their worktrees.

    Builds normally remove their own workspaces; this catches the ones left by
    a build that was killed outright (SIGKILL, a crash, a closed machine).
    After deleting, ``git worktree prune`` drops the registrations in
    *repo_root* whose directories are gone. *include_kept* also removes runs
    kept with ``--keep-workspaces``.

    Returns:
        The run directories that were removed.
    """
    removed = []
    root = WORKSPACE_ROOT.resolve()
    for run_dir in stale_runs(include_kept=include_kept):
        if run_dir.parent != root:
            continue
        shutil.rmtree(run_dir, ignore_errors=True)
        if not run_dir.exists():
            removed.append(run_dir)
    if removed and repo_root is not None:
        _git(Path(repo_root), "worktree", "prune", check=False)
    return removed


def _remove_empty_run_dir(run_dir: Path) -> None:
    """Delete *run_dir* once it holds nothing but the shared hooks directory.

    Every workspace of a run shares the hooks directory, so it is removed with
    the last one. The workspace root itself stays, since other runs may be
    creating directories in it.
    """
    if run_dir == WORKSPACE_ROOT.resolve() or not run_dir.is_dir():
        return
    if any(entry.name != _NO_HOOKS for entry in run_dir.iterdir()):
        return
    for directory in (run_dir / _NO_HOOKS, run_dir):
        try:
            directory.rmdir()
        except OSError:
            pass


def _checked_path(path: Path) -> Path:
    """Return *path* resolved, refusing anything outside :data:`WORKSPACE_ROOT`.

    Workspaces are deleted recursively; confining them to one directory keeps
    a bad argument from ever pointing that deletion at a real checkout.
    """
    resolved = Path(path).resolve()
    root = WORKSPACE_ROOT.resolve()
    if root not in resolved.parents:
        raise ValueError(f"Workspace path must be inside {root}: {resolved}")
    return resolved