cleveragents-core/benchmarks/resource_dag_bench.py

"""ASV benchmarks for resource DAG link and auto_discover."""

from __future__ import annotations

import contextlib
import typing
from datetime import UTC, datetime
from typing import Any


def _setup_db() -> Any:
    """Create in-memory database and return session factory."""
    from sqlalchemy import create_engine, event
    from sqlalchemy.orm import sessionmaker

    from cleveragents.infrastructure.database.models import Base

    engine = create_engine("sqlite:///:memory:")

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

    Base.metadata.create_all(engine)
    return sessionmaker(bind=engine)


_CB32 = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"
_BENCH_CTR = 0


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


def _seed_types(
    rt_repo: Any,
    session_factory: Any,
    parent_name: str = "bench/dag-parent",
    child_name: str = "bench/dag-child",
) -> None:
    """Seed parent and child resource types."""
    from cleveragents.domain.models.core.resource_type import (
        ResourceKind,
        ResourceTypeSpec,
        SandboxStrategy,
    )

    parent_spec = ResourceTypeSpec(
        name=parent_name,
        description="Bench parent type",
        resource_kind=ResourceKind.PHYSICAL,
        sandbox_strategy=SandboxStrategy.NONE,
        user_addable=True,
        cli_args=[],
        parent_types=[],
        child_types=[child_name],
        auto_discovery={
            "enabled": True,
            "rules": [{"type": child_name, "pattern": "*"}],
        },
        equivalence=None,
        handler=None,
        capabilities={
            "read": True,
            "write": True,
            "sandbox": True,
            "checkpoint": False,
        },
        built_in=False,
    )
    child_spec = ResourceTypeSpec(
        name=child_name,
        description="Bench child type",
        resource_kind=ResourceKind.PHYSICAL,
        sandbox_strategy=SandboxStrategy.NONE,
        user_addable=True,
        cli_args=[],
        parent_types=[],
        child_types=[],
        auto_discovery=None,
        equivalence=None,
        handler=None,
        capabilities={
            "read": True,
            "write": True,
            "sandbox": True,
            "checkpoint": False,
        },
        built_in=False,
    )
    rt_repo.create(parent_spec)
    if child_name != parent_name:
        rt_repo.create(child_spec)
    # Commit so types survive across session boundaries (StaticPool).
    session_factory().commit()


def _seed_single_type(
    rt_repo: Any,
    type_name: str = "bench/chain-type",
) -> None:
    """Seed a single self-linking resource type for DAG chain benchmarks."""
    from cleveragents.domain.models.core.resource_type import (
        ResourceKind,
        ResourceTypeSpec,
        SandboxStrategy,
    )

    spec = ResourceTypeSpec(
        name=type_name,
        description="Bench chain type",
        resource_kind=ResourceKind.PHYSICAL,
        sandbox_strategy=SandboxStrategy.NONE,
        user_addable=True,
        cli_args=[],
        parent_types=[],
        child_types=[type_name],
        auto_discovery=None,
        equivalence=None,
        handler=None,
        capabilities={
            "read": True,
            "write": True,
            "sandbox": True,
            "checkpoint": False,
        },
        built_in=False,
    )
    rt_repo.create(spec)


def _make_resource(
    res_repo: Any,
    type_name: str,
    session_factory: Any,
) -> str:
    """Create a resource and return its ID."""
    from cleveragents.domain.models.core.resource import (
        PhysVirt,
        Resource,
        ResourceCapabilities,
    )

    rid = _bench_ulid()
    res = Resource(
        resource_id=rid,
        name=None,
        resource_type_name=type_name,
        classification=PhysVirt.PHYSICAL,
        description="Bench resource",
        properties={},
        location=None,
        content_hash=None,
        sandbox_strategy=None,
        capabilities=ResourceCapabilities(),
        created_at=datetime.now(tz=UTC),
        updated_at=datetime.now(tz=UTC),
    )
    res_repo.create(res)
    # Commit so the resource survives across session boundaries.
    session_factory().commit()
    return rid


class TimeLinkChild:
    """Benchmark ResourceRepository.link_child()."""

    def setup(self) -> None:
        from cleveragents.infrastructure.database.repositories import (
            ResourceRepository,
            ResourceTypeRepository,
        )

        self.factory = _setup_db()
        self.rt_repo = ResourceTypeRepository(self.factory)
        self.res_repo = ResourceRepository(self.factory)
        _seed_types(self.rt_repo, self.factory)
        self.parent_id = _make_resource(self.res_repo, "bench/dag-parent", self.factory)
        self._child_ctr = 0

    def time_link_child(self) -> None:
        child_id = _make_resource(self.res_repo, "bench/dag-child", self.factory)
        self.res_repo.link_child(self.parent_id, child_id)


class TimeUnlinkChild:
    """Benchmark ResourceRepository.unlink_child()."""

    def setup(self) -> None:
        from cleveragents.infrastructure.database.repositories import (
            ResourceRepository,
            ResourceTypeRepository,
        )

        self.factory = _setup_db()
        self.rt_repo = ResourceTypeRepository(self.factory)
        self.res_repo = ResourceRepository(self.factory)
        _seed_types(self.rt_repo, self.factory)
        self.parent_id = _make_resource(self.res_repo, "bench/dag-parent", self.factory)
        self.child_id = _make_resource(self.res_repo, "bench/dag-child", self.factory)
        self.res_repo.link_child(self.parent_id, self.child_id)
        self.factory().commit()

    def time_unlink_child(self) -> None:
        self.res_repo.unlink_child(self.parent_id, self.child_id)
        # Re-link so the benchmark can run multiple times
        self.res_repo.link_child(self.parent_id, self.child_id)


class TimeAutoDiscoverChildren:
    """Benchmark ResourceRepository.auto_discover_children()."""

    def setup(self) -> None:
        from cleveragents.infrastructure.database.repositories import (
            ResourceRepository,
            ResourceTypeRepository,
        )

        self.factory = _setup_db()
        self.rt_repo = ResourceTypeRepository(self.factory)
        self.res_repo = ResourceRepository(self.factory)
        _seed_types(self.rt_repo, self.factory)
        self.parent_id = _make_resource(self.res_repo, "bench/dag-parent", self.factory)

    def time_auto_discover_children(self) -> None:
        self.res_repo.auto_discover_children(self.parent_id)


class TimeCycleDetection:
    """Benchmark cycle detection with deep DAG chains."""

    params: typing.ClassVar[list[int]] = [5, 10, 50]
    param_names: typing.ClassVar[list[str]] = ["chain_depth"]

    def setup(self, chain_depth: int) -> None:
        from cleveragents.infrastructure.database.repositories import (
            ResourceRepository,
            ResourceTypeRepository,
        )

        self.factory = _setup_db()
        self.rt_repo = ResourceTypeRepository(self.factory)
        self.res_repo = ResourceRepository(self.factory)
        _seed_single_type(self.rt_repo, "bench/chain-type")

        # Build a chain: r0 -> r1 -> ... -> rN
        self.chain_ids: list[str] = []
        for _ in range(chain_depth + 1):
            rid = _make_resource(self.res_repo, "bench/chain-type", self.factory)
            self.chain_ids.append(rid)

        for i in range(chain_depth):
            self.res_repo.link_child(
                self.chain_ids[i],
                self.chain_ids[i + 1],
            )
        self.factory().commit()

    def time_cycle_detection(self, chain_depth: int) -> None:
        """Attempt to close the cycle (should raise)."""
        with contextlib.suppress(Exception):
            self.res_repo.link_child(
                self.chain_ids[-1],
                self.chain_ids[0],
            )