cleveragents-core/benchmarks/checkpoint_rollback_bench.py

"""Airspeed Velocity benchmarks for checkpoint and rollback operations.

Measures checkpoint creation, listing, rollback, pruning, and metadata
overhead across varying checkpoint counts.
"""

from __future__ import annotations

import shutil
import subprocess
import tempfile
from typing import ClassVar

from cleveragents.application.services.checkpoint_service import CheckpointService
from cleveragents.domain.models.core.checkpoint import CheckpointRetentionPolicy

_PLAN_ID = "01ARZ3NDEKTSV4RRFFQ69G5FAV"


# ---------------------------------------------------------------------------
# Checkpoint creation benchmarks
# ---------------------------------------------------------------------------


class CheckpointCreateSuite:
    """Benchmark checkpoint creation at varying counts."""

    def setup(self) -> None:
        """Prepare a fresh service instance."""
        self.svc = CheckpointService()

    def time_create_single(self) -> None:
        """Time creating a single checkpoint."""
        self.svc.create_checkpoint(_PLAN_ID, "commit-abc")

    def time_create_ten(self) -> None:
        """Time creating 10 checkpoints."""
        svc = CheckpointService()
        for i in range(10):
            svc.create_checkpoint(_PLAN_ID, f"commit-{i}")

    def time_create_hundred(self) -> None:
        """Time creating 100 checkpoints."""
        svc = CheckpointService()
        for i in range(100):
            svc.create_checkpoint(_PLAN_ID, f"commit-{i}")


# ---------------------------------------------------------------------------
# Checkpoint listing benchmarks
# ---------------------------------------------------------------------------


class CheckpointListSuite:
    """Benchmark listing checkpoints at varying plan sizes."""

    params: ClassVar[list[int]] = [10, 50, 100]
    param_names: ClassVar[list[str]] = ["checkpoint_count"]

    def setup(self, checkpoint_count: int) -> None:
        """Populate service with checkpoints."""
        self.svc = CheckpointService()
        for i in range(checkpoint_count):
            self.svc.create_checkpoint(_PLAN_ID, f"commit-{i}")

    def time_list(self, checkpoint_count: int) -> None:
        """Time listing all checkpoints for a plan."""
        self.svc.list_checkpoints(_PLAN_ID)


# ---------------------------------------------------------------------------
# Rollback benchmarks
# ---------------------------------------------------------------------------


class RollbackSuite:
    """Benchmark rollback operations."""

    def setup(self) -> None:
        """Prepare a service with a real git sandbox and checkpoints."""
        self._tmpdir = tempfile.mkdtemp(prefix="asv-rollback-")
        subprocess.run(
            ["git", "init", self._tmpdir],
            check=True,
            capture_output=True,
        )
        subprocess.run(
            ["git", "commit", "--allow-empty", "-m", "initial"],
            check=True,
            capture_output=True,
            cwd=self._tmpdir,
        )
        self.svc = CheckpointService()
        self.svc.register_sandbox(_PLAN_ID, self._tmpdir)
        self.cp = self.svc.create_checkpoint(_PLAN_ID, "commit-abc")

    def teardown(self) -> None:
        """Remove temporary sandbox directory."""
        shutil.rmtree(self._tmpdir, ignore_errors=True)

    def time_rollback(self) -> None:
        """Time a single rollback operation."""
        self.svc.rollback_to_checkpoint(_PLAN_ID, self.cp.checkpoint_id)


# ---------------------------------------------------------------------------
# Prune benchmarks
# ---------------------------------------------------------------------------


class PruneSuite:
    """Benchmark checkpoint pruning at varying counts."""

    params: ClassVar[list[int]] = [20, 50, 100]
    param_names: ClassVar[list[str]] = ["checkpoint_count"]

    def setup(self, checkpoint_count: int) -> None:
        """Populate service with checkpoints."""
        self.svc = CheckpointService()
        for i in range(checkpoint_count):
            self.svc.create_checkpoint(_PLAN_ID, f"commit-{i}")

    def time_prune_to_five(self, checkpoint_count: int) -> None:
        """Time pruning down to 5 checkpoints."""
        policy = CheckpointRetentionPolicy(max_checkpoints=5, auto_prune=True)
        self.svc.prune_checkpoints(_PLAN_ID, policy)


# ---------------------------------------------------------------------------
# Metadata benchmarks
# ---------------------------------------------------------------------------


class MetadataSuite:
    """Benchmark checkpoint creation with metadata."""

    def setup(self) -> None:
        """Prepare service."""
        self.svc = CheckpointService()

    def time_create_with_metadata(self) -> None:
        """Time creating a checkpoint with full metadata."""
        self.svc.create_checkpoint(
            _PLAN_ID,
            "commit-xyz",
            reason="pre-execution snapshot",
            source_tool="lint-check",
            phase="execute",
        )