rns-net 0.9.0

//! Coalescing buffer for complete, on-wire interface frames.

use std::collections::VecDeque;
use std::io::{self, Write};
use std::time::Duration;

pub(crate) const COALESCE_TARGET: usize = 65_536;
pub(crate) const EGRESS_MID_WATERMARK: usize = 128 * 1024;
pub(crate) const EGRESS_HIGH_WATERMARK: usize = 4 * 1024 * 1024;
const EGRESS_INTERVAL: Duration = Duration::from_secs(1);
const EGRESS_STALL_TICKS: u8 = 3;
const EGRESS_MAX_ETA_SECS: f64 = 10.0;
const EGRESS_RELEASE_ETA_SECS: f64 = 5.0;
pub(crate) const EGRESS_DEAD_TIME: Duration = Duration::from_secs(12);

#[derive(Debug, Default)]
pub(crate) struct EgressController {
    previous_sent: usize,
    zero_ticks: u8,
    last_evaluation: Duration,
    last_drain: Duration,
    stalled: bool,
}

impl EgressController {
    pub(crate) fn stalled(&self) -> bool {
        self.stalled
    }

    /// Evaluate one interval. Returns `true` when the peer made no progress
    /// for the dead-peer deadline and should be torn down.
    pub(crate) fn evaluate(
        &mut self,
        now: Duration,
        buffered: usize,
        sendable: usize,
        total_sent: usize,
    ) -> bool {
        if now.saturating_sub(self.last_evaluation) < EGRESS_INTERVAL {
            return false;
        }
        self.last_evaluation = now;
        let drained = total_sent.saturating_sub(self.previous_sent);
        self.previous_sent = total_sent;

        if buffered == 0 || sendable == 0 {
            self.zero_ticks = 0;
            self.last_drain = now;
            self.stalled = false;
            return false;
        }
        if now.saturating_sub(self.last_drain) >= EGRESS_DEAD_TIME {
            return true;
        }

        if drained > 0 {
            self.last_drain = now;
            self.zero_ticks = 0;
            let clear_eta = buffered as f64 / drained as f64;
            if buffered > EGRESS_MID_WATERMARK && clear_eta > EGRESS_MAX_ETA_SECS {
                self.stalled = true;
            } else if buffered <= EGRESS_MID_WATERMARK || clear_eta < EGRESS_RELEASE_ETA_SECS {
                self.stalled = false;
            }
        } else if buffered > EGRESS_MID_WATERMARK {
            self.zero_ticks = self.zero_ticks.saturating_add(1);
            if self.zero_ticks >= EGRESS_STALL_TICKS {
                self.stalled = true;
            }
        } else {
            self.zero_ticks = 0;
            self.stalled = false;
        }
        false
    }
}

struct Chunk {
    data: Vec<u8>,
    frames: usize,
}

/// Single-owner transmit buffer with bounded coalescing chunks.
///
/// The producer appends finalized frames. The consumer drains immutable
/// chunks, resuming at `head_offset` after a partial write.
pub(crate) struct TransmitBuffer {
    chunks: VecDeque<Chunk>,
    tail: Vec<u8>,
    tail_frames: usize,
    head_offset: usize,
    total_bytes: usize,
    visible_bytes: usize,
    total_frames: usize,
    sent_bytes: usize,
    sent_frames: usize,
}

impl TransmitBuffer {
    pub(crate) fn new() -> Self {
        Self {
            chunks: VecDeque::new(),
            tail: Vec::new(),
            tail_frames: 0,
            head_offset: 0,
            total_bytes: 0,
            visible_bytes: 0,
            total_frames: 0,
            sent_bytes: 0,
            sent_frames: 0,
        }
    }

    /// Append one finalized on-wire frame.
    #[allow(dead_code)]
    pub(crate) fn append(&mut self, frame: Vec<u8>) -> bool {
        self.append_with_limit(frame, None)
    }

    /// Append one frame unless doing so would exceed `limit` buffered bytes.
    pub(crate) fn append_with_limit(&mut self, frame: Vec<u8>, limit: Option<usize>) -> bool {
        let frame_len = frame.len();
        if limit.is_some_and(|limit| {
            self.buffered_bytes()
                .checked_add(frame_len)
                .is_none_or(|total| total > limit)
        }) {
            return false;
        }
        self.total_bytes = self.total_bytes.saturating_add(frame_len);
        self.total_frames = self.total_frames.saturating_add(1);

        if frame_len >= COALESCE_TARGET {
            self.flush_tail();
            self.visible_bytes = self.visible_bytes.saturating_add(frame_len);
            self.chunks.push_back(Chunk {
                data: frame,
                frames: 1,
            });
            return true;
        }

        if !self.tail.is_empty() && self.tail.len() + frame_len > COALESCE_TARGET {
            self.flush_tail();
        }
        self.tail.extend_from_slice(&frame);
        self.tail_frames = self.tail_frames.saturating_add(1);

        // Sparse traffic must become visible immediately.
        if self.chunks.is_empty() {
            self.flush_tail();
        }
        true
    }

