Skip to main content

reifydb_profiler/
visit.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::fmt;
5
6use tracing::field::{Field, Visit};
7
8#[derive(Default, Clone, Debug)]
9pub struct FlowApplyFields {
10	pub node_id: String,
11	pub node_type: String,
12	pub input_rows: u64,
13	pub output_rows: u64,
14	pub apply_time_us: u64,
15	pub lock_wait_us: u64,
16}
17
18impl Visit for FlowApplyFields {
19	fn record_u64(&mut self, field: &Field, value: u64) {
20		match field.name() {
21			"input_rows" => self.input_rows = value,
22			"output_rows" => self.output_rows = value,
23			"apply_time_us" => self.apply_time_us = value,
24			"lock_wait_us" => self.lock_wait_us = value,
25			_ => {}
26		}
27	}
28
29	fn record_i64(&mut self, field: &Field, value: i64) {
30		if value >= 0 {
31			self.record_u64(field, value as u64);
32		}
33	}
34
35	fn record_str(&mut self, field: &Field, value: &str) {
36		match field.name() {
37			"node_id" => self.node_id.replace_range(.., value),
38			"node_type" => self.node_type.replace_range(.., value),
39			_ => {}
40		}
41	}
42
43	fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
44		match field.name() {
45			"node_id" => {
46				self.node_id.clear();
47				self.node_id.push_str(&format!("{:?}", value));
48			}
49			"node_type" => {
50				self.node_type.clear();
51				self.node_type.push_str(format!("{:?}", value).trim_matches('"'));
52			}
53			_ => {}
54		}
55	}
56}
57
58#[cfg(test)]
59mod tests {
60	use std::sync::Arc;
61
62	use reifydb_runtime::sync::mutex::Mutex;
63	use tracing::{
64		Subscriber, debug_span,
65		span::{Attributes, Id},
66		subscriber::with_default,
67	};
68	use tracing_subscriber::{
69		Layer, Registry,
70		layer::{Context, SubscriberExt},
71		registry::LookupSpan,
72	};
73
74	use super::*;
75
76	struct CaptureLayer {
77		captured: Arc<Mutex<Option<FlowApplyFields>>>,
78	}
79
80	impl<S> Layer<S> for CaptureLayer
81	where
82		S: Subscriber + for<'a> LookupSpan<'a>,
83	{
84		fn on_new_span(&self, attrs: &Attributes<'_>, _id: &Id, _ctx: Context<'_, S>) {
85			let mut v = FlowApplyFields::default();
86			attrs.record(&mut v);
87			*self.captured.lock() = Some(v);
88		}
89	}
90
91	#[test]
92	fn extracts_flow_apply_fields() {
93		let captured = Arc::new(Mutex::new(None));
94		let layer = CaptureLayer {
95			captured: captured.clone(),
96		};
97		let subscriber = Registry::default().with(layer);
98		with_default(subscriber, || {
99			let _span = debug_span!(
100				"flow::engine::apply",
101				node_id = "n1",
102				node_type = "map",
103				input_rows = 10u64,
104				output_rows = 7u64,
105				apply_time_us = 250u64,
106				lock_wait_us = 5u64,
107			);
108		});
109		let captured = captured.lock().clone().unwrap();
110		assert_eq!(captured.node_id, "n1");
111		assert_eq!(captured.node_type, "map");
112		assert_eq!(captured.input_rows, 10);
113		assert_eq!(captured.output_rows, 7);
114		assert_eq!(captured.apply_time_us, 250);
115		assert_eq!(captured.lock_wait_us, 5);
116	}
117}