openrtc_file_transfer/

flow_control.rs

//! Shared pacing of application body bytes. One controller serves all bulk
//! executors in a runtime; admission, ACKs and lifecycle control bypass it.
//! Call `acquire` before starting transport I/O deadlines. Dropping the future
//! cancels a queued slice without leaving a peer slot or timer behind.
use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::{Arc, Mutex};
use tokio::sync::{oneshot, Notify};
use tokio::time::{Duration, Instant};

pub const MAX_BULK_BURST_BYTES: u64 = 64 * 1024;
const MAX_PENDING: usize = 4096;
const MAX_SAFE_INTEGER: u64 = 9_007_199_254_740_991;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BulkDirection {
    Upload,
    Download,
}
impl BulkDirection {
    fn index(self) -> usize {
        match self {
            Self::Upload => 0,
            Self::Download => 1,
        }
    }
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct BulkRates {
    pub upload_bytes_per_second: Option<u64>,
    pub download_bytes_per_second: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BulkFlowError {
    InvalidRate,
    InvalidSlice,
    QueueFull,
    OwnerClosed,
}
impl std::fmt::Display for BulkFlowError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            Self::InvalidRate => "Bulk rate must be unset or a positive safe integer",
            Self::InvalidSlice => "Invalid bulk body slice",
            Self::QueueFull => "receiver_busy: bulk flow queue is full",
            Self::OwnerClosed => "execution_suspended: bulk flow owner closed",
        })
    }
}
impl std::error::Error for BulkFlowError {}
struct Pending {
    id: u64,
    bytes: u64,
    scope: Option<String>,
    ready: oneshot::Sender<()>,
}
struct ScopedBudget {
    rate: u64,
    tokens: f64,
    updated: Instant,
}
impl ScopedBudget {
    fn refill(&mut self) {
        let now = Instant::now();
        self.tokens = (self.tokens
            + now.saturating_duration_since(self.updated).as_secs_f64() * self.rate as f64)
            .min(MAX_BULK_BURST_BYTES as f64);
        self.updated = now;
    }
}
struct DirectionState {
    rate: Option<u64>,
    tokens: f64,
    updated: Instant,
    peers: HashMap<String, VecDeque<Pending>>,
    order: VecDeque<String>,
    scopes: HashMap<String, ScopedBudget>,
}
impl DirectionState {
    fn new() -> Self {
        Self {
            rate: None,
            tokens: MAX_BULK_BURST_BYTES as f64,
            updated: Instant::now(),
            peers: HashMap::new(),
            order: VecDeque::new(),
            scopes: HashMap::new(),
        }
    }
    fn refill(&mut self) {
        let now = Instant::now();
        if let Some(rate) = self.rate {
            self.tokens = (self.tokens
                + now.saturating_duration_since(self.updated).as_secs_f64() * rate as f64)
                .min(MAX_BULK_BURST_BYTES as f64);
        }
        self.updated = now;
    }
    fn admit_scope(&mut self, key: &str, rate: u64) -> Result<(), BulkFlowError> {
        if let Some(scope) = self.scopes.get(key) {
            return if scope.rate == rate {
                Ok(())
            } else {
                Err(BulkFlowError::InvalidRate)
            };
        }
        // A queued slice owns its scope. Other scopes may retire only after
        // their full burst could refill, preserving debt across short retries.
        let active: HashSet<_> = self
            .peers
            .values()
            .flat_map(|queue| queue.iter())
            .filter_map(|slice| slice.scope.clone())
            .collect();
        self.scopes.retain(|key, scope| {
            active.contains(key)
                || scope.updated.elapsed().as_secs_f64() * (scope.rate as f64)
                    < MAX_BULK_BURST_BYTES as f64
        });
        if self.scopes.len() >= MAX_PENDING {
            return Err(BulkFlowError::QueueFull);
        }
        self.scopes.insert(
            key.into(),
            ScopedBudget {
                rate,
                tokens: MAX_BULK_BURST_BYTES as f64,
                updated: Instant::now(),
            },
        );
        Ok(())
    }
    fn pump(&mut self) -> Option<Duration> {
        self.refill();
        let mut blocked = 0;
        let mut wait: Option<Duration> = None;
        while !self.order.is_empty() && blocked < self.order.len() {
            let peer = self.order.pop_front().expect("flow peer order");
            let queue = self.peers.get_mut(&peer).expect("flow peer order");
            while queue.front().is_some_and(|slice| slice.ready.is_closed()) {
                queue.pop_front();
            }
            if queue.is_empty() {
                self.peers.remove(&peer);
                continue;
            }
            let front = queue.front().expect("pending flow slice");
            let bytes = front.bytes;
            let scope_key = front.scope.clone();
            let global_delay = self
                .rate
                .filter(|_| self.tokens < bytes as f64)
                .map(|rate| (bytes as f64 - self.tokens) / rate as f64);
            let mut scoped_delay = None;
            if let Some(key) = &scope_key {
                let scope = self.scopes.get_mut(key).expect("queued flow scope");
                scope.refill();
                if scope.tokens < bytes as f64 {
                    let scoped = (bytes as f64 - scope.tokens) / scope.rate as f64;
                    scoped_delay = Some(scoped);
                }
            }
            // An aggregate-ready peer keeps its turn while global tokens
            // refill. Skipping it would let smaller slices starve a large one.
            if scoped_delay.is_none() {
                if let Some(delay) = global_delay {
                    let delay = Duration::from_secs_f64(delay.max(0.000_000_001));
                    self.order.push_front(peer);
                    return Some(wait.map_or(delay, |prior| prior.min(delay)));
                }
            }
            let delay = match (global_delay, scoped_delay) {
                (Some(global), Some(scoped)) => Some(global.max(scoped)),
                (global, scoped) => global.or(scoped),
            };
            if let Some(delay) = delay {
                // Do not reserve aggregate tokens for a scope that cannot run.
                // Other ready peers retain their fair opportunity this turn.
                let delay = Duration::from_secs_f64(delay.max(0.000_000_001));
                wait = Some(wait.map_or(delay, |prior| prior.min(delay)));
                self.order.push_back(peer);
                blocked += 1;
                continue;
            }
            if self.rate.is_some() {
                self.tokens -= bytes as f64;
            }
            if let Some(key) = scope_key {
                self.scopes.get_mut(&key).expect("queued flow scope").tokens -= bytes as f64;
            }
            let slice = queue.pop_front().expect("pending flow slice");
            if queue.is_empty() {
                self.peers.remove(&peer);
            } else {
                self.order.push_back(peer);
            }
            let _ = slice.ready.send(());
            blocked = 0;
        }
        wait
    }
}
struct FlowState {
    directions: [DirectionState; 2],
    next_id: u64,
}
struct FlowInner {
    state: Mutex<FlowState>,
    changed: Notify,
}
/// Clone this handle, rather than constructing a controller per transfer.
#[derive(Clone)]
pub struct BulkFlowController {
    inner: Arc<FlowInner>,
}
struct WaitingSlice {
    owner: Arc<FlowInner>,
    direction: usize,
    peer: String,
    id: u64,
}
impl Drop for WaitingSlice {
    fn drop(&mut self) {
        let mut state = self
            .owner
            .state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let direction = &mut state.directions[self.direction];
        if let Some(queue) = direction.peers.get_mut(&self.peer) {
            queue.retain(|slice| slice.id != self.id);
            if queue.is_empty() {
                direction.peers.remove(&self.peer);
                direction.order.retain(|peer| peer != &self.peer);
            }
        }
        direction.pump();
        drop(state);
        self.owner.changed.notify_waiters();
    }
}
impl BulkFlowController {
    pub fn new(rates: BulkRates) -> Result<Self, BulkFlowError> {
        let result = Self {
            inner: Arc::new(FlowInner {
                state: Mutex::new(FlowState {
                    directions: [DirectionState::new(), DirectionState::new()],
                    next_id: 0,
                }),
                changed: Notify::new(),
            }),
        };
        result.configure(rates)?;
        Ok(result)
    }
    pub fn configure(&self, rates: BulkRates) -> Result<(), BulkFlowError> {
        let rates = [
            rates.upload_bytes_per_second,
            rates.download_bytes_per_second,
        ];
        if rates
            .into_iter()
            .flatten()
            .any(|rate| rate == 0 || rate > MAX_SAFE_INTEGER)
        {
            return Err(BulkFlowError::InvalidRate);
        }
        let mut state = self
            .inner
            .state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        for (direction, rate) in state.directions.iter_mut().zip(rates) {
            direction.refill();
            direction.rate = rate;
            direction.pump();
        }
        drop(state);
        self.inner.changed.notify_waiters();
        Ok(())
    }
    pub async fn acquire(
        &self,
        direction: BulkDirection,
        peer: &str,
        bytes: u64,
    ) -> Result<(), BulkFlowError> {
        self.acquire_inner(direction, peer, None, bytes).await
    }
    /// Admit a body slice against both aggregate and immutable scoped rates
    /// in one decision. Reuse a stable scope across attempts; controls bypass it.
    pub async fn acquire_scoped(
        &self,
        direction: BulkDirection,
        peer: &str,
        scope: &str,
        bytes_per_second: Option<u64>,
        bytes: u64,
    ) -> Result<(), BulkFlowError> {
        if scope.is_empty()
            || scope.len() > 128
            || scope.chars().any(char::is_control)
            || bytes_per_second.is_some_and(|rate| rate == 0 || rate > MAX_SAFE_INTEGER)
        {
            return Err(BulkFlowError::InvalidRate);
        }
        self.acquire_inner(
            direction,
            peer,
            bytes_per_second.map(|rate| (scope, rate)),
            bytes,
        )
        .await
    }
    async fn acquire_inner(
        &self,
        direction: BulkDirection,
        peer: &str,
        scope: Option<(&str, u64)>,
        bytes: u64,
    ) -> Result<(), BulkFlowError> {
        if peer.is_empty()
            || peer.len() > 512
            || peer.chars().any(char::is_control)
            || bytes > MAX_BULK_BURST_BYTES
        {
            return Err(BulkFlowError::InvalidSlice);
        }
        if bytes == 0 {
            return Ok(());
        }
        let direction = direction.index();
        let (send, mut receive) = oneshot::channel();
        let id = {
            let mut state = self
                .inner
                .state
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            let pending: usize = state
                .directions
                .iter()
                .flat_map(|state| state.peers.values())
                .map(VecDeque::len)
                .sum();
            if pending >= MAX_PENDING {
                return Err(BulkFlowError::QueueFull);
            }
            if let Some((key, rate)) = scope {
                state.directions[direction].admit_scope(key, rate)?;
            }
            let id = state
                .next_id
                .checked_add(1)
                .ok_or(BulkFlowError::QueueFull)?;
            state.next_id = id;
            let state = &mut state.directions[direction];
            if !state.peers.contains_key(peer) {
                state.order.push_back(peer.into());
            }
            state
                .peers
                .entry(peer.into())
                .or_default()
                .push_back(Pending {
                    id,
                    bytes,
                    scope: scope.map(|(key, _)| key.into()),
                    ready: send,
                });
            state.pump();
            id
        };
        let _lease = WaitingSlice {
            owner: self.inner.clone(),
            direction,
            peer: peer.into(),
            id,
        };
        loop {
            // Register before inspecting state so policy updates/cancellation
            // cannot be missed between choosing a deadline and sleeping.
            let changed = self.inner.changed.notified();
            tokio::pin!(changed);
            changed.as_mut().enable();
            let wait = {
                let mut state = self
                    .inner
                    .state
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner);
                state.directions[direction].pump()
            };
            tokio::select! {
                result = &mut receive => return result.map_err(|_| BulkFlowError::OwnerClosed),
                _ = &mut changed => {},
                _ = async { match wait { Some(delay) => tokio::time::sleep(delay).await,
                    None => std::future::pending().await } } => {},
            }
        }
    }
}

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

    #[tokio::test(start_paused = true)]
    async fn aggregate_ready_large_slice_keeps_its_turn_ahead_of_later_small_slices() {
        let flow = BulkFlowController::new(BulkRates {
            upload_bytes_per_second: Some(1024),
            ..BulkRates::default()
        })
        .unwrap();
        flow.acquire(BulkDirection::Upload, "initial", MAX_BULK_BURST_BYTES)
            .await
            .unwrap();
        let large = {
            let flow = flow.clone();
            tokio::spawn(async move {
                flow.acquire(BulkDirection::Upload, "large", MAX_BULK_BURST_BYTES)
                    .await
            })
        };
        tokio::task::yield_now().await;
        let small = {
            let flow = flow.clone();
            tokio::spawn(async move { flow.acquire(BulkDirection::Upload, "small", 1024).await })
        };
        tokio::task::yield_now().await;
        tokio::time::advance(Duration::from_secs(1)).await;
        tokio::task::yield_now().await;
        tokio::task::yield_now().await;
        assert!(
            !small.is_finished(),
            "Later small slices consumed credit owed to an aggregate-ready large peer"
        );
        assert!(!large.is_finished());
        tokio::time::advance(Duration::from_secs(63)).await;
        tokio::task::yield_now().await;
        tokio::task::yield_now().await;
        assert!(large.is_finished());
        large.await.unwrap().unwrap();
        assert!(!small.is_finished());
        tokio::time::advance(Duration::from_secs(1)).await;
        small.await.unwrap().unwrap();
    }

    // Native timers and dropped futures have different ownership from the JS
    // controller. Assert public grants against independently measured time.
    #[tokio::test(start_paused = true)]
    async fn parallel_clones_share_one_burst_and_rotate_peers_before_second_slices() {
        let flow = BulkFlowController::new(BulkRates {
            upload_bytes_per_second: Some(1024),
            download_bytes_per_second: Some(2048),
        })
        .unwrap();
        let start = Instant::now();
        flow.acquire(BulkDirection::Upload, "a", MAX_BULK_BURST_BYTES)
            .await
            .unwrap();
        let observed = Arc::new(Mutex::new(Vec::new()));
        let mut tasks = Vec::new();
        for (name, peer) in [("a1", "a"), ("a2", "a"), ("b1", "b")] {
            let (flow, observed) = (flow.clone(), observed.clone());
            tasks.push(tokio::spawn(async move {
                flow.acquire(BulkDirection::Upload, peer, MAX_BULK_BURST_BYTES)
                    .await
                    .unwrap();
                observed
                    .lock()
                    .unwrap()
                    .push((name, Instant::now().duration_since(start)));
            }));
            tokio::task::yield_now().await;
        }
        assert!(
            observed.lock().unwrap().is_empty(),
            "clones must not each get a fresh burst"
        );
        // Upload consumption cannot consume the independent download burst.
        flow.acquire(BulkDirection::Download, "b", MAX_BULK_BURST_BYTES)
            .await
            .unwrap();
        for _ in 0..3 {
            tokio::time::advance(Duration::from_secs(64)).await;
            tokio::task::yield_now().await;
            tokio::task::yield_now().await;
        }
        for task in tasks {
            task.await.unwrap();
        }
        let events = observed.lock().unwrap();
        assert_eq!(
            events.iter().map(|(name, _)| *name).collect::<Vec<_>>(),
            ["a1", "b1", "a2"]
        );
        for (index, (_, elapsed)) in events.iter().enumerate() {
            let attempted = (index as u64 + 2) * MAX_BULK_BURST_BYTES;
            assert!(attempted <= MAX_BULK_BURST_BYTES + elapsed.as_secs() * 1024);
        }
    }

    #[tokio::test(start_paused = true)]
    async fn cancelled_head_releases_waiters_and_live_limits_wake_only_their_direction() {
        let flow = BulkFlowController::new(BulkRates {
            upload_bytes_per_second: Some(1024),
            download_bytes_per_second: Some(1024),
        })
        .unwrap();
        for direction in [BulkDirection::Upload, BulkDirection::Download] {
            flow.acquire(direction, "initial", MAX_BULK_BURST_BYTES)
                .await
                .unwrap();
        }
        let head = {
            let flow = flow.clone();
            tokio::spawn(async move {
                flow.acquire(BulkDirection::Upload, "cancelled", MAX_BULK_BURST_BYTES)
                    .await
            })
        };
        tokio::task::yield_now().await;
        let smaller = {
            let flow = flow.clone();
            tokio::spawn(async move { flow.acquire(BulkDirection::Upload, "other", 1024).await })
        };
        tokio::task::yield_now().await;
        head.abort();
        let _ = head.await;
        tokio::time::advance(Duration::from_secs(1)).await;
        tokio::task::yield_now().await;
        assert!(
            smaller.is_finished(),
            "cancelled head cannot leave a 64-second deadline behind"
        );
        smaller.await.unwrap().unwrap();
        let upload = {
            let flow = flow.clone();
            tokio::spawn(async move {
                flow.acquire(BulkDirection::Upload, "a", MAX_BULK_BURST_BYTES)
                    .await
            })
        };
        let download = {
            let flow = flow.clone();
            tokio::spawn(async move {
                flow.acquire(BulkDirection::Download, "a", MAX_BULK_BURST_BYTES)
                    .await
            })
        };
        tokio::task::yield_now().await;
        assert!(!upload.is_finished());
        assert!(!download.is_finished());
        flow.configure(BulkRates {
            upload_bytes_per_second: None,
            download_bytes_per_second: Some(1024),
        })
        .unwrap();
        tokio::task::yield_now().await;
        assert!(upload.is_finished());
        assert!(!download.is_finished());
        upload.await.unwrap().unwrap();
        download.abort();
        let _ = download.await;
        assert_eq!(
            flow.acquire(BulkDirection::Upload, "a", MAX_BULK_BURST_BYTES + 1)
                .await,
            Err(BulkFlowError::InvalidSlice)
        );
        assert_eq!(
            flow.configure(BulkRates {
                upload_bytes_per_second: Some(0),
                download_bytes_per_second: None
            }),
            Err(BulkFlowError::InvalidRate)
        );
    }
}