feat: use instrument hooks

Author: not-matthias

src/pytest_codspeed/instruments/valgrind.py

@@ -6,36 +6,30 @@
 
 from pytest_codspeed import __semver_version__
 from pytest_codspeed.instruments import Instrument
-from pytest_codspeed.instruments.valgrind._wrapper import get_lib
+from pytest_codspeed.instruments.hooks import InstrumentHooks
 
 if TYPE_CHECKING:
     from typing import Any, Callable
 
     from pytest import Session
 
     from pytest_codspeed.instruments import P, T
-    from pytest_codspeed.instruments.valgrind._wrapper import LibType
     from pytest_codspeed.plugin import CodSpeedConfig
 
 SUPPORTS_PERF_TRAMPOLINE = sys.version_info >= (3, 12)
 
 
 class ValgrindInstrument(Instrument):
     instrument = "valgrind"
-    lib: LibType | None
 
     def __init__(self, config: CodSpeedConfig) -> None:
         self.benchmark_count = 0
         self.should_measure = os.environ.get("CODSPEED_ENV") is not None
         if self.should_measure:
-            self.lib = get_lib()
-            self.lib.dump_stats_at(
-                f"Metadata: pytest-codspeed {__semver_version__}".encode("ascii")
-            )
+            InstrumentHooks.set_integration("pytest-codspeed", __semver_version__)
+
             if SUPPORTS_PERF_TRAMPOLINE:
                 sys.activate_stack_trampoline("perf")  # type: ignore
-        else:
-            self.lib = None
 
     def get_instrument_config_str_and_warns(self) -> tuple[str, list[str]]:
         config = (
@@ -61,7 +55,8 @@ def measure(
         **kwargs: P.kwargs,
     ) -> T:
         self.benchmark_count += 1
-        if self.lib is None:  # Thus should_measure is False
+
+        if not InstrumentHooks.is_instrumented():  # Thus should_measure is False
             return fn(*args, **kwargs)
 
         def __codspeed_root_frame__() -> T:
@@ -71,14 +66,13 @@ def __codspeed_root_frame__() -> T:
             # Warmup CPython performance map cache
             __codspeed_root_frame__()
 
-        self.lib.zero_stats()
-        self.lib.start_instrumentation()
+        InstrumentHooks.start_benchmark()
         try:
             return __codspeed_root_frame__()
         finally:
             # Ensure instrumentation is stopped even if the test failed
-            self.lib.stop_instrumentation()
-            self.lib.dump_stats_at(uri.encode("ascii"))
+            InstrumentHooks.stop_benchmark()
+            InstrumentHooks.set_current_benchmark(uri)
 
     def report(self, session: Session) -> None:
         reporter = session.config.pluginmanager.get_plugin("terminalreporter")

src/pytest_codspeed/instruments/walltime.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import sys
 from dataclasses import asdict, dataclass
 from math import ceil
 from statistics import mean, quantiles, stdev
@@ -11,7 +12,9 @@
 from rich.table import Table
 from rich.text import Text
 
+from pytest_codspeed import __semver_version__
 from pytest_codspeed.instruments import Instrument
+from pytest_codspeed.instruments.hooks import InstrumentHooks
 
 if TYPE_CHECKING:
     from typing import Any, Callable
@@ -133,15 +136,18 @@ class Benchmark:
 def run_benchmark(
     name: str, uri: str, fn: Callable[P, T], args, kwargs, config: BenchmarkConfig
 ) -> tuple[Benchmark, T]:
+    def __codspeed_root_frame__() -> T:
+        return fn(*args, **kwargs)
+
     # Compute the actual result of the function
-    out = fn(*args, **kwargs)
+    out = __codspeed_root_frame__()
 
     # Warmup
     times_per_round_ns: list[float] = []
     warmup_start = start = perf_counter_ns()
     while True:
         start = perf_counter_ns()
-        fn(*args, **kwargs)
+        __codspeed_root_frame__()
         end = perf_counter_ns()
         times_per_round_ns.append(end - start)
         if end - warmup_start > config.warmup_time_ns:
@@ -166,16 +172,19 @@ def run_benchmark(
     # Benchmark
     iter_range = range(iter_per_round)
     run_start = perf_counter_ns()
+    InstrumentHooks.start_benchmark()
     for _ in range(rounds):
         start = perf_counter_ns()
         for _ in iter_range:
-            fn(*args, **kwargs)
+            __codspeed_root_frame__()
         end = perf_counter_ns()
         times_per_round_ns.append(end - start)
 
         if end - run_start > config.max_time_ns:
             # TODO: log something
             break
+    InstrumentHooks.stop_benchmark()
+    InstrumentHooks.set_current_benchmark(uri)
     benchmark_end = perf_counter_ns()
     total_time = (benchmark_end - run_start) / 1e9
 
@@ -196,6 +205,8 @@ class WallTimeInstrument(Instrument):
     def __init__(self, config: CodSpeedConfig) -> None:
         self.config = config
         self.benchmarks: list[Benchmark] = []
+        sys.activate_stack_trampoline("perf")  # type: ignore
+        InstrumentHooks.set_integration("pytest-codspeed", __semver_version__)
 
     def get_instrument_config_str_and_warns(self) -> tuple[str, list[str]]:
         return f"mode: walltime, timer_resolution: {TIMER_RESOLUTION_NS:.1f}ns", []