cleveragents-core/benchmarks/concurrency_lock_bench.py

"""ASV benchmarks for concurrency lock overhead.

Measures the cost of lock acquire, release, renew, and cleanup
operations against an in-memory SQLite database.
"""

from __future__ import annotations

from datetime import UTC, datetime, timedelta
from typing import Any

from sqlalchemy import create_engine, event
from sqlalchemy.orm import Session, sessionmaker

from cleveragents.application.services.lock_service import LockService
from cleveragents.infrastructure.database.models import Base, LockModel

_bench_counter = 9000


def _next_id() -> str:
    global _bench_counter
    _bench_counter += 1
    return f"bench-res-{_bench_counter}"


def _build_lock_service() -> tuple[LockService, sessionmaker[Session]]:
    engine = create_engine(
        "sqlite:///:memory:",
        echo=False,
        future=True,
        connect_args={"check_same_thread": False},
    )

    @event.listens_for(engine, "connect")
    def _fk(dbapi_conn: Any, _rec: Any) -> None:
        cursor = dbapi_conn.cursor()
        cursor.execute("PRAGMA foreign_keys=ON")
        cursor.close()

    Base.metadata.create_all(engine)
    sf: sessionmaker[Session] = sessionmaker(
        bind=engine,
        expire_on_commit=False,
        autoflush=False,
        autocommit=False,
        class_=Session,
    )
    return LockService(session_factory=sf), sf


class TimeAcquireLock:
    """Benchmark acquiring a new lock."""

    timeout = 30

    def setup(self) -> None:
        self.svc, self._sf = _build_lock_service()

    def time_acquire(self) -> None:
        rid = _next_id()
        self.svc.acquire("owner-bench", "plan", rid)


class TimeReentrantAcquire:
    """Benchmark re-entrant lock acquisition."""

    timeout = 30

    def setup(self) -> None:
        self.svc, self._sf = _build_lock_service()
        self.svc.acquire("owner-bench", "plan", "reentrant-res")

    def time_reentrant_acquire(self) -> None:
        self.svc.acquire("owner-bench", "plan", "reentrant-res")


class TimeReleaseLock:
    """Benchmark releasing a lock."""

    timeout = 30

    def setup(self) -> None:
        self.svc, self._sf = _build_lock_service()
        self._rid = _next_id()
        self.svc.acquire("owner-bench", "plan", self._rid)

    def time_release(self) -> None:
        self.svc.release("owner-bench", "plan", self._rid)
        self.svc.acquire("owner-bench", "plan", self._rid)


class TimeRenewLock:
    """Benchmark renewing a lock."""

    timeout = 30

    def setup(self) -> None:
        self.svc, self._sf = _build_lock_service()
        self.svc.acquire("owner-bench", "plan", "renew-res")

    def time_renew(self) -> None:
        self.svc.renew("owner-bench", "plan", "renew-res")


class TimeCleanupExpired:
    """Benchmark expired lock cleanup."""

    timeout = 30

    def setup(self) -> None:
        self.svc, self._sf = _build_lock_service()
        now = datetime.now(tz=UTC)
        expired = (now - timedelta(seconds=60)).isoformat()
        acquired = (now - timedelta(seconds=120)).isoformat()
        session = self._sf()
        for i in range(50):
            session.add(
                LockModel(
                    owner_id=f"expired-{i}",
                    resource_type="plan",
                    resource_id=f"cleanup-{i}",
                    acquired_at=acquired,
                    expires_at=expired,
                )
            )
        session.commit()
        session.close()

    def time_cleanup(self) -> None:
        self.svc.cleanup_expired()


class TimeIsLocked:
    """Benchmark is_locked check."""

    timeout = 30

    def setup(self) -> None:
        self.svc, self._sf = _build_lock_service()
        self.svc.acquire("owner-bench", "plan", "locked-res")

    def time_is_locked(self) -> None:
        self.svc.is_locked("plan", "locked-res")