Switch to encode-inline + raw-copy flush path

Thread-local buffers now encode all events directly via
Encoder<Vec<u8>> and flush raw encoded bytes to the central
collector. The flush thread raw-copies these bytes into the output
file (each batch's header acts as a reset frame for decoders).

Key changes:
- buffer.rs: encode all events directly via write_infallible/
  intern_string_infallible, size-based flush (64KB default),
  remove Vec<RawEvent> path entirely.
- collector.rs: Batch now only contains encoded_bytes.
- recorder/mod.rs: flush_once only processes encoded batches,
  skips header-only batches.
- writer.rs: RotatingWriter::write_encoded_batch does raw-copy +
  reset_state + trailing header. Default trait impl returns error.
  Add Box/NullWriter delegation.
- event_writer.rs: write_raw_event gated behind #[cfg(test)],
  write_transcoded_batch increments event counter.
- tests/common.rs: CapturingWriter decodes encoded batches back
  to RawEvent for test inspection.
- threadlocal_encoding.rs: E2E tests for encoder reset, raw-copy
  concatenation, timestamp preservation, different schemas, empty
  batches, and fallible iteration.
- benches/threadlocal_encode.rs: benchmark comparing direct encode
  vs thread-local encode + raw-copy.

All 294 tests pass. Clippy clean.

Author: rcoh

dial9-tokio-telemetry/Cargo.toml

@@ -109,3 +109,7 @@ required-features = ["cpu-profiling"]
 [[bench]]
 name = "writer_encode"
 harness = false
+
+[[bench]]
+name = "threadlocal_encode"
+harness = false

dial9-tokio-telemetry/benches/threadlocal_encode.rs

@@ -0,0 +1,119 @@
+//! Benchmark comparing direct encode vs thread-local encode + raw-copy.
+//!
+//! Measures the end-to-end cost of:
+//! - Direct: encode events directly through a single Encoder<Vec<u8>>
+//! - Thread-local: encode through a thread-local Encoder, reset(), then
+//!   raw-copy into a central Vec<u8> (reset-frame pattern)
+//!
+//! Usage:
+//!   cargo bench --bench threadlocal_encode
+
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use dial9_tokio_telemetry::telemetry::format::{
+    PollEndEvent, PollStartEvent, WorkerParkEvent, WorkerUnparkEvent,
+};
+use dial9_tokio_telemetry::telemetry::{TaskId, WorkerId};
+use dial9_trace_format::encoder::Encoder;
+
+fn make_batch(worker: usize) -> Vec<(u64, WorkerId, TaskId)> {
+    let wid = WorkerId::from(worker);
+    let task = TaskId::from_u32(1);
+    let mut events = Vec::with_capacity(350);
+
+    for cycle in 0..3u64 {
+        let base = cycle * 10_000;
+        events.push((base + 100, wid, task));
+        for i in 0..170u64 {
+            events.push((base + 200 + i * 10, wid, task));
+        }
+        events.push((base + 3000, wid, task));
+    }
+    events
+}
+
+fn encode_batch(encoder: &mut Encoder<Vec<u8>>, batch: &[(u64, WorkerId, TaskId)]) {
+    let spawn_loc = encoder.intern_string_infallible("test");
+    for &(ts, wid, task) in batch {
+        encoder.write_infallible(&PollStartEvent {
+            timestamp_ns: ts,
+            worker_id: wid,
+            local_queue: 3,
+            task_id: task,
+            spawn_loc,
+        });
+        encoder.write_infallible(&PollEndEvent {
+            timestamp_ns: ts + 5,
+            worker_id: wid,
+        });
+    }
+    encoder.write_infallible(&WorkerParkEvent {
+        timestamp_ns: batch.last().unwrap().0 + 100,
+        worker_id: batch[0].1,
+        local_queue: 0,
+        cpu_time_ns: 600_000,
+    });
+    encoder.write_infallible(&WorkerUnparkEvent {
+        timestamp_ns: batch[0].0,
+        worker_id: batch[0].1,
+        local_queue: 5,
+        cpu_time_ns: 500_000,
+        sched_wait_ns: 1_000,
+    });
+}
+
+fn bench_direct_encode(c: &mut Criterion) {
+    let mut group = c.benchmark_group("direct_encode");
+
+    for num_batches in [1, 10, 100] {
+        let batches: Vec<_> = (0..num_batches).map(|i| make_batch(i % 8)).collect();
+        let total_events: usize = batches.iter().map(|b| b.len() * 2).sum();
+        group.throughput(criterion::Throughput::Elements(total_events as u64));
+
+        group.bench_with_input(
+            BenchmarkId::new("batches", num_batches),
+            &batches,
+            |b, batches| {
+                b.iter(|| {
+                    let mut encoder = Encoder::new();
+                    for batch in batches {
+                        encode_batch(&mut encoder, batch);
+                    }
+                    encoder.finish()
+                });
+            },
+        );
+    }
+    group.finish();
+}
+
+fn bench_threadlocal_rawcopy(c: &mut Criterion) {
+    let mut group = c.benchmark_group("threadlocal_rawcopy");
+
+    for num_batches in [1, 10, 100] {
+        let batches: Vec<_> = (0..num_batches).map(|i| make_batch(i % 8)).collect();
+        let total_events: usize = batches.iter().map(|b| b.len() * 2).sum();
+        group.throughput(criterion::Throughput::Elements(total_events as u64));
+
+        group.bench_with_input(
+            BenchmarkId::new("batches", num_batches),
+            &batches,
+            |b, batches| {
+                b.iter(|| {
+                    let mut output = Vec::new();
+                    dial9_trace_format::codec::encode_header(&mut output).unwrap();
+                    for batch in batches {
+                        let mut local = Encoder::new();
+                        encode_batch(&mut local, batch);
+                        let bytes = local.reset_to(Vec::new());
+                        output.extend_from_slice(&bytes);
+                    }
+                    output
+                });
+            },
+        );
+    }
+    group.finish();
+}
+
+criterion_group!(benches, bench_direct_encode, bench_threadlocal_rawcopy);
+criterion_main!(benches);

dial9-tokio-telemetry/src/telemetry/buffer.rs

@@ -1,17 +1,27 @@
 //! `ThreadLocalBuffer` is the entrypoint for almost all dial9 events
 //!
-//! The TL buffer is created lazily the first time an event is sent. Events are buffered into a fixed-size Vec (currently 1024 items)
-//! before being flushed to the central collector.
+//! The TL buffer is created lazily the first time an event is sent. Events are encoded directly
+//! into a thread-local `Encoder<Vec<u8>>` and flushed to the central collector when the encoded
+//! batch reaches the configured batch size (default 1 MB).
 use crate::telemetry::collector::CentralCollector;
 use crate::telemetry::events::RawEvent;
+use crate::telemetry::format::*;
+use dial9_trace_format::InternedString;
+use dial9_trace_format::encoder::Encoder;
 use std::cell::RefCell;
+use std::collections::HashMap;
+use std::panic::Location;
 use std::sync::Arc;
 
-const BUFFER_CAPACITY: usize = 1024;
+/// Default maximum encoded batch size before flushing (64KB).
+const DEFAULT_BATCH_SIZE: usize = 63 * 1024;
 
 pub(crate) struct ThreadLocalBuffer {
-    pub(crate) events: Vec<RawEvent>,
+    encoder: Encoder<Vec<u8>>,
+    event_count: usize,
+    batch_size: usize,
     collector: Option<Arc<CentralCollector>>,
+    location_cache: HashMap<&'static Location<'static>, String>,
 }
 
 impl Default for ThreadLocalBuffer {
@@ -22,9 +32,19 @@ impl Default for ThreadLocalBuffer {
 
 impl ThreadLocalBuffer {
     fn new() -> Self {
+        Self::with_batch_size(DEFAULT_BATCH_SIZE)
+    }
+
+    fn with_batch_size(batch_size: usize) -> Self {
         Self {
-            events: Vec::with_capacity(BUFFER_CAPACITY),
+            // make the Vec 1KB bigger to reduce the risk of reallocating
+            // TODO: harden this code, we should ensure we never have to re-allocate this buffer.
+            encoder: Encoder::new_to(Vec::with_capacity(batch_size + 1024))
+                .expect("Vec::write_all cannot fail"),
+            event_count: 0,
+            batch_size,
             collector: None,
+            location_cache: HashMap::new(),
         }
     }
 
@@ -36,28 +56,143 @@ impl ThreadLocalBuffer {
         }
     }
 
+    // todo: this is now really "encode tokio event"
+    fn encode_event(&mut self, event: &RawEvent) {
+        match event {
+            RawEvent::PollStart {
+                timestamp_nanos,
+                worker_id,
+                worker_local_queue_depth,
+                task_id,
+                location,
+            } => {
+                let spawn_loc = self.intern_location(location);
+                self.encoder.write_infallible(&PollStartEvent {
+                    timestamp_ns: *timestamp_nanos,
+                    worker_id: *worker_id,
+                    local_queue: *worker_local_queue_depth as u8,
+                    task_id: *task_id,
+                    spawn_loc,
+                });
+            }
+            RawEvent::PollEnd {
+                timestamp_nanos,
+                worker_id,
+            } => self.encoder.write_infallible(&PollEndEvent {
+                timestamp_ns: *timestamp_nanos,
+                worker_id: *worker_id,
+            }),
+            RawEvent::WorkerPark {
+                timestamp_nanos,
+                worker_id,
+                worker_local_queue_depth,
+                cpu_time_nanos,
+            } => self.encoder.write_infallible(&WorkerParkEvent {
+                timestamp_ns: *timestamp_nanos,
+                worker_id: *worker_id,
+                local_queue: *worker_local_queue_depth as u8,
+                cpu_time_ns: *cpu_time_nanos,
+            }),
+            RawEvent::WorkerUnpark {
+                timestamp_nanos,
+                worker_id,
+                worker_local_queue_depth,
+                cpu_time_nanos,
+                sched_wait_delta_nanos,
+            } => self.encoder.write_infallible(&WorkerUnparkEvent {
+                timestamp_ns: *timestamp_nanos,
+                worker_id: *worker_id,
+                local_queue: *worker_local_queue_depth as u8,
+                cpu_time_ns: *cpu_time_nanos,
+                sched_wait_ns: *sched_wait_delta_nanos,
+            }),
+            RawEvent::QueueSample {
+                timestamp_nanos,
+                global_queue_depth,
+            } => self.encoder.write_infallible(&QueueSampleEvent {
+                timestamp_ns: *timestamp_nanos,
+                global_queue: *global_queue_depth as u8,
+            }),
+            RawEvent::TaskSpawn {
+                timestamp_nanos,
+                task_id,
+                location,
+            } => {
+                let spawn_loc = self.intern_location(location);
+                self.encoder.write_infallible(&TaskSpawnEvent {
+                    timestamp_ns: *timestamp_nanos,
+                    task_id: *task_id,
+                    spawn_loc,
+                });
+            }
+            RawEvent::TaskTerminate {
+                timestamp_nanos,
+                task_id,
+            } => self.encoder.write_infallible(&TaskTerminateEvent {
+                timestamp_ns: *timestamp_nanos,
+                task_id: *task_id,
+            }),
+            RawEvent::WakeEvent {
+                timestamp_nanos,
+                waker_task_id,
+                woken_task_id,
+                target_worker,
+            } => self.encoder.write_infallible(&WakeEventEvent {
+                timestamp_ns: *timestamp_nanos,
+                waker_task_id: *waker_task_id,
+                woken_task_id: *woken_task_id,
+                target_worker: *target_worker,
+            }),
+            RawEvent::CpuSample(data) => {
+                let thread_name = match &data.thread_name {
+                    Some(name) => self.encoder.intern_string_infallible(name.as_str()),
+                    None => self.encoder.intern_string_infallible("<no thread name>"),
+                };
+                self.encoder.write_infallible(&CpuSampleEvent {
+                    timestamp_ns: data.timestamp_nanos,
+                    worker_id: data.worker_id,
+                    tid: data.tid,
+                    source: data.source,
+                    thread_name,
+                    callchain: dial9_trace_format::StackFrames(data.callchain.clone()),
+                });
+            }
+        }
+    }
+
+    fn intern_location(&mut self, location: &'static Location<'static>) -> InternedString {
+        let s = self
+            .location_cache
+            .entry(location)
+            .or_insert_with(|| location.to_string());
+        self.encoder.intern_string_infallible(s)
+    }
+
     fn record_event(&mut self, event: RawEvent) {
-        self.events.push(event);
+        self.encode_event(&event);
+        self.event_count += 1;
     }
 
     fn should_flush(&self) -> bool {
-        self.events.len() >= BUFFER_CAPACITY
+        self.encoder.bytes_written() as usize >= self.batch_size
     }
 
-    fn flush(&mut self) -> Vec<RawEvent> {
-        std::mem::replace(&mut self.events, Vec::with_capacity(BUFFER_CAPACITY))
+    fn flush(&mut self) -> crate::telemetry::collector::Batch {
+        let encoded_bytes = self.encoder.reset_to(Vec::with_capacity(self.batch_size));
+        self.event_count = 0;
+        crate::telemetry::collector::Batch { encoded_bytes }
     }
 }
 
 impl Drop for ThreadLocalBuffer {
     fn drop(&mut self) {
-        if !self.events.is_empty() {
+        if self.event_count > 0 {
             if let Some(collector) = self.collector.take() {
-                collector.accept_flush(std::mem::take(&mut self.events));
+                collector.accept_flush(self.flush());
             } else {
                 tracing::warn!(
                     "dial9-tokio-telemetry: dropping {} unflushed events (no collector registered on this thread)",
-                    self.events.len()
+                    self.event_count
                 );
             }
         }
@@ -88,9 +223,8 @@ pub(crate) fn record_event(event: RawEvent, collector: &Arc<CentralCollector>) {
 pub(crate) fn drain_to_collector(collector: &CentralCollector) {
     BUFFER.with(|buf| {
         let mut buf = buf.borrow_mut();
-        let events = buf.flush();
-        if !events.is_empty() {
-            collector.accept_flush(events);
+        if buf.event_count > 0 {
+            collector.accept_flush(buf.flush());
         }
     });
 }
@@ -108,34 +242,42 @@ mod tests {
     #[test]
     fn test_buffer_creation() {
         let buffer = ThreadLocalBuffer::new();
-        assert_eq!(buffer.events.len(), 0);
-        assert_eq!(buffer.events.capacity(), BUFFER_CAPACITY);
+        assert_eq!(buffer.event_count, 0);
+        assert_eq!(buffer.batch_size, DEFAULT_BATCH_SIZE);
     }
 
     #[test]
     fn test_record_event() {
         let mut buffer = ThreadLocalBuffer::new();
         buffer.record_event(poll_end_event());
-        assert_eq!(buffer.events.len(), 1);
+        assert_eq!(buffer.event_count, 1);
+        assert!(buffer.encoder.bytes_written() > 0);
     }
 
     #[test]
-    fn test_should_flush() {
-        let mut buffer = ThreadLocalBuffer::new();
+    fn test_should_flush_respects_batch_size() {
+        // Use a tiny batch size so a single event triggers flush.
+        let mut buffer = ThreadLocalBuffer::with_batch_size(1);
         assert!(!buffer.should_flush());
-        for _ in 0..BUFFER_CAPACITY {
-            buffer.record_event(poll_end_event());
-        }
+        buffer.record_event(poll_end_event());
         assert!(buffer.should_flush());
     }
 
+    #[test]
+    fn test_should_flush_default_batch_size() {
+        let mut buffer = ThreadLocalBuffer::new();
+        assert!(!buffer.should_flush());
+        buffer.record_event(poll_end_event());
+        // A single small event should not exceed 1 MB.
+        assert!(!buffer.should_flush());
+    }
+
     #[test]
     fn test_flush() {
         let mut buffer = ThreadLocalBuffer::new();
         buffer.record_event(poll_end_event());
-        let flushed = buffer.flush();
-        assert_eq!(flushed.len(), 1);
-        assert_eq!(buffer.events.len(), 0);
-        assert_eq!(buffer.events.capacity(), BUFFER_CAPACITY);
+        let batch = buffer.flush();
+        assert!(!batch.encoded_bytes.is_empty());
+        assert_eq!(buffer.event_count, 0);
     }
 }