feat: use instrument hooks

Author: not-matthias

src/pytest_codspeed/instruments/valgrind.py

@@ -1,41 +1,36 @@
 from __future__ import annotations
 
-import os
 import sys
 from typing import TYPE_CHECKING
 
 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
+    instrument_hooks: InstrumentHooks | 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")
-            )
-            if SUPPORTS_PERF_TRAMPOLINE:
-                sys.activate_stack_trampoline("perf")  # type: ignore
-        else:
-            self.lib = None
+        try:
+            self.instrument_hooks = InstrumentHooks()
+            self.instrument_hooks.set_integration("pytest-codspeed", __semver_version__)
+        except RuntimeError:
+            self.instrument_hooks = None
+
+        self.should_measure = self.instrument_hooks is not None
 
     def get_instrument_config_str_and_warns(self) -> tuple[str, list[str]]:
         config = (
@@ -61,7 +56,8 @@ def measure(
         **kwargs: P.kwargs,
     ) -> T:
         self.benchmark_count += 1
-        if self.lib is None:  # Thus should_measure is False
+
+        if not self.instrument_hooks:
             return fn(*args, **kwargs)
 
         def __codspeed_root_frame__() -> T:
@@ -71,14 +67,15 @@ def __codspeed_root_frame__() -> T:
             # Warmup CPython performance map cache
             __codspeed_root_frame__()
 
-        self.lib.zero_stats()
-        self.lib.start_instrumentation()
+        # Manually call the library function to avoid an extra stack frame. Also
+        # call the callgrind markers directly to avoid extra overhead.
+        self.instrument_hooks.lib.callgrind_start_instrumentation()
         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"))
+            self.instrument_hooks.lib.callgrind_stop_instrumentation()
+            self.instrument_hooks.set_executed_benchmark(uri)
 
     def report(self, session: Session) -> None:
         reporter = session.config.pluginmanager.get_plugin("terminalreporter")

src/pytest_codspeed/instruments/walltime.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import os
+import warnings
 from dataclasses import asdict, dataclass
 from math import ceil
 from statistics import mean, quantiles, stdev
@@ -11,7 +13,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
@@ -131,17 +135,26 @@ class Benchmark:
 
 
 def run_benchmark(
-    name: str, uri: str, fn: Callable[P, T], args, kwargs, config: BenchmarkConfig
+    instrument_hooks: InstrumentHooks | None,
+    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 +179,21 @@ def run_benchmark(
     # Benchmark
     iter_range = range(iter_per_round)
     run_start = perf_counter_ns()
+    if instrument_hooks:
+        instrument_hooks.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
+    if instrument_hooks:
+        instrument_hooks.stop_benchmark()
+        instrument_hooks.set_executed_benchmark(uri)
     benchmark_end = perf_counter_ns()
     total_time = (benchmark_end - run_start) / 1e9
 
@@ -192,8 +210,19 @@ def run_benchmark(
 
 class WallTimeInstrument(Instrument):
     instrument = "walltime"
+    instrument_hooks: InstrumentHooks | None
 
     def __init__(self, config: CodSpeedConfig) -> None:
+        try:
+            self.instrument_hooks = InstrumentHooks()
+            self.instrument_hooks.set_integration("pytest-codspeed", __semver_version__)
+        except RuntimeError as e:
+            if os.environ.get("CODSPEED_ENV") is not None:
+                warnings.warn(
+                    f"Failed to initialize instrument hooks: {e}", RuntimeWarning
+                )
+            self.instrument_hooks = None
+
         self.config = config
         self.benchmarks: list[Benchmark] = []
 
@@ -209,6 +238,7 @@ def measure(
         **kwargs: P.kwargs,
     ) -> T:
         bench, out = run_benchmark(
+            instrument_hooks=self.instrument_hooks,
             name=name,
             uri=uri,
             fn=fn,