Keep trace capture responsive during daemon recovery

bt-daemon/src/journal.rs

@@ -333,6 +333,26 @@ pub struct JournalReader {
 }
 
 impl JournalReader {
+    pub(crate) fn position(&self) -> u64 {
+        self.position
+    }
+
+    /// Rebuild decoder context before continuing at a saved record boundary.
+    pub(crate) async fn resume(path: &Path, position: u64, through: u64) -> anyhow::Result<Self> {
+        anyhow::ensure!(position <= through, "invalid recovery range");
+        let mut reader = Self::open(path, through)
+            .await?
+            .ok_or_else(|| anyhow::anyhow!("recovery journal is missing"))?;
+        while reader.position < position {
+            anyhow::ensure!(
+                reader.next_record().await?.is_some(),
+                "recovery journal ended before its boundary"
+            );
+        }
+        anyhow::ensure!(reader.position == position, "invalid recovery boundary");
+        Ok(reader)
+    }
+
     /// Open a journal for streaming, reading at most `through` bytes.
     ///
     /// That bound is what keeps replay from consuming the very event that
@@ -385,49 +405,19 @@ impl JournalReader {
             }
             self.position += read as u64;
             self.line_no += 1;
-            let line = std::str::from_utf8(&bytes)
-                .map_err(|error| {
-                    anyhow::anyhow!("journal {}:{}: {error}", self.path.display(), self.line_no)
-                })?
-                .trim();
-            if line.is_empty() {
-                continue;
+            let path = self.path.clone();
+            let line_no = self.line_no;
+            let position = self.position;
+            let mut pi_context = std::mem::take(&mut self.pi_context);
+            let (pi_context, record) = tokio::task::spawn_blocking(move || {
+                let record = decode_record(&bytes, &path, line_no, position, &mut pi_context);
+                (pi_context, record)
+            })
+            .await?;
+            self.pi_context = pi_context;
+            if let Some(record) = record? {
+                return Ok(Some(record));
             }
-            let value: serde_json::Value = serde_json::from_str(line).map_err(|error| {
-                anyhow::anyhow!("journal {}:{}: {error}", self.path.display(), self.line_no)
-            })?;
-            let record = if value
-                .get("_bt_record_type")
-                .and_then(serde_json::Value::as_str)
-                == Some("delivery_checkpoint")
-            {
-                let checkpoint: DeliveryCheckpointRecord =
-                    serde_json::from_value(value).map_err(|error| {
-                        anyhow::anyhow!("journal {}:{}: {error}", self.path.display(), self.line_no)
-                    })?;
-                JournalRecord::DeliveryCheckpoint {
-                    route: checkpoint.route,
-                    through: checkpoint.through,
-                }
-            } else if value
-                .get("_bt_record_type")
-                .and_then(serde_json::Value::as_str)
-                == Some("pipeline_control")
-            {
-                JournalRecord::Control(serde_json::from_value(value)?)
-            } else {
-                let mut event = serde_json::from_value(value).map_err(|error| {
-                    anyhow::anyhow!("journal {}:{}: {error}", self.path.display(), self.line_no)
-                })?;
-                expand_pi_payload(&mut event, &mut self.pi_context).map_err(|error| {
-                    anyhow::anyhow!("journal {}:{}: {error}", self.path.display(), self.line_no)
-                })?;
-                JournalRecord::Event(event)
-            };
-            return Ok(Some(JournalRecordEntry {
-                record,
-                through: self.position,
-            }));
         }
     }
 
@@ -461,6 +451,51 @@ impl JournalReader {
     }
 }
 
