zelos_trace/
source.rs

1use std::{collections::HashMap, sync::Arc};
2
3use anyhow::{Result, anyhow};
4use parking_lot::RwLock;
5use uuid::Uuid;
6use zelos_trace_types::{
7    Value,
8    ipc::{
9        IpcMessage, IpcMessageWithId, Sender, TraceEvent, TraceEventFieldMetadata,
10        TraceEventFieldNamedValues, TraceEventSchema, TraceSegmentEnd, TraceSegmentStart,
11    },
12};
13
14use crate::time::now_time_ns;
15
16/// TraceSourceEvent is a child of a TraceSource, which contains the schema for the event, as well as helpers for building and emitting new events for that schema.
17#[derive(Debug)]
18pub struct TraceSourceEvent {
19    id: Uuid,
20    source_name: String,
21    sender: Sender,
22    pub name: String,
23    pub schema: Vec<TraceEventFieldMetadata>,
24}
25
26impl TraceSourceEvent {
27    pub fn emit(&self, time_ns: i64, fields: impl Iterator<Item = (String, Value)>) -> Result<()> {
28        let evt = TraceEvent {
29            time_ns,
30            name: self.name.clone(),
31            fields: fields.collect(),
32        };
33
34        self.sender.send(IpcMessageWithId {
35            segment_id: self.id,
36            source_name: self.source_name.clone(),
37            msg: IpcMessage::TraceEvent(evt),
38        })?;
39
40        Ok(())
41    }
42
43    pub async fn emit_async(
44        &self,
45        time_ns: i64,
46        fields: impl Iterator<Item = (String, Value)>,
47    ) -> Result<()> {
48        let evt = TraceEvent {
49            time_ns,
50            name: self.name.clone(),
51            fields: fields.collect(),
52        };
53
54        self.sender
55            .send_async(IpcMessageWithId {
56                segment_id: self.id,
57                source_name: self.source_name.clone(),
58                msg: IpcMessage::TraceEvent(evt),
59            })
60            .await?;
61
62        Ok(())
63    }
64
65    pub fn build(&self) -> builder::EventBuilder<'_> {
66        builder::EventBuilder::new(self)
67    }
68}
69
70/// TraceSource is the main interface to emitting Zelos trace events. It provides convencience methods for building new
71/// event schemas and emitting events from them.
72#[derive(Debug)]
73pub struct TraceSource {
74    pub id: Uuid,
75    pub source_name: String,
76    sender: Sender,
77    events: RwLock<HashMap<String, Arc<TraceSourceEvent>>>,
78}
79
80impl TraceSource {
81    pub fn new(source_name: &str, sender: Sender) -> Self {
82        let id = Uuid::now_v7();
83        let src = TraceSource {
84            id,
85            source_name: source_name.to_string(),
86            sender,
87            events: RwLock::new(HashMap::new()),
88        };
89
90        tracing::debug!(?id, ?source_name, "TraceSource::new");
91
92        if let Err(e) = src.emit_start() {
93            tracing::error!("Error emitting trace segment start: {}", e);
94        }
95
96        src
97    }
98
99    fn emit(&self, msg: IpcMessage) -> Result<()> {
100        self.sender.send(IpcMessageWithId {
101            segment_id: self.id,
102            source_name: self.source_name.clone(),
103            msg,
104        })?;
105
106        Ok(())
107    }
108
109    async fn emit_async(&self, msg: IpcMessage) -> Result<()> {
110        self.sender
111            .send_async(IpcMessageWithId {
112                segment_id: self.id,
113                source_name: self.source_name.clone(),
114                msg,
115            })
116            .await?;
117
118        Ok(())
119    }
120
121    pub fn emit_start(&self) -> Result<()> {
122        self.emit(IpcMessage::TraceSegmentStart(TraceSegmentStart {
123            time_ns: now_time_ns(),
124            source_name: self.source_name.clone(),
125        }))
126    }
127
128    pub fn emit_end(&self) -> Result<()> {
129        self.emit(IpcMessage::TraceSegmentEnd(TraceSegmentEnd {
130            time_ns: now_time_ns(),
131        }))
132    }
133
134    pub fn add_value_table(
135        &self,
136        name: &str,
137        field_name: &str,
138        values: impl Iterator<Item = (Value, String)>,
139    ) -> Result<()> {
140        self.emit(IpcMessage::TraceEventFieldNamedValues(
141            TraceEventFieldNamedValues {
142                event_name: name.to_string(),
143                field_name: field_name.to_string(),
144                values: values.collect(),
145            },
146        ))
147    }
148
149    pub fn add_event(
150        &self,
151        name: &str,
152        schema: impl Iterator<Item = TraceEventFieldMetadata>,
153    ) -> Result<Arc<TraceSourceEvent>> {
154        if self.events.read().contains_key(name) {
155            return Err(anyhow!("Event={} already exists", name));
156        }
157
158        let msg = Arc::new(TraceSourceEvent {
159            id: self.id,
160            source_name: self.source_name.clone(),
161            sender: self.sender.clone(),
162            name: name.to_string(),
163            schema: schema.collect(),
164        });
165
166        // Emit the event to the router
167        self.emit(IpcMessage::TraceEventSchema(TraceEventSchema {
168            name: name.to_string(),
169            fields: msg.schema.clone(),
170        }))?;
171
172        // Insert the event into our metadata store
173        self.events.write().insert(name.to_string(), msg.clone());
174
175        Ok(msg)
176    }
177
178    pub async fn add_event_async(
179        &self,
180        name: &str,
181        schema: impl Iterator<Item = TraceEventFieldMetadata>,
182    ) -> Result<Arc<TraceSourceEvent>> {
183        if self.events.read().contains_key(name) {
184            return Err(anyhow!("Event={} already exists", name));
185        }
186
187        let msg = Arc::new(TraceSourceEvent {
188            id: self.id,
189            source_name: self.source_name.clone(),
190            sender: self.sender.clone(),
191            name: name.to_string(),
192            schema: schema.collect(),
193        });
194
195        // Emit the event to the router
196        self.emit_async(IpcMessage::TraceEventSchema(TraceEventSchema {
197            name: name.to_string(),
198            fields: msg.schema.clone(),
199        }))
200        .await?;
201
202        // Insert the event into our metadata store
203        self.events.write().insert(name.to_string(), msg.clone());
204
205        Ok(msg)
206    }
207
208    pub fn get_event(&self, name: &str) -> Result<Arc<TraceSourceEvent>> {
209        self.events
210            .read()
211            .get(name)
212            .cloned()
213            .ok_or_else(|| anyhow!("Event not found"))
214    }
215
216    pub fn build_event<'a>(&'a self, name: &'a str) -> builder::TraceSourceEventBuilder<'a> {
217        builder::TraceSourceEventBuilder::new(self, name)
218    }
219}
220
221impl Drop for TraceSource {
222    fn drop(&mut self) {
223        if let Err(e) = self.emit_end() {
224            tracing::debug!("Error emitting trace segment end: {}", e);
225        }
226    }
227}
228
229pub mod builder {
230    use zelos_trace_types::DataType;
231
232    use super::*;
233
234    #[must_use]
235    pub struct EventBuilder<'a> {
236        parent: &'a TraceSourceEvent,
237        data: HashMap<String, Value>,
238    }
239
240    impl<'a> EventBuilder<'a> {
241        pub(crate) fn new(parent: &'a TraceSourceEvent) -> Self {
242            EventBuilder {
243                parent,
244                data: HashMap::new(),
245            }
246        }
247
248        /// Emit the event at the current time.
249        pub fn emit(mut self) -> Result<()> {
250            self.parent.emit(now_time_ns(), self.data.drain())
251        }
252
253        /// Emit the event at a specific time.
254        pub fn emit_at(mut self, time_ns: i64) -> Result<()> {
255            self.parent.emit(time_ns, self.data.drain())
256        }
257
258        /// Emit the event at the current time via async
259        pub async fn emit_async(mut self) -> Result<()> {
260            self.parent
261                .emit_async(now_time_ns(), self.data.drain())
262                .await
263        }
264
265        /// Emit the event at a specific time via async
266        pub async fn emit_at_async(mut self, time_ns: i64) -> Result<()> {
267            self.parent.emit_async(time_ns, self.data.drain()).await
268        }
269
270        /// Attempt to insert a value into the event, returning an error if the field is not found or the type does not match.
271        pub fn try_insert(&mut self, name: &str, value: Value) -> Result<()> {
272            // Find the field in the schema
273            let field = self
274                .parent
275                .schema
276                .iter()
277                .find(|field| field.name == name)
278                .ok_or_else(|| anyhow!("Field '{}' not found in schema", name))?;
279
280            // Check if our value matches the field type
281            if field.data_type != value.data_type() {
282                return Err(anyhow!(
283                    "Type mismatch for field '{}': expected {:?}, found {:?}",
284                    name,
285                    field.data_type,
286                    value.data_type()
287                ));
288            }
289
290            // Insert the value into our event
291            self.data.insert(name.to_string(), value);
292            Ok(())
293        }
294
295        pub fn try_insert_i8(mut self, name: &str, value: i8) -> Result<Self> {
296            self.try_insert(name, Value::Int8(value))?;
297            Ok(self)
298        }
299
300        pub fn try_insert_i16(mut self, name: &str, value: i16) -> Result<Self> {
301            self.try_insert(name, Value::Int16(value))?;
302            Ok(self)
303        }
304
305        pub fn try_insert_i32(mut self, name: &str, value: i32) -> Result<Self> {
306            self.try_insert(name, Value::Int32(value))?;
307            Ok(self)
308        }
309
310        pub fn try_insert_i64(mut self, name: &str, value: i64) -> Result<Self> {
311            self.try_insert(name, Value::Int64(value))?;
312            Ok(self)
313        }
314
315        pub fn try_insert_u8(mut self, name: &str, value: u8) -> Result<Self> {
316            self.try_insert(name, Value::UInt8(value))?;
317            Ok(self)
318        }
319
320        pub fn try_insert_u16(mut self, name: &str, value: u16) -> Result<Self> {
321            self.try_insert(name, Value::UInt16(value))?;
322            Ok(self)
323        }
324
325        pub fn try_insert_u32(mut self, name: &str, value: u32) -> Result<Self> {
326            self.try_insert(name, Value::UInt32(value))?;
327            Ok(self)
328        }
329
330        pub fn try_insert_u64(mut self, name: &str, value: u64) -> Result<Self> {
331            self.try_insert(name, Value::UInt64(value))?;
332            Ok(self)
333        }
334
335        pub fn try_insert_f32(mut self, name: &str, value: f32) -> Result<Self> {
336            self.try_insert(name, Value::Float32(value))?;
337            Ok(self)
338        }
339
340        pub fn try_insert_f64(mut self, name: &str, value: f64) -> Result<Self> {
341            self.try_insert(name, Value::Float64(value))?;
342            Ok(self)
343        }
344
345        pub fn try_insert_timestamp_ns(mut self, name: &str, value: i64) -> Result<Self> {
346            self.try_insert(name, Value::TimestampNs(value))?;
347            Ok(self)
348        }
349
350        pub fn try_insert_binary(mut self, name: &str, value: Vec<u8>) -> Result<Self> {
351            self.try_insert(name, Value::Binary(value))?;
352            Ok(self)
353        }
354
355        pub fn try_insert_string(mut self, name: &str, value: String) -> Result<Self> {
356            self.try_insert(name, Value::String(value))?;
357            Ok(self)
358        }
359
360        pub fn try_insert_bool(mut self, name: &str, value: bool) -> Result<Self> {
361            self.try_insert(name, Value::Boolean(value))?;
362            Ok(self)
363        }
364    }
365
366    #[must_use]
367    pub struct TraceSourceEventBuilder<'a> {
368        source: &'a TraceSource,
369        name: &'a str,
370        schema: HashMap<String, TraceEventFieldMetadata>,
371    }
372
373    impl<'a> TraceSourceEventBuilder<'a> {
374        pub(crate) fn new(source: &'a TraceSource, name: &'a str) -> Self {
375            TraceSourceEventBuilder {
376                source,
377                name,
378                schema: HashMap::new(),
379            }
380        }
381
382        /// Build the event and add it to the source.
383        pub fn build(self) -> Result<Arc<TraceSourceEvent>> {
384            self.source.add_event(self.name, self.schema.into_values())
385        }
386
387        /// Build the event and add it to the source via async.
388        pub async fn build_async(self) -> Result<Arc<TraceSourceEvent>> {
389            self.source
390                .add_event_async(self.name, self.schema.into_values())
391                .await
392        }
393
394        pub fn add_field(mut self, name: &str, data_type: DataType, unit: Option<String>) -> Self {
395            self.schema.insert(
396                name.to_string(),
397                TraceEventFieldMetadata {
398                    name: name.to_string(),
399                    data_type,
400                    unit,
401                },
402            );
403            self
404        }
405
406        pub fn add_i8_field(self, name: &str, unit: Option<String>) -> Self {
407            self.add_field(name, DataType::Int8, unit)
408        }
409
410        pub fn add_i16_field(self, name: &str, unit: Option<String>) -> Self {
411            self.add_field(name, DataType::Int16, unit)
412        }
413
414        pub fn add_i32_field(self, name: &str, unit: Option<String>) -> Self {
415            self.add_field(name, DataType::Int32, unit)
416        }
417
418        pub fn add_i64_field(self, name: &str, unit: Option<String>) -> Self {
419            self.add_field(name, DataType::Int64, unit)
420        }
421
422        pub fn add_u8_field(self, name: &str, unit: Option<String>) -> Self {
423            self.add_field(name, DataType::UInt8, unit)
424        }
425
426        pub fn add_u16_field(self, name: &str, unit: Option<String>) -> Self {
427            self.add_field(name, DataType::UInt16, unit)
428        }
429
430        pub fn add_u32_field(self, name: &str, unit: Option<String>) -> Self {
431            self.add_field(name, DataType::UInt32, unit)
432        }
433
434        pub fn add_u64_field(self, name: &str, unit: Option<String>) -> Self {
435            self.add_field(name, DataType::UInt64, unit)
436        }
437
438        pub fn add_f32_field(self, name: &str, unit: Option<String>) -> Self {
439            self.add_field(name, DataType::Float32, unit)
440        }
441
442        pub fn add_f64_field(self, name: &str, unit: Option<String>) -> Self {
443            self.add_field(name, DataType::Float64, unit)
444        }
445
446        pub fn add_timestamp_ns_field(self, name: &str, unit: Option<String>) -> Self {
447            self.add_field(name, DataType::TimestampNs, unit)
448        }
449
450        pub fn add_binary_field(self, name: &str, unit: Option<String>) -> Self {
451            self.add_field(name, DataType::Binary, unit)
452        }
453
454        pub fn add_string_field(self, name: &str, unit: Option<String>) -> Self {
455            self.add_field(name, DataType::String, unit)
456        }
457
458        pub fn add_bool_field(self, name: &str, unit: Option<String>) -> Self {
459            self.add_field(name, DataType::Boolean, unit)
460        }
461    }
462}
463
464#[cfg(test)]
465mod test {
466    use zelos_trace_types::DataType;
467
468    use super::*;
469
470    #[test]
471    fn test_source_basic() -> Result<()> {
472        // Create our flume channel for messages
473        let (sender, receiver) = flume::unbounded::<IpcMessageWithId>();
474        // Create our source
475        let src = TraceSource::new("src", sender.clone());
476        let id = src.id.clone();
477
478        // Check we get a start event on creation
479        {
480            let m = receiver.recv()?;
481            assert_eq!(m.segment_id, id);
482            assert!(matches!(m.msg, IpcMessage::TraceSegmentStart(_)));
483        }
484
485        // Create our event
486        let evt = src
487            .build_event("hello")
488            .add_i32_field("sig", None)
489            .build()?;
490
491        // Check specific contents of the TraceEventSchema
492        {
493            let m = receiver.recv()?;
494            assert_eq!(m.segment_id, id);
495            if let IpcMessage::TraceEventSchema(schema) = &m.msg {
496                assert_eq!(schema.name, "hello");
497                assert_eq!(schema.fields.len(), 1);
498                assert_eq!(schema.fields[0].name, "sig");
499                assert_eq!(schema.fields[0].data_type, DataType::Int32);
500            } else {
501                panic!("Expected TraceEventSchema");
502            }
503        }
504
505        // Use our event handle to insert
506        evt.build().try_insert_i32("sig", 10)?.emit()?;
507
508        // Check specific contents of the first TraceEvent
509        {
510            let m = receiver.recv()?;
511            assert_eq!(m.segment_id, id);
512            if let IpcMessage::TraceEvent(event) = &m.msg {
513                assert_eq!(event.name, "hello");
514                assert_eq!(event.fields.len(), 1);
515                let data = event.fields.iter().collect::<Vec<_>>();
516                assert_eq!(data[0].0, "sig");
517                assert_eq!(*data[0].1, Value::Int32(10));
518            } else {
519                panic!("Expected TraceEvent");
520            }
521        }
522
523        // Get the message and insert using that
524        src.get_event("hello")?
525            .build()
526            .try_insert_i32("sig", 20)?
527            .emit()?;
528
529        // Check specific contents of the second TraceEvent
530        {
531            let m = receiver.recv()?;
532            assert_eq!(m.segment_id, id);
533            if let IpcMessage::TraceEvent(event) = &m.msg {
534                assert_eq!(event.name, "hello");
535                assert_eq!(event.fields.len(), 1);
536                let data = event.fields.iter().collect::<Vec<_>>();
537                assert_eq!(data[0].0, "sig");
538                assert_eq!(*data[0].1, Value::Int32(20));
539            } else {
540                panic!("Expected TraceEvent");
541            }
542        }
543
544        // Drop the source to trigger the end event
545        drop(src);
546
547        // Check we get an end event on drop
548        {
549            let m = receiver.recv()?;
550            assert_eq!(m.segment_id, id);
551            assert!(matches!(m.msg, IpcMessage::TraceSegmentEnd(_)));
552        }
553
554        Ok(())
555    }
556}
zelos_trace/source.rs

zelos_trace/
source.rs
