handle stale keyframe

Author: lherman-cs

pulsebeam-runtime/src/collections/ring.rs

@@ -30,6 +30,15 @@ impl<T: Clone> RingBuffer<T> {
         self.tail
     }
 
+    /// Returns the number of sequence number slots between the tail and the head.
+    /// This represents the span of the buffer, including any empty (lost) packets.
+    pub fn len(&self) -> u64 {
+        if !self.initialized {
+            return 0;
+        }
+        self.head.wrapping_sub(self.tail)
+    }
+
     fn initialize(&mut self, seq: u64) {
         self.head = seq.wrapping_add(1);
         self.tail = seq;
@@ -479,4 +488,33 @@ mod tests {
         assert_eq!(buffer.tail(), 10);
         assert_eq!(buffer.head(), 13);
     }
+
+    #[test]
+    fn test_len() {
+        let mut buffer = RingBuffer::<i32>::new(10);
+        assert_eq!(buffer.len(), 0);
+
+        buffer.insert(100, 1); // tail=100, head=101
+        assert_eq!(buffer.len(), 1);
+
+        buffer.insert(101, 2); // tail=100, head=102
+        assert_eq!(buffer.len(), 2);
+
+        buffer.insert(105, 6); // tail=100, head=106
+        assert_eq!(buffer.len(), 6); // Length is the span, not the count of items
+
+        let _ = buffer.pop_front(); // tail=101, head=106
+        assert_eq!(buffer.len(), 5);
+
+        // Test wrap-around
+        let mut buffer = RingBuffer::<i32>::new(10);
+        let start_seq = u64::MAX - 1;
+        buffer.insert(start_seq, 1);
+        buffer.insert(start_seq.wrapping_add(1), 2); // u64::MAX
+        buffer.insert(start_seq.wrapping_add(2), 3); // 0
+
+        assert_eq!(buffer.tail(), start_seq);
+        assert_eq!(buffer.head(), start_seq.wrapping_add(3));
+        assert_eq!(buffer.len(), 3);
+    }
 }

pulsebeam/src/rtp/buffer.rs

@@ -1,6 +1,7 @@
 use crate::rtp::RtpPacket;
 use pulsebeam_runtime::collections::ring::RingBuffer;
 use str0m::rtp::SeqNo;
+use tokio::time::Instant;
 
 const KEYFRAME_BUFFER_CAPACITY: usize = 256;
 
