coven_storage/
transfer_budget.rs

1//! The handle's transfer limits and shared leases for response bodies.
2
3use coven_foundation::locks::lock_unpoisoned;
4use std::{
5    num::NonZeroUsize,
6    sync::{Arc, Mutex},
7};
8use tokio::sync::Notify;
9
10/// Limits shared by the open store's transfer consumers (E5).
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12pub struct TransferLimits {
13    /// Maximum file uploads running together in each lane, a part one
14    /// sends beside its first counting as one more.
15    pub uploads: NonZeroUsize,
16    /// Maximum requests for file bytes in flight together, each holding a
17    /// lease while its body is read; one file's pin takes as many as it can.
18    /// The sequence's own reads don't count against it.
19    pub downloads: NonZeroUsize,
20}
21impl Default for TransferLimits {
22    fn default() -> Self {
23        Self { uploads: NonZeroUsize::MIN, downloads: NonZeroUsize::MIN }
24    }
25}
26
27/// Compose once at the handle's root and pass to every connection and file
28/// owner. Existing response bodies retain their leases when a limit changes.
29pub struct TransferBudget {
30    state: Mutex<BudgetState>,
31    changed: Notify,
32    /// How many times a download has found the limit full and waited, for
33    /// a test.
34    #[cfg(any(test, feature = "test-utils"))]
35    waited: tokio::sync::watch::Sender<usize>,
36}
37struct BudgetState {
38    limits: TransferLimits,
39    downloads: usize,
40}
41impl TransferBudget {
42    /// An unused budget with the app's configured limits.
43    pub fn new(limits: TransferLimits) -> Self {
44        Self {
45            state: Mutex::new(BudgetState { limits, downloads: 0 }),
46            changed: Notify::new(),
47            #[cfg(any(test, feature = "test-utils"))]
48            waited: tokio::sync::watch::channel(0).0,
49        }
50    }
51    /// The current limits; the next upload to start checks the upload limit.
52    pub fn limits(&self) -> TransferLimits {
53        lock_unpoisoned(&self.state).limits
54    }
55    /// Change admission for subsequent downloads without cancelling an active
56    /// response. Reducing the limit waits for enough existing bodies to end.
57    pub fn set_limits(&self, limits: TransferLimits) {
58        lock_unpoisoned(&self.state).limits = limits;
59        self.changed.notify_waiters();
60    }
61    pub(crate) async fn download(self: &Arc<Self>) -> DownloadLease {
62        loop {
63            let changed = self.changed.notified();
64            tokio::pin!(changed);
65            changed.as_mut().enable();
66            {
67                let mut state = lock_unpoisoned(&self.state);
68                if state.downloads < state.limits.downloads.get() {
69                    state.downloads += 1;
70                    return DownloadLease(self.clone());
71                }
72            }
73            #[cfg(any(test, feature = "test-utils"))]
74            self.waited.send_modify(|waited| *waited += 1);
75            changed.await;
76        }
77    }
78
79    /// Wait until downloads have found the limit full and waited for a body
80    /// to end `count` more times than before this call.
81    #[cfg(any(test, feature = "test-utils"))]
82    pub fn downloads_waiting_from_now(
83        &self,
84        count: usize,
85    ) -> impl std::future::Future<Output = ()> + Send + 'static {
86        let mut waited = self.waited.subscribe();
87        let now = *waited.borrow_and_update();
88        async move {
89            coven_foundation::clock::within(waited.wait_for(|waited| *waited >= now + count))
90                .await
91                .expect("the downloads wait")
92                .expect("the budget is kept");
93        }
94    }
95}
96
97pub(crate) struct DownloadLease(Arc<TransferBudget>);
98impl Drop for DownloadLease {
99    fn drop(&mut self) {
100        lock_unpoisoned(&self.0.state).downloads -= 1;
101        self.0.changed.notify_waiters();
102    }
103}
104
105#[cfg(test)]
106#[path = "transfer_budget_tests.rs"]
107mod tests;
coven_storage/transfer_budget.rs

coven_storage/
transfer_budget.rs
