Gitea-Tools/issue_lock_store.py

"""Atomic per-issue lock store (#438).

Distinct issues use separate lock files under ``GITEA_ISSUE_LOCK_DIR`` so
concurrent author sessions do not clobber each other. Acquisition is
serialized per issue with ``fcntl.flock`` and fail-closed stale handling.
"""

from __future__ import annotations

import errno
import fcntl
import json
import os
import tempfile
import uuid
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from typing import Any

DEFAULT_LOCK_DIR = os.environ.get("GITEA_ISSUE_LOCK_DIR", "/tmp/gitea_issue_locks")
LEGACY_LOCK_FILE = os.environ.get("GITEA_ISSUE_LOCK_FILE", "/tmp/gitea_issue_lock.json")
LOCK_VERSION = 1


class LockContentionError(RuntimeError):
    """Raised when an exclusive per-issue lock cannot be acquired."""


def _sanitize(value: str) -> str:
    text = (value or "").strip()
    if not text:
        return "unknown"
    return "".join(c if c.isalnum() or c in "-_" else "-" for c in text)


def lock_scope_key(remote: str, org: str, repo: str, issue_number: int) -> str:
    return "_".join(
        (
            _sanitize(remote),
            _sanitize(org),
            _sanitize(repo),
            f"issue-{issue_number}",
        )
    )


def issue_lock_path(
    remote: str,
    org: str,
    repo: str,
    issue_number: int,
    *,
    lock_dir: str | None = None,
) -> str:
    base = (lock_dir or DEFAULT_LOCK_DIR).strip()
    return os.path.join(base, f"{lock_scope_key(remote, org, repo, issue_number)}.json")


def flock_path(json_path: str) -> str:
    return f"{json_path}.lock"


def _iso(value: datetime) -> str:
    return value.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def parse_timestamp(value: str | None) -> datetime | None:
    text = (value or "").strip()
    if not text:
        return None
    try:
        return datetime.fromisoformat(text.replace("Z", "+00:00")).astimezone(timezone.utc)
    except ValueError:
        return None


def is_process_alive(pid: int | None) -> bool:
    if not pid or pid <= 0:
        return False
    try:
        os.kill(int(pid), 0)
        return True
    except OSError as exc:
        return exc.errno != errno.ESRCH
    except (TypeError, ValueError):
        return False


def read_lock_file(path: str) -> dict[str, Any] | None:
    if not os.path.exists(path):
        return None
    try:
        with open(path, encoding="utf-8") as handle:
            data = json.load(handle)
        return data if isinstance(data, dict) else None
    except Exception:
        return None


def atomic_write_json(path: str, data: dict[str, Any]) -> None:
    directory = os.path.dirname(path) or "."
    os.makedirs(directory, exist_ok=True)
    fd, tmp_path = tempfile.mkstemp(dir=directory, prefix=".lock-", suffix=".tmp")
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            json.dump(data, handle)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(tmp_path, path)
    finally:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)


@contextmanager
def _exclusive_file_lock(lock_path: str):
    os.makedirs(os.path.dirname(lock_path) or ".", exist_ok=True)
    fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o600)
    try:
        try:
            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError as exc:
            raise LockContentionError(
                f"could not acquire exclusive lock on '{lock_path}'"
            ) from exc
        yield fd
    finally:
        try:
            fcntl.flock(fd, fcntl.LOCK_UN)
        finally:
            os.close(fd)


def _lease_expires_at(lock_data: dict[str, Any] | None) -> datetime | None:
    if not lock_data:
        return None
    lease = lock_data.get("work_lease")
    if isinstance(lease, dict):
        return parse_timestamp(lease.get("expires_at"))
    return parse_timestamp(lock_data.get("expires_at"))


