Skip to content

Commit aa3f613

Browse files
tzanko-matevtzanko-matev
committed
Stage 1
1 parent b19c1ed commit aa3f613

File tree

12 files changed

+154
-33
lines changed

12 files changed

+154
-33
lines changed

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -737,9 +737,7 @@ mod tests {
737737
module
738738
.setattr("__package__", py.None())
739739
.expect("set __package__");
740-
module
741-
.setattr("__spec__", py.None())
742-
.expect("set __spec__");
740+
module.setattr("__spec__", py.None()).expect("set __spec__");
743741
modules
744742
.set_item("__main__", module)
745743
.expect("register __main__");

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

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -108,11 +108,7 @@ pub trait Tracer: Send + Any {
108108
}
109109

110110
/// Provide the process exit status ahead of tracer teardown.
111-
fn set_exit_status(
112-
&mut self,
113-
_py: Python<'_>,
114-
_exit_code: Option<i32>,
115-
) -> PyResult<()> {
111+
fn set_exit_status(&mut self, _py: Python<'_>, _exit_code: Option<i32>) -> PyResult<()> {
116112
Ok(())
117113
}
118114

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

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,7 @@ pub(crate) mod install;
1010
pub mod tracer;
1111

1212
pub use api::Tracer;
13-
pub use install::{
14-
flush_installed_tracer,
15-
install_tracer,
16-
uninstall_tracer,
17-
update_exit_status,
18-
};
13+
pub use install::{flush_installed_tracer, install_tracer, uninstall_tracer, update_exit_status};
1914

2015
const MONITORING_TOOL_NAME: &str = "codetracer";
2116

‎codetracer-python-recorder/src/policy/ffi.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -169,7 +169,7 @@ mod tests {
169169
None,
170170
None,
171171
)
172-
.expect_err("invalid variant should error");
172+
.expect_err("invalid variant should error");
173173
// Ensure the error maps through map_recorder_error by checking the display text.
174174
let message = Python::with_gil(|py| err.value(py).to_string());
175175
assert!(

‎codetracer-python-recorder/src/runtime/tracer/events.rs‎

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,29 @@ impl Tracer for RuntimeTracer {
176176
code: &CodeObjectWrapper,
177177
_offset: i32,
178178
) -> CallbackResult {
179+
if self.module_name_from_globals {
180+
if let Ok(qualname) = code.qualname(py) {
181+
if qualname == "<module>" {
182+
let globals_name = match capture_frame(py, code) {
183+
Ok(snapshot) => {
184+
let mapping = snapshot.globals().unwrap_or_else(|| snapshot.locals());
185+
mapping
186+
.get_item("__name__")
187+
.ok()
188+
.flatten()
189+
.and_then(|value| value.extract::<String>().ok())
190+
.map(|name| name.trim().to_string())
191+
.filter(|name| !name.is_empty())
192+
}
193+
Err(_) => None,
194+
};
195+
196+
self.filter
197+
.set_module_name_hint(code.id(), globals_name.clone());
198+
}
199+
}
200+
}
201+
179202
if let Some(outcome) = self.evaluate_gate(py, code, true) {
180203
return Ok(outcome);
181204
}
@@ -398,11 +421,7 @@ impl Tracer for RuntimeTracer {
398421
)
399422
}
400423

401-
fn set_exit_status(
402-
&mut self,
403-
_py: Python<'_>,
404-
exit_code: Option<i32>,
405-
) -> PyResult<()> {
424+
fn set_exit_status(&mut self, _py: Python<'_>, exit_code: Option<i32>) -> PyResult<()> {
406425
self.record_exit_status(exit_code);
407426
Ok(())
408427
}
@@ -458,7 +477,9 @@ impl Tracer for RuntimeTracer {
458477

459478
if self.lifecycle.encountered_failure() {
460479
if policy.keep_partial_trace {
461-
if let Err(err) = self.lifecycle.finalise(&mut self.writer, &self.filter, &exit_summary)
480+
if let Err(err) =
481+
self.lifecycle
482+
.finalise(&mut self.writer, &self.filter, &exit_summary)
462483
{
463484
with_error_code(ErrorCode::TraceIncomplete, || {
464485
log::warn!(

‎codetracer-python-recorder/src/runtime/tracer/filtering.rs‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ pub(crate) struct FilterCoordinator {
2323
engine: Option<Arc<TraceFilterEngine>>,
2424
ignored_code_ids: HashSet<usize>,
2525
scope_cache: HashMap<usize, Arc<ScopeResolution>>,
26+
module_name_hints: HashMap<usize, String>,
2627
stats: FilterStats,
2728
}
2829

@@ -32,6 +33,7 @@ impl FilterCoordinator {
3233
engine,
3334
ignored_code_ids: HashSet::new(),
3435
scope_cache: HashMap::new(),
36+
module_name_hints: HashMap::new(),
3537
stats: FilterStats::default(),
3638
}
3739
}
@@ -52,6 +54,21 @@ impl FilterCoordinator {
5254
self.stats.values_mut()
5355
}
5456

57+
pub(crate) fn set_module_name_hint(&mut self, code_id: usize, hint: Option<String>) {
58+
match hint {
59+
Some(value) => {
60+
self.module_name_hints.insert(code_id, value);
61+
}
62+
None => {
63+
self.module_name_hints.remove(&code_id);
64+
}
65+
}
66+
}
67+
68+
pub(crate) fn module_name_hint(&self, code_id: usize) -> Option<String> {
69+
self.module_name_hints.get(&code_id).cloned()
70+
}
71+
5572
pub(crate) fn clear_caches(&mut self) {
5673
self.ignored_code_ids.clear();
5774
self.scope_cache.clear();
@@ -113,6 +130,17 @@ impl FilterCoordinator {
113130

114131
match engine.resolve(py, code) {
115132
Ok(resolution) => {
133+
let hint = self.module_name_hints.get(&code_id).cloned();
134+
let resolution = if let Some(hint) = hint {
135+
if resolution.module_name() == Some(hint.as_str()) {
136+
resolution
137+
} else {
138+
Arc::new(resolution.clone_with_module_name(Some(hint)))
139+
}
140+
} else {
141+
resolution
142+
};
143+
116144
if resolution.exec() == ExecDecision::Trace {
117145
self.scope_cache.insert(code_id, Arc::clone(&resolution));
118146
} else {
@@ -140,6 +168,7 @@ impl FilterCoordinator {
140168
fn mark_ignored(&mut self, code_id: usize) {
141169
self.scope_cache.remove(&code_id);
142170
self.ignored_code_ids.insert(code_id);
171+
self.module_name_hints.remove(&code_id);
143172
}
144173
}
145174

‎codetracer-python-recorder/src/runtime/tracer/lifecycle.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,10 @@ impl LifecycleController {
121121
enverr!(ErrorCode::Io, "failed to finalise trace events")
122122
.with_context("source", err.to_string())
123123
})?;
124-
debug!("[Lifecycle] writing exit metadata: code={:?}, label={:?}", exit_summary.code, exit_summary.label);
124+
debug!(
125+
"[Lifecycle] writing exit metadata: code={:?}, label={:?}",
126+
exit_summary.code, exit_summary.label
127+
);
125128
self.append_filter_metadata(filter)?;
126129
self.append_exit_metadata(exit_summary)?;
127130
Ok(())

‎codetracer-python-recorder/src/runtime/tracer/runtime_tracer.rs‎

Lines changed: 50 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ pub struct RuntimeTracer {
125125
pub(super) io: IoCoordinator,
126126
pub(super) filter: FilterCoordinator,
127127
pub(super) module_names: ModuleIdentityCache,
128+
pub(super) module_name_from_globals: bool,
128129
session_exit: SessionExitState,
129130
}
130131

@@ -135,6 +136,7 @@ impl RuntimeTracer {
135136
format: TraceEventsFileFormat,
136137
activation_path: Option<&Path>,
137138
trace_filter: Option<Arc<TraceFilterEngine>>,
139+
module_name_from_globals: bool,
138140
) -> Self {
139141
let mut writer = NonStreamingTraceWriter::new(program, args);
140142
writer.set_format(format);
@@ -147,6 +149,7 @@ impl RuntimeTracer {
147149
io: IoCoordinator::new(),
148150
filter: FilterCoordinator::new(trace_filter),
149151
module_names: ModuleIdentityCache::new(),
152+
module_name_from_globals,
150153
session_exit: SessionExitState::default(),
151154
}
152155
}
@@ -231,7 +234,10 @@ impl RuntimeTracer {
231234
code: &CodeObjectWrapper,
232235
allow_disable: bool,
233236
) -> Option<CallbackOutcome> {
234-
let is_active = self.lifecycle.activation_mut().should_process_event(py, code);
237+
let is_active = self
238+
.lifecycle
239+
.activation_mut()
240+
.should_process_event(py, code);
235241
if matches!(
236242
self.should_trace_code(py, code),
237243
TraceDecision::SkipAndDisable
@@ -290,6 +296,12 @@ impl RuntimeTracer {
290296
}
291297

292298
fn derive_module_name(&self, py: Python<'_>, code: &CodeObjectWrapper) -> Option<String> {
299+
if self.module_name_from_globals {
300+
if let Some(name) = self.filter.module_name_hint(code.id()) {
301+
return Some(name);
302+
}
303+
}
304+
293305
let resolution = self.filter.cached_resolution(code.id());
294306
if let Some(resolution) = resolution.as_ref() {
295307
let hints = ModuleNameHints {
@@ -388,8 +400,14 @@ mod tests {
388400
#[test]
389401
fn skips_synthetic_filename_events() {
390402
Python::with_gil(|py| {
391-
let mut tracer =
392-
RuntimeTracer::new("test.py", &[], TraceEventsFileFormat::Json, None, None);
403+
let mut tracer = RuntimeTracer::new(
404+
"test.py",
405+
&[],
406+
TraceEventsFileFormat::Json,
407+
None,
408+
None,
409+
false,
410+
);
393411
ensure_test_module(py);
394412
let script = format!("{PRELUDE}\nsnapshot()\n");
395413
{
@@ -485,6 +503,7 @@ result = compute()\n"
485503
TraceEventsFileFormat::Json,
486504
Some(script_path.as_path()),
487505
None,
506+
false,
488507
);
489508

490509
{
@@ -536,6 +555,7 @@ result = compute()\n"
536555
TraceEventsFileFormat::Json,
537556
None,
538557
None,
558+
false,
539559
);
540560
let store = tracer.line_snapshot_store();
541561

@@ -612,6 +632,7 @@ result = compute()\n"
612632
TraceEventsFileFormat::Json,
613633
None,
614634
None,
635+
false,
615636
);
616637
let outputs = TraceOutputPaths::new(tmp.path(), TraceEventsFileFormat::Json);
617638
tracer.begin(&outputs, 1).expect("begin tracer");
@@ -705,6 +726,7 @@ result = compute()\n"
705726
TraceEventsFileFormat::Json,
706727
None,
707728
None,
729+
false,
708730
);
709731
let outputs = TraceOutputPaths::new(tmp.path(), TraceEventsFileFormat::Json);
710732
tracer.begin(&outputs, 1).expect("begin tracer");
@@ -812,6 +834,7 @@ result = compute()\n"
812834
TraceEventsFileFormat::Json,
813835
None,
814836
None,
837+
false,
815838
);
816839
let outputs = TraceOutputPaths::new(tmp.path(), TraceEventsFileFormat::Json);
817840
tracer.begin(&outputs, 1).expect("begin tracer");
@@ -1066,8 +1089,14 @@ def start_call():
10661089

10671090
fn run_traced_script(body: &str) -> Vec<Snapshot> {
10681091
Python::with_gil(|py| {
1069-
let mut tracer =
1070-
RuntimeTracer::new("test.py", &[], TraceEventsFileFormat::Json, None, None);
1092+
let mut tracer = RuntimeTracer::new(
1093+
"test.py",
1094+
&[],
1095+
TraceEventsFileFormat::Json,
1096+
None,
1097+
None,
1098+
false,
1099+
);
10711100
ensure_test_module(py);
10721101
let tmp = tempfile::tempdir().expect("create temp dir");
10731102
let script_path = tmp.path().join("script.py");
@@ -1191,6 +1220,7 @@ sensitive("s3cr3t")
11911220
TraceEventsFileFormat::Json,
11921221
None,
11931222
Some(engine),
1223+
false,
11941224
);
11951225

11961226
{
@@ -1297,8 +1327,14 @@ sensitive("s3cr3t")
12971327
.call_method1("insert", (0, pkg_root.to_string_lossy().as_ref()))
12981328
.expect("insert temp root");
12991329

1300-
let tracer =
1301-
RuntimeTracer::new("runner.py", &[], TraceEventsFileFormat::Json, None, None);
1330+
let tracer = RuntimeTracer::new(
1331+
"runner.py",
1332+
&[],
1333+
TraceEventsFileFormat::Json,
1334+
None,
1335+
None,
1336+
false,
1337+
);
13021338

13031339
let builtins = py.import("builtins").expect("builtins");
13041340
let compile = builtins.getattr("compile").expect("compile builtin");
@@ -1361,6 +1397,7 @@ dropper()
13611397
TraceEventsFileFormat::Json,
13621398
None,
13631399
Some(engine),
1400+
false,
13641401
);
13651402

13661403
{
@@ -1448,6 +1485,7 @@ initializer("omega")
14481485
TraceEventsFileFormat::Json,
14491486
None,
14501487
Some(engine),
1488+
false,
14511489
);
14521490

14531491
{
@@ -1501,6 +1539,7 @@ initializer("omega")
15011539
TraceEventsFileFormat::Json,
15021540
None,
15031541
None,
1542+
false,
15041543
);
15051544
tracer.begin(&outputs, 1).expect("begin tracer");
15061545
tracer.record_exit_status(Some(7));
@@ -1607,6 +1646,7 @@ sensitive("s3cr3t")
16071646
TraceEventsFileFormat::Json,
16081647
None,
16091648
Some(engine),
1649+
false,
16101650
);
16111651
tracer.begin(&outputs, 1).expect("begin tracer");
16121652

@@ -2163,6 +2203,7 @@ snapshot()
21632203
TraceEventsFileFormat::Json,
21642204
None,
21652205
None,
2206+
false,
21662207
);
21672208
tracer.begin(&outputs, 1).expect("begin tracer");
21682209

@@ -2197,6 +2238,7 @@ snapshot()
21972238
TraceEventsFileFormat::Json,
21982239
None,
21992240
None,
2241+
false,
22002242
);
22012243
tracer.begin(&outputs, 1).expect("begin tracer");
22022244
tracer.mark_failure();
@@ -2241,6 +2283,7 @@ snapshot()
22412283
TraceEventsFileFormat::Json,
22422284
None,
22432285
None,
2286+
false,
22442287
);
22452288
tracer.begin(&outputs, 1).expect("begin tracer");
22462289
tracer.mark_failure();

0 commit comments

Comments
 (0)