refactor to always poll per mutation

Author: lherman-cs

pulsebeam-runtime/src/net/mod.rs

@@ -29,6 +29,8 @@ pub struct RecvPacketBatch {
     pub stride: usize,
     pub len: usize,
     pub transport: Transport,
+
+    offset: usize,
 }
 
 impl RecvPacketBatch {
@@ -38,42 +40,16 @@ impl RecvPacketBatch {
         &self.buf
     }
 
-    pub fn iter(&self) -> RecvPacketBatchIter<'_> {
-        RecvPacketBatchIter {
-            batch: self,
-            offset: 0,
-        }
-    }
-}
-
-impl<'a> IntoIterator for &'a RecvPacketBatch {
-    type Item = &'a [u8];
-    type IntoIter = RecvPacketBatchIter<'a>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        self.iter()
-    }
-}
-
-pub struct RecvPacketBatchIter<'a> {
-    batch: &'a RecvPacketBatch,
-    offset: usize,
-}
-
-impl<'a> Iterator for RecvPacketBatchIter<'a> {
-    type Item = &'a [u8];
-
-    fn next(&mut self) -> Option<Self::Item> {
-        assert!(
-            self.batch.stride != 0,
-            "stride must not be zero in iteration"
-        );
-        let tail = self.batch.len.min(self.offset + self.batch.stride);
+    // https://stackoverflow.com/questions/68606470/how-to-return-a-reference-when-implementing-an-iterator
+    // GAT is not stabilized, so it's not possible to return a reference in an iterator yet.
+    pub fn next_packet(&mut self) -> Option<&[u8]> {
+        assert!(self.stride != 0, "stride must not be zero in iteration");
+        let tail = self.len.min(self.offset + self.stride);
         if self.offset >= tail {
             return None;
         }
 
-        let buf = &self.batch.buf[self.offset..tail];
+        let buf = &self.buf[self.offset..tail];
         self.offset = tail;
         Some(buf)
     }
@@ -201,38 +177,38 @@ mod tests {
 
     #[test]
     fn recv_packet_batch_iter_yields_multiple_segments_without_off_by_one() {
-        let batch = RecvPacketBatch {
+        let mut batch = RecvPacketBatch {
             src: test_addr(),
             dst: test_addr(),
             buf: (0u8..20).collect(),
             stride: 6,
             len: 20,
             transport: Transport::Udp(UdpMode::Scalar),
+            offset: 0,
         };
 
-        let chunks: Vec<&[u8]> = batch.iter().collect();
-        assert_eq!(chunks.len(), 4);
-        assert_eq!(chunks[0], &[0, 1, 2, 3, 4, 5]);
-        assert_eq!(chunks[1], &[6, 7, 8, 9, 10, 11]);
-        assert_eq!(chunks[2], &[12, 13, 14, 15, 16, 17]);
-        assert_eq!(chunks[3], &[18, 19]);
+        assert_eq!(batch.next_packet(), Some(&[0, 1, 2, 3, 4, 5][..]));
+        assert_eq!(batch.next_packet(), Some(&[6, 7, 8, 9, 10, 11][..]));
+        assert_eq!(batch.next_packet(), Some(&[12, 13, 14, 15, 16, 17][..]));
+        assert_eq!(batch.next_packet(), Some(&[18, 19][..]));
+        assert_eq!(batch.next_packet(), None);
     }
 
     #[test]
     fn recv_packet_batch_iter_yields_multiple_segments_with_exact_len() {
-        let batch = RecvPacketBatch {
+        let mut batch = RecvPacketBatch {
             src: test_addr(),
             dst: test_addr(),
             buf: (0u8..20).collect(),
             stride: 6,
             len: 18,
             transport: Transport::Udp(UdpMode::Scalar),
+            offset: 0,
         };
 
-        let chunks: Vec<&[u8]> = batch.iter().collect();
-        assert_eq!(chunks.len(), 3);
-        assert_eq!(chunks[0], &[0, 1, 2, 3, 4, 5]);
-        assert_eq!(chunks[1], &[6, 7, 8, 9, 10, 11]);
-        assert_eq!(chunks[2], &[12, 13, 14, 15, 16, 17]);
+        assert_eq!(batch.next_packet(), Some(&[0, 1, 2, 3, 4, 5][..]));
+        assert_eq!(batch.next_packet(), Some(&[6, 7, 8, 9, 10, 11][..]));
+        assert_eq!(batch.next_packet(), Some(&[12, 13, 14, 15, 16, 17][..]));
+        assert_eq!(batch.next_packet(), None);
     }
 }

pulsebeam-runtime/src/net/tcp.rs

@@ -455,6 +455,7 @@ impl TcpTransport {
                     stride: len,
                     len,
                     transport: Transport::Tcp,
+                    offset: 0,
                 });
             }
             TcpEvent::Closed(src) => {

pulsebeam-runtime/src/net/udp.rs

@@ -158,6 +158,7 @@ impl UdpTransportReader {
                             stride: m.stride,  // Downstream will use this to skip through buf
                             len: m.len,        // Total length of the GRO batch
                             transport: Transport::Udp(UdpMode::Batch),
+                            offset: 0,
                         });
                     }
                     Ok(out.len() - prev_len)

pulsebeam-runtime/src/net/udp_scalar.rs

@@ -100,6 +100,7 @@ impl UdpTransportReader {
                     buf: slot,
                     stride: n,
                     len: n,
+                    offset: 0,
                 });
             }
             Err(err) => {

pulsebeam/src/participant/core.rs

@@ -280,72 +280,54 @@ impl ParticipantCore {
     }
 
     pub fn poll(&mut self, now: Instant, events: &mut EventQueue) {
-        loop {
-            let next_deadline = self.poll_until_deadline(now, events);
-            match next_deadline {
-                None => {
-                    events.exit();
-                    return;
-                }
-                Some(deadline) if deadline > now => {
-                    events.update_deadline(deadline);
-                    return;
-                }
-                Some(_) => {
-                    // Deadline is immediate, feed another Timeout and keep draining
-                    // https://discordapp.com/channels/1436725650302959829/1436796781462687817
-                    let _ = self.rtc.handle_input(Input::Timeout(now.into()));
-                }
-            }
-        }
-    }
+        'drain: loop {
+            let Some(rtc_deadline) = self.poll_rtc(events) else {
+                events.exit();
+                return;
+            };
 
-    fn poll_until_deadline(&mut self, now: Instant, events: &mut EventQueue) -> Option<Instant> {
-        if now >= self.last_slow_poll + SLOW_POLL_INTERVAL {
-            self.poll_slow(now, events);
-            self.last_slow_poll = now;
-        }
+            if now >= self.last_slow_poll + SLOW_POLL_INTERVAL {
+                self.poll_slow(now, events);
+                self.last_slow_poll = now;
+                continue;
+            }
 
-        while let Some(batch) = self.pending_ingress.pop_front() {
-            let transport = match batch.transport {
-                Transport::Udp(_) => str0m::net::Protocol::Udp,
-                Transport::Tcp => str0m::net::Protocol::Tcp,
-            };
+            while let Some(batch) = self.pending_ingress.front_mut() {
+                let transport = match batch.transport {
+                    Transport::Udp(_) => str0m::net::Protocol::Udp,
+                    Transport::Tcp => str0m::net::Protocol::Tcp,
+                };
 
-            for pkt in batch.into_iter() {
-                if let Ok(contents) = (*pkt).try_into() {
-                    let recv = str0m::net::Receive {
-                        proto: transport,
-                        source: batch.src,
-                        destination: batch.dst,
-                        contents,
-                    };
-                    if self
-                        .rtc
-                        .handle_input(Input::Receive(now.into(), recv))
-                        .is_err()
-                    {
-                        continue;
-                    }
+                let src = batch.src;
+                let dst = batch.dst;
+                let Some(pkt) = batch.next_packet() else {
+                    self.pending_ingress.pop_front();
+                    continue;
+                };
 
-                    self.poll_rtc(events)?;
-                } else {
+                let Ok(contents) = (*pkt).try_into() else {
                     tracing::warn!(src = %batch.src, "Dropping malformed UDP packet");
-                }
+                    // no point iterating the batch, this is already malicous
+                    self.pending_ingress.pop_front();
+                    continue;
+                };
+
+                let recv = str0m::net::Receive {
+                    proto: transport,
+                    source: src,
+                    destination: dst,
+                    contents,
+                };
+                let _ = self.rtc.handle_input(Input::Receive(now.into(), recv));
+                continue 'drain;
             }
-        }
 
-        loop {
-            let rtc_deadline = self.poll_rtc(events)?;
             let did_work = self.signaling.poll(&mut self.rtc, &self.downstream);
             if did_work {
-                // Signaling wrote data. The RTC engine is now "dirty" (has output to send).
-                // We loop back to `poll_rtc` immediately to flush `Output::Transmit`.
                 continue;
             }
 
             if self.downstream.dirty_allocation {
-                // Make sure rtc is updated with new allocations
                 let assignments_changed = self.downstream.update_allocations(&mut self.rtc.bwe());
                 if assignments_changed {
                     self.signaling.mark_assignments_dirty();
@@ -355,9 +337,14 @@ impl ParticipantCore {
             }
 
             let next_slow_poll = self.last_slow_poll + SLOW_POLL_INTERVAL;
+            let deadline = rtc_deadline.min(next_slow_poll);
+
+            if deadline > now {
+                events.update_deadline(deadline);
+                return;
+            }
 
-            // No new work generated. We are synced. Return the RTC deadline.
-            return Some(rtc_deadline.min(next_slow_poll));
+            let _ = self.rtc.handle_input(Input::Timeout(now.into()));
         }
     }