    /// Make the current coalescing tail visible to the consumer.
    pub(crate) fn flush(&mut self) {
        self.flush_tail();
    }

    fn flush_tail(&mut self) {
        if self.tail.is_empty() {
            return;
        }
        let data = std::mem::take(&mut self.tail);
        let frames = std::mem::take(&mut self.tail_frames);
        self.visible_bytes = self.visible_bytes.saturating_add(data.len());
        self.chunks.push_back(Chunk { data, frames });
    }

    /// Drain all currently writable chunks, stopping cleanly on backpressure.
    pub(crate) fn drain_to(&mut self, writer: &mut impl Write) -> io::Result<usize> {
        let mut total_written = 0;
        loop {
            if self.chunks.is_empty() {
                self.flush_tail();
                if self.chunks.is_empty() {
                    break;
                }
            }

            let write_result = {
                let chunk = &self.chunks.front().expect("chunk queue is non-empty").data;
                if self.head_offset >= chunk.len() {
                    self.pop_head();
                    continue;
                }
                writer.write(&chunk[self.head_offset..])
            };

            match write_result {
                Ok(0) => break,
                Ok(written) => {
                    total_written += written;
                    self.sent_bytes = self.sent_bytes.saturating_add(written);
                    self.head_offset += written;
                    let chunk_len = self
                        .chunks
                        .front()
                        .expect("chunk remains queued")
                        .data
                        .len();
                    if self.head_offset >= chunk_len {
                        self.pop_head();
                    } else {
                        break;
                    }
                }
                Err(error) if error.kind() == io::ErrorKind::WouldBlock => break,
                Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
                Err(error) => return Err(error),
            }
        }
        Ok(total_written)
    }

    fn pop_head(&mut self) {
        if let Some(chunk) = self.chunks.pop_front() {
            self.sent_frames = self.sent_frames.saturating_add(chunk.frames);
        }
        self.head_offset = 0;
    }

    pub(crate) fn buffered_bytes(&self) -> usize {
        self.total_bytes.saturating_sub(self.sent_bytes)
    }

    pub(crate) fn sendable_bytes(&self) -> usize {
        self.visible_bytes.saturating_sub(self.sent_bytes)
    }

    pub(crate) fn total_sent_bytes(&self) -> usize {
        self.sent_bytes
    }

    #[allow(dead_code)]
    pub(crate) fn buffered_frames(&self) -> usize {
        self.total_frames.saturating_sub(self.sent_frames)
    }

    #[allow(dead_code)]
    pub(crate) fn buffered_chunks(&self) -> usize {
        self.chunks.len()
    }
}

impl Default for TransmitBuffer {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[derive(Default)]
    struct StepWriter {
        output: Vec<u8>,
        max_write: usize,
        block_next: bool,
    }

    impl Write for StepWriter {
        fn write(&mut self, data: &[u8]) -> io::Result<usize> {
            if self.block_next {
                self.block_next = false;
                return Err(io::Error::from(io::ErrorKind::WouldBlock));
            }
            let written = data.len().min(self.max_write.max(1));
            self.output.extend_from_slice(&data[..written]);
            Ok(written)
        }

        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    #[test]
    fn empty_and_sparse_state_is_immediately_sendable() {
        let mut buffer = TransmitBuffer::new();
        assert_eq!(buffer.buffered_bytes(), 0);
        assert_eq!(buffer.sendable_bytes(), 0);
        assert_eq!(buffer.buffered_frames(), 0);
        assert_eq!(buffer.buffered_chunks(), 0);

        buffer.append(vec![0x42; 200]);
        assert_eq!(buffer.buffered_bytes(), 200);
        assert_eq!(buffer.sendable_bytes(), 200);
        assert_eq!(buffer.buffered_frames(), 1);
        assert_eq!(buffer.buffered_chunks(), 1);
    }

    #[test]
    fn small_frames_coalesce_into_target_bounded_chunks() {
        let mut buffer = TransmitBuffer::new();
        for _ in 0..2_000 {
            buffer.append(vec![0x31; 200]);
        }

        assert_eq!(buffer.buffered_bytes(), 400_000);
        assert_eq!(buffer.buffered_frames(), 2_000);
        assert!(buffer.chunks.iter().all(|chunk| {
            !chunk.data.is_empty() && chunk.data.len() <= COALESCE_TARGET && chunk.frames > 0
        }));
        assert_eq!(
            buffer.sendable_bytes() + buffer.tail.len(),
            buffer.buffered_bytes()
        );
    }

