Markk116/teleprof
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix recursion issue, better macros

Browse Source
This commit is contained in:
Mark Kalsbeek
2025-12-11 08:03:49 +01:00
parent 06b8cac896
commit 11d9ebe2b6
4 changed files with 220 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
examples/bouncing_ball.rs
View File

@@ -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<Mutex<Ball>>,
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<Mutex<Ball>>) {
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<Mutex<Ball>>, framebuffer: &mut [u32]) {
clear_background(framebuffer);
draw_ball(ball, framebuffer);

2
teleprof-macros/Cargo.toml
View File

@@ -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"

99
teleprof-macros/src/lib.rs
View File

@@ -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(),
}
}

165
teleprof/src/lib.rs
View File

@@ -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<String>) {
}
/// Create a profiling span
#[allow(dead_code)]
pub struct SpanGuard {
span_id: u64,
previous_parent: Option<u64>,
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<Option<u64>> = Cell::new(None);
static PARENT_SPAN_NAME: Cell<Option<&'static str>> = Cell::new(None);
}
static EVENT_SENDER: Lazy<Sender<Event>> = 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<u64, Vec<&CompletedSpan>> = HashMap::new();
let mut depths: HashMap<u64, usize> = 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<u64> = 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<usize, Vec<&CompletedSpan>> = 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<u64, Vec<&CompletedSpan>>,
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(