def assess_lock_freshness(
    lock_data: dict[str, Any] | None,
    *,
    now: datetime | None = None,
) -> dict[str, Any]:
    """Classify a lock as live, expired, stale, or absent."""
    current = now or datetime.now(timezone.utc)
    if not lock_data:
        return {
            "status": "absent",
            "live": False,
            "stale": False,
            "reason": "no lock record",
        }

    expires_at = _lease_expires_at(lock_data)
    heartbeat_at = parse_timestamp(lock_data.get("last_heartbeat_at"))
    if heartbeat_at is None and isinstance(lock_data.get("work_lease"), dict):
        heartbeat_at = parse_timestamp(lock_data["work_lease"].get("last_heartbeat_at"))

    pid = lock_data.get("pid")
    pid_alive = is_process_alive(pid) if pid is not None else False

    if expires_at and expires_at <= current:
        return {
            "status": "expired",
            "live": False,
            "stale": True,
            "reason": f"lease expired at {expires_at.isoformat()}",
            "pid_alive": pid_alive,
        }

    if pid is not None and not pid_alive:
        return {
            "status": "stale",
            "live": False,
            "stale": True,
            "reason": f"owner pid {pid} is not alive",
            "pid_alive": False,
        }

    return {
        "status": "live",
        "live": True,
        "stale": False,
        "reason": "lock heartbeat and lease are fresh",
        "pid_alive": pid_alive,
        "heartbeat_at": _iso(heartbeat_at) if heartbeat_at else None,
        "expires_at": _iso(expires_at) if expires_at else None,
    }


def _same_owner(
    existing: dict[str, Any],
    *,
    branch_name: str,
    worktree_path: str,
    claimant: dict[str, Any] | None,
) -> bool:
    same_branch = existing.get("branch_name") == branch_name
    same_worktree = os.path.realpath(str(existing.get("worktree_path") or "")) == os.path.realpath(
        worktree_path
    )
    if not (same_branch and same_worktree):
        return False
    if claimant and isinstance(existing.get("claimant"), dict):
        return (
            existing["claimant"].get("profile") == claimant.get("profile")
            and existing["claimant"].get("username") == claimant.get("username")
        )
    return same_branch and same_worktree


def update_legacy_session_pointer(record: dict[str, Any]) -> dict[str, Any]:
    """Update the legacy global pointer without clobbering another live session."""
    existing = read_lock_file(LEGACY_LOCK_FILE)
    if existing:
        freshness = assess_lock_freshness(existing)
        if freshness["status"] == "live":
            different_issue = existing.get("issue_number") != record.get("issue_number")
            different_pid = existing.get("pid") != os.getpid()
            if different_issue and different_pid and is_process_alive(existing.get("pid")):
                return {
                    "updated": False,
                    "reason": (
                        "retained global pointer for live issue "
                        f"#{existing.get('issue_number')} pid={existing.get('pid')}; "
                        f"per-issue lock at '{record.get('lock_path')}' is authoritative"
                    ),
                }
    atomic_write_json(LEGACY_LOCK_FILE, record)
    return {"updated": True, "path": LEGACY_LOCK_FILE}


