stonks-oracle/services/intelligence_pipeline_v3/orchestrator/leases.py

"""Lease management for pipeline stage workers.

Leases ensure exactly-once processing semantics. A worker must acquire
a lease before processing a stage. Expired leases allow re-processing
by another worker.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from uuid import UUID, uuid4


class LeaseExpiredError(Exception):
    """Raised when an operation is attempted on an expired lease."""

    def __init__(self, lease_id: UUID, expired_at: datetime) -> None:
        self.lease_id = lease_id
        self.expired_at = expired_at
        super().__init__(
            f"Lease {lease_id} expired at {expired_at.isoformat()}"
        )


@dataclass
class Lease:
    """A time-bounded processing lease for a pipeline stage."""

    lease_id: UUID
    run_id: UUID
    stage: str
    worker_id: str
    acquired_at: datetime
    expires_at: datetime
    released: bool = False
    renewed_count: int = 0

    @property
    def is_expired(self) -> bool:
        """Check if the lease has passed its expiry time."""
        return datetime.now(timezone.utc) >= self.expires_at

    @property
    def is_active(self) -> bool:
        """Check if the lease is currently active."""
        return not self.released and not self.is_expired

    def renew(self, extension: timedelta) -> None:
        """Extend the lease expiry.

        Raises LeaseExpiredError if already expired.
        """
        if self.is_expired:
            raise LeaseExpiredError(self.lease_id, self.expires_at)
        if self.released:
            raise LeaseExpiredError(self.lease_id, self.expires_at)
        self.expires_at = datetime.now(timezone.utc) + extension
        self.renewed_count += 1

    def release(self) -> None:
        """Mark the lease as released (work completed or abandoned)."""
        self.released = True


@dataclass
class LeaseManager:
    """Manages leases for pipeline stage workers.

    In production, this would use Redis or database-backed distributed locks.
    This implementation provides the lease lifecycle logic for testing.
    """

    default_ttl: timedelta = field(default_factory=lambda: timedelta(seconds=120))
    _active_leases: dict[tuple[UUID, str], Lease] = field(default_factory=dict)
    _all_leases: list[Lease] = field(default_factory=list)

    def acquire(
        self,
        run_id: UUID,
        stage: str,
        worker_id: str,
        ttl: timedelta | None = None,
    ) -> Lease | None:
        """Attempt to acquire a lease for a (run_id, stage) pair.

        Returns None if an active lease already exists for that pair.
        Expired leases are cleaned up and allow re-acquisition.
        """
        key = (run_id, stage)
        existing = self._active_leases.get(key)

        if existing is not None:
            if existing.is_active:
                return None  # Already leased
            # Expired — clean up
            del self._active_leases[key]

        lease = Lease(
            lease_id=uuid4(),
            run_id=run_id,
            stage=stage,
            worker_id=worker_id,
            acquired_at=datetime.now(timezone.utc),
            expires_at=datetime.now(timezone.utc) + (ttl or self.default_ttl),
        )
        self._active_leases[key] = lease
        self._all_leases.append(lease)
        return lease

    def release(self, lease: Lease) -> None:
        """Release a lease, making the slot available."""
        lease.release()
        key = (lease.run_id, lease.stage)
        if key in self._active_leases and self._active_leases[key] is lease:
            del self._active_leases[key]

    def renew(self, lease: Lease, extension: timedelta | None = None) -> None:
        """Renew an active lease. Raises LeaseExpiredError if expired."""
        lease.renew(extension or self.default_ttl)

    def is_leased(self, run_id: UUID, stage: str) -> bool:
        """Check if a (run_id, stage) pair has an active lease."""
        key = (run_id, stage)
        existing = self._active_leases.get(key)
        if existing is None:
            return False
        if not existing.is_active:
            del self._active_leases[key]
            return False
        return True

    def active_count(self) -> int:
        """Number of currently active leases."""
        # Clean up expired
        expired_keys = [
            k for k, v in self._active_leases.items() if not v.is_active
        ]
        for k in expired_keys:
            del self._active_leases[k]
        return len(self._active_leases)

    def get_expired(self) -> list[Lease]:
        """Get all expired but unreleased leases (for recovery)."""
        return [
            lease
            for lease in self._active_leases.values()
            if lease.is_expired and not lease.released
        ]