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

dial9_core/
encoder.rs
