Skip to main content

syncweb_core/
transfer_limits.rs

1//! Rolling-window transfer quotas and a queue of content waiting for budget.
2//!
3//! iroh-blobs has no per-request rate limiter and iroh-docs downloads eagerly,
4//! so instead of trying to throttle a transfer in flight this module implements
5//! a *download policy* in the spirit of qBittorrent: transfers are queued and
6//! only start when their full size fits inside every configured per-hour,
7//! per-day, and per-month byte quota. Sharing (upload) is bounded by pausing
8//! the folder's live session when an upload quota would be exceeded.
9
10use std::{
11    collections::VecDeque,
12    path::PathBuf,
13    sync::{Mutex, RwLock},
14    time::{Duration, Instant},
15};
16
17use iroh_blobs::Hash;
18use iroh_docs::NamespaceId;
19use serde::{Deserialize, Serialize};
20
21use crate::error::Result;
22use crate::schedule::parse_byte_size;
23
24/// Rolling window for the per-hour quota.
25pub const WINDOW_HOUR: Duration = Duration::from_hours(1);
26/// Rolling window for the per-day quota.
27pub const WINDOW_DAY: Duration = Duration::from_hours(24);
28/// Rolling window for the per-month quota (30 days).
29pub const WINDOW_MONTH: Duration = Duration::from_hours(720);
30
31/// Which side of a transfer a quota applies to.
32#[derive(Clone, Copy, Debug, Eq, PartialEq)]
33#[non_exhaustive]
34pub enum Direction {
35    Download,
36    Upload,
37}
38
39/// Parsed byte quotas for one direction. `None` means unlimited.
40#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
41#[non_exhaustive]
42pub struct TransferQuota {
43    pub per_hour: Option<u64>,
44    pub per_day: Option<u64>,
45    pub per_month: Option<u64>,
46}
47
48impl TransferQuota {
49    #[must_use]
50    pub const fn is_unlimited(&self) -> bool {
51        self.per_hour.is_none() && self.per_day.is_none() && self.per_month.is_none()
52    }
53
54    /// The most restrictive configured cap, if any.
55    #[must_use]
56    pub fn tightest_cap(&self) -> Option<u64> {
57        [self.per_hour, self.per_day, self.per_month]
58            .into_iter()
59            .flatten()
60            .min()
61    }
62}
63
64/// Serializable `[transfer_limits]` configuration.
65#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
66#[serde(deny_unknown_fields)]
67#[non_exhaustive]
68pub struct TransferLimitsConfig {
69    #[serde(default, skip_serializing_if = "Option::is_none")]
70    pub max_download_per_hour: Option<String>,
71    #[serde(default, skip_serializing_if = "Option::is_none")]
72    pub max_download_per_day: Option<String>,
73    #[serde(default, skip_serializing_if = "Option::is_none")]
74    pub max_download_per_month: Option<String>,
75    #[serde(default, skip_serializing_if = "Option::is_none")]
76    pub max_upload_per_hour: Option<String>,
77    #[serde(default, skip_serializing_if = "Option::is_none")]
78    pub max_upload_per_day: Option<String>,
79    #[serde(default, skip_serializing_if = "Option::is_none")]
80    pub max_upload_per_month: Option<String>,
81}
82
83impl TransferLimitsConfig {
84    #[must_use]
85    pub const fn is_empty(&self) -> bool {
86        self.max_download_per_hour.is_none()
87            && self.max_download_per_day.is_none()
88            && self.max_download_per_month.is_none()
89            && self.max_upload_per_hour.is_none()
90            && self.max_upload_per_day.is_none()
91            && self.max_upload_per_month.is_none()
92    }
93
94    /// Validate and store one `transfer_limits.*` key.
95    ///
96    /// # Errors
97    ///
98    /// Returns an error if the key is not a recognised transfer-limit key or
99    /// the value is not a valid byte count.
100    pub fn set(&mut self, key: &str, value: &str) -> Result<bool> {
101        let slot = match key {
102            "transfer_limits.max_download_per_hour" => &mut self.max_download_per_hour,
103            "transfer_limits.max_download_per_day" => &mut self.max_download_per_day,
104            "transfer_limits.max_download_per_month" => &mut self.max_download_per_month,
105            "transfer_limits.max_upload_per_hour" => &mut self.max_upload_per_hour,
106            "transfer_limits.max_upload_per_day" => &mut self.max_upload_per_day,
107            "transfer_limits.max_upload_per_month" => &mut self.max_upload_per_month,
108            _ => return Ok(false),
109        };
110        parse_byte_size(value)?;
111        *slot = Some(value.to_owned());
112        Ok(true)
113    }
114}
115
116/// Parsed per-direction quotas.
117#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
118#[non_exhaustive]
119pub struct TransferLimits {
120    pub download: TransferQuota,
121    pub upload: TransferQuota,
122}
123
124impl TransferLimits {
125    /// Parse a configuration into quotas.
126    ///
127    /// # Errors
128    ///
129    /// Returns an error if any configured value is not a valid byte count.
130    pub fn from_config(config: &TransferLimitsConfig) -> Result<Self> {
131        let quota = |hour: &Option<String>, day: &Option<String>, month: &Option<String>| -> Result<TransferQuota> {
132            Ok(TransferQuota {
133                per_hour: hour.as_deref().map(parse_byte_size).transpose()?.flatten(),
134                per_day: day.as_deref().map(parse_byte_size).transpose()?.flatten(),
135                per_month: month.as_deref().map(parse_byte_size).transpose()?.flatten(),
136            })
137        };
138        Ok(Self {
139            download: quota(
140                &config.max_download_per_hour,
141                &config.max_download_per_day,
142                &config.max_download_per_month,
143            )?,
144            upload: quota(
145                &config.max_upload_per_hour,
146                &config.max_upload_per_day,
147                &config.max_upload_per_month,
148            )?,
149        })
150    }
151
152    #[must_use]
153    pub const fn is_unlimited(&self) -> bool {
154        self.download.is_unlimited() && self.upload.is_unlimited()
155    }
156
157    #[must_use]
158    pub const fn quota(&self, direction: Direction) -> TransferQuota {
159        match direction {
160            Direction::Download => self.download,
161            Direction::Upload => self.upload,
162        }
163    }
164}
165
166#[derive(Clone, Copy, Debug)]
167struct Sample {
168    at: Instant,
169    direction: Direction,
170    bytes: u64,
171}
172
173/// Rolling byte accounting against a set of [`TransferLimits`].
174#[derive(Debug)]
175pub struct TransferBudget {
176    limits: RwLock<TransferLimits>,
177    samples: Mutex<VecDeque<Sample>>,
178}
179
180impl TransferBudget {
181    #[must_use]
182    pub const fn new(limits: TransferLimits) -> Self {
183        Self {
184            limits: RwLock::new(limits),
185            samples: Mutex::new(VecDeque::new()),
186        }
187    }
188
189    #[must_use]
190    pub fn limits(&self) -> TransferLimits {
191        *self.limits.read().unwrap_or_else(std::sync::PoisonError::into_inner)
192    }
193
194    /// Replace the active quotas without discarding recorded usage.
195    pub fn set_limits(&self, limits: TransferLimits) {
196        *self.limits.write().unwrap_or_else(std::sync::PoisonError::into_inner) = limits;
197    }
198
199    /// Bytes that may still be transferred for `direction` right now.
200    #[must_use]
201    pub fn available(&self, direction: Direction) -> u64 {
202        self.available_at(direction, Instant::now())
203    }
204
205    /// Bytes that may still be transferred for `direction` at `now`.
206    ///
207    /// The result is the smallest remaining allowance across the configured
208    /// rolling windows, or [`u64::MAX`] when the direction is unlimited.
209    #[must_use]
210    pub fn available_at(&self, direction: Direction, now: Instant) -> u64 {
211        let quota = self.limits().quota(direction);
212        let mut available = u64::MAX;
213        for (cap, window) in [
214            (quota.per_hour, WINDOW_HOUR),
215            (quota.per_day, WINDOW_DAY),
216            (quota.per_month, WINDOW_MONTH),
217        ] {
218            if let Some(limit) = cap {
219                let used = self.used_at(direction, window, now);
220                available = available.min(limit.saturating_sub(used));
221            }
222        }
223        available
224    }
225
226    /// Whether `bytes` can be transferred for `direction` without breaking a quota.
227    #[must_use]
228    pub fn can_transfer(&self, direction: Direction, bytes: u64) -> bool {
229        bytes <= self.available(direction)
230    }
231
232    /// Record transferred bytes.
233    pub fn record(&self, direction: Direction, bytes: u64) {
234        self.record_at(direction, bytes, Instant::now());
235    }
236
237    /// Record transferred bytes observed at `now`.
238    pub fn record_at(&self, direction: Direction, bytes: u64, now: Instant) {
239        let cutoff = now.checked_sub(WINDOW_MONTH).unwrap_or(now);
240        let mut samples = self.samples.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
241        samples.push_back(Sample {
242            at: now,
243            direction,
244            bytes,
245        });
246        while samples.front().is_some_and(|sample| sample.at < cutoff) {
247            samples.pop_front();
248        }
249        drop(samples);
250    }
251
252    /// Bytes transferred for `direction` within `window` ending at `now`.
253    #[must_use]
254    pub fn used_at(&self, direction: Direction, window: Duration, now: Instant) -> u64 {
255        let samples = self.samples.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
256        samples
257            .iter()
258            .filter(|sample| sample.direction == direction)
259            .filter(|sample| now.saturating_duration_since(sample.at) < window)
260            .fold(0_u64, |total, sample| total.saturating_add(sample.bytes))
261    }
262
263    /// Total bytes recorded in all windows (diagnostics).
264    #[must_use]
265    pub fn recorded_total(&self) -> u64 {
266        let samples = self.samples.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
267        samples
268            .iter()
269            .fold(0_u64, |total, sample| total.saturating_add(sample.bytes))
270    }
271}
272
273/// A blob waiting for download budget.
274#[derive(Clone, Debug, Eq, PartialEq)]
275#[non_exhaustive]
276pub struct PendingTransfer {
277    pub namespace: NamespaceId,
278    pub hash: Hash,
279    pub size: u64,
280    pub path: PathBuf,
281}
282
283impl PendingTransfer {
284    #[must_use]
285    pub const fn new(namespace: NamespaceId, hash: Hash, size: u64, path: PathBuf) -> Self {
286        Self {
287            namespace,
288            hash,
289            size,
290            path,
291        }
292    }
293}
294
295/// Size-prioritized queue of blobs waiting for their download quota.
296///
297/// Shorter blobs are started first when they fit the remaining budget. This
298/// improves time-to-first-content for mixed workloads while leaving the quota
299/// accounting unchanged.
300#[derive(Debug, Default)]
301pub struct TransferQueue {
302    inner: Mutex<VecDeque<PendingTransfer>>,
303}
304
305impl TransferQueue {
306    #[must_use]
307    pub fn new() -> Self {
308        Self::default()
309    }
310
311    /// Add a transfer unless the same `(namespace, hash)` is already queued.
312    ///
313    /// Returns `true` when the item was newly added.
314    pub fn enqueue(&self, item: PendingTransfer) -> bool {
315        let mut queue = self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
316        if queue
317            .iter()
318            .any(|queued| queued.namespace == item.namespace && queued.hash == item.hash)
319        {
320            return false;
321        }
322        queue.push_back(item);
323        true
324    }
325
326    /// Remove and return the smallest item that fits the download budget.
327    pub fn pop_fitting(&self, budget: &TransferBudget) -> Option<PendingTransfer> {
328        let mut queue = self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
329        let available = budget.available(Direction::Download);
330        let position = queue
331            .iter()
332            .enumerate()
333            .filter(|(_, item)| item.size <= available)
334            .min_by_key(|(_, item)| item.size)
335            .map(|(position, _)| position)?;
336        queue.remove(position)
337    }
338
339    /// Drop everything queued for a namespace (e.g. after leaving a folder).
340    pub fn remove_namespace(&self, namespace: NamespaceId) {
341        let mut queue = self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
342        queue.retain(|item| item.namespace != namespace);
343    }
344
345    #[must_use]
346    pub fn len(&self) -> usize {
347        self.inner
348            .lock()
349            .unwrap_or_else(std::sync::PoisonError::into_inner)
350            .len()
351    }
352
353    #[must_use]
354    pub fn is_empty(&self) -> bool {
355        self.len() == 0
356    }
357
358    #[must_use]
359    pub fn pending_bytes(&self) -> u64 {
360        self.inner
361            .lock()
362            .unwrap_or_else(std::sync::PoisonError::into_inner)
363            .iter()
364            .fold(0_u64, |total, item| total.saturating_add(item.size))
365    }
366
367    #[must_use]
368    pub fn snapshot(&self) -> Vec<PendingTransfer> {
369        self.inner
370            .lock()
371            .unwrap_or_else(std::sync::PoisonError::into_inner)
372            .iter()
373            .cloned()
374            .collect()
375    }
376}
377
378#[cfg(test)]
379mod tests {
380    use super::*;
381
382    fn limits(
383        download: (Option<u64>, Option<u64>, Option<u64>),
384        upload: (Option<u64>, Option<u64>, Option<u64>),
385    ) -> TransferLimits {
386        TransferLimits {
387            download: TransferQuota {
388                per_hour: download.0,
389                per_day: download.1,
390                per_month: download.2,
391            },
392            upload: TransferQuota {
393                per_hour: upload.0,
394                per_day: upload.1,
395                per_month: upload.2,
396            },
397        }
398    }
399
400    #[test]
401    fn unlimited_quota_has_no_ceiling() {
402        let budget = TransferBudget::new(limits((None, None, None), (None, None, None)));
403        assert_eq!(budget.available(Direction::Download), u64::MAX);
404        assert!(budget.can_transfer(Direction::Download, u64::MAX));
405    }
406
407    #[test]
408    fn hourly_cap_is_enforced() {
409        let budget = TransferBudget::new(limits((Some(1_000), None, None), (None, None, None)));
410        let now = Instant::now();
411        budget.record_at(Direction::Download, 400, now);
412        assert_eq!(budget.available_at(Direction::Download, now), 600);
413        budget.record_at(Direction::Download, 600, now);
414        assert!(!budget.can_transfer(Direction::Download, 1));
415    }
416
417    #[test]
418    fn tightest_window_wins() {
419        let budget = TransferBudget::new(limits((Some(1_000), Some(1_500), None), (None, None, None)));
420        let now = Instant::now();
421        budget.record_at(Direction::Download, 900, now);
422        // hour remaining 100, day remaining 600 -> 100
423        assert_eq!(budget.available_at(Direction::Download, now), 100);
424    }
425
426    #[test]
427    fn expired_samples_free_budget() {
428        let budget = TransferBudget::new(limits((Some(1_000), None, None), (None, None, None)));
429        let now = Instant::now();
430        let earlier = now.checked_sub(Duration::from_secs(4_000)).expect("underflow");
431        budget.record_at(Direction::Download, 1_000, earlier);
432        assert_eq!(budget.available_at(Direction::Download, now), 1_000);
433    }
434
435    #[test]
436    fn directions_are_accounted_separately() {
437        let budget = TransferBudget::new(limits((Some(500), None, None), (Some(500), None, None)));
438        let now = Instant::now();
439        budget.record_at(Direction::Upload, 500, now);
440        assert_eq!(budget.available_at(Direction::Upload, now), 0);
441        assert_eq!(budget.available_at(Direction::Download, now), 500);
442    }
443
444    #[test]
445    fn queue_deduplicates_and_respects_budget() {
446        let namespace = NamespaceId::from([7_u8; 32]);
447        let queue = TransferQueue::new();
448        assert!(queue.enqueue(PendingTransfer::new(
449            namespace,
450            Hash::from_bytes([1; 32]),
451            400,
452            PathBuf::from("a"),
453        )));
454        assert!(!queue.enqueue(PendingTransfer::new(
455            namespace,
456            Hash::from_bytes([1; 32]),
457            400,
458            PathBuf::from("a"),
459        )));
460        assert!(queue.enqueue(PendingTransfer::new(
461            namespace,
462            Hash::from_bytes([2; 32]),
463            900,
464            PathBuf::from("b"),
465        )));
466        assert_eq!(queue.len(), 2);
467
468        // A 500-byte budget skips the 900-byte item and pops the 400-byte item.
469        let tight = TransferBudget::new(limits((Some(500), None, None), (None, None, None)));
470        let popped = queue.pop_fitting(&tight).expect("400-byte item fits");
471        assert_eq!(popped.size, 400);
472        assert_eq!(queue.len(), 1);
473        assert!(queue.pop_fitting(&tight).is_none(), "900 bytes must not fit");
474        assert_eq!(queue.pending_bytes(), 900);
475    }
476
477    #[test]
478    fn queue_prioritizes_shorter_fitting_transfers() {
479        let namespace = NamespaceId::from([8_u8; 32]);
480        let queue = TransferQueue::new();
481        queue.enqueue(PendingTransfer::new(
482            namespace,
483            Hash::from_bytes([1; 32]),
484            900,
485            PathBuf::from("large"),
486        ));
487        queue.enqueue(PendingTransfer::new(
488            namespace,
489            Hash::from_bytes([2; 32]),
490            400,
491            PathBuf::from("small"),
492        ));
493
494        let budget = TransferBudget::new(limits((None, None, None), (None, None, None)));
495        let popped = queue.pop_fitting(&budget).expect("a queued transfer fits");
496        assert_eq!(popped.path, PathBuf::from("small"));
497        assert_eq!(queue.len(), 1);
498    }
499
500    #[test]
501    fn config_parses_byte_suffixes() {
502        let mut config = TransferLimitsConfig::default();
503        assert!(config.set("transfer_limits.max_download_per_hour", "2G").unwrap());
504        assert!(config.set("transfer_limits.max_upload_per_month", "500M").unwrap());
505        assert!(!config.set("transfer_limits.unknown", "1").unwrap());
506        let limits = TransferLimits::from_config(&config).unwrap();
507        assert_eq!(limits.download.per_hour, Some(2_000_000_000));
508        assert_eq!(limits.upload.per_month, Some(500_000_000));
509        assert!(!limits.is_unlimited());
510    }
511}