feat: timestamp-based queue

Lethe10137 · Lethe10137 · commit 372bfb740d8a · 2026-01-09T08:31:53.000Z
diff --git a/rattan-core/Cargo.toml b/rattan-core/Cargo.toml
@@ -30,7 +30,7 @@ bytesize = "2.0.1"
 camellia = { git = "https://github.com/minhuw/camellia.git", optional = true }
 clap = { workspace = true }
 csv = { workspace = true }
-derive_more = { workspace = true}
+derive_more = { workspace = true }
 dyn-clone = "1.0"
 etherparse = "0.18.0"
 figment = { workspace = true }
diff --git a/rattan-core/src/cells/bandwidth/mod.rs b/rattan-core/src/cells/bandwidth/mod.rs
@@ -6,13 +6,15 @@ use netem_trace::{model::BwTraceConfig, Bandwidth, BwTrace, Delay};
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
 use tokio::sync::mpsc;
+use tokio::sync::mpsc::error::TryRecvError;
 use tokio::time::Instant;
 use tracing::{debug, info, trace};
 
 use super::{
     AtomicCellState, Cell, CellState, ControlInterface, CurrentConfig, Egress, Ingress, Packet,
     LARGE_DURATION, TRACE_START_INSTANT,
 };
+use crate::cells::bandwidth::queue::AQM;
 use crate::error::Error;
 use crate::metal::timer::Timer;
 
@@ -91,7 +93,7 @@ where
     egress: mpsc::UnboundedReceiver<P>,
     bw_type: BwType,
     bandwidth: Bandwidth,
-    packet_queue: Q,
+    packet_queue: AQM<Q, P>,
     next_available: Instant,
     config_rx: mpsc::UnboundedReceiver<BwCellConfig<P, Q>>,
     timer: Timer,
@@ -143,6 +145,7 @@ where
     /// queue and thus being returnd, this function returns `Some`. When this happens, the caller
     /// should send out the packet immediately.
     #[inline(always)]
+    #[must_use]
     fn enqueue_packet(&mut self, new_packet: P) -> Option<P> {
         if let Some(packet_to_drop) = self.packet_queue.enqueue(new_packet) {
             let timestamp = packet_to_drop.get_timestamp();
@@ -155,7 +158,6 @@ where
         None
     }
 }
-
 #[async_trait]
 impl<P, Q> Egress<P> for BwCellEgress<P, Q>
 where
@@ -187,7 +189,29 @@ where
             }
         }
 
-        let mut packet = self.packet_queue.dequeue();
+        while self.packet_queue.need_more_packets(self.next_available) {
+            match self.egress.try_recv() {
+                Ok(new_packet) => {
+                    let new_packet = crate::check_cell_state!(self.state, new_packet);
+                    if let Some(packet) = self.enqueue_packet(new_packet) {
+                        return Some(packet);
+                    }
+                }
+                Err(TryRecvError::Empty) => {
+                    break;
+                }
+                Err(TryRecvError::Disconnected) => {
+                    return None;
+                }
+            }
+        }
+        // Here, either:
+        //    1) No more packets can be retrived from egress, or
+        //    2) A packet, that should enter the queue after `self.next_available` is seen.
+        // Thus the `dequeue_at()` see a correct queue, containing any packet that should
+        // enter the AQM at `self.next_available`.
+        let mut packet = self.packet_queue.dequeue_at(self.next_available);
+
         while packet.is_none() {
             // the queue is empty, wait for the next packet
             tokio::select! {
@@ -198,10 +222,11 @@ where
                 // `new_packet` can be None only if `self.egress` is closed.
                 new_packet = self.egress.recv() => {
                     let new_packet = crate::check_cell_state!(self.state, new_packet?);
+                    let timestamp = new_packet.get_timestamp();
                     if let Some(packet) = self.enqueue_packet(new_packet){
                         return Some(packet);
                     }
-                    packet = self.packet_queue.dequeue();
+                    packet = self.packet_queue.dequeue_at(timestamp);
                 }
             }
         }
@@ -377,7 +402,7 @@ where
                 egress: tx,
                 bw_type: bw_type.into().unwrap_or_default(),
                 bandwidth: bandwidth.into().unwrap_or(MAX_BANDWIDTH),
-                packet_queue,
+                packet_queue: AQM::new(packet_queue),
                 next_available: Instant::now(),
                 config_rx,
                 timer: Timer::new()?,
@@ -400,7 +425,7 @@ where
     egress: mpsc::UnboundedReceiver<P>,
     bw_type: BwType,
     trace: Box<dyn BwTrace>,
-    packet_queue: Q,
+    packet_queue: AQM<Q, P>,
     current_bandwidth: CurrentConfig<Bandwidth>,
     next_available: Instant,
     next_change: Instant,