def acquire_issue_lock(
    *,
    issue_number: int,
    branch_name: str,
    remote: str,
    org: str,
    repo: str,
    worktree_path: str,
    work_lease: dict[str, Any],
    claimant: dict[str, Any] | None = None,
    lock_provenance: dict[str, Any] | None = None,
    lock_dir: str | None = None,
    allow_stale_recovery: bool = False,
) -> dict[str, Any]:
    """Acquire an atomic per-issue lock; fail closed on live competing ownership."""
    path = issue_lock_path(remote, org, repo, issue_number, lock_dir=lock_dir)
    sentinel = flock_path(path)
    now = datetime.now(timezone.utc)
    session_id = str(uuid.uuid4())
    resolved_worktree = os.path.realpath(worktree_path)

    try:
        with _exclusive_file_lock(sentinel):
            existing = read_lock_file(path)
            freshness = assess_lock_freshness(existing, now=now)
            if existing and freshness["live"]:
                if not _same_owner(
                    existing,
                    branch_name=branch_name,
                    worktree_path=resolved_worktree,
                    claimant=claimant,
                ):
                    owner = existing.get("claimant") or {}
                    raise RuntimeError(
                        f"Issue #{issue_number} already has an active lock owned by "
                        f"pid={existing.get('pid')} session={existing.get('session_id')} "
                        f"profile={owner.get('profile')} on branch "
                        f"'{existing.get('branch_name')}' (fail closed)"
                    )
            elif existing and freshness["stale"]:
                if not allow_stale_recovery:
                    raise RuntimeError(
                        f"Issue #{issue_number} has a stale lock ({freshness['reason']}). "
                        "Recovery review is required before takeover (fail closed)"
                    )

            record: dict[str, Any] = {
                "lock_version": LOCK_VERSION,
                "issue_number": issue_number,
                "branch_name": branch_name,
                "remote": remote,
                "org": org,
                "repo": repo,
                "worktree_path": resolved_worktree,
                "work_lease": work_lease,
                "pid": os.getpid(),
                "session_id": session_id,
                "claimant": claimant or {},
                "created_at": _iso(now),
                "last_heartbeat_at": _iso(now),
                "lock_path": path,
            }
            if lock_provenance:
                record["lock_provenance"] = lock_provenance
            atomic_write_json(path, record)
            pointer = update_legacy_session_pointer(record)
    except LockContentionError as exc:
        competing = read_lock_file(path)
        if competing:
            freshness = assess_lock_freshness(competing, now=now)
            owner = competing.get("claimant") or {}
            raise RuntimeError(
                f"Issue #{issue_number} lock contention: {exc}; competing owner "
                f"pid={competing.get('pid')} session={competing.get('session_id')} "
                f"profile={owner.get('profile')} status={freshness['status']} (fail closed)"
            ) from exc
        raise RuntimeError(f"Issue #{issue_number} lock contention: {exc} (fail closed)") from exc

    freshness = assess_lock_freshness(record, now=now)
    return {
        "acquired": True,
        "lock_path": path,
        "session_id": session_id,
        "freshness": freshness,
        "legacy_pointer": pointer,
        "lock_proof": format_lock_proof(record, freshness=freshness),
        "record": record,
    }


def resolve_lock_for_issue(
    *,
    issue_number: int,
    remote: str | None = None,
    org: str | None = None,
    repo: str | None = None,
    lock_dir: str | None = None,
) -> dict[str, Any] | None:
    """Load the per-issue lock, falling back to the legacy global pointer."""
    if remote and org and repo:
        scoped = read_lock_file(
            issue_lock_path(remote, org, repo, issue_number, lock_dir=lock_dir)
        )
        if scoped:
            return scoped
    legacy = read_lock_file(LEGACY_LOCK_FILE)
    if legacy and legacy.get("issue_number") == issue_number:
        return legacy
    return None


def resolve_lock_for_branch(
    branch_name: str,
    *,
    lock_dir: str | None = None,
) -> dict[str, Any] | None:
    """Resolve a lock using the issue number embedded in a branch name."""
    import re

    match = re.search(r"issue-(\d+)", branch_name or "", re.IGNORECASE)
    if not match:
        legacy = read_lock_file(LEGACY_LOCK_FILE)
        if legacy and legacy.get("branch_name") == branch_name:
            return legacy
        return None

    issue_number = int(match.group(1))
    base = (lock_dir or DEFAULT_LOCK_DIR).strip()
    if os.path.isdir(base):
        suffix = f"_issue-{issue_number}.json"
        for name in os.listdir(base):
            if name.endswith(suffix):
                record = read_lock_file(os.path.join(base, name))
                if record and record.get("branch_name") == branch_name:
                    return record
    legacy = read_lock_file(LEGACY_LOCK_FILE)
    if legacy and legacy.get("issue_number") == issue_number:
        return legacy
    return None


