cleveragents-core/benchmarks/uow_lifecycle_bench.py

"""ASV benchmarks for UnitOfWork lifecycle repository wiring.

Measures overhead of creating UnitOfWorkContext and accessing
actions / lifecycle_plans repositories.
"""

from __future__ import annotations

from datetime import datetime

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

from cleveragents.domain.models.core.action import Action, ActionState
from cleveragents.domain.models.core.plan import (
    NamespacedName,
    PlanIdentity,
    PlanPhase,
    PlanTimestamps,
    ProcessingState,
)
from cleveragents.domain.models.core.plan import Plan as V3Plan
from cleveragents.infrastructure.database.models import Base
from cleveragents.infrastructure.database.unit_of_work import UnitOfWorkContext

_CB32 = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"
_bench_counter = 5000


def _bench_ulid() -> str:
    global _bench_counter
    _bench_counter += 1
    n = _bench_counter
    suffix = ""
    for _ in range(8):
        suffix = _CB32[n % 32] + suffix
        n //= 32
    return f"01HGZ6FE0AQDYTR4BX{suffix}"


def _make_bench_action(name: str = "local/bench-uow-action") -> Action:
    parts = name.split("/", 1)
    namespace = parts[0] if len(parts) == 2 else "local"
    short_name = parts[1] if len(parts) == 2 else parts[0]
    now = datetime.now()
    return Action(
        namespaced_name=NamespacedName(namespace=namespace, name=short_name),
        description=f"Bench action {short_name}",
        long_description=None,
        definition_of_done="Done",
        strategy_actor="local/s",
        execution_actor="local/e",
        estimation_actor=None,
        review_actor=None,
        arguments=[],
        reusable=True,
        read_only=False,
        state=ActionState("available"),
        created_at=now,
        updated_at=now,
        created_by=None,
        tags=[],
    )


def _make_bench_plan(
    plan_id: str | None = None,
    action_name: str = "local/bench-uow-action",
    plan_name: str | None = None,
) -> V3Plan:
    now = datetime.now()
    name = plan_name or f"bench-uow-plan-{_bench_ulid()}"
    return V3Plan(
        identity=PlanIdentity(plan_id=plan_id or _bench_ulid(), attempt=1),
        namespaced_name=NamespacedName(namespace="local", name=name),
        action_name=action_name,
        description="Bench plan",
        definition_of_done="Done",
        strategy_actor="local/s",
        execution_actor="local/e",
        phase=PlanPhase("action"),
        processing_state=ProcessingState("queued"),
        timestamps=PlanTimestamps(created_at=now, updated_at=now),
        project_links=[],
        arguments={},
        arguments_order=[],
        invariants=[],
        reusable=True,
        read_only=False,
        created_by=None,
        tags=[],
    )


class TimeUowContextCreation:
    """Benchmark UnitOfWorkContext instantiation."""

    timeout = 30.0

    def setup(self) -> None:
        self.engine = create_engine(
            "sqlite:///:memory:",
            echo=False,
            future=True,
            connect_args={"check_same_thread": False},
        )

        @event.listens_for(self.engine, "connect")
        def _fk(dbapi_conn, _rec):  # type: ignore[no-untyped-def]
            cursor = dbapi_conn.cursor()
            cursor.execute("PRAGMA foreign_keys=ON")
            cursor.close()

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

    def time_create_context(self) -> None:
        session = self.sf()
        _ = UnitOfWorkContext(session)
        session.close()

    def time_access_actions_repo(self) -> None:
        session = self.sf()
        ctx = UnitOfWorkContext(session)
        _ = ctx.actions
        session.close()

    def time_access_lifecycle_plans_repo(self) -> None:
        session = self.sf()
        ctx = UnitOfWorkContext(session)
        _ = ctx.lifecycle_plans
        session.close()


class TimeUowActionPlanRoundTrip:
    """Benchmark creating action + plan via UoW and retrieving."""

    timeout = 30.0

    def setup(self) -> None:
        self.engine = create_engine(
            "sqlite:///:memory:",
            echo=False,
            future=True,
            connect_args={"check_same_thread": False},
        )

        @event.listens_for(self.engine, "connect")
        def _fk(dbapi_conn, _rec):  # type: ignore[no-untyped-def]
            cursor = dbapi_conn.cursor()
            cursor.execute("PRAGMA foreign_keys=ON")
            cursor.close()

        Base.metadata.create_all(self.engine)
        self.sf: sessionmaker[Session] = sessionmaker(
            bind=self.engine,
            expire_on_commit=False,
            autoflush=False,
            autocommit=False,
            class_=Session,
        )
        # Seed an action for plan creation benchmarks
        session = self.sf()
        ctx = UnitOfWorkContext(session)
        ctx.actions.create(_make_bench_action("local/bench-uow-action"))
        session.commit()
        session.close()

    def time_create_action_via_uow(self) -> None:
        session = self.sf()
        ctx = UnitOfWorkContext(session)
        uid = _bench_ulid()
        ctx.actions.create(_make_bench_action(f"local/bench-action-{uid}"))
        session.commit()
        session.close()

    def time_create_plan_via_uow(self) -> None:
        session = self.sf()
        ctx = UnitOfWorkContext(session)
        ctx.lifecycle_plans.create(_make_bench_plan())
        session.commit()
        session.close()

    def time_create_action_and_plan_via_uow(self) -> None:
        uid = _bench_ulid()
        # NamespacedName.validate_name lowercases the name, so the action
        # is stored with a lowercased namespaced_name PK. The plan's
        # action_name (a plain-string FK) must match exactly.
        action_name = f"local/bench-combo-{uid}".lower()
        # Commit the action first so the FK reference is durable before
        # the plan repository (which has its own retry/rollback logic)
        # attempts the insert.
        session = self.sf()
        ctx = UnitOfWorkContext(session)
        ctx.actions.create(_make_bench_action(action_name))
        session.commit()
        session.close()

        session = self.sf()
        ctx = UnitOfWorkContext(session)
        ctx.lifecycle_plans.create(_make_bench_plan(action_name=action_name))
        session.commit()
        session.close()