Skip to content

Commit 97f34de

Browse files
tzanko-matevtzanko-matev
committed
Step 4
1 parent 2305f35 commit 97f34de

File tree

6 files changed

+171
-163
lines changed

6 files changed

+171
-163
lines changed

‎codetracer-python-recorder/src/lib.rs‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,15 +6,17 @@
66
77
pub mod code_object;
88
mod logging;
9+
pub mod monitoring;
910
mod runtime;
1011
mod session;
11-
pub mod tracer;
1212

1313
pub use crate::code_object::{CodeObjectRegistry, CodeObjectWrapper};
14-
pub use crate::session::{flush_tracing, is_tracing, start_tracing, stop_tracing};
15-
pub use crate::tracer::{
16-
install_tracer, uninstall_tracer, CallbackOutcome, CallbackResult, EventSet, Tracer,
14+
pub use crate::monitoring as tracer;
15+
pub use crate::monitoring::{
16+
flush_installed_tracer, install_tracer, uninstall_tracer, CallbackOutcome, CallbackResult,
17+
EventSet, Tracer,
1718
};
19+
pub use crate::session::{flush_tracing, is_tracing, start_tracing, stop_tracing};
1820

1921
use pyo3::prelude::*;
2022

‎codetracer-python-recorder/src/monitoring/mod.rs‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
use pyo3::prelude::*;
2+
use pyo3::types::PyCFunction;
3+
use std::sync::OnceLock;
4+
5+
mod tracer;
6+
7+
pub use tracer::{flush_installed_tracer, install_tracer, uninstall_tracer, Tracer};
8+
9+
const MONITORING_TOOL_NAME: &str = "codetracer";
10+
11+
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12+
#[repr(transparent)]
13+
pub struct EventId(pub i32);
14+
15+
#[allow(non_snake_case)]
16+
#[derive(Clone, Copy, Debug)]
17+
pub struct MonitoringEvents {
18+
pub BRANCH: EventId,
19+
pub CALL: EventId,
20+
pub C_RAISE: EventId,
21+
pub C_RETURN: EventId,
22+
pub EXCEPTION_HANDLED: EventId,
23+
pub INSTRUCTION: EventId,
24+
pub JUMP: EventId,
25+
pub LINE: EventId,
26+
pub PY_RESUME: EventId,
27+
pub PY_RETURN: EventId,
28+
pub PY_START: EventId,
29+
pub PY_THROW: EventId,
30+
pub PY_UNWIND: EventId,
31+
pub PY_YIELD: EventId,
32+
pub RAISE: EventId,
33+
pub RERAISE: EventId,
34+
//pub STOP_ITERATION: EventId, //See comment in Tracer trait
35+
}
36+
37+
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
38+
pub struct ToolId {
39+
pub id: u8,
40+
}
41+
42+
pub type CallbackFn<'py> = Bound<'py, PyCFunction>;
43+
44+
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
45+
pub struct EventSet(pub i32);
46+
47+
pub const NO_EVENTS: EventSet = EventSet(0);
48+
49+
/// Outcome returned by tracer callbacks to control CPython monitoring.
50+
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
51+
pub enum CallbackOutcome {
52+
/// Continue receiving events for the current location.
53+
Continue,
54+
/// Disable future events for the current location by returning
55+
/// `sys.monitoring.DISABLE`.
56+
DisableLocation,
57+
}
58+
59+
/// Result type shared by tracer callbacks.
60+
pub type CallbackResult = PyResult<CallbackOutcome>;
61+
62+
static MONITORING_EVENTS: OnceLock<MonitoringEvents> = OnceLock::new();
63+
64+
impl EventSet {
65+
pub const fn empty() -> Self {
66+
NO_EVENTS
67+
}
68+
69+
pub fn contains(&self, ev: &EventId) -> bool {
70+
(self.0 & ev.0) != 0
71+
}
72+
}
73+
74+
pub fn acquire_tool_id(py: Python<'_>) -> PyResult<ToolId> {
75+
let monitoring = py.import("sys")?.getattr("monitoring")?;
76+
const FALLBACK_ID: u8 = 5;
77+
monitoring.call_method1("use_tool_id", (FALLBACK_ID, MONITORING_TOOL_NAME))?;
78+
Ok(ToolId { id: FALLBACK_ID })
79+
}
80+
81+
pub fn load_monitoring_events(py: Python<'_>) -> PyResult<MonitoringEvents> {
82+
let monitoring = py.import("sys")?.getattr("monitoring")?;
83+
let events = monitoring.getattr("events")?;
84+
Ok(MonitoringEvents {
85+
BRANCH: EventId(events.getattr("BRANCH")?.extract()?),
86+
CALL: EventId(events.getattr("CALL")?.extract()?),
87+
C_RAISE: EventId(events.getattr("C_RAISE")?.extract()?),
88+
C_RETURN: EventId(events.getattr("C_RETURN")?.extract()?),
89+
EXCEPTION_HANDLED: EventId(events.getattr("EXCEPTION_HANDLED")?.extract()?),
90+
INSTRUCTION: EventId(events.getattr("INSTRUCTION")?.extract()?),
91+
JUMP: EventId(events.getattr("JUMP")?.extract()?),
92+
LINE: EventId(events.getattr("LINE")?.extract()?),
93+
PY_RESUME: EventId(events.getattr("PY_RESUME")?.extract()?),
94+
PY_RETURN: EventId(events.getattr("PY_RETURN")?.extract()?),
95+
PY_START: EventId(events.getattr("PY_START")?.extract()?),
96+
PY_THROW: EventId(events.getattr("PY_THROW")?.extract()?),
97+
PY_UNWIND: EventId(events.getattr("PY_UNWIND")?.extract()?),
98+
PY_YIELD: EventId(events.getattr("PY_YIELD")?.extract()?),
99+
RAISE: EventId(events.getattr("RAISE")?.extract()?),
100+
RERAISE: EventId(events.getattr("RERAISE")?.extract()?),
101+
//STOP_ITERATION: EventId(events.getattr("STOP_ITERATION")?.extract()?), //See comment in Tracer trait
102+
})
103+
}
104+
105+
pub fn monitoring_events(py: Python<'_>) -> PyResult<&'static MonitoringEvents> {
106+
if let Some(ev) = MONITORING_EVENTS.get() {
107+
return Ok(ev);
108+
}
109+
let ev = load_monitoring_events(py)?;
110+
let _ = MONITORING_EVENTS.set(ev);
111+
Ok(MONITORING_EVENTS.get().unwrap())
112+
}
113+
114+
pub fn register_callback(
115+
py: Python<'_>,
116+
tool: &ToolId,
117+
event: &EventId,
118+
cb: Option<&CallbackFn<'_>>,
119+
) -> PyResult<()> {
120+
let monitoring = py.import("sys")?.getattr("monitoring")?;
121+
match cb {
122+
Some(cb) => {
123+
monitoring.call_method("register_callback", (tool.id, event.0, cb), None)?;
124+
}
125+
None => {
126+
monitoring.call_method("register_callback", (tool.id, event.0, py.None()), None)?;
127+
}
128+
}
129+
Ok(())
130+
}
131+
132+
pub fn events_union(ids: &[EventId]) -> EventSet {
133+
let mut bits = 0i32;
134+
for id in ids {
135+
bits |= id.0;
136+
}
137+
EventSet(bits)
138+
}
139+
140+
pub fn set_events(py: Python<'_>, tool: &ToolId, set: EventSet) -> PyResult<()> {
141+
let monitoring = py.import("sys")?.getattr("monitoring")?;
142+
monitoring.call_method1("set_events", (tool.id, set.0))?;
143+
Ok(())
144+
}
145+
146+
pub fn free_tool_id(py: Python<'_>, tool: &ToolId) -> PyResult<()> {
147+
let monitoring = py.import("sys")?.getattr("monitoring")?;
148+
monitoring.call_method1("free_tool_id", (tool.id,))?;
149+
Ok(())
150+
}

‎codetracer-python-recorder/src/tracer.rs‎ renamed to ‎codetracer-python-recorder/src/monitoring/tracer.rs‎

Lines changed: 9 additions & 154 deletions
Original file line numberDiff line numberDiff line change
@@ -1,153 +1,17 @@
1+
use std::any::Any;
2+
use std::sync::Mutex;
3+
14
use crate::code_object::{CodeObjectRegistry, CodeObjectWrapper};
25
use pyo3::{
36
exceptions::PyRuntimeError,
47
prelude::*,
5-
types::{PyAny, PyCFunction, PyCode, PyModule},
8+
types::{PyAny, PyCode, PyModule},
69
};
7-
use std::any::Any;
8-
use std::sync::{Mutex, OnceLock};
9-
10-
const MONITORING_TOOL_NAME: &str = "codetracer";
11-
12-
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
13-
#[repr(transparent)]
14-
pub struct EventId(pub i32);
15-
16-
#[allow(non_snake_case)]
17-
#[derive(Clone, Copy, Debug)]
18-
pub struct MonitoringEvents {
19-
pub BRANCH: EventId,
20-
pub CALL: EventId,
21-
pub C_RAISE: EventId,
22-
pub C_RETURN: EventId,
23-
pub EXCEPTION_HANDLED: EventId,
24-
pub INSTRUCTION: EventId,
25-
pub JUMP: EventId,
26-
pub LINE: EventId,
27-
pub PY_RESUME: EventId,
28-
pub PY_RETURN: EventId,
29-
pub PY_START: EventId,
30-
pub PY_THROW: EventId,
31-
pub PY_UNWIND: EventId,
32-
pub PY_YIELD: EventId,
33-
pub RAISE: EventId,
34-
pub RERAISE: EventId,
35-
//pub STOP_ITERATION: EventId, //See comment in Tracer trait
36-
}
37-
38-
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
39-
pub struct ToolId {
40-
pub id: u8,
41-
}
42-
43-
pub type CallbackFn<'py> = Bound<'py, PyCFunction>;
44-
45-
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
46-
pub struct EventSet(pub i32);
47-
48-
pub const NO_EVENTS: EventSet = EventSet(0);
49-
50-
/// Outcome returned by tracer callbacks to control CPython monitoring.
51-
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
52-
pub enum CallbackOutcome {
53-
/// Continue receiving events for the current location.
54-
Continue,
55-
/// Disable future events for the current location by returning
56-
/// `sys.monitoring.DISABLE`.
57-
DisableLocation,
58-
}
59-
60-
/// Result type shared by tracer callbacks.
61-
pub type CallbackResult = PyResult<CallbackOutcome>;
62-
63-
impl EventSet {
64-
pub const fn empty() -> Self {
65-
NO_EVENTS
66-
}
67-
pub fn contains(&self, ev: &EventId) -> bool {
68-
(self.0 & ev.0) != 0
69-
}
70-
}
71-
72-
pub fn acquire_tool_id(py: Python<'_>) -> PyResult<ToolId> {
73-
let monitoring = py.import("sys")?.getattr("monitoring")?;
74-
const FALLBACK_ID: u8 = 5;
75-
monitoring.call_method1("use_tool_id", (FALLBACK_ID, MONITORING_TOOL_NAME))?;
76-
Ok(ToolId { id: FALLBACK_ID })
77-
}
78-
79-
pub fn load_monitoring_events(py: Python<'_>) -> PyResult<MonitoringEvents> {
80-
let monitoring = py.import("sys")?.getattr("monitoring")?;
81-
let events = monitoring.getattr("events")?;
82-
Ok(MonitoringEvents {
83-
BRANCH: EventId(events.getattr("BRANCH")?.extract()?),
84-
CALL: EventId(events.getattr("CALL")?.extract()?),
85-
C_RAISE: EventId(events.getattr("C_RAISE")?.extract()?),
86-
C_RETURN: EventId(events.getattr("C_RETURN")?.extract()?),
87-
EXCEPTION_HANDLED: EventId(events.getattr("EXCEPTION_HANDLED")?.extract()?),
88-
INSTRUCTION: EventId(events.getattr("INSTRUCTION")?.extract()?),
89-
JUMP: EventId(events.getattr("JUMP")?.extract()?),
90-
LINE: EventId(events.getattr("LINE")?.extract()?),
91-
PY_RESUME: EventId(events.getattr("PY_RESUME")?.extract()?),
92-
PY_RETURN: EventId(events.getattr("PY_RETURN")?.extract()?),
93-
PY_START: EventId(events.getattr("PY_START")?.extract()?),
94-
PY_THROW: EventId(events.getattr("PY_THROW")?.extract()?),
95-
PY_UNWIND: EventId(events.getattr("PY_UNWIND")?.extract()?),
96-
PY_YIELD: EventId(events.getattr("PY_YIELD")?.extract()?),
97-
RAISE: EventId(events.getattr("RAISE")?.extract()?),
98-
RERAISE: EventId(events.getattr("RERAISE")?.extract()?),
99-
//STOP_ITERATION: EventId(events.getattr("STOP_ITERATION")?.extract()?), //See comment in Tracer trait
100-
})
101-
}
102-
103-
static MONITORING_EVENTS: OnceLock<MonitoringEvents> = OnceLock::new();
10410

105-
pub fn monitoring_events(py: Python<'_>) -> PyResult<&'static MonitoringEvents> {
106-
if let Some(ev) = MONITORING_EVENTS.get() {
107-
return Ok(ev);
108-
}
109-
let ev = load_monitoring_events(py)?;
110-
let _ = MONITORING_EVENTS.set(ev);
111-
Ok(MONITORING_EVENTS.get().unwrap())
112-
}
113-
114-
pub fn register_callback(
115-
py: Python<'_>,
116-
tool: &ToolId,
117-
event: &EventId,
118-
cb: Option<&CallbackFn<'_>>,
119-
) -> PyResult<()> {
120-
let monitoring = py.import("sys")?.getattr("monitoring")?;
121-
match cb {
122-
Some(cb) => {
123-
monitoring.call_method("register_callback", (tool.id, event.0, cb), None)?;
124-
}
125-
None => {
126-
monitoring.call_method("register_callback", (tool.id, event.0, py.None()), None)?;
127-
}
128-
}
129-
Ok(())
130-
}
131-
132-
pub fn events_union(ids: &[EventId]) -> EventSet {
133-
let mut bits = 0i32;
134-
for id in ids {
135-
bits |= id.0;
136-
}
137-
EventSet(bits)
138-
}
139-
140-
pub fn set_events(py: Python<'_>, tool: &ToolId, set: EventSet) -> PyResult<()> {
141-
let monitoring = py.import("sys")?.getattr("monitoring")?;
142-
monitoring.call_method1("set_events", (tool.id, set.0))?;
143-
Ok(())
144-
}
145-
146-
pub fn free_tool_id(py: Python<'_>, tool: &ToolId) -> PyResult<()> {
147-
let monitoring = py.import("sys")?.getattr("monitoring")?;
148-
monitoring.call_method1("free_tool_id", (tool.id,))?;
149-
Ok(())
150-
}
11+
use super::{
12+
acquire_tool_id, free_tool_id, monitoring_events, register_callback, set_events,
13+
CallbackOutcome, CallbackResult, EventSet, MonitoringEvents, ToolId, NO_EVENTS,
14+
};
15115

15216
/// Trait implemented by tracing backends.
15317
///
@@ -545,7 +409,6 @@ pub fn uninstall_tracer(py: Python<'_>) -> PyResult<()> {
545409
if global.mask.contains(&events.EXCEPTION_HANDLED) {
546410
register_callback(py, &global.tool, &events.EXCEPTION_HANDLED, None)?;
547411
}
548-
// See comment in tracer trait
549412
// if global.mask.contains(&events.STOP_ITERATION) {
550413
// register_callback(py, &global.tool, &events.STOP_ITERATION, None)?;
551414
// }
@@ -556,7 +419,6 @@ pub fn uninstall_tracer(py: Python<'_>) -> PyResult<()> {
556419
register_callback(py, &global.tool, &events.C_RAISE, None)?;
557420
}
558421

559-
global.registry.clear();
560422
set_events(py, &global.tool, NO_EVENTS)?;
561423
free_tool_id(py, &global.tool)?;
562424
}
@@ -658,14 +520,7 @@ fn callback_py_start(
658520
if let Some(global) = GLOBAL.lock().unwrap().as_mut() {
659521
let wrapper = global.registry.get_or_insert(py, &code);
660522
let result = global.tracer.on_py_start(py, &wrapper, instruction_offset);
661-
return match result {
662-
Ok(outcome) => global.handle_callback(py, Ok(outcome)),
663-
Err(err) => {
664-
let _ = set_events(py, &global.tool, NO_EVENTS);
665-
log::error!("Event monitoring turned off due to exception. No new events will be recorded! {}", err);
666-
Err(err)
667-
}
668-
};
523+
return global.handle_callback(py, result);
669524
}
670525
Ok(py.None())
671526
}

‎codetracer-python-recorder/src/runtime/mod.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ use runtime_tracing::NonStreamingTraceWriter;
1818
use runtime_tracing::{Line, TraceEventsFileFormat, TraceWriter};
1919

2020
use crate::code_object::CodeObjectWrapper;
21-
use crate::tracer::{
21+
use crate::monitoring::{
2222
events_union, CallbackOutcome, CallbackResult, EventSet, MonitoringEvents, Tracer,
2323
};
2424

@@ -449,7 +449,7 @@ impl Tracer for RuntimeTracer {
449449
#[cfg(test)]
450450
mod tests {
451451
use super::*;
452-
use crate::tracer::CallbackOutcome;
452+
use crate::monitoring::CallbackOutcome;
453453
use pyo3::types::{PyCode, PyModule};
454454
use pyo3::wrap_pyfunction;
455455
use runtime_tracing::{FullValueRecord, TraceLowLevelEvent, ValueRecord};

‎codetracer-python-recorder/src/session.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@ use pyo3::exceptions::PyRuntimeError;
66
use pyo3::prelude::*;
77

88
use crate::logging::init_rust_logging_with_default;
9+
use crate::monitoring::{flush_installed_tracer, install_tracer, uninstall_tracer};
910
use crate::runtime::{RuntimeTracer, TraceOutputPaths};
10-
use crate::tracer::{flush_installed_tracer, install_tracer, uninstall_tracer};
1111

1212
/// Global flag tracking whether tracing is active.
1313
static ACTIVE: AtomicBool = AtomicBool::new(false);

0 commit comments

Comments
 (0)