use std::collections::VecDeque;
pub const DEFAULT_MAX_BUFFERED_DATAGRAMS: usize = 16;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "camelCase")]
pub struct PeerDatagramStats {
pub received: u64,
pub sent: u64,
pub bytes_received: u64,
pub bytes_sent: u64,
pub dropped_incoming: u64,
pub dropped_outgoing: u64,
pub expired_incoming: u64,
pub expired_outgoing: u64,
pub send_failures: u64,
pub incoming_buffered_datagrams: usize,
pub outgoing_buffered_datagrams: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OutgoingPeerDatagram {
pub payload: Vec<u8>,
pub remaining_max_age_ms: Option<u64>,
}
#[derive(Debug, Clone)]
struct QueuedDatagram {
payload: Vec<u8>,
queued_at_ms: u64,
}
#[derive(Debug)]
pub struct PeerDatagramPolicy {
max_datagram_size: usize,
incoming_max_age_ms: Option<u64>,
outgoing_max_age_ms: Option<u64>,
incoming_max_buffered: usize,
outgoing_max_buffered: usize,
incoming: VecDeque<QueuedDatagram>,
outgoing: VecDeque<QueuedDatagram>,
outgoing_in_flight: usize,
stats: PeerDatagramStats,
}
impl PeerDatagramPolicy {
pub fn new(max_datagram_size: usize) -> Self {
Self {
max_datagram_size: max_datagram_size.max(1),
incoming_max_age_ms: None,
outgoing_max_age_ms: None,
incoming_max_buffered: DEFAULT_MAX_BUFFERED_DATAGRAMS,
outgoing_max_buffered: DEFAULT_MAX_BUFFERED_DATAGRAMS,
incoming: VecDeque::new(),
outgoing: VecDeque::new(),
outgoing_in_flight: 0,
stats: PeerDatagramStats::default(),
}
}
pub fn set_incoming_max_age_ms(&mut self, value: Option<u64>, now_ms: u64) {
self.incoming_max_age_ms = value.filter(|value| *value > 0);
self.drop_expired_incoming(now_ms);
}
pub fn set_outgoing_max_age_ms(&mut self, value: Option<u64>) {
self.outgoing_max_age_ms = value.filter(|value| *value > 0);
}
pub fn set_incoming_max_buffered(&mut self, value: usize) -> anyhow::Result<()> {
if value == 0 {
anyhow::bail!("incoming datagram buffer limit must be positive");
}
self.incoming_max_buffered = value;
while self.incoming.len() > value {
self.incoming.pop_front();
self.stats.dropped_incoming = self.stats.dropped_incoming.saturating_add(1);
}
Ok(())
}
pub fn set_outgoing_max_buffered(&mut self, value: usize) -> anyhow::Result<()> {
if value == 0 {
anyhow::bail!("outgoing datagram buffer limit must be positive");
}
self.outgoing_max_buffered = value;
Ok(())
}
pub fn push_incoming(&mut self, payload: Vec<u8>, now_ms: u64) -> anyhow::Result<()> {
self.validate_payload(&payload)?;
self.stats.received = self.stats.received.saturating_add(1);
self.stats.bytes_received = self
.stats
.bytes_received
.saturating_add(payload.len() as u64);
self.drop_expired_incoming(now_ms);
while self.incoming.len() >= self.incoming_max_buffered {
self.incoming.pop_front();
self.stats.dropped_incoming = self.stats.dropped_incoming.saturating_add(1);
}
self.incoming.push_back(QueuedDatagram {
payload,
queued_at_ms: now_ms,
});
Ok(())
}
pub fn pop_incoming(&mut self, now_ms: u64) -> Option<Vec<u8>> {
self.drop_expired_incoming(now_ms);
self.incoming.pop_front().map(|entry| entry.payload)
}
pub fn push_outgoing(&mut self, payload: Vec<u8>, now_ms: u64) -> anyhow::Result<bool> {
self.validate_payload(&payload)?;
if self.outgoing.len() >= self.outgoing_max_buffered {
return Ok(false);
}
self.outgoing.push_back(QueuedDatagram {
payload,
queued_at_ms: now_ms,
});
Ok(true)
}
pub fn pop_outgoing(&mut self, now_ms: u64) -> Option<OutgoingPeerDatagram> {
while let Some(entry) = self.outgoing.pop_front() {
let age_ms = now_ms.saturating_sub(entry.queued_at_ms);
if self
.outgoing_max_age_ms
.is_some_and(|max_age_ms| age_ms > max_age_ms)
{
self.stats.expired_outgoing = self.stats.expired_outgoing.saturating_add(1);
continue;
}
self.outgoing_in_flight = self.outgoing_in_flight.saturating_add(1);
return Some(OutgoingPeerDatagram {
payload: entry.payload,
remaining_max_age_ms: self
.outgoing_max_age_ms
.map(|max_age_ms| max_age_ms.saturating_sub(age_ms)),
});
}
None
}
pub fn record_sent(&mut self, bytes: usize) {
self.outgoing_in_flight = self.outgoing_in_flight.saturating_sub(1);
self.stats.sent = self.stats.sent.saturating_add(1);
self.stats.bytes_sent = self.stats.bytes_sent.saturating_add(bytes as u64);
}
pub fn record_expired_outgoing(&mut self) {
self.outgoing_in_flight = self.outgoing_in_flight.saturating_sub(1);
self.stats.expired_outgoing = self.stats.expired_outgoing.saturating_add(1);
}
pub fn record_send_failure(&mut self) {
self.outgoing_in_flight = self.outgoing_in_flight.saturating_sub(1);
self.stats.send_failures = self.stats.send_failures.saturating_add(1);
}
pub fn stats(&mut self, now_ms: u64) -> PeerDatagramStats {
self.drop_expired_incoming(now_ms);
PeerDatagramStats {
incoming_buffered_datagrams: self.incoming.len(),
outgoing_buffered_datagrams: self.outgoing.len() + self.outgoing_in_flight,
..self.stats
}
}
pub fn close_incoming(&mut self) {
self.stats.dropped_incoming = self
.stats
.dropped_incoming
.saturating_add(self.incoming.len() as u64);
self.incoming.clear();
}
pub fn close_outgoing(&mut self) {
self.stats.dropped_outgoing = self
.stats
.dropped_outgoing
.saturating_add(self.outgoing.len() as u64);
self.outgoing.clear();
}
pub fn close(&mut self) {
self.close_incoming();
self.close_outgoing();
}
fn validate_payload(&self, payload: &[u8]) -> anyhow::Result<()> {
if payload.len() > self.max_datagram_size {
anyhow::bail!(
"peer datagram payload exceeds {} bytes",
self.max_datagram_size
);
}
Ok(())
}
fn drop_expired_incoming(&mut self, now_ms: u64) {
let Some(max_age_ms) = self.incoming_max_age_ms else {
return;
};
while self
.incoming
.front()
.is_some_and(|entry| now_ms.saturating_sub(entry.queued_at_ms) > max_age_ms)
{
self.incoming.pop_front();
self.stats.expired_incoming = self.stats.expired_incoming.saturating_add(1);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn bounds_and_expires_incoming_datagrams() {
let mut policy = PeerDatagramPolicy::new(4);
policy.set_incoming_max_buffered(2).unwrap();
policy.set_incoming_max_age_ms(Some(10), 0);
policy.push_incoming(vec![1], 0).unwrap();
policy.push_incoming(vec![2], 0).unwrap();
policy.push_incoming(vec![3], 0).unwrap();
assert_eq!(policy.stats(0).dropped_incoming, 1);
assert_eq!(policy.stats(0).incoming_buffered_datagrams, 2);
assert_eq!(policy.stats(11).expired_incoming, 2);
assert_eq!(policy.stats(11).incoming_buffered_datagrams, 0);
}
#[test]
fn applies_outgoing_backpressure_and_freshness() {
let mut policy = PeerDatagramPolicy::new(4);
policy.set_outgoing_max_buffered(1).unwrap();
policy.set_outgoing_max_age_ms(Some(10));
assert!(policy.push_outgoing(vec![1], 0).unwrap());
assert!(!policy.push_outgoing(vec![2], 0).unwrap());
assert!(policy.pop_outgoing(11).is_none());
assert_eq!(policy.stats(11).expired_outgoing, 1);
assert_eq!(policy.stats(11).outgoing_buffered_datagrams, 0);
}
#[test]
fn reports_in_flight_send_outcomes() {
let mut policy = PeerDatagramPolicy::new(4);
policy.push_outgoing(vec![1, 2], 0).unwrap();
let outgoing = policy.pop_outgoing(0).unwrap();
assert_eq!(outgoing.remaining_max_age_ms, None);
assert_eq!(policy.stats(0).outgoing_buffered_datagrams, 1);
policy.record_sent(outgoing.payload.len());
assert_eq!(policy.stats(0).sent, 1);
assert_eq!(policy.stats(0).bytes_sent, 2);
assert_eq!(policy.stats(0).outgoing_buffered_datagrams, 0);
}
}