feat: add monitoring tracer trait and hooks

Author: tzanko-matev

Justfile

@@ -34,6 +34,7 @@ dev:
 # Run unit tests of dev build
 test:
     uv run --group dev --group test pytest
+    uv run --group dev --group test cargo test --manifest-path codetracer-python-recorder/Cargo.toml
 
 # Run tests only on the pure recorder
 test-pure:

codetracer-python-recorder/Cargo.lock

@@ -8,8 +8,9 @@ repository = "https://github.com/metacraft-labs/codetracer-python-recorder"
 
 [lib]
 name = "codetracer_python_recorder"
-crate-type = ["cdylib"]
+crate-type = ["cdylib", "rlib"]
 
 [dependencies]
-pyo3 = { version = "0.25.1", features = ["extension-module"] }
-runtime_tracing = "0.14.0"
+pyo3 = { version = "0.25.1", features = ["extension-module", "abi3-py312"] }
+runtime_tracing = "0.14.0"
+bitflags = "2.4"

codetracer-python-recorder/Cargo.toml

@@ -3,6 +3,9 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 
+pub mod tracer;
+pub use crate::tracer::{install_tracer, uninstall_tracer, EventMask, Tracer};
+
 /// Global flag tracking whether tracing is active.
 static ACTIVE: AtomicBool = AtomicBool::new(false);
 

codetracer-python-recorder/src/lib.rs