@@ -478,11 +503,12 @@ where
         if let Some((bandwidth, duration)) = self.trace.next_bw() {
             self.change_bandwidth(bandwidth, change_time);
             self.next_change = change_time + duration;
-            #[cfg(test)]
-            trace!(
-                "Bandwidth changed to {:?}, next change after {:?}",
+            // #[cfg(test)]
+            eprintln!(
+                "Bandwidth changed to {:?}, next change after {:?}. now {:?}",
                 bandwidth,
-                self.next_change - Instant::now()
+                self.next_change - Instant::now(),
+                Instant::now(),
             );
             true
         } else {
@@ -556,7 +582,29 @@ where
             }
         }
 
-        let mut packet = self.packet_queue.dequeue();
+        while self.packet_queue.need_more_packets(self.next_available) {
+            match self.egress.try_recv() {
+                Ok(new_packet) => {
+                    let new_packet = crate::check_cell_state!(self.state, new_packet);
+                    if let Some(packet) = self.enqueue_packet(new_packet) {
+                        return Some(packet);
+                    }
+                }
+                Err(TryRecvError::Empty) => {
+                    break;
+                }
+                Err(TryRecvError::Disconnected) => {
+                    return None;
+                }
+            }
+        }
+
+        // Here, either:
+        //    1) No more packets can be retrived from egress, or
+        //    2) A packet, that should enter the queue after `self.next_available` is seen.
+        // Thus the `dequeue_at()` see a correct queue, containing any packet that should
+        // enter the AQM at `self.next_available`.
+        let mut packet = self.packet_queue.dequeue_at(self.next_available);
 
         while packet.is_none() {
             // the queue is empty, wait for the next packet
@@ -571,10 +619,11 @@ where
                 // `new_packet` can be None only if `self.egress` is closed.
                 new_packet = self.egress.recv() => {
                     let new_packet = crate::check_cell_state!(self.state, new_packet?);
+                    let timestamp = new_packet.get_timestamp();
                     if let Some(packet) = self.enqueue_packet(new_packet){
                         return Some(packet);
                     }
-                    packet = self.packet_queue.dequeue();
+                    packet = self.packet_queue.dequeue_at(timestamp);
                 }
             }
         }