@@ -32,7 +33,7 @@ pub struct KeyframeBuffer {
     ring: RingBuffer<RtpPacket>,
     state: State,
     /// The sequence number of the keyframe that we are currently buffering or draining.
-    keyframe_start_seq: Option<SeqNo>,
+    keyframe_segment: Option<(SeqNo, Instant)>,
     /// The next sequence number we expect to pop. This is used to handle out-of-order
     /// packets and to detect and skip gaps.
     next_pop_seq: Option<SeqNo>,
@@ -49,14 +50,19 @@ impl KeyframeBuffer {
         Self {
             ring: RingBuffer::new(KEYFRAME_BUFFER_CAPACITY),
             state: State::Waiting,
-            keyframe_start_seq: None,
+            keyframe_segment: None,
             next_pop_seq: None,
         }
     }
 
     /// Returns `true` if a keyframe has been received and the buffer is ready to be drained.
-    pub fn is_ready(&self) -> bool {
-        self.state == State::Buffering || self.state == State::Draining
+    pub fn is_ready(&self, target_playout: Instant) -> bool {
+        self.state == State::Buffering
+            && self
+                .keyframe_segment
+                .map(|s| s.1 >= target_playout)
+                .unwrap_or_default()
+            || self.state == State::Draining
     }
 
     /// Adds an RTP packet to the buffer, applying state-specific logic.
@@ -65,12 +71,15 @@ impl KeyframeBuffer {
             // A new keyframe has arrived. We must handle it based on our current state.
             match self.state {
                 State::Waiting | State::Buffering => {
-                    // If we were waiting for a keyframe, or buffering an old one that
-                    // the consumer hasn't started reading yet, we reset to this new keyframe.
-                    self.ring.advance_tail_to(*pkt.seq_no); // Evict any older packets.
-                    self.keyframe_start_seq = Some(pkt.seq_no);
-                    self.next_pop_seq = Some(pkt.seq_no);
-                    self.state = State::Buffering;
+                    if *pkt.seq_no > self.ring.head() {
+                        // If we were waiting for a keyframe, or buffering an old one that
+                        // the consumer hasn't started reading yet, we reset to this new keyframe.
+                        self.ring.advance_tail_to(*pkt.seq_no); // Evict any older packets.
+                        self.keyframe_segment
+                            .replace((pkt.seq_no, pkt.playout_time));
+                        self.next_pop_seq = Some(pkt.seq_no);
+                        self.state = State::Buffering;
+                    }
                 }
                 State::Draining => {
                     // We are already draining a keyframe. We will buffer this new one,
@@ -113,8 +122,8 @@ impl KeyframeBuffer {
         let mut next_seq = self.next_pop_seq?;
 
         // Search for the next available packet, starting from `next_pop_seq`.
-        // We limit the search to the buffer's capacity to prevent infinite loops.
-        for _ in 0..KEYFRAME_BUFFER_CAPACITY {
+        // We limit the search to the buffer's len to prevent infinite loops.
+        for _ in 0..self.ring.len() {
             if let Some(pkt) = self.ring.remove(*next_seq) {
                 // Found the next packet. Update our expected sequence and return it.
                 self.next_pop_seq = Some(next_seq.wrapping_add(1).into());

pulsebeam/src/rtp/switcher.rs

@@ -1,4 +1,5 @@
 use str0m::media::Frequency;
+use tokio::time::Instant;
 
 use crate::rtp::RtpPacket;
 use crate::rtp::buffer::KeyframeBuffer;
@@ -16,6 +17,8 @@ pub struct Switcher {
     /// The state for the *new* stream we are switching to.
     /// This is `Some` only when a switch is in progress.
     staging: Option<KeyframeBuffer>,
+
+    latest_playout: Instant,
 }
 
 impl Switcher {
@@ -24,49 +27,54 @@ impl Switcher {
             timeline: Timeline::new(clock_rate),
             pending: None,
             staging: None,
+            latest_playout: Instant::now(),
         }
     }
 
     /// Pushes a packet from the **old/current** stream.
     /// This is typically used to forward the stream that is already playing out.
     pub fn push(&mut self, pkt: RtpPacket) {
-        // rewrite active stream here so we don't race with the new stream
-        let pkt = self.timeline.rewrite(pkt);
         self.pending.replace(pkt);
     }
 
     /// Pushes a packet for the **new** stream we are preparing to switch to.
     /// The first call to this method will initiate the switching process.
     pub fn stage(&mut self, pkt: RtpPacket) {
         let staging = self.staging.get_or_insert_default();
-
-        if pkt.is_keyframe_start {
-            self.timeline.rebase(&pkt);
-        }
         staging.push(pkt);
     }
 
     /// Returns true if the new stream has received a keyframe and is ready to be popped.
     pub fn is_ready(&self) -> bool {
-        self.staging.as_ref().map(|s| s.is_ready()).unwrap_or(false)
+        self.staging
+            .as_ref()
+            .map(|s| s.is_ready(self.latest_playout))
+            .unwrap_or(false)
     }
 
     /// Pops the next available packet, prioritizing the old stream to ensure a smooth drain.
     pub fn pop(&mut self) -> Option<RtpPacket> {
         // --- Priority 1: Drain the pending packet from the OLD stream. ---
         if let Some(pending_pkt) = self.pending.take() {
-            return Some(pending_pkt);
+            if pending_pkt.playout_time > self.latest_playout {
+                self.latest_playout = pending_pkt.playout_time;
+            }
+            return Some(self.timeline.rewrite(pending_pkt));
+        }
+
+        if !self.is_ready() {
+            return None;
         }
 
-        // TODO: deal with older frames here.
         // --- Priority 2: Pop packets from the NEW stream if a switch is in progress. ---
         if let Some(staging) = &mut self.staging {
             if let Some(staged_pkt) = staging.pop() {
+                if staged_pkt.is_keyframe_start {
+                    self.timeline.rebase(&staged_pkt);
+                }
                 return Some(self.timeline.rewrite(staged_pkt));
             } else {
-                if staging.is_ready() {
-                    self.staging = None;
-                }
+                self.staging = None;
                 return None;
             }
         }