reifydb_profiler/
visit.rs1use 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}