Skip to content

Commit 3c04842

Browse files
authored
feat: add monitoring tracer trait and hooks (#22)
## Summary - add `Tracer` trait with event mask and dispatcher for sys.monitoring - hook tracer to Python monitoring and expose install/uninstall helpers - run Rust tests via `just test` and document tracer abstraction design ## Testing - `just test` ------ https://chatgpt.com/codex/tasks/task_e_68a317f93b1c8323b654132620bf3c36
2 parents e4a9105 + 1b92837 commit 3c04842

File tree

7 files changed

+321
-5
lines changed

7 files changed

+321
-5
lines changed

Justfile

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,15 @@ dev:
3232
uv run --directory codetracer-python-recorder maturin develop --uv
3333

3434
# Run unit tests of dev build
35-
test:
36-
uv run --group dev --group test pytest
35+
test: cargo-test py-test
36+
37+
# Run Rust unit tests without default features to link Python C library
38+
cargo-test:
39+
uv run cargo test --manifest-path codetracer-python-recorder/Cargo.toml --no-default-features
3740

41+
py-test:
42+
uv run --group dev --group test pytest
43+
3844
# Run tests only on the pure recorder
3945
test-pure:
4046
uv run --group dev --group test pytest codetracer-pure-python-recorder

codetracer-python-recorder/Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

codetracer-python-recorder/Cargo.toml

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,16 @@ repository = "https://github.com/metacraft-labs/codetracer-python-recorder"
88

99
[lib]
1010
name = "codetracer_python_recorder"
11-
crate-type = ["cdylib"]
11+
crate-type = ["cdylib", "rlib"]
12+
13+
[features]
14+
extension-module = ["pyo3/extension-module"]
15+
default = ["extension-module"]
1216

1317
[dependencies]
14-
pyo3 = { version = "0.25.1", features = ["extension-module"] }
15-
runtime_tracing = "0.14.0"
18+
pyo3 = { version = "0.25.1" }
19+
runtime_tracing = "0.14.0"
20+
bitflags = "2.4"
21+
22+
[dev-dependencies]
23+
pyo3 = { version = "0.25.1", features = ["auto-initialize"] }

codetracer-python-recorder/src/lib.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@ use std::sync::atomic::{AtomicBool, Ordering};
33
use pyo3::exceptions::PyRuntimeError;
44
use pyo3::prelude::*;
55

6+
pub mod tracer;
7+
pub use crate::tracer::{install_tracer, uninstall_tracer, EventSet, Tracer};
8+
69
/// Global flag tracking whether tracing is active.
710
static ACTIVE: AtomicBool = AtomicBool::new(false);
811

Lines changed: 244 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,244 @@
1+
use std::sync::{Mutex, OnceLock};
2+
use pyo3::{
3+
exceptions::PyRuntimeError,
4+
prelude::*,
5+
types::{PyAny, PyCFunction, PyModule},
6+
};
7+
8+
const MONITORING_TOOL_NAME: &str = "codetracer";
9+
10+
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11+
#[repr(transparent)]
12+
pub struct EventId(pub i32);
13+
14+
#[allow(non_snake_case)]
15+
#[derive(Clone, Copy, Debug)]
16+
pub struct MonitoringEvents {
17+
pub BRANCH: EventId,
18+
pub CALL: EventId,
19+
pub C_RAISE: EventId,
20+
pub C_RETURN: EventId,
21+
pub EXCEPTION_HANDLED: EventId,
22+
pub INSTRUCTION: EventId,
23+
pub JUMP: EventId,
24+
pub LINE: EventId,
25+
pub PY_RESUME: EventId,
26+
pub PY_RETURN: EventId,
27+
pub PY_START: EventId,
28+
pub PY_THROW: EventId,
29+
pub PY_UNWIND: EventId,
30+
pub PY_YIELD: EventId,
31+
pub RAISE: EventId,
32+
pub RERAISE: EventId,
33+
pub STOP_ITERATION: EventId,
34+
}
35+
36+
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
37+
pub struct ToolId {
38+
pub id: u8,
39+
}
40+
41+
pub type CallbackFn<'py> = Bound<'py, PyCFunction>;
42+
43+
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
44+
pub struct EventSet(pub i32);
45+
46+
pub const NO_EVENTS: EventSet = EventSet(0);
47+
48+
impl EventSet {
49+
pub const fn empty() -> Self {
50+
NO_EVENTS
51+
}
52+
pub fn contains(&self, ev: &EventId) -> bool {
53+
(self.0 & ev.0) != 0
54+
}
55+
}
56+
57+
pub fn acquire_tool_id(py: Python<'_>) -> PyResult<ToolId> {
58+
let monitoring = py.import("sys")?.getattr("monitoring")?;
59+
const FALLBACK_ID: u8 = 5;
60+
monitoring.call_method1("use_tool_id", (FALLBACK_ID, MONITORING_TOOL_NAME))?;
61+
Ok(ToolId { id: FALLBACK_ID })
62+
}
63+
64+
pub fn load_monitoring_events(py: Python<'_>) -> PyResult<MonitoringEvents> {
65+
let monitoring = py.import("sys")?.getattr("monitoring")?;
66+
let events = monitoring.getattr("events")?;
67+
Ok(MonitoringEvents {
68+
BRANCH: EventId(events.getattr("BRANCH")?.extract()?),
69+
CALL: EventId(events.getattr("CALL")?.extract()?),
70+
C_RAISE: EventId(events.getattr("C_RAISE")?.extract()?),
71+
C_RETURN: EventId(events.getattr("C_RETURN")?.extract()?),
72+
EXCEPTION_HANDLED: EventId(events.getattr("EXCEPTION_HANDLED")?.extract()?),
73+
INSTRUCTION: EventId(events.getattr("INSTRUCTION")?.extract()?),
74+
JUMP: EventId(events.getattr("JUMP")?.extract()?),
75+
LINE: EventId(events.getattr("LINE")?.extract()?),
76+
PY_RESUME: EventId(events.getattr("PY_RESUME")?.extract()?),
77+
PY_RETURN: EventId(events.getattr("PY_RETURN")?.extract()?),
78+
PY_START: EventId(events.getattr("PY_START")?.extract()?),
79+
PY_THROW: EventId(events.getattr("PY_THROW")?.extract()?),
80+
PY_UNWIND: EventId(events.getattr("PY_UNWIND")?.extract()?),
81+
PY_YIELD: EventId(events.getattr("PY_YIELD")?.extract()?),
82+
RAISE: EventId(events.getattr("RAISE")?.extract()?),
83+
RERAISE: EventId(events.getattr("RERAISE")?.extract()?),
84+
STOP_ITERATION: EventId(events.getattr("STOP_ITERATION")?.extract()?),
85+
})
86+
}
87+
88+
static MONITORING_EVENTS: OnceLock<MonitoringEvents> = OnceLock::new();
89+
90+
pub fn monitoring_events(py: Python<'_>) -> PyResult<&'static MonitoringEvents> {
91+
if let Some(ev) = MONITORING_EVENTS.get() {
92+
return Ok(ev);
93+
}
94+
let ev = load_monitoring_events(py)?;
95+
let _ = MONITORING_EVENTS.set(ev);
96+
Ok(MONITORING_EVENTS.get().unwrap())
97+
}
98+
99+
pub fn register_callback(
100+
py: Python<'_>,
101+
tool: &ToolId,
102+
event: &EventId,
103+
cb: Option<&CallbackFn<'_>>,
104+
) -> PyResult<()> {
105+
let monitoring = py.import("sys")?.getattr("monitoring")?;
106+
match cb {
107+
Some(cb) => {
108+
monitoring.call_method("register_callback", (tool.id, event.0, cb), None)?;
109+
}
110+
None => {
111+
monitoring.call_method("register_callback", (tool.id, event.0, py.None()), None)?;
112+
}
113+
}
114+
Ok(())
115+
}
116+
117+
pub fn events_union(ids: &[EventId]) -> EventSet {
118+
let mut bits = 0i32;
119+
for id in ids {
120+
bits |= id.0;
121+
}
122+
EventSet(bits)
123+
}
124+
125+
pub fn set_events(py: Python<'_>, tool: &ToolId, set: EventSet) -> PyResult<()> {
126+
let monitoring = py.import("sys")?.getattr("monitoring")?;
127+
monitoring.call_method1("set_events", (tool.id, set.0))?;
128+
Ok(())
129+
}
130+
131+
pub fn free_tool_id(py: Python<'_>, tool: &ToolId) -> PyResult<()> {
132+
let monitoring = py.import("sys")?.getattr("monitoring")?;
133+
monitoring.call_method1("free_tool_id", (tool.id,))?;
134+
Ok(())
135+
}
136+
137+
138+
/// Trait implemented by tracing backends.
139+
///
140+
/// Each method corresponds to an event from `sys.monitoring`. Default
141+
/// implementations allow implementers to only handle the events they care
142+
/// about.
143+
pub trait Tracer: Send {
144+
/// Return the set of events the tracer wants to receive.
145+
fn interest(&self, _events: &MonitoringEvents) -> EventSet {
146+
NO_EVENTS
147+
}
148+
149+
/// Called on Python function calls.
150+
fn on_call(
151+
&mut self,
152+
_py: Python<'_>,
153+
_code: &Bound<'_, PyAny>,
154+
_offset: i32,
155+
_callable: &Bound<'_, PyAny>,
156+
_arg0: Option<&Bound<'_, PyAny>>,
157+
) {
158+
}
159+
160+
/// Called on line execution.
161+
fn on_line(&mut self, _py: Python<'_>, _code: &Bound<'_, PyAny>, _lineno: u32) {}
162+
}
163+
164+
struct Global {
165+
tracer: Box<dyn Tracer>,
166+
mask: EventSet,
167+
tool: ToolId,
168+
}
169+
170+
static GLOBAL: Mutex<Option<Global>> = Mutex::new(None);
171+
172+
/// Install a tracer and hook it into Python's `sys.monitoring`.
173+
pub fn install_tracer(py: Python<'_>, tracer: Box<dyn Tracer>) -> PyResult<()> {
174+
let mut guard = GLOBAL.lock().unwrap();
175+
if guard.is_some() {
176+
return Err(PyRuntimeError::new_err("tracer already installed"));
177+
}
178+
179+
let tool = acquire_tool_id(py)?;
180+
let events = monitoring_events(py)?;
181+
182+
let module = PyModule::new(py, "_codetracer_callbacks")?;
183+
184+
let mask = tracer.interest(events);
185+
186+
if mask.contains(&events.CALL) {
187+
let cb = wrap_pyfunction!(callback_call, &module)?;
188+
189+
register_callback(py, &tool, &events.CALL, Some(&cb))?;
190+
191+
}
192+
if mask.contains(&events.LINE) {
193+
let cb = wrap_pyfunction!(callback_line, &module)?;
194+
register_callback(py, &tool, &events.LINE, Some(&cb))?;
195+
}
196+
set_events(py, &tool, mask)?;
197+
198+
199+
*guard = Some(Global {
200+
tracer,
201+
mask,
202+
tool,
203+
});
204+
Ok(())
205+
}
206+
207+
/// Remove the installed tracer if any.
208+
pub fn uninstall_tracer(py: Python<'_>) -> PyResult<()> {
209+
let mut guard = GLOBAL.lock().unwrap();
210+
if let Some(global) = guard.take() {
211+
let events = monitoring_events(py)?;
212+
if global.mask.contains(&events.CALL) {
213+
register_callback(py, &global.tool, &events.CALL, None)?;
214+
}
215+
if global.mask.contains(&events.LINE) {
216+
register_callback(py, &global.tool, &events.LINE, None)?;
217+
}
218+
set_events(py, &global.tool, NO_EVENTS)?;
219+
free_tool_id(py, &global.tool)?;
220+
}
221+
Ok(())
222+
}
223+
224+
#[pyfunction]
225+
fn callback_call(
226+
py: Python<'_>,
227+
code: Bound<'_, PyAny>,
228+
offset: i32,
229+
callable: Bound<'_, PyAny>,
230+
arg0: Option<Bound<'_, PyAny>>,
231+
) -> PyResult<()> {
232+
if let Some(global) = GLOBAL.lock().unwrap().as_mut() {
233+
global.tracer.on_call(py, &code, offset, &callable, arg0.as_ref());
234+
}
235+
Ok(())
236+
}
237+
238+
#[pyfunction]
239+
fn callback_line(py: Python<'_>, code: Bound<'_, PyAny>, lineno: u32) -> PyResult<()> {
240+
if let Some(global) = GLOBAL.lock().unwrap().as_mut() {
241+
global.tracer.on_line(py, &code, lineno);
242+
}
243+
Ok(())
244+
}
Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
use codetracer_python_recorder::{install_tracer, uninstall_tracer, EventSet, Tracer};
2+
use codetracer_python_recorder::tracer::{MonitoringEvents, events_union};
3+
use pyo3::prelude::*;
4+
use std::ffi::CString;
5+
use std::sync::atomic::{AtomicUsize, Ordering};
6+
7+
static CALL_COUNT: AtomicUsize = AtomicUsize::new(0);
8+
9+
struct PrintTracer;
10+
11+
impl Tracer for PrintTracer {
12+
fn interest(&self, events:&MonitoringEvents) -> EventSet {
13+
events_union(&[events.CALL])
14+
}
15+
16+
fn on_call(
17+
&mut self,
18+
_py: Python<'_>,
19+
_code: &pyo3::Bound<'_, pyo3::types::PyAny>,
20+
_offset: i32,
21+
_callable: &pyo3::Bound<'_, pyo3::types::PyAny>,
22+
_arg0: Option<&pyo3::Bound<'_, pyo3::types::PyAny>>,
23+
) {
24+
CALL_COUNT.fetch_add(1, Ordering::SeqCst);
25+
}
26+
}
27+
28+
#[test]
29+
fn tracer_prints_on_call() {
30+
Python::with_gil(|py| {
31+
CALL_COUNT.store(0, Ordering::SeqCst);
32+
if let Err(e) = install_tracer(py, Box::new(PrintTracer)) {
33+
e.print(py);
34+
panic!("Install Tracer failed");
35+
}
36+
let code = CString::new("def foo():\n return 1\nfoo()").expect("CString::new failed");
37+
if let Err(e) = py.run(code.as_c_str(), None, None) {
38+
e.print(py);
39+
uninstall_tracer(py).ok();
40+
panic!("Python raised an exception");
41+
}
42+
uninstall_tracer(py).unwrap();
43+
let count = CALL_COUNT.load(Ordering::SeqCst);
44+
assert!(count >= 1, "expected at least one CALL event, got {}", count);
45+
});
46+
}
47+

design-docs/design-001.md

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,13 @@ The tracer collects `sys.monitoring` events, converts them to `runtime_tracing`
88

99
## Architecture
1010

11+
### Tracer Abstraction
12+
Rust code exposes a `Tracer` trait representing callbacks for Python
13+
`sys.monitoring` events. Implementations advertise their desired events via an
14+
`EventMask` bit flag returned from `interest`. A `Dispatcher` wraps a trait
15+
object and forwards events only when the mask contains the corresponding flag,
16+
allowing tracers to implement just the methods they care about.
17+
1118
### Tool Initialization
1219
- Acquire a tool identifier via `sys.monitoring.use_tool_id`; store it for the lifetime of the tracer.
1320
```rs

0 commit comments

Comments
 (0)