    #[test]
    fn target_sized_and_larger_frames_get_dedicated_chunks() {
        let mut buffer = TransmitBuffer::new();
        let below = vec![0x31; COALESCE_TARGET - 1];
        let at = vec![0x32; COALESCE_TARGET];
        let above = vec![0x33; COALESCE_TARGET + 1];

        buffer.append(below.clone());
        buffer.append(at.clone());
        buffer.append(above.clone());

        let chunks: Vec<_> = buffer.chunks.iter().collect();
        assert_eq!(chunks.len(), 3);
        assert_eq!(chunks[0].data, below);
        assert_eq!(chunks[1].data, at);
        assert_eq!(chunks[2].data, above);
        assert!(chunks.iter().all(|chunk| chunk.frames == 1));
    }

    #[test]
    fn partial_head_resumes_and_would_block_is_not_an_error() {
        let expected = vec![0x55; COALESCE_TARGET * 2];
        let mut buffer = TransmitBuffer::new();
        buffer.append(expected.clone());
        let mut writer = StepWriter {
            max_write: 4_096,
            block_next: false,
            ..StepWriter::default()
        };

        assert_eq!(buffer.drain_to(&mut writer).unwrap(), 4_096);
        assert_eq!(buffer.head_offset, 4_096);
        writer.block_next = true;
        assert_eq!(buffer.drain_to(&mut writer).unwrap(), 0);
        while buffer.sendable_bytes() > 0 {
            buffer.drain_to(&mut writer).unwrap();
        }

        assert_eq!(writer.output, expected);
        assert_eq!(buffer.buffered_bytes(), 0);
        assert_eq!(buffer.buffered_frames(), 0);
    }

    #[test]
    fn explicit_and_automatic_flush_release_the_tail() {
        let mut buffer = TransmitBuffer::new();
        buffer.append(vec![1; 200]);
        buffer.append(vec![2; 200]);
        buffer.append(vec![3; 200]);
        assert_eq!(buffer.sendable_bytes(), 200);
        assert_eq!(buffer.buffered_bytes(), 600);

        buffer.flush();
        assert_eq!(buffer.sendable_bytes(), 600);
        assert_eq!(buffer.buffered_chunks(), 2);

        let mut second = TransmitBuffer::new();
        second.append(vec![1; 200]);
        second.append(vec![2; 200]);
        second.append(vec![3; 200]);
        let mut writer = StepWriter {
            max_write: usize::MAX,
            ..StepWriter::default()
        };
        assert_eq!(second.drain_to(&mut writer).unwrap(), 600);
        assert_eq!(
            writer.output,
            [vec![1; 200], vec![2; 200], vec![3; 200]].concat()
        );
    }

    #[test]
    fn fifty_thousand_frame_burst_preserves_order_and_accounting() {
        const FRAME_COUNT: usize = 50_000;
        let mut buffer = TransmitBuffer::new();
        let mut expected = Vec::with_capacity(FRAME_COUNT * 8);
        for index in 0..FRAME_COUNT {
            let frame = (index as u64).to_be_bytes().to_vec();
            expected.extend_from_slice(&frame);
            buffer.append(frame);
        }
        buffer.flush();

        let mut writer = StepWriter {
            max_write: usize::MAX,
            ..StepWriter::default()
        };
        assert_eq!(buffer.drain_to(&mut writer).unwrap(), expected.len());
        assert_eq!(writer.output, expected);
        assert_eq!(buffer.buffered_bytes(), 0);
        assert_eq!(buffer.sendable_bytes(), 0);
        assert_eq!(buffer.buffered_frames(), 0);
        assert_eq!(buffer.buffered_chunks(), 0);
    }

    #[test]
    fn fully_drained_accounting_reopens_capacity_for_the_next_burst() {
        let mut buffer = TransmitBuffer::new();
        let mut writer = StepWriter {
            max_write: usize::MAX,
            ..StepWriter::default()
        };

        for _ in 0..10 {
            buffer.append(vec![0x41; 202]);
        }
        buffer.flush();
        assert_eq!(buffer.buffered_bytes(), 2_020);
        assert_eq!(buffer.drain_to(&mut writer).unwrap(), 2_020);
        assert_eq!(buffer.buffered_bytes(), 0);
        assert_eq!(buffer.buffered_frames(), 0);

        for _ in 0..10 {
            buffer.append(vec![0x42; 202]);
        }
        assert_eq!(buffer.buffered_bytes(), 2_020);
        assert_eq!(buffer.buffered_frames(), 10);
    }