+fn decode_record(
+    bytes: &[u8],
+    path: &Path,
+    line_no: usize,
+    position: u64,
+    pi_context: &mut Vec<serde_json::Value>,
+) -> anyhow::Result<Option<JournalRecordEntry>> {
+    let line = std::str::from_utf8(bytes)
+        .map_err(|error| anyhow::anyhow!("journal {}:{}: {error}", path.display(), line_no))?
+        .trim();
+    if line.is_empty() {
+        return Ok(None);
+    }
+    let value: serde_json::Value = serde_json::from_str(line)
+        .map_err(|error| anyhow::anyhow!("journal {}:{}: {error}", path.display(), line_no))?;
+    let record = if value
+        .get("_bt_record_type")
+        .and_then(serde_json::Value::as_str)
+        == Some("delivery_checkpoint")
+    {
+        let checkpoint: DeliveryCheckpointRecord = serde_json::from_value(value)
+            .map_err(|error| anyhow::anyhow!("journal {}:{}: {error}", path.display(), line_no))?;
+        JournalRecord::DeliveryCheckpoint {
+            route: checkpoint.route,
+            through: checkpoint.through,
+        }
+    } else if value
+        .get("_bt_record_type")
+        .and_then(serde_json::Value::as_str)
+        == Some("pipeline_control")
+    {
+        JournalRecord::Control(serde_json::from_value(value)?)
+    } else {
+        let mut event = serde_json::from_value(value)
+            .map_err(|error| anyhow::anyhow!("journal {}:{}: {error}", path.display(), line_no))?;
+        expand_pi_payload(&mut event, pi_context)
+            .map_err(|error| anyhow::anyhow!("journal {}:{}: {error}", path.display(), line_no))?;
+        JournalRecord::Event(event)
+    };
+    Ok(Some(JournalRecordEntry {
+        record,
+        through: position,
+    }))
+}
+
 const PI_MESSAGES_DELTA: &str = "_bt_messages_delta";
 
 fn compact_pi_payload(event: &mut RedactedEnvelope, previous: &mut Vec<serde_json::Value>) {

bt-daemon/src/pipeline/input.rs

@@ -0,0 +1,121 @@
+//! Stage input shared by direct capture and recovery. Storage supplies only
+//! byte positions; original live payloads arrive from capture.
+pub use crate::journal::PipelineControl;
+use crate::wire::Envelope;
+use parking_lot::Mutex;
+use std::{collections::VecDeque, sync::Arc};
+
+const RECORD_BUDGET: usize = 64;
+const BYTE_BUDGET: u64 = 8 * 1024 * 1024;
+
+/// Durable byte range returned by the journal append.
+#[derive(Clone)]
+pub(crate) struct SourcePosition {
+    pub before: u64,
+    pub through: u64,
+}
+
+pub(crate) enum SourceEvent {
+    Event(Box<Envelope>),
+    Control(PipelineControl),
+    LegacyDeliveryCheckpoint,
+}
+
+pub(crate) struct SourceRecord {
+    pub event: SourceEvent,
+    pub position: SourcePosition,
+}
+
+impl SourceRecord {
+    fn bytes(&self) -> u64 {
+        self.position.through.saturating_sub(self.position.before)
+    }
+}
+
+#[derive(Default)]
+struct Buffer {
+    records: VecDeque<SourceRecord>,
+    bytes: u64,
+}
+
+#[derive(Clone, Default)]
+pub(crate) struct SourceInput(Arc<Mutex<Buffer>>);
+
+impl SourceInput {
+    pub(crate) fn push_event(&self, event: Envelope, position: SourcePosition) {
+        self.push(SourceRecord {
+            event: SourceEvent::Event(Box::new(event)),
+            position,
+        });
+    }
+    pub(crate) fn is_empty(&self) -> bool {
+        self.0.lock().records.is_empty()
+    }
+    pub(crate) fn push(&self, record: SourceRecord) {
+        let mut buffer = self.0.lock();
+        let bytes = record.bytes();
+        // An oversized event travels alone, just as it does on the SpanOp feed.
+        while !buffer.records.is_empty()
+            && (buffer.records.len() >= RECORD_BUDGET
+                || buffer.bytes.saturating_add(bytes) > BYTE_BUDGET)
+        {
+            if let Some(evicted) = buffer.records.pop_front() {
+                buffer.bytes -= evicted.bytes();
+            }
+        }
+        buffer.bytes = buffer.bytes.saturating_add(bytes);
+        buffer.records.push_back(record);
+    }
+
+    pub(crate) fn take(&self, position: u64, through: u64) -> Option<SourceRecord> {
+        let mut buffer = self.0.lock();
+        // Recovery may already have consumed records whose live copies remained.
+        while buffer
+            .records
+            .front()
+            .is_some_and(|record| record.position.through <= position)
+        {
+            if let Some(consumed) = buffer.records.pop_front() {
+                buffer.bytes -= consumed.bytes();
+            }
+        }
+        let first = buffer.records.front()?;
+        if first.position.before != position || first.position.through > through {
+            return None;
+        }
+        let record = buffer.records.pop_front()?;
+        buffer.bytes -= record.bytes();
+        Some(record)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[test]
+    fn oversized_events_flow_alone_and_eviction_exposes_a_recovery_gap() {
+        let input = SourceInput::default();
+        let event: Envelope = serde_json::from_value(json!({"source":"debug","session_id":"direct","event":"native","ts_ms":1,"payload":{"history":"x".repeat(9 * 1024 * 1024)}})).unwrap();
+        let stride = 9 * 1024 * 1024;
+        let position = |before| SourcePosition {
+            before,
+            through: before + stride,
+        };
+        input.push_event(event.clone(), position(0));
+        assert!(input.take(0, 0).is_none());
+        let SourceEvent::Event(actual) = input.take(0, stride).unwrap().event else {
+            panic!("expected original event");
+        };
+        assert_eq!(actual.payload, event.payload);
+        input.push_event(event.clone(), position(stride));
+        input.push_event(event, position(2 * stride));
+        assert!(
+            input.take(stride, 3 * stride).is_none(),
+            "eviction requires recovery of the gap"
+        );
+        assert!(input.take(2 * stride, 3 * stride).is_some());
+        assert!(input.is_empty());
+    }
+}