1use 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
24pub const WINDOW_HOUR: Duration = Duration::from_hours(1);
26pub const WINDOW_DAY: Duration = Duration::from_hours(24);
28pub const WINDOW_MONTH: Duration = Duration::from_hours(720);
30
31#[derive(Clone, Copy, Debug, Eq, PartialEq)]
33#[non_exhaustive]
34pub enum Direction {
35 Download,
36 Upload,
37}
38
39#[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 #[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#[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 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#[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 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#[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 pub fn set_limits(&self, limits: TransferLimits) {
196 *self.limits.write().unwrap_or_else(std::sync::PoisonError::into_inner) = limits;
197 }
198
199 #[must_use]
201 pub fn available(&self, direction: Direction) -> u64 {
202 self.available_at(direction, Instant::now())
203 }
204
205 #[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 #[must_use]
228 pub fn can_transfer(&self, direction: Direction, bytes: u64) -> bool {
229 bytes <= self.available(direction)
230 }
231
232 pub fn record(&self, direction: Direction, bytes: u64) {
234 self.record_at(direction, bytes, Instant::now());
235 }
236
237 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 #[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 #[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#[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#[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 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 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 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 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 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}