dial9_core/
encoder.rs

1//! `ThreadLocalBuffer` is the entrypoint for almost all dial9 events
2//!
3//! The TL buffer is created lazily the first time an event is sent. Events are encoded directly
4//! into a thread-local `Encoder<Vec<u8>>` and flushed to the central collector when the encoded
5//! batch reaches the configured batch size (default 1 MB).
6//!
7//! Each buffer is wrapped in `Arc<Mutex<…>>` so the flush thread can intrusively
8//! drain idle/silent threads via `TlBufferHandle`s registered in `SharedState`.
9use crate::collector::CentralCollector;
10use crate::primitives::sync::atomic::{AtomicU64, Ordering};
11use crate::primitives::sync::{Arc, Mutex, Weak};
12use dial9_trace_format::encoder::{Encoder, FxHashMap};
13use dial9_trace_format::{InternedStackFrames, InternedString};
14use std::panic::Location;
15use std::time::Duration;
16
17// ── Public API types ────────────────────────────────────────────────────────
18
19/// Scoped encoder for writing events into the thread-local trace buffer.
20///
21/// Provides access to string interning and event encoding. The borrow lifetime
22/// ensures that [`InternedString`] handles created via [`intern_string`](Self::intern_string)
23/// are used within the same batch — they become invalid after the buffer flushes.
24///
25/// You don't construct this directly; it's passed to [`Encodable::encode`].
26pub struct ThreadLocalEncoder<'a> {
27    encoder: &'a mut Encoder<Vec<u8>>,
28    location_cache: &'a mut FxHashMap<&'static Location<'static>, String>,
29    /// Events written through this handle; feeds the buffer's batch count so
30    /// callers that decline to write (e.g. the metrique sink dropping an
31    /// entry) are not counted.
32    events_written: &'a mut usize,
33}
34
35impl std::fmt::Debug for ThreadLocalEncoder<'_> {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        f.debug_struct("ThreadLocalEncoder").finish_non_exhaustive()
38    }
39}
40
41impl ThreadLocalEncoder<'_> {
42    /// Intern a string into the trace's string pool, returning a compact handle.
43    ///
44    /// If the string was already interned in this batch, returns the existing handle
45    /// (no duplicate wire data). The returned [`InternedString`] is only valid for
46    /// encoding within this same [`Encodable::encode`] call.
47    pub fn intern_string(&mut self, s: &str) -> InternedString {
48        self.encoder.intern_string_infallible(s)
49    }
50
51    /// Intern a stack-frame vector into the trace's stack pool, returning a compact handle.
52    ///
53    /// If the stack was already interned in this batch, returns the existing handle
54    /// (no duplicate wire data). The returned [`InternedStackFrames`] is only valid for
55    /// encoding within this same [`Encodable::encode`] call.
56    pub fn intern_stack_frames(&mut self, frames: &[u64]) -> InternedStackFrames {
57        self.encoder.intern_stack_frames_infallible(frames)
58    }
59
60    /// Encode a [`TraceEvent`](dial9_trace_format::TraceEvent) struct into the buffer.
61    pub fn encode(&mut self, event: &impl dial9_trace_format::TraceEvent) {
62        self.encoder.write_infallible(event);
63        *self.events_written += 1;
64    }
65
66    /// Write an event with a dynamically-registered schema.
67    ///
68    /// `timestamp_ns` is the event's monotonic clock timestamp.
69    /// `values` must match the schema's field definitions in order (excluding the timestamp).
70    /// The schema is auto-registered on first use per buffer flush cycle.
71    ///
72    /// # Example
73    ///
74    /// ```ignore
75    /// use dial9_trace_format::types::FieldValue;
76    ///
77    /// enc.write_event(&schema, timestamp_ns, &[
78    ///     FieldValue::Varint(worker_id),
79    ///     FieldValue::PooledString(enc.intern_string("hello")),
80    /// ]);
81    /// ```
82    #[doc(hidden)]
83    /// Errors on validation failure; the event is dropped and the calling
84    /// thread is never panicked. The caller decides whether and how to
85    /// report (it knows the event's provenance; rate-limit in loops).
86    #[must_use = "a validation failure means the event was dropped"]
87    pub fn write_event(
88        &mut self,
89        schema: &dial9_trace_format::encoder::Schema,
90        timestamp_ns: u64,
91        values: &[dial9_trace_format::types::FieldValue],
92    ) -> std::io::Result<()> {
93        self.encoder.write_event(schema, timestamp_ns, values)?;
94        *self.events_written += 1;
95        Ok(())
96    }
97
98    /// Intern a `&'static Location` (caching the `to_string()` result).
99    #[doc(hidden)]
100    pub fn intern_location(&mut self, location: &'static Location<'static>) -> InternedString {
101        let s = self
102            .location_cache
103            .entry(location)
104            .or_insert_with(|| location.to_string());
105        self.encoder.intern_string_infallible(s)
106    }
107}
108
109/// Trait for types that can be encoded into a dial9 trace.
110///
111/// # Simple case — `#[derive(TraceEvent)]`
112///
113/// Any type implementing [`TraceEvent`](dial9_trace_format::TraceEvent) automatically
114/// implements `Encodable` via a blanket impl, so you can pass it directly to
115/// `Dial9Handle::record_event`:
116///
117/// ```ignore
118/// #[derive(TraceEvent)]
119/// struct MyEvent {
120///     #[traceevent(timestamp)]
121///     timestamp_ns: u64,
122///     request_count: u32,
123/// }
124/// handle.record_event(MyEvent { timestamp_ns: now, request_count: 42 });
125/// ```
126///
127/// # Advanced case — string interning
128///
129/// Implement `Encodable` manually when you need [`InternedString`] fields
130/// for efficient repeated-string encoding:
131///
132/// ```ignore
133/// struct HttpRequest { timestamp_ns: u64, method: String, status: u32 }
134///
135/// impl Encodable for HttpRequest {
136///     fn encode(&self, enc: &mut ThreadLocalEncoder<'_>) {
137///         let method = enc.intern_string(&self.method);
138///         enc.encode(&HttpRequestWire {
139///             timestamp_ns: self.timestamp_ns,
140///             method,
141///             status: self.status,
142///         });
143///     }
144/// }
145/// ```
146///
147/// # Wire event naming
148///
149/// The event name in the trace comes from the struct passed to
150/// [`ThreadLocalEncoder::encode`], not from the type implementing `Encodable`.
151/// In the example above, the trace will contain events named `"HttpRequestWire"`,
152/// not `"HttpRequest"`.
153///
154pub trait Encodable {
155    /// Encode this event into the thread-local trace buffer.
156    ///
157    /// Implementations should call [`ThreadLocalEncoder::encode`] exactly once.
158    /// Each `encode` call is counted as one event for buffer flush decisions;
159    /// calling `encode` multiple times will produce multiple wire events but
160    /// only one event will be counted.
161    fn encode(&self, encoder: &mut ThreadLocalEncoder<'_>);
162}
163
164impl<T: dial9_trace_format::TraceEvent> Encodable for T {
165    fn encode(&self, encoder: &mut ThreadLocalEncoder<'_>) {
166        encoder.encode(self);
167    }
168}
169
170// ── Thread-local buffer internals ───────────────────────────────────────────
171
172/// Tracks the last drain epoch at which a particular thread-local buffer
173/// was flushed. The flush thread reads this (relaxed) to skip buffers
174/// that have self-flushed recently, avoiding contention with busy workers.
175#[derive(Clone)]
176pub(crate) struct FlushEpoch(Arc<AtomicU64>);
177
178impl FlushEpoch {
179    pub(crate) fn new() -> Self {
180        Self(Arc::new(AtomicU64::new(0)))
181    }
182
183    pub(crate) fn store(&self, epoch: u64) {
184        self.0.store(epoch, Ordering::Relaxed);
185    }
186
187    pub(crate) fn load(&self) -> u64 {
188        self.0.load(Ordering::Relaxed)
189    }
190}
191
192/// Default maximum encoded batch size before flushing (~1MB).
193const DEFAULT_BATCH_SIZE: usize = 1023 * 1024;
194
195pub(crate) struct ThreadLocalBuffer {
196    encoder: Encoder<Vec<u8>>,
197    event_count: usize,
198    batch_size: usize,
199    collector: Option<Arc<CentralCollector>>,
200    /// Caches `Location::to_string()` to avoid re-formatting on every event.
201    /// Bounded by the number of `#[track_caller]` call sites in the program,
202    /// which is fixed at compile time, so this does not grow unboundedly.
203    location_cache: FxHashMap<&'static Location<'static>, String>,
204    /// Last drain epoch at which this buffer was flushed. Shared with the
205    /// flush thread via `TlBufferHandle` so it can skip busy workers.
206    pub(crate) flush_epoch: FlushEpoch,
207}
208
209impl Default for ThreadLocalBuffer {
210    fn default() -> Self {
211        Self::new()
212    }
213}
214
215impl ThreadLocalBuffer {
216    fn new() -> Self {
217        Self::with_batch_size(DEFAULT_BATCH_SIZE)
218    }
219
220    fn with_batch_size(batch_size: usize) -> Self {
221        Self {
222            // Allocate 1KB extra headroom so typical events never trigger a realloc.
223            encoder: Encoder::new_to(Vec::with_capacity(batch_size + 1024))
224                .expect("Vec::write_all cannot fail"),
225            event_count: 0,
226            batch_size,
227            collector: None,
228            location_cache: FxHashMap::default(),
229            flush_epoch: FlushEpoch::new(),
230        }
231    }
232
233    /// Ensure the collector reference is set. Called on every record_event;
234    /// only the first call per thread actually stores the Arc.
235    /// Returns `true` on the first call (when the collector was not yet set).
236    fn set_collector(&mut self, collector: &Arc<CentralCollector>) -> bool {
237        if self.collector.is_none() {
238            self.collector = Some(Arc::clone(collector));
239            return true;
240        }
241        false
242    }
243
244    fn thread_local_encoder(&mut self) -> ThreadLocalEncoder<'_> {
245        ThreadLocalEncoder {
246            encoder: &mut self.encoder,
247            location_cache: &mut self.location_cache,
248            events_written: &mut self.event_count,
249        }
250    }
251
252    // Only reached via `encode_single` (test-util) and core's own tests.
253    #[cfg_attr(not(feature = "test-util"), allow(dead_code))]
254    fn record_encodable(&mut self, event: &dyn Encodable) {
255        event.encode(&mut self.thread_local_encoder());
256    }
257
258    fn should_flush(&self) -> bool {
259        self.encoder.bytes_written() as usize >= self.batch_size
260    }
261
262    pub(crate) fn flush(&mut self) -> crate::collector::Batch {
263        let event_count = self.event_count as u64;
264        let encoded_bytes = self
265            .encoder
266            .reset_to_infallible(Vec::with_capacity(self.batch_size));
267        self.event_count = 0;
268        crate::collector::Batch::new(encoded_bytes, event_count)
269    }
270
271    pub(crate) fn has_pending_events(&self) -> bool {
272        self.event_count > 0
273    }
274}
275
276crate::test_util_pub! {
277/// Encode a single event into a self-contained batch (header + event).
278fn encode_single(event: &dyn Encodable) -> Vec<u8> {
279    let mut buf = ThreadLocalBuffer::with_batch_size(1024);
280    buf.record_encodable(event);
281    buf.flush().into_encoded_bytes()
282}
283}
284
285impl Drop for ThreadLocalBuffer {
286    fn drop(&mut self) {
287        if self.event_count > 0 {
288            if let Some(collector) = self.collector.take() {
289                collector.accept_flush(self.flush());
290            } else {
291                crate::rate_limit::rate_limited!(Duration::from_secs(60), {
292                    tracing::warn!(
293                        "dial9-tokio-telemetry: dropping {} unflushed events (no collector registered on this thread)",
294                        self.event_count
295                    );
296                });
297            }
298        }
299    }
300}
301
302/// A handle to a thread-local buffer, held by `SharedState` so the flush
303/// thread can intrusively drain idle/silent buffers.
304pub(crate) struct TlBufferHandle {
305    pub(crate) buffer: Weak<Mutex<ThreadLocalBuffer>>,
306    pub(crate) flush_epoch: FlushEpoch,
307}
308
309crate::primitives::thread_local! {
310    static BUFFER: Arc<Mutex<ThreadLocalBuffer>> = Arc::new(Mutex::new(ThreadLocalBuffer::new()));
311}
312
313/// Drain the current thread's buffer into `collector`, even if not full.
314/// Used at shutdown and before flush cycles to avoid losing events.
315pub(crate) fn drain_to_collector(collector: &CentralCollector) {
316    BUFFER.with(|buf| {
317        let mut buf = match buf.lock() {
318            Ok(guard) => guard,
319            Err(_) => {
320                crate::rate_limit::rate_limited!(Duration::from_secs(60), {
321                    tracing::error!("dial9: thread-local buffer mutex poisoned in drain_to_collector; skipping drain");
322                });
323                return;
324            }
325        };
326        if buf.event_count > 0 {
327            collector.accept_flush(buf.flush());
328        }
329    });
330}
331
332pub(crate) fn record_encodable_event(
333    event: &dyn Encodable,
334    collector: &Arc<CentralCollector>,
335    drain_epoch: &AtomicU64,
336) -> Option<TlBufferHandle> {
337    with_encoder(|enc| event.encode(enc), collector, drain_epoch)
338}
339
340pub(crate) fn with_encoder(
341    f: impl FnOnce(&mut ThreadLocalEncoder<'_>),
342    collector: &Arc<CentralCollector>,
343    drain_epoch: &AtomicU64,
344) -> Option<TlBufferHandle> {
345    BUFFER.with(|arc| {
346        let mut buf = match arc.lock() {
347            Ok(guard) => guard,
348            Err(_) => {
349                crate::rate_limit::rate_limited!(Duration::from_secs(60), {
350                    tracing::error!("dial9: thread-local buffer mutex poisoned in with_encoder; dropping events for this thread");
351                });
352                return None;
353            }
354        };
355        let first_call = buf.set_collector(collector);
356        f(&mut buf.thread_local_encoder());
357        let current_epoch = drain_epoch.load(Ordering::Relaxed);
358        if buf.should_flush() || buf.flush_epoch.load() < current_epoch {
359            collector.accept_flush(buf.flush());
360            buf.flush_epoch.store(current_epoch);
361        }
362        if first_call {
363            Some(TlBufferHandle {
364                buffer: Arc::downgrade(arc),
365                flush_epoch: buf.flush_epoch.clone(),
366            })
367        } else {
368            None
369        }
370    })
371}
372
373#[cfg(test)]
374mod tests {
375    use super::*;
376
377    fn sample_event() -> crate::format::ClockSyncEvent {
378        crate::format::ClockSyncEvent {
379            timestamp_ns: 1000,
380            realtime_ns: 2000,
381        }
382    }
383
384    #[derive(dial9_trace_format::TraceEvent)]
385    struct BorrowedEvent<'a> {
386        #[traceevent(timestamp)]
387        timestamp_ns: u64,
388        value: &'a str,
389    }
390
391    #[test]
392    fn test_buffer_creation() {
393        let buffer = ThreadLocalBuffer::new();
394        assert_eq!(buffer.event_count, 0);
395        assert_eq!(buffer.batch_size, DEFAULT_BATCH_SIZE);
396    }
397
398    #[test]
399    fn test_record_event() {
400        let mut buffer = ThreadLocalBuffer::new();
401        buffer.record_encodable(&sample_event());
402        assert_eq!(buffer.event_count, 1);
403        assert!(buffer.encoder.bytes_written() > 0);
404    }
405
406    #[test]
407    fn borrowed_trace_event_is_encodable() {
408        let value = String::from("borrowed");
409        let event = BorrowedEvent {
410            timestamp_ns: 1000,
411            value: &value,
412        };
413        let mut buffer = ThreadLocalBuffer::new();
414        buffer.record_encodable(&event);
415        assert_eq!(buffer.event_count, 1);
416        assert!(buffer.encoder.bytes_written() > 0);
417    }
418
419    #[test]
420    fn test_should_flush_respects_batch_size() {
421        // Use a tiny batch size so a single event triggers flush.
422        let mut buffer = ThreadLocalBuffer::with_batch_size(1);
423        assert!(!buffer.should_flush());
424        buffer.record_encodable(&sample_event());
425        assert!(buffer.should_flush());
426    }
427
428    #[test]
429    fn test_should_flush_default_batch_size() {
430        let mut buffer = ThreadLocalBuffer::new();
431        assert!(!buffer.should_flush());
432        buffer.record_encodable(&sample_event());
433        // A single small event should not exceed 1 MB.
434        assert!(!buffer.should_flush());
435    }
436
437    #[test]
438    fn test_flush() {
439        let mut buffer = ThreadLocalBuffer::new();
440        buffer.record_encodable(&sample_event());
441        let batch = buffer.flush();
442        assert!(!batch.encoded_bytes().is_empty());
443        assert_eq!(buffer.event_count, 0);
444    }
445
446    #[test]
447    fn test_flush_epoch_store_load() {
448        let epoch = FlushEpoch::new();
449        assert_eq!(epoch.load(), 0);
450        epoch.store(42);
451        assert_eq!(epoch.load(), 42);
452    }
453
454    #[test]
455    fn test_flush_epoch_shared_across_threads() {
456        let epoch = FlushEpoch::new();
457        let epoch_clone = epoch.clone();
458        let handle = std::thread::spawn(move || {
459            epoch_clone.store(7);
460        });
461        handle.join().unwrap();
462        assert_eq!(epoch.load(), 7);
463    }
464
465    #[test]
466    fn test_flush_epoch_stamped_on_self_flush() {
467        let collector = Arc::new(CentralCollector::new());
468        let drain_epoch = AtomicU64::new(5);
469        // Use a tiny batch size so a single event triggers self-flush.
470        // We can't use record_event (thread-local) easily, so test the
471        // logic directly: flush + stamp.
472        let mut buffer = ThreadLocalBuffer::with_batch_size(1);
473        buffer.set_collector(&collector);
474        buffer.record_encodable(&sample_event());
475        assert!(buffer.should_flush());
476        buffer
477            .flush_epoch
478            .store(drain_epoch.load(Ordering::Relaxed));
479        collector.accept_flush(buffer.flush());
480        assert_eq!(buffer.flush_epoch.load(), 5);
481    }
482
483    #[test]
484    fn test_mutex_accessible_from_another_thread() {
485        let buf = Arc::new(Mutex::new(ThreadLocalBuffer::new()));
486        let buf_clone = Arc::clone(&buf);
487        // Write an event from a different thread.
488        let handle = std::thread::spawn(move || {
489            let mut guard = buf_clone.lock().unwrap();
490            guard.record_encodable(&sample_event());
491            assert_eq!(guard.event_count, 1);
492        });
493        handle.join().unwrap();
494        // Main thread can also access it.
495        let guard = buf.lock().unwrap();
496        assert_eq!(guard.event_count, 1);
497    }
498}
dial9_core/encoder.rs

dial9_core/
encoder.rs
