hyperactor_telemetry/
trace_dispatcher.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9//! Unified telemetry layer that captures trace events once and fans out to multiple exporters
10//! on a background thread, eliminating redundant capture and moving work off the application
11//! thread.
12
13use std::cell::Cell;
14use std::sync::Arc;
15use std::sync::atomic::AtomicU64;
16use std::sync::atomic::Ordering;
17use std::sync::mpsc;
18use std::thread::JoinHandle;
19use std::time::Duration;
20use std::time::SystemTime;
21
22use smallvec::SmallVec;
23use tracing::Id;
24use tracing::Subscriber;
25use tracing::level_filters::LevelFilter;
26use tracing_subscriber::filter::Targets;
27use tracing_subscriber::layer::Context;
28use tracing_subscriber::layer::Layer;
29use tracing_subscriber::registry::LookupSpan;
30
31use crate::EntityEvent;
32
33const QUEUE_CAPACITY: usize = 100_000;
34
35/// Type alias for trace event fields
36/// We expect that most trace events have fewer than 4 fields.
37pub(crate) type TraceFields = SmallVec<[(&'static str, FieldValue); 4]>;
38
39#[inline]
40pub(crate) fn get_field<'a>(fields: &'a TraceFields, key: &str) -> Option<&'a FieldValue> {
41    fields.iter().find(|(k, _)| *k == key).map(|(_, v)| v)
42}
43
44/// Unified representation of a trace event captured from the tracing layer.
45/// This is captured once on the application thread, then sent to the background
46/// worker for fan-out to multiple exporters.
47#[derive(Debug, Clone)]
48pub enum TraceEvent {
49    /// A new span was created (on_new_span)
50    NewSpan {
51        id: u64,
52        name: &'static str,
53        target: &'static str,
54        level: tracing::Level,
55        fields: TraceFields,
56        timestamp: SystemTime,
57        parent_id: Option<u64>,
58        thread_name: &'static str,
59        file: Option<&'static str>,
60        line: Option<u32>,
61    },
62    /// A span was entered (on_enter)
63    SpanEnter {
64        id: u64,
65        timestamp: SystemTime,
66        thread_name: &'static str,
67    },
68    /// A span was exited (on_exit)
69    SpanExit {
70        id: u64,
71        timestamp: SystemTime,
72        thread_name: &'static str,
73    },
74    /// A span was closed (dropped)
75    SpanClose { id: u64, timestamp: SystemTime },
76    /// A tracing event occurred (e.g., tracing::info!())
77    Event {
78        name: &'static str,
79        target: &'static str,
80        level: tracing::Level,
81        fields: TraceFields,
82        timestamp: SystemTime,
83        parent_span: Option<u64>,
84        thread_id: &'static str,
85        thread_name: &'static str,
86        module_path: Option<&'static str>,
87        file: Option<&'static str>,
88        line: Option<u32>,
89    },
90    /// An entity lifecycle event emitted outside the tracing subscriber.
91    Entity(EntityEvent),
92}
93
94/// Simplified field value representation for trace events
95#[derive(Debug, Clone)]
96pub enum FieldValue {
97    Bool(bool),
98    I64(i64),
99    U64(u64),
100    F64(f64),
101    Str(String),
102    Debug(String),
103}
104
105/// Trait for sinks that receive trace events from the dispatcher.
106/// Implementations run on the background worker thread and can perform
107/// expensive I/O operations without blocking the application.
108pub trait TraceEventSink: Send + 'static {
109    /// Consume a single event. Called on background thread.
110    fn consume(&mut self, event: &TraceEvent) -> Result<(), anyhow::Error>;
111
112    /// Optional target/level filter for this sink.
113    ///
114    /// The worker loop automatically applies this filter before calling `consume()`,
115    /// so sinks don't need to check target/level in their consume implementation.
116    /// Only `NewSpan` and `Event` are filtered by target/level; other event types
117    /// are always passed through.
118    ///
119    /// # Returns
120    /// - `None` - No filtering, all events are consumed (default)
121    /// - `Some(Targets)` - Only consume events matching the target/level filter
122    ///
123    /// # Example
124    /// ```ignore
125    /// fn target_filter(&self) -> Option<&Targets> {
126    ///     Some(Targets::new()
127    ///         .with_target("opentelemetry", LevelFilter::OFF)
128    ///         .with_default(LevelFilter::DEBUG))
129    /// }
130    /// ```
131    fn target_filter(&self) -> Option<&Targets> {
132        None
133    }
134
135    /// Flush any buffered events to the backend.
136    /// Called periodically and on shutdown.
137    fn flush(&mut self) -> Result<(), anyhow::Error>;
138
139    /// Optional: return name for debugging/logging
140    fn name(&self) -> &str {
141        std::any::type_name::<Self>()
142    }
143}
144
145thread_local! {
146    /// Cached thread info (thread_name, thread_id) for minimal overhead.
147    /// Strings are leaked once per thread to get &'static str - threads are long-lived so this is fine.
148    /// Uses Cell since (&'static str, &'static str) is Copy.
149    static CACHED_THREAD_INFO: Cell<Option<(&'static str, &'static str)>> = const { Cell::new(None) };
150}
151
152#[inline(always)]
153fn get_thread_info() -> (&'static str, &'static str) {
154    CACHED_THREAD_INFO.with(|cache| {
155        if let Some(info) = cache.get() {
156            return info;
157        }
158
159        let thread_name: &'static str = Box::leak(
160            std::thread::current()
161                .name()
162                .unwrap_or("")
163                .to_string()
164                .into_boxed_str(),
165        );
166
167        #[cfg(target_os = "linux")]
168        let thread_id: &'static str = {
169            // SAFETY: syscall(SYS_gettid) is always safe to call - it's a read-only
170            // syscall that returns the current thread's kernel thread ID (TID).
171            // The cast to u64 is safe because gettid() returns a positive pid_t.
172            let tid = unsafe { libc::syscall(libc::SYS_gettid) as u64 };
173            Box::leak(tid.to_string().into_boxed_str())
174        };
175        #[cfg(not(target_os = "linux"))]
176        let thread_id: &'static str = {
177            let tid_num = std::thread::current().id().as_u64().get();
178            Box::leak(tid_num.to_string().into_boxed_str())
179        };
180
181        cache.set(Some((thread_name, thread_id)));
182        (thread_name, thread_id)
183    })
184}
185
186/// Control messages for the dispatcher (e.g., adding sinks dynamically)
187pub enum DispatcherControl {
188    /// Add a new sink to receive events
189    AddSink(Box<dyn TraceEventSink>),
190}
191
192/// The trace event dispatcher that captures events once and dispatches to multiple sinks
193/// on a background thread.
194pub struct TraceEventDispatcher {
195    sender: Option<mpsc::SyncSender<TraceEvent>>,
196    /// Separate channel so we are always notified of when the main queue is full and events are being dropped.
197    dropped_sender: Option<mpsc::Sender<TraceEvent>>,
198    _worker_handle: WorkerHandle,
199    max_level: Option<LevelFilter>,
200    dropped_events: Arc<AtomicU64>,
201}
202
203struct WorkerHandle {
204    join_handle: Option<JoinHandle<()>>,
205}
206
207thread_local! {
208    static IN_SEND: Cell<bool> = const { Cell::new(false) };
209}
210
211struct InSendGuard;
212
213impl Drop for InSendGuard {
214    fn drop(&mut self) {
215        IN_SEND.with(|f| f.set(false));
216    }
217}
218
219impl TraceEventDispatcher {
220    /// Create a new trace event dispatcher with the given sinks.
221    /// Uses a bounded channel (capacity QUEUE_CAPACITY) to ensure telemetry never blocks
222    /// the application. Events are dropped with a warning if the queue is full.
223    /// A separate unbounded priority channel guarantees delivery of critical events
224    /// like drop notifications (safe because drop events are rate-limited).
225    ///
226    /// Takes the global control receiver for dynamic sink registration. Sinks registered
227    /// via `register_sink()` before or after this call will be added to the dispatcher.
228    ///
229    /// # Arguments
230    /// * `sinks` - List of sinks to dispatch events to.
231    pub(crate) fn new(sinks: Vec<Box<dyn TraceEventSink>>) -> Self {
232        let max_level = Self::derive_max_level(&sinks);
233
234        let (sender, receiver) = mpsc::sync_channel(QUEUE_CAPACITY);
235        let (dropped_sender, dropped_receiver) = mpsc::channel();
236        // Take the global control receiver - sinks registered via register_sink() will be received here
237        let control_receiver = crate::take_sink_control_receiver();
238        let dropped_events = Arc::new(AtomicU64::new(0));
239        let dropped_events_worker = Arc::clone(&dropped_events);
240
241        let worker_handle = std::thread::Builder::new()
242            .name("telemetry-worker".into())
243            .spawn(move || {
244                worker_loop(
245                    receiver,
246                    dropped_receiver,
247                    control_receiver,
248                    sinks,
249                    dropped_events_worker,
250                );
251            })
252            .expect("failed to spawn telemetry worker thread");
253
254        Self {
255            sender: Some(sender),
256            dropped_sender: Some(dropped_sender),
257            _worker_handle: WorkerHandle {
258                join_handle: Some(worker_handle),
259            },
260            max_level,
261            dropped_events,
262        }
263    }
264
265    fn derive_max_level(sinks: &[Box<dyn TraceEventSink>]) -> Option<LevelFilter> {
266        let mut max_level: Option<LevelFilter> = None;
267
268        for sink in sinks {
269            let sink_max = match sink.target_filter() {
270                None => LevelFilter::TRACE,
271                Some(targets) => {
272                    let levels = [
273                        (tracing::Level::TRACE, LevelFilter::TRACE),
274                        (tracing::Level::DEBUG, LevelFilter::DEBUG),
275                        (tracing::Level::INFO, LevelFilter::INFO),
276                        (tracing::Level::WARN, LevelFilter::WARN),
277                        (tracing::Level::ERROR, LevelFilter::ERROR),
278                    ];
279                    let mut result = LevelFilter::OFF;
280                    for (level, filter) in levels {
281                        if targets.would_enable("", &level) {
282                            result = filter;
283                            break;
284                        }
285                    }
286                    result
287                }
288            };
289
290            max_level = Some(match max_level {
291                None => sink_max,
292                Some(current) => std::cmp::max(current, sink_max),
293            });
294        }
295
296        max_level
297    }
298
299    fn send_event(&self, event: TraceEvent) {
300        // Re-entrancy guard. A `Layer` callback may emit `tracing` events
301        // through code it touches—notably std's mpmc channel, which is itself
302        // instrumented—and those events loop back through this subscriber.
303        // Without this guard, the recursion exhausts the stack and SIGSEGVs.
304        if IN_SEND.with(|f| f.replace(true)) {
305            return;
306        }
307        let _reset = InSendGuard;
308
309        if let Some(sender) = &self.sender
310            && let Err(mpsc::TrySendError::Full(_)) = sender.try_send(event)
311        {
312            let dropped = self.dropped_events.fetch_add(1, Ordering::Relaxed) + 1;
313
314            if dropped == 1 || dropped.is_multiple_of(1000) {
315                eprintln!(
316                    "[telemetry]: {}  events and log lines dropped que to full queue (capacity: {})",
317                    dropped, QUEUE_CAPACITY
318                );
319                self.send_drop_event(dropped);
320            }
321        }
322    }
323
324    pub(crate) fn sender(&self) -> mpsc::SyncSender<TraceEvent> {
325        self.sender
326            .as_ref()
327            .expect("trace event dispatcher sender should exist during initialization")
328            .clone()
329    }
330
331    fn send_drop_event(&self, total_dropped: u64) {
332        if let Some(dropped_sender) = &self.dropped_sender {
333            let (thread_name, thread_id) = get_thread_info();
334
335            let mut fields = TraceFields::new();
336            fields.push((
337                "message",
338                FieldValue::Str(format!(
339                    "Telemetry events and log lines dropped due to full queue (capacity: {}). Worker may be falling behind.",
340                    QUEUE_CAPACITY
341                )),
342            ));
343            fields.push(("dropped_count", FieldValue::U64(total_dropped)));
344
345            // We want to just directly construct and send a `TraceEvent::Event` here so we don't need to
346            // reason very hard about whether or not we are creating a DoS loop
347            let drop_event = TraceEvent::Event {
348                name: "dropped events",
349                target: module_path!(),
350                level: tracing::Level::ERROR,
351                fields,
352                timestamp: SystemTime::now(),
353                parent_span: None,
354                thread_id,
355                thread_name,
356                module_path: Some(module_path!()),
357                file: Some(file!()),
358                line: Some(line!()),
359            };
360
361            if dropped_sender.send(drop_event).is_err() {
362                // Last resort
363                eprintln!(
364                    "[telemetry] CRITICAL: {} events and log lines dropped and unable to log to telemetry \
365                     (worker thread may have died). Telemetry system offline.",
366                    total_dropped
367                );
368            }
369        }
370    }
371}
372
373impl Drop for TraceEventDispatcher {
374    fn drop(&mut self) {
375        // Explicitly drop both senders to close the channels.
376        // The next field to be dropped is `worker_handle` which
377        // will run its own drop impl to join the thread and flush
378        drop(self.sender.take());
379        drop(self.dropped_sender.take());
380    }
381}
382
383impl<S> Layer<S> for TraceEventDispatcher
384where
385    S: Subscriber + for<'a> LookupSpan<'a>,
386{
387    fn on_new_span(&self, attrs: &tracing::span::Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
388        let metadata = attrs.metadata();
389        let mut fields = TraceFields::new();
390
391        let mut visitor = FieldVisitor(&mut fields);
392        attrs.record(&mut visitor);
393
394        let parent_id = if let Some(parent) = attrs.parent() {
395            Some(parent.into_u64())
396        } else {
397            ctx.current_span().id().map(|id| id.into_u64())
398        };
399
400        let (thread_name, _) = get_thread_info();
401
402        let event = TraceEvent::NewSpan {
403            id: id.into_u64(),
404            name: metadata.name(),
405            target: metadata.target(),
406            level: *metadata.level(),
407            fields,
408            timestamp: SystemTime::now(),
409            parent_id,
410            thread_name,
411            file: metadata.file(),
412            line: metadata.line(),
413        };
414
415        self.send_event(event);
416    }
417
418    fn on_enter(&self, id: &Id, _ctx: Context<'_, S>) {
419        let (thread_name, _) = get_thread_info();
420        let event = TraceEvent::SpanEnter {
421            id: id.into_u64(),
422            timestamp: SystemTime::now(),
423            thread_name,
424        };
425
426        self.send_event(event);
427    }
428
429    fn on_exit(&self, id: &Id, _ctx: Context<'_, S>) {
430        let (thread_name, _) = get_thread_info();
431        let event = TraceEvent::SpanExit {
432            id: id.into_u64(),
433            timestamp: SystemTime::now(),
434            thread_name,
435        };
436
437        self.send_event(event);
438    }
439
440    fn on_event(&self, event: &tracing::Event<'_>, ctx: Context<'_, S>) {
441        let metadata = event.metadata();
442        let mut fields = TraceFields::new();
443        let mut visitor = FieldVisitor(&mut fields);
444        event.record(&mut visitor);
445
446        let parent_span = ctx.event_span(event).map(|span| span.id().into_u64());
447
448        let (thread_name, thread_id) = get_thread_info();
449
450        let trace_event = TraceEvent::Event {
451            name: metadata.name(),
452            target: metadata.target(),
453            level: *metadata.level(),
454            fields,
455            timestamp: SystemTime::now(),
456            parent_span,
457            thread_id,
458            thread_name,
459            module_path: metadata.module_path(),
460            file: metadata.file(),
461            line: metadata.line(),
462        };
463
464        self.send_event(trace_event);
465    }
466
467    fn on_close(&self, id: Id, _ctx: Context<'_, S>) {
468        let event = TraceEvent::SpanClose {
469            id: id.into_u64(),
470            timestamp: SystemTime::now(),
471        };
472
473        self.send_event(event);
474    }
475
476    fn max_level_hint(&self) -> Option<LevelFilter> {
477        self.max_level
478    }
479}
480
481struct FieldVisitor<'a>(&'a mut TraceFields);
482
483impl tracing::field::Visit for FieldVisitor<'_> {
484    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
485        self.0.push((field.name(), FieldValue::Bool(value)));
486    }
487
488    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
489        self.0.push((field.name(), FieldValue::I64(value)));
490    }
491
492    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
493        self.0.push((field.name(), FieldValue::U64(value)));
494    }
495
496    fn record_f64(&mut self, field: &tracing::field::Field, value: f64) {
497        self.0.push((field.name(), FieldValue::F64(value)));
498    }
499
500    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
501        self.0
502            .push((field.name(), FieldValue::Str(value.to_string())));
503    }
504
505    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
506        self.0
507            .push((field.name(), FieldValue::Debug(format!("{:?}", value))));
508    }
509}
510
511/// Background worker loop that receives events from both regular and priority channels,
512/// and dispatches them to sinks. Priority events are processed first.
513/// Runs until both senders are dropped.
514fn worker_loop(
515    receiver: mpsc::Receiver<TraceEvent>,
516    dropped_receiver: mpsc::Receiver<TraceEvent>,
517    control_receiver: Option<mpsc::Receiver<DispatcherControl>>,
518    mut sinks: Vec<Box<dyn TraceEventSink>>,
519    dropped_events: Arc<AtomicU64>,
520) {
521    const FLUSH_INTERVAL: Duration = Duration::from_millis(100);
522    const FLUSH_EVENT_COUNT: usize = 1000;
523    let mut last_flush = std::time::Instant::now();
524    let mut events_since_flush = 0;
525
526    fn flush_sinks(sinks: &mut [Box<dyn TraceEventSink>]) {
527        for sink in sinks {
528            if let Err(e) = sink.flush() {
529                eprintln!("[telemetry] sink {} failed to flush: {}", sink.name(), e);
530            }
531        }
532    }
533
534    fn process_control_messages(
535        control_receiver: Option<&mpsc::Receiver<DispatcherControl>>,
536        sinks: &mut Vec<Box<dyn TraceEventSink>>,
537    ) {
538        if let Some(ctrl_rx) = control_receiver {
539            while let Ok(control) = ctrl_rx.try_recv() {
540                match control {
541                    DispatcherControl::AddSink(sink) => {
542                        sinks.push(sink);
543                    }
544                }
545            }
546        }
547    }
548
549    fn dispatch_to_sinks(sinks: &mut [Box<dyn TraceEventSink>], event: TraceEvent) {
550        for sink in sinks {
551            if match &event {
552                TraceEvent::NewSpan { target, level, .. }
553                | TraceEvent::Event { target, level, .. } => match sink.target_filter() {
554                    Some(targets) => targets.would_enable(target, level),
555                    None => true,
556                },
557                // Target filters are tracing/log filters. Variants without
558                // target/level metadata, including semantic entity table rows,
559                // must reach sinks so each sink can decide whether to consume
560                // or ignore them.
561                _ => true,
562            } && let Err(e) = sink.consume(&event)
563            {
564                eprintln!(
565                    "[telemetry] sink {} failed to consume event: {}",
566                    sink.name(),
567                    e
568                );
569            }
570        }
571    }
572
573    loop {
574        while let Ok(event) = dropped_receiver.try_recv() {
575            dispatch_to_sinks(&mut sinks, event);
576            events_since_flush += 1;
577        }
578
579        match receiver.recv_timeout(FLUSH_INTERVAL) {
580            Ok(event) => {
581                // A control message may have arrived while we were blocked in
582                // `recv_timeout`. Drain it before dispatching the event that
583                // woke us so dynamically registered sinks see subsequent
584                // replayed events.
585                process_control_messages(control_receiver.as_ref(), &mut sinks);
586                dispatch_to_sinks(&mut sinks, event);
587                events_since_flush += 1;
588
589                if events_since_flush >= FLUSH_EVENT_COUNT || last_flush.elapsed() >= FLUSH_INTERVAL
590                {
591                    flush_sinks(&mut sinks);
592                    last_flush = std::time::Instant::now();
593                    events_since_flush = 0;
594                }
595            }
596            Err(mpsc::RecvTimeoutError::Timeout) => {
597                flush_sinks(&mut sinks);
598                last_flush = std::time::Instant::now();
599                events_since_flush = 0;
600            }
601            Err(mpsc::RecvTimeoutError::Disconnected) => {
602                break;
603            }
604        }
605    }
606
607    // The event queues are closing, but the control queue is independent.
608    // Apply any sink registrations already queued before draining telemetry
609    // events so shutdown delivery follows the same ordering as the live loop.
610    process_control_messages(control_receiver.as_ref(), &mut sinks);
611
612    while let Ok(event) = dropped_receiver.try_recv() {
613        dispatch_to_sinks(&mut sinks, event);
614    }
615    while let Ok(event) = receiver.try_recv() {
616        dispatch_to_sinks(&mut sinks, event);
617    }
618
619    flush_sinks(&mut sinks);
620
621    let total_dropped = dropped_events.load(Ordering::Relaxed);
622    if total_dropped > 0 {
623        eprintln!(
624            "[telemetry] Telemetry worker shutting down. Total events dropped during session: {}",
625            total_dropped
626        );
627    }
628}
629
630impl Drop for WorkerHandle {
631    fn drop(&mut self) {
632        if let Some(handle) = self.join_handle.take()
633            && let Err(e) = handle.join()
634        {
635            eprintln!("[telemetry] worker thread panicked: {:?}", e);
636        }
637    }
638}
639
640#[cfg(test)]
641mod tests {
642    use std::io::Read;
643    use std::os::unix::net::UnixListener;
644    use std::sync::Arc;
645    use std::sync::Mutex;
646    use std::sync::atomic::AtomicU64;
647    use std::sync::atomic::Ordering;
648    use std::sync::mpsc;
649
650    use super::*;
651
652    static TEST_SEQ: AtomicU64 = AtomicU64::new(0);
653
654    #[derive(Default)]
655    struct RecordingSink {
656        events: Arc<Mutex<Vec<TraceEvent>>>,
657    }
658
659    impl TraceEventSink for RecordingSink {
660        fn consume(&mut self, event: &TraceEvent) -> Result<(), anyhow::Error> {
661            self.events.lock().unwrap().push(event.clone());
662            Ok(())
663        }
664
665        fn flush(&mut self) -> Result<(), anyhow::Error> {
666            Ok(())
667        }
668    }
669
670    struct CountingSink {
671        entity_events: Arc<AtomicU64>,
672        target_filter: Option<Targets>,
673    }
674
675    impl TraceEventSink for CountingSink {
676        fn consume(&mut self, event: &TraceEvent) -> Result<(), anyhow::Error> {
677            if matches!(event, TraceEvent::Entity(_)) {
678                self.entity_events.fetch_add(1, Ordering::Relaxed);
679            }
680            Ok(())
681        }
682
683        fn target_filter(&self) -> Option<&Targets> {
684            self.target_filter.as_ref()
685        }
686
687        fn flush(&mut self) -> Result<(), anyhow::Error> {
688            Ok(())
689        }
690    }
691
692    fn span_close(id: u64) -> TraceEvent {
693        TraceEvent::SpanClose {
694            id,
695            timestamp: SystemTime::now(),
696        }
697    }
698
699    fn event() -> TraceEvent {
700        TraceEvent::Event {
701            name: "test_event",
702            target: "test",
703            level: tracing::Level::INFO,
704            fields: TraceFields::new(),
705            timestamp: SystemTime::now(),
706            parent_span: None,
707            thread_id: "1",
708            thread_name: "test",
709            module_path: Some("test"),
710            file: Some("test.rs"),
711            line: Some(1),
712        }
713    }
714
715    fn entity_event() -> TraceEvent {
716        TraceEvent::Entity(EntityEvent::Mesh(crate::MeshEvent {
717            id: 11,
718            timestamp: SystemTime::now(),
719            class: "Host".to_string(),
720            given_name: "test_mesh".to_string(),
721            full_name: "test_mesh".to_string(),
722            shape_json: "{}".to_string(),
723            parent_mesh_id: None,
724            parent_view_json: None,
725        }))
726    }
727
728    fn socket_path(name: &str) -> std::path::PathBuf {
729        let seq = TEST_SEQ.fetch_add(1, Ordering::Relaxed);
730        let dir = std::env::temp_dir().join(format!(
731            "monarch_trace_dispatcher_{}_{}",
732            std::process::id(),
733            seq
734        ));
735        std::fs::create_dir_all(&dir).unwrap();
736        dir.join(name)
737    }
738
739    fn read_frame_table(listener: UnixListener) -> String {
740        let (mut stream, _addr) = listener.accept().unwrap();
741        let mut name_len_bytes = [0; 2];
742        stream.read_exact(&mut name_len_bytes).unwrap();
743        let name_len = u16::from_be_bytes(name_len_bytes) as usize;
744        let mut name_bytes = vec![0; name_len];
745        stream.read_exact(&mut name_bytes).unwrap();
746        String::from_utf8(name_bytes).unwrap()
747    }
748
749    #[test]
750    fn send_event_delivers_repeatedly() {
751        let sink = RecordingSink::default();
752        let recorded = Arc::clone(&sink.events);
753        let dispatcher = TraceEventDispatcher::new(vec![Box::new(sink)]);
754
755        dispatcher.send_event(span_close(1));
756        dispatcher.send_event(span_close(2));
757        dispatcher.send_event(span_close(3));
758
759        drop(dispatcher);
760        assert_eq!(recorded.lock().unwrap().len(), 3);
761    }
762
763    #[test]
764    fn send_event_drops_on_reentrance() {
765        let sink = RecordingSink::default();
766        let recorded = Arc::clone(&sink.events);
767        let dispatcher = TraceEventDispatcher::new(vec![Box::new(sink)]);
768
769        // Simulate that this thread is already inside `send_event`. The nested
770        // call must short-circuit; otherwise a `Layer` callback that re-enters
771        // the subscriber would recurse without bound.
772        IN_SEND.with(|f| f.set(true));
773        dispatcher.send_event(span_close(1));
774        IN_SEND.with(|f| f.set(false));
775
776        drop(dispatcher);
777        assert!(recorded.lock().unwrap().is_empty());
778    }
779
780    #[test]
781    fn entity_event_bypasses_target_filter_and_reaches_sink() {
782        let entity_events = Arc::new(AtomicU64::new(0));
783        let sink = CountingSink {
784            entity_events: Arc::clone(&entity_events),
785            target_filter: Some(Targets::new().with_default(LevelFilter::OFF)),
786        };
787        let dispatcher = TraceEventDispatcher::new(vec![Box::new(sink)]);
788
789        dispatcher.send_event(entity_event());
790
791        drop(dispatcher);
792        assert_eq!(entity_events.load(Ordering::Relaxed), 1);
793    }
794
795    #[test]
796    fn shutdown_drain_processes_pending_sink_registration() {
797        let (sender, receiver) = mpsc::sync_channel(1);
798        let (dropped_sender, dropped_receiver) = mpsc::channel();
799        let (control_sender, control_receiver) = mpsc::channel();
800        let recorded = Arc::new(Mutex::new(Vec::new()));
801        let dropped_events = Arc::new(AtomicU64::new(0));
802
803        let worker = std::thread::spawn(move || {
804            worker_loop(
805                receiver,
806                dropped_receiver,
807                Some(control_receiver),
808                Vec::new(),
809                dropped_events,
810            );
811        });
812
813        // Let the worker pass its top-of-loop dropped-event drain and block on
814        // the main event receiver. The queued dropped event below then reaches
815        // the shutdown drain path, where pending sink registrations must be
816        // applied first.
817        std::thread::sleep(Duration::from_millis(200));
818
819        control_sender
820            .send(DispatcherControl::AddSink(Box::new(RecordingSink {
821                events: Arc::clone(&recorded),
822            })))
823            .unwrap();
824        dropped_sender.send(span_close(7)).unwrap();
825
826        drop(sender);
827        drop(dropped_sender);
828        drop(control_sender);
829        worker.join().unwrap();
830
831        assert_eq!(recorded.lock().unwrap().len(), 1);
832    }
833
834    #[test]
835    fn unix_socket_sink_receives_dispatched_event() {
836        let path = socket_path("telemetry.sock");
837        let listener = UnixListener::bind(&path).unwrap();
838        let (sender, receiver) = mpsc::channel();
839        let read_handle = std::thread::spawn(move || {
840            sender.send(read_frame_table(listener)).unwrap();
841        });
842
843        let sink = Arc::new(crate::unix_sink::UnixSocketSink::new());
844        sink.set_path(path).unwrap();
845        let dispatcher =
846            TraceEventDispatcher::new(vec![crate::unix_sink::adapter_for_test(Arc::clone(&sink))]);
847
848        dispatcher.send_event(event());
849
850        drop(dispatcher);
851        let table_name = receiver.recv_timeout(Duration::from_secs(5)).unwrap();
852        read_handle.join().unwrap();
853        assert_eq!(table_name, monarch_telemetry_schema::trace_tables::EVENTS);
854    }
855}
hyperactor_telemetry/trace_dispatcher.rs

hyperactor_telemetry/
trace_dispatcher.rs
