Add reset-frame support + try_for_each_event to decoder (#134)

rcoh · jlizen · web-flow · commit 58642ca73a47 · 2026-03-26T13:49:52.000Z
* Add reset-frame support + try_for_each_event to decoder

Decoder changes for the thread-local encoding pipeline:

- Reset frames: when a valid TRC\0 header appears mid-stream, the
  decoder resets schemas, string pool, and timestamp base, then
  continues. This enables raw-copy concatenation of independently
  encoded batches. Handled in next_frame, next_frame_ref, and
  try_for_each_event.
- try_for_each_event: like for_each_event but the callback may
  return Err to stop early. Includes TryForEachError enum.
- for_each_event simplified to delegate to try_for_each_event.
- schema_cache switched from HashMap&lt;WireTypeId, _&gt; to Vec&lt;Option&lt;_&gt;&gt;
  indexed by WireTypeId.0 for faster lookup.
- JS decoder (decode.js): mid-stream header detection and state reset.

All 115 existing tests pass.

* Update dial9-trace-format/src/decoder.rs

Co-authored-by: Jess Izen &lt;44884346+jlizen@users.noreply.github.com&gt;

---------

Co-authored-by: Jess Izen &lt;44884346+jlizen@users.noreply.github.com&gt;
diff --git a/dial9-trace-format/js/decode.js b/dial9-trace-format/js/decode.js
@@ -105,6 +105,20 @@ class TraceDecoder {
     if (this._pos >= this._view.byteLength) return null;
     try {
       const tag = this._view.getUint8(this._pos);
+      // Mid-stream header = reset frame (concatenated thread-local batch)
+      if (tag === MAGIC[0] && this._pos + 5 <= this._view.byteLength) {
+        let isHeader = true;
+        for (let i = 1; i < 4; i++) {
+          if (this._view.getUint8(this._pos + i) !== MAGIC[i]) { isHeader = false; break; }
+        }
+        if (isHeader) {
+          this.schemas = new Map();
+          this.stringPool = new Map();
+          this._timestampBaseNs = 0n;
+          this._pos += 5; // skip header
+          return this.nextFrame();
+        }
+      }
       this._pos++;
       switch (tag) {
         case TAG_SCHEMA: return this._decodeSchema();
diff --git a/dial9-trace-format/src/decoder.rs b/dial9-trace-format/src/decoder.rs
@@ -25,6 +25,24 @@ impl fmt::Display for DecodeError {
 
 impl std::error::Error for DecodeError {}
 
+/// Error returned by [`Decoder::try_for_each_event`].
+#[derive(Debug)]
+pub enum TryForEachError<E> {
+    Decode(DecodeError),
+    User(E),
+}
+
+impl<E: fmt::Display> fmt::Display for TryForEachError<E> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            TryForEachError::Decode(e) => write!(f, "{e}"),
+            TryForEachError::User(e) => write!(f, "{e}"),
+        }
+    }
+}
+
+impl<E: fmt::Display + fmt::Debug> std::error::Error for TryForEachError<E> {}
+
 /// A decoded event passed to [`Decoder::for_each_event`].
 ///
 /// `'a` is the lifetime of the input data buffer (strings, stack frames borrow from it).
@@ -48,7 +66,7 @@ pub struct StringPool(pub(crate) HashMap<InternedString, String>);
 
 impl StringPool {
     pub(crate) fn new() -> Self {
-        Self(HashMap::new())
+        Self(HashMap::default())
     }
 
     pub(crate) fn insert(&mut self, id: InternedString, value: String) {
@@ -108,7 +126,7 @@ pub struct Decoder<'a> {
     data: &'a [u8],
     pos: usize,
     registry: SchemaRegistry,
-    schema_cache: HashMap<WireTypeId, SchemaCache>,
+    schema_cache: Vec<Option<SchemaCache>>,
     string_pool: StringPool,
     version: u8,
     timestamp_base_ns: u64,
@@ -121,7 +139,7 @@ impl<'a> Decoder<'a> {
             data,
             pos: HEADER_SIZE,
             registry: SchemaRegistry::new(),
-            schema_cache: HashMap::new(),
+            schema_cache: Vec::new(),
             string_pool: StringPool::new(),
             version,
             timestamp_base_ns: 0,
@@ -140,6 +158,30 @@ impl<'a> Decoder<'a> {
         &self.string_pool
     }
 
+    /// Reset decoder state (schemas, string pool, timestamp base) as if
+    /// starting a fresh stream. Used when a mid-stream header is encountered
+    /// (the "reset frame" pattern for concatenated thread-local batches).
+    fn reset_state(&mut self) {
+        self.registry = SchemaRegistry::new();
+        self.schema_cache.clear();
+        self.string_pool = StringPool::new();
+        self.timestamp_base_ns = 0;
+    }
+
+    /// If the current position starts with a valid header, reset state and
+    /// skip past it, returning true.
+    fn try_consume_reset_header(&mut self) -> bool {
+        if self.pos + HEADER_SIZE <= self.data.len()
+            && codec::decode_header(&self.data[self.pos..]).is_some()
+        {
+            self.reset_state();
+            self.pos += HEADER_SIZE;
+            true
+        } else {
+            false
+        }
+    }
+
     /// Consume this decoder and create an [`Encoder`] that appends to the
     /// decoded trace. The encoder inherits the string pool, schema registry,
     /// and timestamp base so new frames are compatible with the existing data.
@@ -155,22 +197,26 @@ impl<'a> Decoder<'a> {
         )
     }
 
-    fn schema_info(&self, type_id: WireTypeId) -> Option<SchemaInfo<'_>> {
-        self.schema_cache.get(&type_id).map(|c| SchemaInfo {
-            field_types: &c.field_types,
-            has_timestamp: c.has_timestamp,
-        })
+    pub(crate) fn schema_info(&self, type_id: WireTypeId) -> Option<SchemaInfo<'_>> {
+        self.schema_cache
+            .get(type_id.0 as usize)
+            .and_then(|s| s.as_ref())
+            .map(|c| SchemaInfo {
+                field_types: &c.field_types,
+                has_timestamp: c.has_timestamp,
+            })
     }
 
     fn register_schema(&mut self, type_id: WireTypeId, entry: SchemaEntry) -> Result<(), String> {
-        self.schema_cache.insert(
-            type_id,
-            SchemaCache {
-                name: entry.name.clone(),
-                field_types: entry.fields.iter().map(|f| f.field_type).collect(),
-                has_timestamp: entry.has_timestamp,
-            },
-        );
+        let idx = type_id.0 as usize;
+        if idx >= self.schema_cache.len() {
+            self.schema_cache.resize_with(idx + 1, || None);
+        }
+        self.schema_cache[idx] = Some(SchemaCache {
+            name: entry.name.clone(),
+            field_types: entry.fields.iter().map(|f| f.field_type).collect(),
+            has_timestamp: entry.has_timestamp,
+        });
         self.registry.register(type_id, entry)
     }
 
@@ -181,6 +227,9 @@ impl<'a> Decoder<'a> {
         if self.pos >= self.data.len() {
             return Ok(None);
         }
+        if self.try_consume_reset_header() {
+            return self.next_frame();
+        }
         let remaining = &self.data[self.pos..];
         let base = self.timestamp_base_ns;
         let (frame, consumed) =
@@ -243,6 +292,9 @@ impl<'a> Decoder<'a> {
         if self.pos >= self.data.len() {
             return Ok(None);
         }
+        if self.try_consume_reset_header() {
+            return self.next_frame_ref();
+        }
         let remaining = &self.data[self.pos..];
         let base = self.timestamp_base_ns;
         let (frame, consumed) =
@@ -312,6 +364,22 @@ impl<'a> Decoder<'a> {
         &mut self,
         mut f: impl for<'f> FnMut(RawEvent<'a, 'f>),
     ) -> Result<(), DecodeError> {
+        self.try_for_each_event(|ev| {
+            f(ev);
+            Ok::<(), std::convert::Infallible>(())
+        })
+        .map_err(|e| match e {
+            TryForEachError::Decode(d) => d,
+            TryForEachError::User(inf) => match inf {},
+        })
+    }
+
+    /// Like [`for_each_event`](Self::for_each_event), but the callback may
+    /// return an error to stop iteration early.
+    pub fn try_for_each_event<E>(
+        &mut self,
+        mut f: impl for<'f> FnMut(RawEvent<'a, 'f>) -> Result<(), E>,
+    ) -> Result<(), TryForEachError<E>> {
         let mut values_buf: Vec<FieldValueRef<'a>> = Vec::new();
         while self.pos < self.data.len() {
             let remaining = &self.data[self.pos..];
@@ -328,19 +396,23 @@ impl<'a> Decoder<'a> {
                             WireTypeId(u16::from_le_bytes(b.try_into().unwrap()))
                         }
                         None => {
-                            return Err(DecodeError {
+                            return Err(TryForEachError::Decode(DecodeError {
                                 pos: self.pos,
                                 message: "truncated event frame".into(),
-                            });
+                            }));
                         }
                     };
-                    let cache = match self.schema_cache.get(&type_id) {
+                    let cache = match self
+                        .schema_cache
+                        .get(type_id.0 as usize)
+                        .and_then(|s| s.as_ref())
+                    {
                         Some(c) => c,
                         None => {
-                            return Err(DecodeError {
+                            return Err(TryForEachError::Decode(DecodeError {
                                 pos: self.pos,
                                 message: format!("unknown type_id {type_id:?}"),
-                            });
+                            }));
                         }
                     };
 
@@ -351,10 +423,10 @@ impl<'a> Decoder<'a> {
                                 Some(self.timestamp_base_ns + delta as u64)
                             }
                             None => {
-                                return Err(DecodeError {
+                                return Err(TryForEachError::Decode(DecodeError {
                                     pos: self.pos + pos,
                                     message: "truncated timestamp delta".into(),
-                                });
+                                }));
                             }
                         }
                     } else {
@@ -369,10 +441,10 @@ impl<'a> Decoder<'a> {
                                 pos += consumed;
                             }
                             None => {
-                                return Err(DecodeError {
+                                return Err(TryForEachError::Decode(DecodeError {
                                     pos: self.pos + pos,
                                     message: "truncated field value".into(),
-                                });
+                                }));
                             }
                         }
                     }
@@ -386,35 +458,36 @@ impl<'a> Decoder<'a> {
                         timestamp_ns,
                         fields: &values_buf,
                         string_pool: &self.string_pool,
-                    });
+                    })
+                    .map_err(TryForEachError::User)?;
                 }
                 codec::TAG_TIMESTAMP_RESET => {
-                    // Handle timestamp resets inline to avoid next_frame_ref's
-                    // recursive consumption of the following frame.
                     let ts = match self.data.get(self.pos + 1..self.pos + 9) {
                         Some(b) => u64::from_le_bytes(b.try_into().unwrap()),
                         None => {
-                            return Err(DecodeError {
+                            return Err(TryForEachError::Decode(DecodeError {
                                 pos: self.pos,
                                 message: "truncated timestamp reset".into(),
-                            });
+                            }));
                         }
                     };
                     self.timestamp_base_ns = ts;
                     self.pos += 9;
                 }
                 _ => {
-                    // Use next_frame_ref for non-event frames (schema, pool, symbol table)
-                    // `next_frame_ref` will update the decoder state as we read the frames (e.g. the pooled strings)
+                    // Mid-stream header = reset frame (tag 0x54 = 'T' from TRC\0)
+                    if tag == codec::MAGIC[0] && self.try_consume_reset_header() {
+                        continue;
+                    }
                     match self.next_frame_ref() {
                         Ok(Some(_)) => {}
                         Ok(None) => {
-                            return Err(DecodeError {
+                            return Err(TryForEachError::Decode(DecodeError {
                                 pos: self.pos,
                                 message: format!("failed to decode frame with tag 0x{tag:02x}"),
-                            });
+                            }));
                         }
-                        Err(e) => return Err(e),
+                        Err(e) => return Err(TryForEachError::Decode(e)),
                     }
                 }
             }