@@ -0,0 +1,149 @@
+use std::sync::Mutex;
+
+use bitflags::bitflags;
+use pyo3::{
+    exceptions::PyRuntimeError,
+    prelude::*,
+    types::{PyModule, PyTuple},
+    ffi,
+};
+
+bitflags! {
+    /// Bitmask of monitoring events a tracer can subscribe to.
+    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
+    pub struct EventMask: u32 {
+        /// Python function calls.
+        const CALL = 1 << 0;
+        /// Line execution events.
+        const LINE = 1 << 1;
+    }
+}
+
+#[non_exhaustive]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum Event {
+    Call,
+    Line,
+}
+
+/// Trait implemented by tracing backends.
+///
+/// Each method corresponds to an event from `sys.monitoring`. Default
+/// implementations allow implementers to only handle the events they care
+/// about.
+pub trait Tracer: Send {
+    /// Return the set of events the tracer wants to receive.
+    fn interest(&self) -> EventMask { EventMask::empty() }
+
+    /// Called on Python function calls.
+    fn on_call(&mut self, _py: Python<'_>, _frame: *mut ffi::PyFrameObject) {}
+
+    /// Called on line execution.
+    fn on_line(&mut self, _py: Python<'_>, _frame: *mut ffi::PyFrameObject, _lineno: u32) {}
+}
+
+/// Dispatcher routing events based on the tracer's interest mask.
+pub struct Dispatcher {
+    tracer: Box<dyn Tracer>,
+    mask: EventMask,
+}
+
+impl Dispatcher {
+    pub fn new(tracer: Box<dyn Tracer>) -> Self {
+        let mask = tracer.interest();
+        Self { tracer, mask }
+    }
+
+    pub fn dispatch_call(&mut self, py: Python<'_>, frame: *mut ffi::PyFrameObject) {
+        if self.mask.contains(EventMask::CALL) {
+            self.tracer.on_call(py, frame);
+        }
+    }
+
+    pub fn dispatch_line(&mut self, py: Python<'_>, frame: *mut ffi::PyFrameObject, lineno: u32) {
+        if self.mask.contains(EventMask::LINE) {
+            self.tracer.on_line(py, frame, lineno);
+        }
+    }
+}
+
+struct Global {
+    dispatcher: Dispatcher,
+    tool_id: u8,
+    callbacks: Vec<Py<PyAny>>,
+}
+
+static GLOBAL: Mutex<Option<Global>> = Mutex::new(None);
+
+/// Install a tracer and hook it into Python's `sys.monitoring`.
+pub fn install_tracer(py: Python<'_>, tracer: Box<dyn Tracer>) -> PyResult<()> {
+    let mask = tracer.interest();
+    let mut guard = GLOBAL.lock().unwrap();
+    if guard.is_some() {
+        return Err(PyRuntimeError::new_err("tracer already installed"));
+    }
+    let monitoring = py.import("sys")?.getattr("monitoring")?;
+    let tool_id: u8 = monitoring.call_method1("use_tool_id", ("codetracer",))?.extract()?;
+    let events = monitoring.getattr("events")?;
+    let module = PyModule::new(py, "_codetracer_callbacks")?;
+
+    let mut callbacks = Vec::new();
+    if mask.contains(EventMask::CALL) {
+        module.add_function(wrap_pyfunction!(callback_call, &module)?)?;
+        let cb = module.getattr("callback_call")?;
+        let ev = events.getattr("CALL")?;
+        monitoring.call_method("register_callback", (tool_id, ev, &cb), None)?;
+        callbacks.push(cb.unbind());
+    }
+    if mask.contains(EventMask::LINE) {
+        module.add_function(wrap_pyfunction!(callback_line, &module)?)?;
+        let cb = module.getattr("callback_line")?;
+        let ev = events.getattr("LINE")?;
+        monitoring.call_method("register_callback", (tool_id, ev, &cb), None)?;
+        callbacks.push(cb.unbind());
+    }
+    monitoring.call_method("set_events", (tool_id, mask.bits(), mask.bits()), None)?;
+
+    *guard = Some(Global { dispatcher: Dispatcher::new(tracer), tool_id, callbacks });
+    Ok(())
+}
+
+/// Remove the installed tracer if any.
+pub fn uninstall_tracer(py: Python<'_>) -> PyResult<()> {
+    let mut guard = GLOBAL.lock().unwrap();
+    if let Some(global) = guard.take() {
+        let monitoring = py.import("sys")?.getattr("monitoring")?;
+        let events = monitoring.getattr("events")?;
+        if global.dispatcher.mask.contains(EventMask::CALL) {
+            let ev = events.getattr("CALL")?;
+            monitoring.call_method("register_callback", (global.tool_id, ev, py.None()), None)?;
+        }
+        if global.dispatcher.mask.contains(EventMask::LINE) {
+            let ev = events.getattr("LINE")?;
+            monitoring.call_method("register_callback", (global.tool_id, ev, py.None()), None)?;
+        }
+        monitoring.call_method("set_events", (global.tool_id, 0u32, global.dispatcher.mask.bits()), None)?;
+        monitoring.call_method1("free_tool_id", (global.tool_id,))?;
+    }
+    Ok(())
+}
+
+#[pyfunction]
+fn callback_call(py: Python<'_>, args: &Bound<'_, PyTuple>) -> PyResult<()> {
+    let frame = args.get_item(0)?.as_ptr() as *mut ffi::PyFrameObject;
+    if let Some(global) = GLOBAL.lock().unwrap().as_mut() {
+        global.dispatcher.dispatch_call(py, frame);
+    }
+    Ok(())
+}
+
+#[pyfunction]
+fn callback_line(py: Python<'_>, args: &Bound<'_, PyTuple>) -> PyResult<()> {
+    let frame = args.get_item(0)?.as_ptr() as *mut ffi::PyFrameObject;
+    let lineno: u32 = args.get_item(1)?.extract()?;
+    if let Some(global) = GLOBAL.lock().unwrap().as_mut() {
+        global.dispatcher.dispatch_line(py, frame, lineno);
+    }
+    Ok(())
+}
+

codetracer-python-recorder/src/tracer.rs

@@ -0,0 +1,27 @@
+use codetracer_python_recorder::{install_tracer, uninstall_tracer, EventMask, Tracer};
+use pyo3::prelude::*;
+use pyo3::ffi::PyFrameObject;
+use std::ffi::CString;
+
+struct PrintTracer;
+
+impl Tracer for PrintTracer {
+    fn interest(&self) -> EventMask {
+        EventMask::CALL
+    }
+
+    fn on_call(&mut self, _py: Python<'_>, _frame: *mut PyFrameObject) {
+        println!("call event");
+    }
+}
+
+#[test]
+fn tracer_prints_on_call() {
+    Python::with_gil(|py| {
+        install_tracer(py, Box::new(PrintTracer)).unwrap();
+        let code = CString::new("def foo():\n    return 1\nfoo()").unwrap();
+        py.run(&code, None, None).unwrap();
+        uninstall_tracer(py).unwrap();
+    });
+}
+

codetracer-python-recorder/tests/print_tracer.rs

@@ -8,6 +8,13 @@ The tracer collects `sys.monitoring` events, converts them to `runtime_tracing`
 
 ## Architecture
 
+### Tracer Abstraction
+Rust code exposes a `Tracer` trait representing callbacks for Python
+`sys.monitoring` events. Implementations advertise their desired events via an
+`EventMask` bit flag returned from `interest`. A `Dispatcher` wraps a trait
+object and forwards events only when the mask contains the corresponding flag,
+allowing tracers to implement just the methods they care about.
+
 ### Tool Initialization
 - Acquire a tool identifier via `sys.monitoring.use_tool_id`; store it for the lifetime of the tracer.
   ```rs