    #[test]
    fn byte_limit_is_a_hard_valve_and_reopens_after_drain() {
        let mut buffer = TransmitBuffer::new();
        let mut accepted_bytes = 0;
        let mut rejected = 0;
        for _ in 0..200 {
            for size in [502, 702] {
                if buffer.append_with_limit(vec![0x55; size], Some(4_096)) {
                    accepted_bytes += size;
                    assert!(buffer.buffered_bytes() <= 4_096);
                } else {
                    rejected += 1;
                    assert!(buffer.buffered_bytes() + size > 4_096);
                }
            }
        }
        assert!(rejected > 0);

        let mut writer = StepWriter {
            max_write: usize::MAX,
            ..StepWriter::default()
        };
        buffer.flush();
        assert_eq!(buffer.drain_to(&mut writer).unwrap(), accepted_bytes);
        assert_eq!(buffer.buffered_bytes(), 0);
        assert!(buffer.append_with_limit(vec![0x33; 2_048], Some(2_048)));
    }

    #[test]
    fn byte_limit_rejects_oversized_frame_without_changing_state() {
        let mut buffer = TransmitBuffer::new();
        assert!(buffer.append_with_limit(vec![0x11; 500], Some(2_048)));
        assert!(!buffer.append_with_limit(vec![0x22; 3_000], Some(2_048)));
        assert_eq!(buffer.buffered_bytes(), 500);
        assert_eq!(buffer.buffered_frames(), 1);

        assert!(buffer.append(vec![0x44; 3_000]));
        assert_eq!(buffer.buffered_bytes(), 3_500);
        assert_eq!(buffer.buffered_frames(), 2);
    }

    #[test]
    fn egress_controller_gates_after_three_zero_drain_ticks_and_releases_below_mid() {
        let mut controller = EgressController::default();
        let buffered = EGRESS_MID_WATERMARK + 1;
        assert!(!controller.evaluate(Duration::from_secs(1), buffered, buffered, 0));
        assert!(!controller.stalled());
        assert!(!controller.evaluate(Duration::from_secs(2), buffered, buffered, 0));
        assert!(!controller.stalled());
        assert!(!controller.evaluate(Duration::from_secs(3), buffered, buffered, 0));
        assert!(controller.stalled());

        assert!(!controller.evaluate(
            Duration::from_secs(4),
            EGRESS_MID_WATERMARK,
            EGRESS_MID_WATERMARK,
            1
        ));
        assert!(!controller.stalled());
    }

    #[test]
    fn egress_controller_applies_eta_hysteresis() {
        let mut controller = EgressController::default();
        let buffered = EGRESS_MID_WATERMARK * 2;

        assert!(!controller.evaluate(Duration::from_secs(1), buffered, buffered, 1));
        assert!(controller.stalled());

        // Eight seconds is inside the five-to-ten-second hysteresis band.
        assert!(!controller.evaluate(Duration::from_secs(2), buffered, buffered, 1 + buffered / 8));
        assert!(controller.stalled());

        assert!(!controller.evaluate(
            Duration::from_secs(3),
            buffered,
            buffered,
            1 + buffered / 8 + buffered
        ));
        assert!(!controller.stalled());
    }

    #[test]
    fn egress_controller_resets_empty_and_escalates_dead_peer() {
        let mut controller = EgressController::default();
        let buffered = EGRESS_MID_WATERMARK + 1;
        assert!(!controller.evaluate(Duration::from_secs(1), 0, 0, 0));
        assert!(!controller.stalled());

        assert!(!controller.evaluate(Duration::from_secs(2), buffered, buffered, 0));
        assert!(!controller.evaluate(Duration::from_secs(3), buffered, buffered, 0));
        assert!(!controller.evaluate(Duration::from_secs(4), buffered, buffered, 0));
        assert!(controller.stalled());
        assert!(controller.evaluate(Duration::from_secs(13), buffered, buffered, 0));
    }

    #[test]
    fn fresh_egress_state_does_not_immediately_reescalate_after_teardown() {
        let buffered = EGRESS_MID_WATERMARK + 1;
        {
            let mut stalled = EgressController::default();
            for tick in 1..=4 {
                stalled.evaluate(Duration::from_secs(tick), buffered, buffered, 0);
            }
            assert!(stalled.stalled());
        }

        // Teardown drops the stalled writer; reconnect builds a fresh writer
        // with an empty buffer and an unstalled controller, mirroring upstream's
        // BackboneInterface._dp_ec_reset (rgit `f0875ac4`). The replacement must
        // start clean and must not re-escalate on its first ticks.
        let mut fresh = EgressController::default();
        let fresh_buffer = TransmitBuffer::new();
        assert!(!fresh.stalled());
        assert_eq!(fresh_buffer.buffered_bytes(), 0);
        assert!(!fresh.evaluate(Duration::from_secs(1), 0, 0, 0));
        assert!(!fresh.stalled());
    }
}