Skip to content

Commit 1b19ac8

Browse files
tzanko-matevtzanko-matev
committed
WS5-3
1 parent b5c4fad commit 1b19ac8

File tree

4 files changed

+304
-77
lines changed

4 files changed

+304
-77
lines changed

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

Lines changed: 75 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,12 @@ use std::sync::Mutex;
66
use crate::code_object::{CodeObjectRegistry, CodeObjectWrapper};
77
use crate::ffi;
88
use crate::policy::{self, OnRecorderError};
9+
use log::{error, warn};
910
use pyo3::{
1011
prelude::*,
1112
types::{PyAny, PyCode, PyModule},
1213
};
1314
use recorder_errors::{usage, ErrorCode};
14-
use log::{error, warn};
1515

1616
use super::{
1717
acquire_tool_id, free_tool_id, monitoring_events, register_callback, set_events,
@@ -270,12 +270,10 @@ fn handle_callback_result(
270270
) -> PyResult<Py<PyAny>> {
271271
match result {
272272
Ok(CallbackOutcome::Continue) => Ok(py.None()),
273-
Ok(CallbackOutcome::DisableLocation) => Ok(
274-
guard
275-
.as_ref()
276-
.map(|global| global.disable_sentinel.clone_ref(py))
277-
.unwrap_or_else(|| py.None()),
278-
),
273+
Ok(CallbackOutcome::DisableLocation) => Ok(guard
274+
.as_ref()
275+
.map(|global| global.disable_sentinel.clone_ref(py))
276+
.unwrap_or_else(|| py.None())),
279277
Err(err) => handle_callback_error(py, guard, err),
280278
}
281279
}
@@ -310,62 +308,78 @@ fn handle_callback_error(
310308

311309
fn uninstall_locked(py: Python<'_>, guard: &mut Option<Global>) -> PyResult<()> {
312310
if let Some(mut global) = guard.take() {
313-
let _ = global.tracer.finish(py);
314-
let events = monitoring_events(py)?;
315-
if global.mask.contains(&events.CALL) {
316-
register_callback(py, &global.tool, &events.CALL, None)?;
317-
}
318-
if global.mask.contains(&events.LINE) {
319-
register_callback(py, &global.tool, &events.LINE, None)?;
320-
}
321-
if global.mask.contains(&events.INSTRUCTION) {
322-
register_callback(py, &global.tool, &events.INSTRUCTION, None)?;
323-
}
324-
if global.mask.contains(&events.JUMP) {
325-
register_callback(py, &global.tool, &events.JUMP, None)?;
326-
}
327-
if global.mask.contains(&events.BRANCH) {
328-
register_callback(py, &global.tool, &events.BRANCH, None)?;
329-
}
330-
if global.mask.contains(&events.PY_START) {
331-
register_callback(py, &global.tool, &events.PY_START, None)?;
332-
}
333-
if global.mask.contains(&events.PY_RESUME) {
334-
register_callback(py, &global.tool, &events.PY_RESUME, None)?;
335-
}
336-
if global.mask.contains(&events.PY_RETURN) {
337-
register_callback(py, &global.tool, &events.PY_RETURN, None)?;
338-
}
339-
if global.mask.contains(&events.PY_YIELD) {
340-
register_callback(py, &global.tool, &events.PY_YIELD, None)?;
341-
}
342-
if global.mask.contains(&events.PY_THROW) {
343-
register_callback(py, &global.tool, &events.PY_THROW, None)?;
344-
}
345-
if global.mask.contains(&events.PY_UNWIND) {
346-
register_callback(py, &global.tool, &events.PY_UNWIND, None)?;
347-
}
348-
if global.mask.contains(&events.RAISE) {
349-
register_callback(py, &global.tool, &events.RAISE, None)?;
350-
}
351-
if global.mask.contains(&events.RERAISE) {
352-
register_callback(py, &global.tool, &events.RERAISE, None)?;
353-
}
354-
if global.mask.contains(&events.EXCEPTION_HANDLED) {
355-
register_callback(py, &global.tool, &events.EXCEPTION_HANDLED, None)?;
356-
}
357-
// if global.mask.contains(&events.STOP_ITERATION) {
358-
// register_callback(py, &global.tool, &events.STOP_ITERATION, None)?;
359-
// }
360-
if global.mask.contains(&events.C_RETURN) {
361-
register_callback(py, &global.tool, &events.C_RETURN, None)?;
362-
}
363-
if global.mask.contains(&events.C_RAISE) {
364-
register_callback(py, &global.tool, &events.C_RAISE, None)?;
311+
let finish_result = global.tracer.finish(py);
312+
313+
let cleanup_result = (|| -> PyResult<()> {
314+
let events = monitoring_events(py)?;
315+
if global.mask.contains(&events.CALL) {
316+
register_callback(py, &global.tool, &events.CALL, None)?;
317+
}
318+
if global.mask.contains(&events.LINE) {
319+
register_callback(py, &global.tool, &events.LINE, None)?;
320+
}
321+
if global.mask.contains(&events.INSTRUCTION) {
322+
register_callback(py, &global.tool, &events.INSTRUCTION, None)?;
323+
}
324+
if global.mask.contains(&events.JUMP) {
325+
register_callback(py, &global.tool, &events.JUMP, None)?;
326+
}
327+
if global.mask.contains(&events.BRANCH) {
328+
register_callback(py, &global.tool, &events.BRANCH, None)?;
329+
}
330+
if global.mask.contains(&events.PY_START) {
331+
register_callback(py, &global.tool, &events.PY_START, None)?;
332+
}
333+
if global.mask.contains(&events.PY_RESUME) {
334+
register_callback(py, &global.tool, &events.PY_RESUME, None)?;
335+
}
336+
if global.mask.contains(&events.PY_RETURN) {
337+
register_callback(py, &global.tool, &events.PY_RETURN, None)?;
338+
}
339+
if global.mask.contains(&events.PY_YIELD) {
340+
register_callback(py, &global.tool, &events.PY_YIELD, None)?;
341+
}
342+
if global.mask.contains(&events.PY_THROW) {
343+
register_callback(py, &global.tool, &events.PY_THROW, None)?;
344+
}
345+
if global.mask.contains(&events.PY_UNWIND) {
346+
register_callback(py, &global.tool, &events.PY_UNWIND, None)?;
347+
}
348+
if global.mask.contains(&events.RAISE) {
349+
register_callback(py, &global.tool, &events.RAISE, None)?;
350+
}
351+
if global.mask.contains(&events.RERAISE) {
352+
register_callback(py, &global.tool, &events.RERAISE, None)?;
353+
}
354+
if global.mask.contains(&events.EXCEPTION_HANDLED) {
355+
register_callback(py, &global.tool, &events.EXCEPTION_HANDLED, None)?;
356+
}
357+
// if global.mask.contains(&events.STOP_ITERATION) {
358+
// register_callback(py, &global.tool, &events.STOP_ITERATION, None)?;
359+
// }
360+
if global.mask.contains(&events.C_RETURN) {
361+
register_callback(py, &global.tool, &events.C_RETURN, None)?;
362+
}
363+
if global.mask.contains(&events.C_RAISE) {
364+
register_callback(py, &global.tool, &events.C_RAISE, None)?;
365+
}
366+
367+
set_events(py, &global.tool, NO_EVENTS)?;
368+
free_tool_id(py, &global.tool)?;
369+
Ok(())
370+
})();
371+
372+
if let Err(err) = finish_result {
373+
if let Err(cleanup_err) = cleanup_result {
374+
warn!(
375+
"failed to reset monitoring callbacks after finish error: {}",
376+
cleanup_err
377+
);
378+
}
379+
return Err(err);
365380
}
366381

367-
set_events(py, &global.tool, NO_EVENTS)?;
368-
free_tool_id(py, &global.tool)?;
382+
cleanup_result?;
369383
}
370384
Ok(())
371385
}

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

Lines changed: 124 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,13 @@ use value_capture::{capture_call_arguments, record_return_value, record_visible_
1717
use std::collections::{hash_map::Entry, HashMap, HashSet};
1818
use std::fs;
1919
use std::path::{Path, PathBuf};
20+
use std::sync::atomic::{AtomicBool, Ordering};
21+
use std::sync::OnceLock;
2022

2123
use pyo3::prelude::*;
2224
use pyo3::types::PyAny;
2325

24-
use recorder_errors::{enverr, usage, ErrorCode, RecorderResult};
26+
use recorder_errors::{bug, enverr, usage, ErrorCode, RecorderResult};
2527
use runtime_tracing::NonStreamingTraceWriter;
2628
use runtime_tracing::{Line, TraceEventsFileFormat, TraceWriter};
2729

@@ -54,6 +56,71 @@ enum ShouldTrace {
5456
SkipAndDisable,
5557
}
5658

59+
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
60+
enum FailureStage {
61+
PyStart,
62+
Line,
63+
Finish,
64+
}
65+
66+
impl FailureStage {
67+
fn as_str(self) -> &'static str {
68+
match self {
69+
FailureStage::PyStart => "py_start",
70+
FailureStage::Line => "line",
71+
FailureStage::Finish => "finish",
72+
}
73+
}
74+
}
75+
76+
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
77+
enum FailureMode {
78+
Stage(FailureStage),
79+
SuppressEvents,
80+
}
81+
82+
static FAILURE_MODE: OnceLock<Option<FailureMode>> = OnceLock::new();
83+
static FAILURE_TRIGGERED: AtomicBool = AtomicBool::new(false);
84+
85+
fn configured_failure_mode() -> Option<FailureMode> {
86+
*FAILURE_MODE.get_or_init(|| {
87+
let raw = std::env::var("CODETRACER_TEST_INJECT_FAILURE").ok();
88+
if let Some(value) = raw.as_deref() {
89+
log::debug!("[RuntimeTracer] test failure injection mode: {}", value);
90+
}
91+
raw.and_then(|raw| match raw.trim().to_ascii_lowercase().as_str() {
92+
"py_start" | "py-start" => Some(FailureMode::Stage(FailureStage::PyStart)),
93+
"line" => Some(FailureMode::Stage(FailureStage::Line)),
94+
"finish" => Some(FailureMode::Stage(FailureStage::Finish)),
95+
"suppress-events" | "suppress_events" | "suppress" => Some(FailureMode::SuppressEvents),
96+
_ => None,
97+
})
98+
})
99+
}
100+
101+
fn should_inject_failure(stage: FailureStage) -> bool {
102+
match configured_failure_mode() {
103+
Some(FailureMode::Stage(mode)) if mode == stage => {
104+
!FAILURE_TRIGGERED.swap(true, Ordering::SeqCst)
105+
}
106+
_ => false,
107+
}
108+
}
109+
110+
fn suppress_events() -> bool {
111+
matches!(configured_failure_mode(), Some(FailureMode::SuppressEvents))
112+
}
113+
114+
fn injected_failure_err(stage: FailureStage) -> PyErr {
115+
let err = bug!(
116+
ErrorCode::TraceIncomplete,
117+
"test-injected failure at {}",
118+
stage.as_str()
119+
)
120+
.with_context("injection_stage", stage.as_str().to_string());
121+
ffi::map_recorder_error(err)
122+
}
123+
57124
fn is_real_filename(filename: &str) -> bool {
58125
let trimmed = filename.trim();
59126
!(trimmed.starts_with('<') && trimmed.ends_with('>'))
@@ -97,6 +164,10 @@ impl RuntimeTracer {
97164
}
98165

99166
fn mark_event(&mut self) {
167+
if suppress_events() {
168+
log::debug!("[RuntimeTracer] skipping event mark due to test injection");
169+
return;
170+
}
100171
self.events_recorded = true;
101172
}
102173

@@ -212,6 +283,10 @@ impl Tracer for RuntimeTracer {
212283
return Ok(CallbackOutcome::Continue);
213284
}
214285

286+
if should_inject_failure(FailureStage::PyStart) {
287+
return Err(injected_failure_err(FailureStage::PyStart));
288+
}
289+
215290
log_event(py, code, "on_py_start", None);
216291

217292
if let Ok(fid) = self.ensure_function_id(py, code) {
@@ -247,6 +322,10 @@ impl Tracer for RuntimeTracer {
247322
return Ok(CallbackOutcome::Continue);
248323
}
249324

325+
if should_inject_failure(FailureStage::Line) {
326+
return Err(injected_failure_err(FailureStage::Line));
327+
}
328+
250329
log_event(py, code, "on_line", Some(lineno));
251330

252331
if let Ok(filename) = code.filename(py) {
@@ -320,6 +399,11 @@ impl Tracer for RuntimeTracer {
320399
fn finish(&mut self, _py: Python<'_>) -> PyResult<()> {
321400
// Trace event entry
322401
log::debug!("[RuntimeTracer] finish");
402+
403+
if should_inject_failure(FailureStage::Finish) {
404+
return Err(injected_failure_err(FailureStage::Finish));
405+
}
406+
323407
let policy = policy_snapshot();
324408

325409
if self.encountered_failure {
@@ -337,8 +421,7 @@ impl Tracer for RuntimeTracer {
337421
);
338422
}
339423
} else {
340-
self
341-
.cleanup_partial_outputs()
424+
self.cleanup_partial_outputs()
342425
.map_err(ffi::map_recorder_error)?;
343426
}
344427
self.ignored_code_ids.clear();
@@ -393,8 +476,15 @@ mod tests {
393476
}
394477

395478
fn reset_policy(_py: Python<'_>) {
396-
policy::configure_policy_py(Some("abort"), Some(false), Some(false), None, None, Some(false))
397-
.expect("reset recorder policy");
479+
policy::configure_policy_py(
480+
Some("abort"),
481+
Some(false),
482+
Some(false),
483+
None,
484+
None,
485+
Some(false),
486+
)
487+
.expect("reset recorder policy");
398488
}
399489

400490
#[test]
@@ -1160,8 +1250,15 @@ snapshot()
11601250
#[test]
11611251
fn finish_enforces_require_trace_policy() {
11621252
Python::with_gil(|py| {
1163-
policy::configure_policy_py(Some("abort"), Some(true), Some(false), None, None, Some(false))
1164-
.expect("enable require_trace policy");
1253+
policy::configure_policy_py(
1254+
Some("abort"),
1255+
Some(true),
1256+
Some(false),
1257+
None,
1258+
None,
1259+
Some(false),
1260+
)
1261+
.expect("enable require_trace policy");
11651262

11661263
let script_dir = tempfile::tempdir().expect("script dir");
11671264
let program_path = script_dir.path().join("program.py");
@@ -1178,7 +1275,9 @@ snapshot()
11781275
);
11791276
tracer.begin(&outputs, 1).expect("begin tracer");
11801277

1181-
let err = tracer.finish(py).expect_err("finish should error when require_trace true");
1278+
let err = tracer
1279+
.finish(py)
1280+
.expect_err("finish should error when require_trace true");
11821281
let message = err.to_string();
11831282
assert!(
11841283
message.contains("ERR_TRACE_MISSING"),
@@ -1213,16 +1312,26 @@ snapshot()
12131312
tracer.finish(py).expect("finish after failure");
12141313

12151314
assert!(!outputs.events().exists(), "expected events file removed");
1216-
assert!(!outputs.metadata().exists(), "expected metadata file removed");
1315+
assert!(
1316+
!outputs.metadata().exists(),
1317+
"expected metadata file removed"
1318+
);
12171319
assert!(!outputs.paths().exists(), "expected paths file removed");
12181320
});
12191321
}
12201322

12211323
#[test]
12221324
fn finish_keeps_partial_outputs_when_policy_allows() {
12231325
Python::with_gil(|py| {
1224-
policy::configure_policy_py(Some("abort"), Some(false), Some(true), None, None, Some(false))
1225-
.expect("enable keep_partial policy");
1326+
policy::configure_policy_py(
1327+
Some("abort"),
1328+
Some(false),
1329+
Some(true),
1330+
None,
1331+
None,
1332+
Some(false),
1333+
)
1334+
.expect("enable keep_partial policy");
12261335

12271336
let script_dir = tempfile::tempdir().expect("script dir");
12281337
let program_path = script_dir.path().join("program.py");
@@ -1243,7 +1352,10 @@ snapshot()
12431352
tracer.finish(py).expect("finish after failure");
12441353

12451354
assert!(outputs.events().exists(), "expected events file retained");
1246-
assert!(outputs.metadata().exists(), "expected metadata file retained");
1355+
assert!(
1356+
outputs.metadata().exists(),
1357+
"expected metadata file retained"
1358+
);
12471359
assert!(outputs.paths().exists(), "expected paths file retained");
12481360

12491361
reset_policy(py);

0 commit comments

Comments
 (0)