From 11d9ebe2b63918aff685d9609a4374124526d4b0 Mon Sep 17 00:00:00 2001 From: Mark Kalsbeek Date: Thu, 11 Dec 2025 08:03:49 +0100 Subject: [PATCH] fix recursion issue, better macros --- examples/bouncing_ball.rs | 31 +++---- teleprof-macros/Cargo.toml | 2 +- teleprof-macros/src/lib.rs | 99 +++++++++++++++++++++- teleprof/src/lib.rs | 165 +++++++++++++++++++++++-------------- 4 files changed, 220 insertions(+), 77 deletions(-) diff --git a/examples/bouncing_ball.rs b/examples/bouncing_ball.rs index fa04b85..7331f7b 100644 --- a/examples/bouncing_ball.rs +++ b/examples/bouncing_ball.rs @@ -4,7 +4,7 @@ use std::time::{Duration, Instant}; use std::sync::{Arc, Mutex}; use std::sync::atomic::{AtomicUsize, Ordering}; use rand::Rng; -use teleprof::instrument; +use teleprof::{instrument, instrument_calls}; const WIDTH: usize = 800; const HEIGHT: usize = 600; @@ -39,12 +39,17 @@ fn main() { // Name the main thread teleprof::set_thread_name("Main"); - println!("Bouncing Ball Demo"); + println!("Bouncing Ball Demo - Enhanced Instrumentation"); println!("The ball window should appear alongside the profiler"); println!("Press Space in profiler window to pause"); println!("Mouse wheel to zoom, drag to pan in profiler"); println!("Press Escape in either window to quit"); println!(); + println!("Features demonstrated:"); + println!("- Flame graph rendering (all depth-N spans on row N)"); + println!("- instrument_calls macro for automatic call instrumentation"); + println!("- Runtime deduplication preventing double-wrapping"); + println!(); let mut window = Window::new( "Bouncing Ball", @@ -74,7 +79,10 @@ fn main() { } } -#[instrument] +// Using instrument_calls - automatically instruments all function calls in the body +// This will show update_physics, render, print_status, thread::sleep, Arc::clone, etc. +// Pause handling is now automatic in SpanGuard! +#[instrument_calls] fn main_frame( ball: &Arc>, framebuffer: &mut [u32], @@ -82,16 +90,8 @@ fn main_frame( frame_count: &mut u32, frame_time: Duration, ) { - // Check if paused - if teleprof::PAUSE.try_lock().is_err() { - println!("Paused!"); - while teleprof::PAUSE.try_lock().is_err() { - thread::sleep(Duration::from_millis(100)); - } - println!("Resumed!"); - } - - // Update physics + // Update physics - this is also instrumented with #[instrument] + // Runtime dedup will prevent double-wrapping! let hit_wall = update_physics(ball, frame_time.as_secs_f32()); // If we hit a wall, spawn a thread to pick a new color @@ -104,7 +104,7 @@ fn main_frame( }); } - // Render + // Render - also instrumented, so dedup will handle it render(ball, framebuffer); // Update window @@ -162,7 +162,8 @@ fn pick_new_color(ball: Arc>) { println!(" → New color selected: RGB({}, {}, {})", r, g, b); } -#[instrument] +// Using instrument_calls here too to see the breakdown of rendering +#[instrument_calls] fn render(ball: &Arc>, framebuffer: &mut [u32]) { clear_background(framebuffer); draw_ball(ball, framebuffer); diff --git a/teleprof-macros/Cargo.toml b/teleprof-macros/Cargo.toml index 2ef149b..ce00968 100644 --- a/teleprof-macros/Cargo.toml +++ b/teleprof-macros/Cargo.toml @@ -7,6 +7,6 @@ edition = "2021" proc-macro = true [dependencies] -syn = { version = "2.0", features = ["full"] } +syn = { version = "2.0", features = ["full", "visit-mut"] } quote = "1.0" proc-macro2 = "1.0" \ No newline at end of file diff --git a/teleprof-macros/src/lib.rs b/teleprof-macros/src/lib.rs index 766cff7..6cd318a 100644 --- a/teleprof-macros/src/lib.rs +++ b/teleprof-macros/src/lib.rs @@ -1,6 +1,7 @@ use proc_macro::TokenStream; use quote::quote; -use syn::{parse_macro_input, ItemFn, Lit}; +use syn::{parse_macro_input, ItemFn, Lit, Expr, Stmt, ExprCall, ExprMethodCall}; +use syn::visit_mut::{self, VisitMut}; #[proc_macro_attribute] pub fn instrument(args: TokenStream, input: TokenStream) -> TokenStream { @@ -55,4 +56,100 @@ pub fn instrument(args: TokenStream, input: TokenStream) -> TokenStream { }; TokenStream::from(instrumented) +} + +#[proc_macro_attribute] +pub fn instrument_calls(_args: TokenStream, input: TokenStream) -> TokenStream { + let mut input_fn = parse_macro_input!(input as ItemFn); + + // Transform the function body to wrap all function calls + let mut visitor = CallInstrumenter; + visitor.visit_block_mut(&mut input_fn.block); + + TokenStream::from(quote! { #input_fn }) +} + +struct CallInstrumenter; + +impl VisitMut for CallInstrumenter { + fn visit_expr_mut(&mut self, expr: &mut Expr) { + // First, recurse into children + visit_mut::visit_expr_mut(self, expr); + + // Then wrap this expression if it's a function call + match expr { + Expr::Call(call) => { + *expr = wrap_call(call); + } + Expr::MethodCall(method_call) => { + *expr = wrap_method_call(method_call); + } + _ => {} + } + } + + fn visit_stmt_mut(&mut self, stmt: &mut Stmt) { + // Handle statements that contain expressions + match stmt { + Stmt::Expr(expr, _) => { + self.visit_expr_mut(expr); + } + Stmt::Local(local) => { + if let Some(init) = &mut local.init { + self.visit_expr_mut(&mut init.expr); + } + } + _ => { + // Use default visitor for other statement types + visit_mut::visit_stmt_mut(self, stmt); + } + } + } +} + +fn wrap_call(call: &ExprCall) -> Expr { + // Extract function name from the call expression + let func_name = extract_function_name(&call.func); + + let func = &call.func; + let args = &call.args; + let attrs = &call.attrs; + + syn::parse_quote! { + { + #(#attrs)* + let _guard = ::teleprof::SpanGuard::new(#func_name); + (#func)(#args) + } + } +} + +fn wrap_method_call(method_call: &ExprMethodCall) -> Expr { + let method_name = method_call.method.to_string(); + + let receiver = &method_call.receiver; + let method = &method_call.method; + let args = &method_call.args; + let turbofish = &method_call.turbofish; + let attrs = &method_call.attrs; + + syn::parse_quote! { + { + #(#attrs)* + let _guard = ::teleprof::SpanGuard::new(#method_name); + (#receiver).#method #turbofish(#args) + } + } +} + +fn extract_function_name(func: &Expr) -> String { + match func { + Expr::Path(path) => { + // Get the last segment of the path + path.path.segments.last() + .map(|seg| seg.ident.to_string()) + .unwrap_or_else(|| "unknown".to_string()) + } + _ => "unknown".to_string(), + } } \ No newline at end of file diff --git a/teleprof/src/lib.rs b/teleprof/src/lib.rs index b749761..ab7cf71 100644 --- a/teleprof/src/lib.rs +++ b/teleprof/src/lib.rs @@ -4,8 +4,8 @@ use std::cell::Cell; use std::sync::{Arc, Mutex}; use std::time::Instant; -// Re-export the macro -pub use teleprof_macros::instrument; +// Re-export the macros +pub use teleprof_macros::{instrument, instrument_calls}; // ============================================================================ // Public API @@ -33,17 +33,52 @@ pub fn set_thread_name(name: impl Into) { } /// Create a profiling span +#[allow(dead_code)] pub struct SpanGuard { span_id: u64, + previous_parent: Option, + previous_parent_name: Option<&'static str>, + name: &'static str, + is_duplicate: bool, } impl SpanGuard { pub fn new(name: &'static str) -> Self { - let span_id = next_span_id(); + // Handle pause mechanism BEFORE creating the span + // This way the pause time isn't counted in the span + if let Err(_guard) = PAUSE.try_lock() { + // We're paused - block until unpaused + // The _guard will be dropped when we acquire the lock, releasing it immediately + let _pause_lock = PAUSE.lock().unwrap(); + // Now we're unpaused, continue with span creation + } + let thread_id = std::thread::current().id(); let parent_id = PARENT_SPAN.with(|p| p.get()); + let parent_name = PARENT_SPAN_NAME.with(|p| p.get()); + + // Check if this is a duplicate span (parent has same name) + let is_duplicate = parent_name == Some(name); + + if is_duplicate { + // Don't create a new span, just pass through + return Self { + span_id: 0, // Dummy value, won't be used + previous_parent: parent_id, + previous_parent_name: parent_name, + name, + is_duplicate: true, + }; + } + + let span_id = next_span_id(); + + // Store the previous parent so we can restore it when this span ends + let previous_parent = parent_id; + let previous_parent_name = parent_name; PARENT_SPAN.with(|p| p.set(Some(span_id))); + PARENT_SPAN_NAME.with(|p| p.set(Some(name))); EVENT_SENDER.send(Event::SpanStart { span_id, @@ -53,19 +88,31 @@ impl SpanGuard { timestamp: Instant::now(), }).ok(); - Self { span_id } + Self { + span_id, + previous_parent, + previous_parent_name, + name, + is_duplicate: false, + } } } impl Drop for SpanGuard { fn drop(&mut self) { + if self.is_duplicate { + // Don't send end event or modify thread-local state + return; + } + EVENT_SENDER.send(Event::SpanEnd { span_id: self.span_id, timestamp: Instant::now(), }).ok(); - // Restore parent - PARENT_SPAN.with(|p| p.set(None)); + // Restore the previous parent (grandparent of this span) + PARENT_SPAN.with(|p| p.set(self.previous_parent)); + PARENT_SPAN_NAME.with(|p| p.set(self.previous_parent_name)); } } @@ -102,6 +149,7 @@ pub(crate) enum Event { thread_local! { static PARENT_SPAN: Cell> = Cell::new(None); + static PARENT_SPAN_NAME: Cell> = Cell::new(None); } static EVENT_SENDER: Lazy> = Lazy::new(|| { @@ -719,62 +767,84 @@ mod window { view: &ViewState, font: &fontdue::Font, ) { - // Build tree structure - let mut roots = Vec::new(); + // Build tree structure and compute depths let mut children: HashMap> = HashMap::new(); + let mut depths: HashMap = HashMap::new(); + let mut roots = Vec::new(); for span in spans { if let Some(parent) = span.parent_id { children.entry(parent).or_default().push(span); } else { - roots.push(*span); + roots.push(span.span_id); + depths.insert(span.span_id, 0); } } - // Sort roots by start time - roots.sort_by_key(|s| s.start); + // Compute depths for all spans using BFS + let mut queue: Vec = roots.clone(); + while let Some(span_id) = queue.pop() { + let current_depth = depths[&span_id]; + if let Some(child_list) = children.get(&span_id) { + for child in child_list { + depths.insert(child.span_id, current_depth + 1); + queue.push(child.span_id); + } + } + } - // Render each root and its children recursively + // Group spans by depth level + let mut spans_by_depth: HashMap> = HashMap::new(); + for span in spans { + if let Some(&depth) = depths.get(&span.span_id) { + spans_by_depth.entry(depth).or_default().push(*span); + } + } + + // Find max depth + let max_depth = spans_by_depth.keys().max().copied().unwrap_or(0); + + // Render each depth level as a row (flame graph style) let row_height = 24; - let mut y_offset = view.icicle_scroll_y as usize; + let scroll_offset = view.icicle_scroll_y as i32; - for root in roots { + for depth in 0..=max_depth { + let y = (depth * row_height) as i32 - scroll_offset; + // Only render if in view - if y_offset < height { - y_offset = render_icicle_span( - framebuffer, - width, - height, - root, - &children, - earliest, - y_offset, - row_height, - view, - font, - 0, // depth - ); - } else { - break; + if y + row_height as i32 > 0 && y < height as i32 { + if let Some(depth_spans) = spans_by_depth.get(&depth) { + for span in depth_spans { + render_icicle_span_at_depth( + framebuffer, + width, + height, + span, + earliest, + y.max(0) as usize, + row_height, + view, + font, + ); + } + } } } } - fn render_icicle_span( + fn render_icicle_span_at_depth( framebuffer: &mut [u32], width: usize, height: usize, span: &CompletedSpan, - children: &HashMap>, earliest: Instant, y: usize, row_height: usize, view: &ViewState, font: &fontdue::Font, - depth: usize, - ) -> usize { + ) { if y + row_height > height { - return y; + return; } let start_time = (span.start - earliest).as_secs_f64(); @@ -799,31 +869,6 @@ mod window { draw_text(framebuffer, width, x2.max(0) as usize + 4, y + 4, span.name, font, 14.0, TEXT_COLOR); } } - - // Render children at next depth level - let mut child_y = y + row_height; - if let Some(child_spans) = children.get(&span.span_id) { - for child in child_spans { - if child_y >= height { - break; - } - child_y = render_icicle_span( - framebuffer, - width, - height, - child, - children, - earliest, - child_y, - row_height, - view, - font, - depth + 1, - ); - } - } - - child_y } fn render_timeline(