def verify_lock_for_mutation(
    lock_data: dict[str, Any] | None,
    *,
    issue_number: int | None = None,
    branch_name: str | None = None,
    worktree_path: str | None = None,
) -> dict[str, Any]:
    """Re-check lock ownership immediately before a mutation (#438)."""
    reasons: list[str] = []
    if not lock_data:
        return {"proven": False, "block": True, "reasons": ["issue lock is missing (fail closed)"]}

    freshness = assess_lock_freshness(lock_data)
    if not freshness["live"]:
        reasons.append(f"issue lock is not live: {freshness['reason']} (fail closed)")

    if issue_number is not None and lock_data.get("issue_number") != issue_number:
        reasons.append(
            f"issue lock targets #{lock_data.get('issue_number')}, expected #{issue_number} (fail closed)"
        )

    if branch_name is not None and lock_data.get("branch_name") != branch_name:
        reasons.append(
            f"issue lock branch '{lock_data.get('branch_name')}' does not match "
            f"'{branch_name}' (fail closed)"
        )

    if worktree_path is not None:
        locked = os.path.realpath(str(lock_data.get("worktree_path") or ""))
        declared = os.path.realpath(worktree_path)
        if locked != declared:
            reasons.append(
                f"issue lock worktree '{locked}' does not match declared '{declared}' (fail closed)"
            )

    return {
        "proven": not reasons,
        "block": bool(reasons),
        "reasons": reasons,
        "freshness": freshness,
        "lock_proof": format_lock_proof(lock_data, freshness=freshness),
    }


def list_live_locks(*, lock_dir: str | None = None, now: datetime | None = None) -> list[dict[str, Any]]:
    """Return live per-issue locks for queue visibility."""
    base = (lock_dir or DEFAULT_LOCK_DIR).strip()
    if not os.path.isdir(base):
        return []

    live: list[dict[str, Any]] = []
    for name in sorted(os.listdir(base)):
        if not name.endswith(".json"):
            continue
        record = read_lock_file(os.path.join(base, name))
        if not record:
            continue
        freshness = assess_lock_freshness(record, now=now)
        if freshness["live"]:
            live.append(
                {
                    "issue_number": record.get("issue_number"),
                    "branch_name": record.get("branch_name"),
                    "remote": record.get("remote"),
                    "org": record.get("org"),
                    "repo": record.get("repo"),
                    "worktree_path": record.get("worktree_path"),
                    "pid": record.get("pid"),
                    "session_id": record.get("session_id"),
                    "claimant": record.get("claimant"),
                    "freshness": freshness,
                    "lock_path": record.get("lock_path") or os.path.join(base, name),
                }
            )
    return live


def format_lock_proof(
    lock_data: dict[str, Any] | None,
    *,
    freshness: dict[str, Any] | None = None,
    competing_live_locks: list[dict[str, Any]] | None = None,
    released: bool | None = None,
) -> str:
    """Canonical issue-lock proof string for final reports."""
    if not lock_data:
        return "issue lock proof: not acquired"
    fresh = freshness or assess_lock_freshness(lock_data)
    owner = lock_data.get("claimant") or {}
    parts = [
        "issue lock proof:",
        f"acquired issue #{lock_data.get('issue_number')}",
        f"branch {lock_data.get('branch_name')}",
        f"owner {owner.get('profile') or 'unknown'}",
        f"pid {lock_data.get('pid')}",
        f"session {lock_data.get('session_id')}",
        f"freshness {fresh.get('status')}",
    ]
    if competing_live_locks is not None:
        parts.append(
            "no competing live lock"
            if not competing_live_locks
            else f"competing live locks {len(competing_live_locks)}"
        )
    if released is True:
        parts.append("lock released")
    elif released is False:
        parts.append("lock retained")
    return "; ".join(parts)