|
4 | 4 | //! signal when CPython should disable further monitoring for a location by propagating |
5 | 5 | //! the `sys.monitoring.DISABLE` sentinel. |
6 | 6 |
|
7 | | -use std::fs; |
8 | | -use std::path::Path; |
9 | | -use std::sync::atomic::{AtomicBool, Ordering}; |
10 | | -use std::sync::Once; |
11 | | - |
12 | | -use pyo3::exceptions::PyRuntimeError; |
13 | | -use pyo3::prelude::*; |
14 | 7 | pub mod code_object; |
| 8 | +mod logging; |
15 | 9 | mod runtime_tracer; |
| 10 | +mod session; |
16 | 11 | pub mod tracer; |
| 12 | + |
17 | 13 | pub use crate::code_object::{CodeObjectRegistry, CodeObjectWrapper}; |
| 14 | +pub use crate::session::{flush_tracing, is_tracing, start_tracing, stop_tracing}; |
18 | 15 | pub use crate::tracer::{ |
19 | 16 | install_tracer, uninstall_tracer, CallbackOutcome, CallbackResult, EventSet, Tracer, |
20 | 17 | }; |
21 | 18 |
|
22 | | -/// Global flag tracking whether tracing is active. |
23 | | -static ACTIVE: AtomicBool = AtomicBool::new(false); |
24 | | - |
25 | | -// Initialize Rust logging once per process. Defaults to debug for this crate |
26 | | -// unless overridden by RUST_LOG. This helps surface debug! output during dev. |
27 | | -static INIT_LOGGER: Once = Once::new(); |
28 | | - |
29 | | -fn init_rust_logging_with_default(default_filter: &str) { |
30 | | - INIT_LOGGER.call_once(|| { |
31 | | - let env = env_logger::Env::default().default_filter_or(default_filter); |
32 | | - // Use a compact format with timestamps and targets to aid debugging. |
33 | | - let mut builder = env_logger::Builder::from_env(env); |
34 | | - builder.format_timestamp_micros().format_target(true); |
35 | | - let _ = builder.try_init(); |
36 | | - }); |
37 | | -} |
38 | | - |
39 | | -/// Start tracing using sys.monitoring and runtime_tracing writer. |
40 | | -#[pyfunction] |
41 | | -fn start_tracing(path: &str, format: &str, activation_path: Option<&str>) -> PyResult<()> { |
42 | | - // Ensure logging is ready before any tracer logs might be emitted. |
43 | | - // Default only our crate to debug to avoid excessive verbosity from deps. |
44 | | - init_rust_logging_with_default("codetracer_python_recorder=debug"); |
45 | | - if ACTIVE.load(Ordering::SeqCst) { |
46 | | - return Err(PyRuntimeError::new_err("tracing already active")); |
47 | | - } |
48 | | - |
49 | | - // Interpret `path` as a directory where trace files will be written. |
50 | | - let out_dir = Path::new(path); |
51 | | - if out_dir.exists() && !out_dir.is_dir() { |
52 | | - return Err(PyRuntimeError::new_err( |
53 | | - "trace path exists and is not a directory", |
54 | | - )); |
55 | | - } |
56 | | - if !out_dir.exists() { |
57 | | - // Best-effort create the directory tree |
58 | | - fs::create_dir_all(&out_dir).map_err(|e| { |
59 | | - PyRuntimeError::new_err(format!("failed to create trace directory: {}", e)) |
60 | | - })?; |
61 | | - } |
62 | | - |
63 | | - // Map format string to enum |
64 | | - let fmt = match format.to_lowercase().as_str() { |
65 | | - "json" => runtime_tracing::TraceEventsFileFormat::Json, |
66 | | - // Use BinaryV0 for "binary" to avoid streaming writer here. |
67 | | - "binary" | "binaryv0" | "binary_v0" | "b0" => { |
68 | | - runtime_tracing::TraceEventsFileFormat::BinaryV0 |
69 | | - } |
70 | | - //TODO AI! We need to assert! that the format is among the known values. |
71 | | - other => { |
72 | | - eprintln!("Unknown format '{}', defaulting to binary (v0)", other); |
73 | | - runtime_tracing::TraceEventsFileFormat::BinaryV0 |
74 | | - } |
75 | | - }; |
76 | | - |
77 | | - // Build output file paths inside the directory. |
78 | | - let (events_path, meta_path, paths_path) = match fmt { |
79 | | - runtime_tracing::TraceEventsFileFormat::Json => ( |
80 | | - out_dir.join("trace.json"), |
81 | | - out_dir.join("trace_metadata.json"), |
82 | | - out_dir.join("trace_paths.json"), |
83 | | - ), |
84 | | - _ => ( |
85 | | - out_dir.join("trace.bin"), |
86 | | - out_dir.join("trace_metadata.json"), |
87 | | - out_dir.join("trace_paths.json"), |
88 | | - ), |
89 | | - }; |
90 | | - |
91 | | - // Activation path: when set, tracing starts only after entering it. |
92 | | - let activation_path = activation_path.map(|s| Path::new(s)); |
93 | | - |
94 | | - Python::with_gil(|py| { |
95 | | - // Program and args: keep minimal; Python-side API stores full session info if needed |
96 | | - let sys = py.import("sys")?; |
97 | | - let argv = sys.getattr("argv")?; |
98 | | - let program: String = argv.get_item(0)?.extract::<String>()?; |
99 | | - //TODO: Error-handling. What to do if argv is empty? Does this ever happen? |
100 | | - |
101 | | - let mut tracer = runtime_tracer::RuntimeTracer::new(&program, &[], fmt, activation_path); |
102 | | - |
103 | | - // Start location: prefer activation path, otherwise best-effort argv[0] |
104 | | - let start_path: &Path = activation_path.unwrap_or(Path::new(&program)); |
105 | | - log::debug!("{}", start_path.display()); |
106 | | - tracer.begin(&meta_path, &paths_path, &events_path, start_path, 1)?; |
107 | | - |
108 | | - // Install callbacks |
109 | | - install_tracer(py, Box::new(tracer))?; |
110 | | - ACTIVE.store(true, Ordering::SeqCst); |
111 | | - Ok(()) |
112 | | - }) |
113 | | -} |
114 | | - |
115 | | -/// Stop tracing by resetting the global flag. |
116 | | -#[pyfunction] |
117 | | -fn stop_tracing() -> PyResult<()> { |
118 | | - Python::with_gil(|py| { |
119 | | - // Uninstall triggers finish() on tracer implementation. |
120 | | - uninstall_tracer(py)?; |
121 | | - ACTIVE.store(false, Ordering::SeqCst); |
122 | | - Ok(()) |
123 | | - }) |
124 | | -} |
125 | | - |
126 | | -/// Query whether tracing is currently active. |
127 | | -#[pyfunction] |
128 | | -fn is_tracing() -> PyResult<bool> { |
129 | | - Ok(ACTIVE.load(Ordering::SeqCst)) |
130 | | -} |
131 | | - |
132 | | -/// Flush buffered trace data (best-effort, non-streaming formats only). |
133 | | -#[pyfunction] |
134 | | -fn flush_tracing() -> PyResult<()> { |
135 | | - Python::with_gil(|py| crate::tracer::flush_installed_tracer(py)) |
136 | | -} |
| 19 | +use pyo3::prelude::*; |
137 | 20 |
|
138 | 21 | /// Python module definition. |
139 | 22 | #[pymodule] |
140 | 23 | fn codetracer_python_recorder(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { |
141 | 24 | // Initialize logging on import so users see logs without extra setup. |
142 | 25 | // Respect RUST_LOG if present; otherwise default to debug for this crate. |
143 | | - init_rust_logging_with_default("codetracer_python_recorder=debug"); |
| 26 | + logging::init_rust_logging_with_default("codetracer_python_recorder=debug"); |
144 | 27 | m.add_function(wrap_pyfunction!(start_tracing, m)?)?; |
145 | 28 | m.add_function(wrap_pyfunction!(stop_tracing, m)?)?; |
146 | 29 | m.add_function(wrap_pyfunction!(is_tracing, m)?)?; |
|
0 commit comments