@@ -606,6 +655,7 @@ where
     fn reset(&mut self) {
         self.next_available = *TRACE_START_INSTANT.get_or_init(Instant::now);
         self.next_change = *TRACE_START_INSTANT.get_or_init(Instant::now);
+        eprintln!("Reset to {:?}", self.next_change);
     }
 
     fn change_state(&self, state: CellState) {
@@ -770,7 +820,7 @@ where
                 egress: tx,
                 bw_type: bw_type.into().unwrap_or_default(),
                 trace,
-                packet_queue,
+                packet_queue: AQM::new(packet_queue),
                 current_bandwidth: CurrentConfig::default(),
                 next_available: Instant::now(),
                 next_change: Instant::now(),
diff --git a/rattan-core/src/cells/bandwidth/queue.rs b/rattan-core/src/cells/bandwidth/queue.rs
@@ -3,7 +3,10 @@ use std::fmt::Debug;
 
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
-use tokio::time::{Duration, Instant};
+use tokio::{
+    sync::mpsc::{error::TryRecvError, UnboundedReceiver},
+    time::{Duration, Instant},
+};
 use tracing::{debug, trace};
 
 use super::BwType;
@@ -14,6 +17,83 @@ fn serde_default<T: Default + PartialEq>(t: &T) -> bool {
     *t == Default::default()
 }
 
+pub enum PacketInboundTryReceiveError {
+    Empty,
+    Failed,
+}
+
+pub trait PacketInbound<P> {
+    fn try_receive(&mut self) -> Result<P, PacketInboundTryReceiveError>;
+}
+
+impl<P> PacketInbound<P> for UnboundedReceiver<P> {
+    fn try_receive(&mut self) -> Result<P, PacketInboundTryReceiveError> {
+        self.try_recv().map_err(|e| match e {
+            TryRecvError::Empty => PacketInboundTryReceiveError::Empty,
+            TryRecvError::Disconnected => PacketInboundTryReceiveError::Failed,
+        })
+    }
+}
+
+pub struct AQM<Q, P>
+where
+    Q: PacketQueue<P>,
+    P: Packet,
+{
+    inbound_buffer: VecDeque<P>,
+    queue: Q,
+    latest_enqueue_timestamp: Option<Instant>,
+}
+
+impl<Q, P> AQM<Q, P>
+where
+    Q: PacketQueue<P>,
+    P: Packet,
+{
+    pub fn new(queue: Q) -> Self {
+        Self {
+            inbound_buffer: VecDeque::with_capacity(1024),
+            queue,
+            latest_enqueue_timestamp: None,
+        }
+    }
+
+    pub fn configure(&mut self, config: Q::Config) {
+        self.queue.configure(config);
+    }
+
+    /// If this returns true, the caller should try to enqueue more packets.
+    pub fn need_more_packets(&self, next_available: Instant) -> bool {
+        self.latest_enqueue_timestamp
+            .is_none_or(|t| t <= next_available)
+    }
+
+    /// Return the packet immediately, if the inner queue is zero-buffered.
+    pub fn enqueue(&mut self, packet: P) -> Option<P> {
+        self.latest_enqueue_timestamp = packet.get_timestamp().into();
+        if self.queue.is_zero_buffer() {
+            packet.into()
+        } else {
+            self.inbound_buffer.push_back(packet);
+            None
+        }
+    }
+
+    /// The caller ensures that:
+    ///   1) This function is not called before the `timestamp` here.
+    ///   2) The timestamp should be non-decending.
+    pub fn dequeue_at(&mut self, timestamp: Instant) -> Option<P> {
+        while let Some(head) = self.inbound_buffer.front() {
+            if head.get_timestamp() <= timestamp {
+                self.queue.enqueue(self.inbound_buffer.pop_front().unwrap());
+            } else {
+                break;
+            }
+        }
+        self.queue.dequeue()
+    }
+}
+
 pub trait PacketQueue<P>: Send
 where
     P: Packet,
@@ -25,14 +105,18 @@ where
 
     fn configure(&mut self, config: Self::Config);
 
-    /// Returns the packet if there is not space for it.
-    fn enqueue(&mut self, packet: P) -> Option<P>;
+    fn enqueue(&mut self, packet: P);
 
     // If the queue is empty, return `None`
     fn dequeue(&mut self) -> Option<P>;
 
     fn is_empty(&self) -> bool;
 
+    // Returns if the buffer is zero-sized.
+    fn is_zero_buffer(&self) -> bool {
+        false
+    }
+
     // How this queue measures the size of a packet.
     // Should return 0 if it measures the size of a packet based on its L3 size.
     // Should return 14 if it measures that based on its L2 size (L3 length + 14 bytes L2 header).
@@ -102,9 +186,8 @@ where
 
     fn configure(&mut self, _config: Self::Config) {}
 
-    fn enqueue(&mut self, packet: P) -> Option<P> {
+    fn enqueue(&mut self, packet: P) {
         self.queue.push_back(packet);
-        None
     }
 
     fn dequeue(&mut self) -> Option<P> {
@@ -212,13 +295,12 @@ where
         self.bw_type = config.bw_type;
     }
 
-    fn enqueue(&mut self, packet: P) -> Option<P> {
-        if self.packet_limit.is_some_and(|limit| limit == 0)
+    fn is_zero_buffer(&self) -> bool {
+        self.packet_limit.is_some_and(|limit| limit == 0)
             || self.byte_limit.is_some_and(|limit| limit == 0)
-        {
-            return packet.into();
-        }
+    }
 
+    fn enqueue(&mut self, packet: P) {
         if self
             .packet_limit
             .is_none_or(|limit| self.queue.len() < limit)
@@ -237,7 +319,6 @@ where
                 "Drop packet(l3_len: {}, extra_len: {}) when enqueue", packet.l3_length(), self.bw_type.extra_length()
             );
         }
-        None
     }
 
     fn dequeue(&mut self) -> Option<P> {
@@ -351,13 +432,12 @@ where
         self.bw_type = config.bw_type;
     }
 
-    fn enqueue(&mut self, packet: P) -> Option<P> {
-        if self.packet_limit.is_some_and(|limit| limit == 0)
+    fn is_zero_buffer(&self) -> bool {
+        self.packet_limit.is_some_and(|limit| limit == 0)
             || self.byte_limit.is_some_and(|limit| limit == 0)
-        {
-            return packet.into();
-        }
+    }
 
+    fn enqueue(&mut self, packet: P) {
         self.now_bytes += packet.l3_length() + self.bw_type.extra_length();
         self.queue.push_back(packet);
         while self
@@ -374,7 +454,6 @@ where
                 "Drop packet(l3_len: {}, extra_len: {}) when enqueue another packet", _packet.l3_length(), self.bw_type.extra_length()
             );
         }
-        None
     }
 
     fn dequeue(&mut self) -> Option<P> {
@@ -554,13 +633,12 @@ where
         self.config = config;
     }
 
-    fn enqueue(&mut self, packet: P) -> Option<P> {
-        if self.config.packet_limit.is_some_and(|limit| limit == 0)
+    fn is_zero_buffer(&self) -> bool {
+        self.config.packet_limit.is_some_and(|limit| limit == 0)
             || self.config.byte_limit.is_some_and(|limit| limit == 0)
-        {
-            return packet.into();
-        }
+    }
 
+    fn enqueue(&mut self, packet: P) {
         if self
             .config
             .packet_limit
@@ -582,7 +660,6 @@ where
                 self.config.bw_type.extra_length()
             );
         }
-        None
     }
 
     fn dequeue(&mut self) -> Option<P> {
