1use std::collections::HashMap;
2use std::fmt;
3use std::sync::atomic::{
4 AtomicU64,
5 Ordering,
6};
7use std::sync::{
8 Arc,
9 Mutex,
10 MutexGuard,
11 PoisonError,
12};
13use std::time::Duration;
14
15use chrono::{
16 DateTime,
17 Utc,
18};
19use tokio::sync::watch;
20use tokio::time::Instant;
21
22use crate::domain::{
23 DomainError,
24 FailureKind,
25};
26
27const FAR_FUTURE: Duration = Duration::from_secs(10 * 365 * 24 * 60 * 60);
28const MAX_INTERVAL: Duration = Duration::from_secs(365 * 24 * 60 * 60);
29
30static NEXT_GENERATION: AtomicU64 = AtomicU64::new(1);
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
33pub struct Generation(u64);
34
35impl Generation {
36 pub fn next() -> Self {
37 Self(NEXT_GENERATION.fetch_add(1, Ordering::Relaxed))
38 }
39}
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum ClaimKind {
43 Scheduled,
44 Forced,
45}
46
47#[derive(Debug, Clone)]
48pub struct BackoffPolicy {
49 pub max_backoff: Duration,
50 pub max_rate_limit: Duration,
51 pub outage_floor: Duration,
52 pub forced_min_gap: Duration,
53}
54
55impl Default for BackoffPolicy {
56 fn default() -> Self {
57 Self {
58 max_backoff: Duration::from_secs(10 * 60),
59 max_rate_limit: Duration::from_secs(60 * 60),
60 outage_floor: Duration::from_secs(2 * 60),
61 forced_min_gap: Duration::from_secs(10),
62 }
63 }
64}
65
66#[derive(Debug, Clone, PartialEq, Eq)]
67pub enum Outcome {
68 Success,
69 Failed {
70 kind: FailureKind,
71 retry_after: Option<Duration>,
72 },
73 Cancelled,
74}
75
76#[derive(Debug, Clone)]
77pub enum FetchResult {
78 Committed(Result<usize, DomainError>),
79 Superseded,
80 Cancelled,
81}
82
83#[derive(Debug, Clone, PartialEq, Eq)]
84pub struct ScheduleView {
85 pub failures: u32,
86 pub last_failure: Option<FailureKind>,
87 pub next_attempt_in: Duration,
88 pub in_flight: bool,
89}
90
91type AttemptSender = watch::Sender<Option<FetchResult>>;
92type AttemptReceiver = watch::Receiver<Option<FetchResult>>;
93
94struct QuotaPause {
95 until: DateTime<Utc>,
96 resume: Instant,
97}
98
99struct Entry {
100 generation: Generation,
101 interval: Duration,
102 next_attempt: Instant,
103 last_attempt: Option<Instant>,
104 last_outcome: Option<Outcome>,
105 failures: u32,
106 last_failure: Option<FailureKind>,
107 attempt: Option<AttemptReceiver>,
108 quota: Option<QuotaPause>,
109}
110
111impl Entry {
112 fn new(generation: Generation, interval: Duration, now: Instant) -> Self {
113 Self {
114 generation,
115 interval: interval.min(MAX_INTERVAL),
116 next_attempt: now,
117 last_attempt: None,
118 last_outcome: None,
119 failures: 0,
120 last_failure: None,
121 attempt: None,
122 quota: None,
123 }
124 }
125
126 fn view(&self, now: Instant) -> ScheduleView {
127 ScheduleView {
128 failures: self.failures,
129 last_failure: self.last_failure,
130 next_attempt_in: self.next_attempt.saturating_duration_since(now),
131 in_flight: self.attempt.is_some(),
132 }
133 }
134
135 fn due_at(&self, floor: Option<Duration>) -> Instant {
136 match (floor, self.last_attempt) {
137 (Some(floor), Some(last)) => self.next_attempt.max(after(last, floor)),
138 _ => self.next_attempt,
139 }
140 }
141
142 fn start_attempt(&mut self, now: Instant) -> (AttemptSender, Option<Instant>) {
143 let (sender, receiver) = watch::channel(None);
144 self.attempt = Some(receiver);
145 (sender, self.last_attempt.replace(now))
146 }
147
148 fn release(&mut self, previous_attempt: Option<Instant>) {
149 self.attempt = None;
150 self.last_attempt = previous_attempt;
151 }
152
153 fn finish(&mut self, policy: &BackoffPolicy, now: Instant, outcome: Outcome) {
154 let interval = self.interval;
155 let wait = match &outcome {
156 Outcome::Success => {
157 self.failures = 0;
158 self.last_failure = None;
159 interval
160 }
161 Outcome::Cancelled => interval,
162 Outcome::Failed { kind, retry_after } => {
163 self.last_failure = Some(*kind);
164 if *kind == FailureKind::RateLimited {
165 self.failures = self.failures.saturating_add(1);
166 let wait = (*retry_after).unwrap_or_else(|| doubled(interval, self.failures));
167 interval.max(wait.min(policy.max_rate_limit))
168 } else if kind.is_transient() {
169 self.failures = self.failures.saturating_add(1);
170 let hint = (*retry_after)
171 .unwrap_or_default()
172 .min(policy.max_rate_limit);
173 interval
174 .max(doubled(interval, self.failures).min(policy.max_backoff))
175 .max(hint)
176 } else {
177 interval
178 }
179 }
180 };
181 self.attempt = None;
182 self.last_outcome = Some(outcome);
183 self.next_attempt = after(now, wait);
184 match self.quota.as_ref().map(|quota| quota.resume) {
185 Some(resume) if resume > now => self.next_attempt = self.next_attempt.max(resume),
186 Some(_) => self.quota = None,
187 None => {}
188 }
189 }
190}
191
192fn doubled(interval: Duration, failures: u32) -> Duration {
193 interval.saturating_mul(1u32.checked_shl(failures).unwrap_or(u32::MAX))
194}
195
196fn after(now: Instant, wait: Duration) -> Instant {
197 now.checked_add(wait).unwrap_or_else(|| now + FAR_FUTURE)
198}
199
200pub struct FetchScheduler {
201 entries: Mutex<HashMap<i64, Entry>>,
202 policy: BackoffPolicy,
203}
204
205impl Default for FetchScheduler {
206 fn default() -> Self {
207 Self::new(BackoffPolicy::default())
208 }
209}
210
211impl FetchScheduler {
212 pub fn new(policy: BackoffPolicy) -> Self {
213 Self {
214 entries: Mutex::new(HashMap::new()),
215 policy,
216 }
217 }
218
219 pub fn policy(&self) -> &BackoffPolicy {
220 &self.policy
221 }
222
223 fn entries(&self) -> MutexGuard<'_, HashMap<i64, Entry>> {
224 self.entries.lock().unwrap_or_else(PoisonError::into_inner)
225 }
226
227 pub fn register(&self, id: i64, generation: Generation, interval: Duration) {
228 let mut entries = self.entries();
229 if entries
230 .get(&id)
231 .is_some_and(|entry| entry.generation == generation)
232 {
233 return;
234 }
235 entries.insert(id, Entry::new(generation, interval, Instant::now()));
236 }
237
238 pub fn set_interval(&self, id: i64, generation: Generation, interval: Duration) {
239 let mut entries = self.entries();
240 let Some(entry) = entries
241 .get_mut(&id)
242 .filter(|entry| entry.generation == generation)
243 else {
244 return;
245 };
246 entry.interval = interval.min(MAX_INTERVAL);
247 if entry.failures == 0 && entry.attempt.is_none() {
248 if let Some(last) = entry.last_attempt {
249 entry.next_attempt = after(last, entry.interval);
250 }
251 }
252 }
253
254 pub fn reset(&self, id: i64, generation: Generation, interval: Duration) {
255 self.entries()
256 .insert(id, Entry::new(generation, interval, Instant::now()));
257 }
258
259 pub fn remove(&self, id: i64) {
260 self.entries().remove(&id);
261 }
262
263 pub fn claim_due(
264 self: &Arc<Self>, now: Instant, floors: &HashMap<i64, Duration>,
265 ) -> Vec<FetchClaim> {
266 let mut entries = self.entries();
267 let mut claims: Vec<FetchClaim> = entries
268 .iter_mut()
269 .filter(|(id, entry)| {
270 !entry.interval.is_zero()
271 && entry.attempt.is_none()
272 && now >= entry.due_at(floors.get(*id).copied())
273 })
274 .map(|(id, entry)| {
275 let (result, previous_attempt) = entry.start_attempt(now);
276 FetchClaim {
277 scheduler: Arc::clone(self),
278 id: *id,
279 generation: entry.generation,
280 result,
281 previous_attempt,
282 done: false,
283 }
284 })
285 .collect();
286 claims.sort_by_key(FetchClaim::id);
287
288 claims
289 }
290
291 pub fn claim_forced(
292 self: &Arc<Self>, id: i64, generation: Generation, now: Instant,
293 ) -> ForcedClaim {
294 let mut entries = self.entries();
295 let Some(entry) = entries
296 .get_mut(&id)
297 .filter(|entry| entry.generation == generation)
298 else {
299 return ForcedClaim::Refused;
300 };
301 if let Some(attempt) = &entry.attempt {
302 return ForcedClaim::Joined(AttemptHandle(attempt.clone()));
303 }
304 if let Some(resume) = entry
305 .quota
306 .as_ref()
307 .map(|quota| quota.resume)
308 .filter(|resume| *resume > now)
309 {
310 return ForcedClaim::Recent(Outcome::Failed {
311 kind: FailureKind::RateLimited,
312 retry_after: Some(resume.saturating_duration_since(now)),
313 });
314 }
315 let rate_limited = matches!(
316 entry.last_outcome,
317 Some(Outcome::Failed {
318 kind: FailureKind::RateLimited,
319 ..
320 })
321 ) && now < entry.next_attempt;
322 let too_soon = entry
323 .last_attempt
324 .is_some_and(|last| now < after(last, self.policy.forced_min_gap));
325 if rate_limited || too_soon {
326 return ForcedClaim::Recent(entry.last_outcome.clone().unwrap_or(Outcome::Cancelled));
327 }
328
329 let (result, previous_attempt) = entry.start_attempt(now);
330 ForcedClaim::Claimed(FetchClaim {
331 scheduler: Arc::clone(self),
332 id,
333 generation,
334 result,
335 previous_attempt,
336 done: false,
337 })
338 }
339
340 pub fn view(&self, id: i64, now: Instant) -> Option<ScheduleView> {
341 self.entries().get(&id).map(|entry| entry.view(now))
342 }
343
344 pub fn next_due_in(&self, id: i64, now: Instant, floor: Option<Duration>) -> Option<Duration> {
345 self.entries()
346 .get(&id)
347 .map(|entry| entry.due_at(floor).saturating_duration_since(now))
348 }
349
350 pub fn set_quota_pause(
351 &self, id: i64, generation: Generation, until: Option<DateTime<Utc>>, now: Instant,
352 wall: DateTime<Utc>,
353 ) -> bool {
354 let mut entries = self.entries();
355 let Some(entry) = entries
356 .get_mut(&id)
357 .filter(|entry| entry.generation == generation)
358 else {
359 return false;
360 };
361 let previous = entry.quota.as_ref().map(|quota| quota.until);
362 let pause = until.and_then(|until| {
363 (until - wall)
364 .to_std()
365 .ok()
366 .filter(|wait| !wait.is_zero())
367 .map(|wait| QuotaPause {
368 until,
369 resume: after(now, wait),
370 })
371 });
372 if let Some(pause) = &pause {
373 entry.next_attempt = entry.next_attempt.max(pause.resume);
374 }
375 let current = pause.as_ref().map(|pause| pause.until);
376 entry.quota = pause;
377
378 previous != current
379 }
380
381 pub fn paused_until(&self, id: i64, wall: DateTime<Utc>) -> Option<DateTime<Utc>> {
382 self.entries()
383 .get(&id)
384 .and_then(|entry| entry.quota.as_ref())
385 .map(|quota| quota.until)
386 .filter(|until| *until > wall)
387 }
388
389 fn finish(
390 &self, id: i64, generation: Generation, now: Instant, outcome: Outcome,
391 ) -> Option<ScheduleView> {
392 let mut entries = self.entries();
393 let entry = entries
394 .get_mut(&id)
395 .filter(|entry| entry.generation == generation)?;
396 entry.finish(&self.policy, now, outcome);
397
398 Some(entry.view(now))
399 }
400
401 fn release(&self, id: i64, generation: Generation, previous_attempt: Option<Instant>) {
402 if let Some(entry) = self
403 .entries()
404 .get_mut(&id)
405 .filter(|entry| entry.generation == generation)
406 {
407 entry.release(previous_attempt);
408 }
409 }
410}
411
412#[derive(Debug)]
413pub enum ForcedClaim {
414 Claimed(FetchClaim),
415 Joined(AttemptHandle),
416 Recent(Outcome),
417 Refused,
418}
419
420pub struct FetchClaim {
421 scheduler: Arc<FetchScheduler>,
422 id: i64,
423 generation: Generation,
424 result: AttemptSender,
425 previous_attempt: Option<Instant>,
426 done: bool,
427}
428
429impl fmt::Debug for FetchClaim {
430 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
431 f.debug_struct("FetchClaim")
432 .field("id", &self.id)
433 .field("generation", &self.generation)
434 .finish()
435 }
436}
437
438impl FetchClaim {
439 pub fn id(&self) -> i64 {
440 self.id
441 }
442
443 pub fn generation(&self) -> Generation {
444 self.generation
445 }
446
447 pub fn handle(&self) -> AttemptHandle {
448 AttemptHandle(self.result.subscribe())
449 }
450
451 pub fn complete(
452 mut self, now: Instant, outcome: Outcome, result: FetchResult,
453 ) -> Option<ScheduleView> {
454 self.done = true;
455 self.result.send_replace(Some(result));
456
457 self.scheduler
458 .finish(self.id, self.generation, now, outcome)
459 }
460
461 pub fn release(mut self) {
462 self.done = true;
463 self.result.send_replace(Some(FetchResult::Cancelled));
464 self.scheduler
465 .release(self.id, self.generation, self.previous_attempt);
466 }
467}
468
469impl Drop for FetchClaim {
470 fn drop(&mut self) {
471 if !self.done {
472 self.scheduler
473 .finish(self.id, self.generation, Instant::now(), Outcome::Cancelled);
474 }
475 }
476}
477
478pub struct AttemptHandle(AttemptReceiver);
479
480impl fmt::Debug for AttemptHandle {
481 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
482 f.write_str("AttemptHandle")
483 }
484}
485
486impl AttemptHandle {
487 pub async fn wait(mut self) -> FetchResult {
488 match self.0.wait_for(Option::is_some).await {
489 Ok(result) => result.clone().unwrap_or(FetchResult::Cancelled),
490 Err(_) => FetchResult::Cancelled,
491 }
492 }
493}
494
495#[cfg(test)]
496mod tests {
497 use super::*;
498
499 fn secs(seconds: u64) -> Duration {
500 Duration::from_secs(seconds)
501 }
502
503 fn scheduler_with(entries: &[(i64, u64)]) -> (Arc<FetchScheduler>, HashMap<i64, Generation>) {
504 let scheduler = Arc::new(FetchScheduler::default());
505 let generations = entries
506 .iter()
507 .map(|(id, interval)| {
508 let generation = Generation::next();
509 scheduler.register(*id, generation, secs(*interval));
510 (*id, generation)
511 })
512 .collect();
513
514 (scheduler, generations)
515 }
516
517 fn no_floors() -> HashMap<i64, Duration> {
518 HashMap::new()
519 }
520
521 fn due_ids(scheduler: &Arc<FetchScheduler>, now: Instant) -> Vec<i64> {
522 scheduler
523 .claim_due(now, &no_floors())
524 .iter()
525 .map(FetchClaim::id)
526 .collect()
527 }
528
529 fn only_claim(scheduler: &Arc<FetchScheduler>, now: Instant) -> FetchClaim {
530 let mut claims = scheduler.claim_due(now, &no_floors());
531 assert_eq!(claims.len(), 1);
532 claims.remove(0)
533 }
534
535 fn failed(kind: FailureKind, retry_after: Option<Duration>) -> Outcome {
536 Outcome::Failed { kind, retry_after }
537 }
538
539 fn finish_with(claim: FetchClaim, now: Instant, outcome: Outcome) -> ScheduleView {
540 let result = match &outcome {
541 Outcome::Failed { .. } => {
542 FetchResult::Committed(Err(DomainError::NetworkError("unreachable".to_string())))
543 }
544 _ => FetchResult::Committed(Ok(0)),
545 };
546 claim.complete(now, outcome, result).unwrap()
547 }
548
549 fn due_after(seconds: u64) -> Vec<i64> {
550 let (scheduler, generations) = scheduler_with(&[(1, 10), (2, 600), (3, 0), (4, 30)]);
551 let start = Instant::now();
552 for id in [1, 2] {
553 let ForcedClaim::Claimed(claim) = scheduler.claim_forced(id, generations[&id], start)
554 else {
555 panic!("expected a claim for {id}");
556 };
557 finish_with(claim, start, Outcome::Success);
558 }
559
560 due_ids(&scheduler, start + secs(seconds))
561 }
562
563 #[test]
564 fn each_provider_waits_for_its_own_interval() {
565 assert_eq!(due_after(5), vec![4]);
566 assert_eq!(due_after(60), vec![1, 4]);
567 assert_eq!(due_after(599), vec![1, 4]);
568 assert_eq!(due_after(600), vec![1, 2, 4]);
569 }
570
571 #[test]
572 fn counters_do_not_change_from_view_or_from_claiming() {
573 let (scheduler, _) = scheduler_with(&[(1, 30)]);
574 let now = Instant::now();
575
576 let before = scheduler.view(1, now).unwrap();
577 assert_eq!(scheduler.view(1, now).unwrap(), before);
578 let claim = only_claim(&scheduler, now);
579 let claimed = scheduler.view(1, now).unwrap();
580
581 assert_eq!(claimed.failures, before.failures);
582 assert_eq!(claimed.next_attempt_in, before.next_attempt_in);
583 assert!(claimed.in_flight);
584 drop(claim);
585 }
586
587 #[test]
588 fn transient_failures_double_up_to_the_cap() {
589 let (scheduler, _) = scheduler_with(&[(1, 30)]);
590 let mut now = Instant::now();
591 let mut waits = Vec::new();
592 for _ in 0..7 {
593 let view = finish_with(
594 only_claim(&scheduler, now),
595 now,
596 failed(FailureKind::Network, None),
597 );
598 waits.push(view.next_attempt_in.as_secs());
599 now += view.next_attempt_in;
600 }
601
602 assert_eq!(waits, vec![60, 120, 240, 480, 600, 600, 600]);
603 assert_eq!(scheduler.view(1, now).unwrap().failures, 7);
604 }
605
606 #[test]
607 fn timeouts_and_unavailable_count_like_network_failures() {
608 for kind in [FailureKind::Timeout, FailureKind::Unavailable] {
609 let (scheduler, _) = scheduler_with(&[(1, 30)]);
610 let now = Instant::now();
611
612 let view = finish_with(only_claim(&scheduler, now), now, failed(kind, None));
613
614 assert_eq!(view.failures, 1, "{kind:?}");
615 assert_eq!(view.next_attempt_in, secs(60), "{kind:?}");
616 }
617 }
618
619 #[test]
620 fn a_server_hint_extends_the_wait_but_never_shortens_it() {
621 let (scheduler, _) = scheduler_with(&[(1, 30)]);
622 let now = Instant::now();
623
624 let long = finish_with(
625 only_claim(&scheduler, now),
626 now,
627 failed(FailureKind::Unavailable, Some(secs(300))),
628 );
629 let later = now + long.next_attempt_in;
630 let short = finish_with(
631 only_claim(&scheduler, later),
632 later,
633 failed(FailureKind::Unavailable, Some(secs(5))),
634 );
635
636 assert_eq!(long.next_attempt_in, secs(300));
637 assert_eq!(short.next_attempt_in, secs(120));
638 }
639
640 #[test]
641 fn rate_limits_use_retry_after_capped_at_an_hour() {
642 let (scheduler, _) = scheduler_with(&[(1, 30)]);
643 let mut now = Instant::now();
644 let mut waits = Vec::new();
645 for retry_after in [Some(secs(120)), Some(secs(2 * 3600)), None, Some(secs(5))] {
646 let view = finish_with(
647 only_claim(&scheduler, now),
648 now,
649 failed(FailureKind::RateLimited, retry_after),
650 );
651 waits.push(view.next_attempt_in.as_secs());
652 now += view.next_attempt_in;
653 }
654
655 assert_eq!(waits, vec![120, 3600, 30 * 8, 30]);
656 }
657
658 #[test]
659 fn an_interval_longer_than_the_wait_caps_is_kept() {
660 let (scheduler, _) = scheduler_with(&[(1, 2 * 3600), (2, 2 * 3600)]);
661 let now = Instant::now();
662 let claims = scheduler.claim_due(now, &no_floors());
663 let mut views = HashMap::new();
664 for claim in claims {
665 let id = claim.id();
666 let outcome = if id == 1 {
667 failed(FailureKind::RateLimited, Some(secs(3600)))
668 } else {
669 failed(FailureKind::Network, None)
670 };
671 views.insert(id, finish_with(claim, now, outcome));
672 }
673
674 assert_eq!(views[&1].next_attempt_in, secs(2 * 3600));
675 assert_eq!(views[&2].next_attempt_in, secs(2 * 3600));
676 }
677
678 #[test]
679 fn auth_failures_do_not_count_and_success_resets() {
680 let (scheduler, _) = scheduler_with(&[(1, 30)]);
681 let mut now = Instant::now();
682
683 let network = finish_with(
684 only_claim(&scheduler, now),
685 now,
686 failed(FailureKind::Network, None),
687 );
688 now += network.next_attempt_in;
689 let auth = finish_with(
690 only_claim(&scheduler, now),
691 now,
692 failed(FailureKind::Auth, None),
693 );
694 now += auth.next_attempt_in;
695 let success = finish_with(only_claim(&scheduler, now), now, Outcome::Success);
696
697 assert_eq!(network.failures, 1);
698 assert_eq!(auth.failures, 1);
699 assert_eq!(auth.last_failure, Some(FailureKind::Auth));
700 assert_eq!(auth.next_attempt_in, secs(30));
701 assert_eq!(success.failures, 0);
702 assert_eq!(success.last_failure, None);
703 assert_eq!(success.next_attempt_in, secs(30));
704 }
705
706 #[test]
707 fn the_outage_floor_delays_but_still_lets_a_probe_through() {
708 let (scheduler, _) = scheduler_with(&[(1, 10)]);
709 let floors = HashMap::from([(1, secs(120))]);
710 let start = Instant::now();
711 finish_with(only_claim(&scheduler, start), start, Outcome::Success);
712
713 assert!(scheduler.claim_due(start + secs(10), &floors).is_empty());
714 assert!(scheduler.claim_due(start + secs(119), &floors).is_empty());
715 assert_eq!(
716 scheduler.next_due_in(1, start + secs(10), Some(secs(120))),
717 Some(secs(110))
718 );
719 assert_eq!(
720 scheduler.next_due_in(1, start + secs(10), None),
721 Some(Duration::ZERO)
722 );
723 let probe = scheduler.claim_due(start + secs(120), &floors);
724
725 assert_eq!(probe.len(), 1);
726 let view = finish_with(
727 probe.into_iter().next().unwrap(),
728 start + secs(120),
729 Outcome::Success,
730 );
731 assert_eq!(view.failures, 0);
732 assert_eq!(view.next_attempt_in, secs(10));
733 }
734
735 #[test]
736 fn interval_zero_is_never_scheduled_but_can_be_forced() {
737 let (scheduler, generations) = scheduler_with(&[(1, 0)]);
738 let now = Instant::now();
739
740 assert!(scheduler
741 .claim_due(now + secs(86_400), &no_floors())
742 .is_empty());
743 assert!(matches!(
744 scheduler.claim_forced(1, generations[&1], now),
745 ForcedClaim::Claimed(_)
746 ));
747 }
748
749 #[test]
750 fn forced_claims_bypass_next_attempt_and_the_outage_floor() {
751 let (scheduler, generations) = scheduler_with(&[(1, 600)]);
752 let start = Instant::now();
753 finish_with(only_claim(&scheduler, start), start, Outcome::Success);
754 let later = start + secs(20);
755
756 assert!(scheduler.claim_due(later, &no_floors()).is_empty());
757 assert!(matches!(
758 scheduler.claim_forced(1, generations[&1], later),
759 ForcedClaim::Claimed(_)
760 ));
761 }
762
763 #[test]
764 fn forced_claims_wait_out_a_rate_limit_and_the_minimum_gap() {
765 let (scheduler, generations) = scheduler_with(&[(1, 30), (2, 30)]);
766 let start = Instant::now();
767 let mut claims = scheduler.claim_due(start, &no_floors());
768 let second = claims.pop().unwrap();
769 let first = claims.pop().unwrap();
770 finish_with(
771 first,
772 start,
773 failed(FailureKind::RateLimited, Some(secs(120))),
774 );
775 finish_with(second, start, Outcome::Success);
776
777 let during_rate_limit = scheduler.claim_forced(1, generations[&1], start + secs(60));
778 let within_gap = scheduler.claim_forced(2, generations[&2], start + secs(9));
779 let after_gap = scheduler.claim_forced(2, generations[&2], start + secs(10));
780 let after_rate_limit = scheduler.claim_forced(1, generations[&1], start + secs(120));
781
782 assert!(matches!(
783 during_rate_limit,
784 ForcedClaim::Recent(Outcome::Failed {
785 kind: FailureKind::RateLimited,
786 ..
787 })
788 ));
789 assert!(matches!(within_gap, ForcedClaim::Recent(Outcome::Success)));
790 assert!(matches!(after_gap, ForcedClaim::Claimed(_)));
791 assert!(matches!(after_rate_limit, ForcedClaim::Claimed(_)));
792 }
793
794 #[tokio::test]
795 async fn forced_claims_join_an_attempt_from_claim_time_until_it_completes() {
796 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
797 let now = Instant::now();
798 let claim = only_claim(&scheduler, now);
799
800 let ForcedClaim::Joined(joiner) = scheduler.claim_forced(1, generations[&1], now) else {
801 panic!("expected to join the in-flight attempt");
802 };
803 let waiting = tokio::spawn(joiner.wait());
804 tokio::task::yield_now().await;
805 assert!(!waiting.is_finished());
806 let ForcedClaim::Joined(late) = scheduler.claim_forced(1, generations[&1], now) else {
807 panic!("expected to join the in-flight attempt");
808 };
809 claim.complete(now, Outcome::Success, FetchResult::Committed(Ok(3)));
810
811 assert!(matches!(
812 waiting.await.unwrap(),
813 FetchResult::Committed(Ok(3))
814 ));
815 assert!(matches!(late.wait().await, FetchResult::Committed(Ok(3))));
816 assert!(!scheduler.view(1, now).unwrap().in_flight);
817 }
818
819 #[tokio::test]
820 async fn a_dropped_claim_cancels_its_joiners_without_counting_a_failure() {
821 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
822 let now = Instant::now();
823 let claim = only_claim(&scheduler, now);
824 let ForcedClaim::Joined(joiner) = scheduler.claim_forced(1, generations[&1], now) else {
825 panic!("expected to join the in-flight attempt");
826 };
827
828 drop(claim);
829
830 assert!(matches!(joiner.wait().await, FetchResult::Cancelled));
831 let view = scheduler.view(1, now).unwrap();
832 assert_eq!(view.failures, 0);
833 assert_eq!(view.last_failure, None);
834 assert!(!view.in_flight);
835 let remaining = scheduler.view(1, Instant::now()).unwrap().next_attempt_in;
836 assert!(remaining <= secs(30));
837 assert!(remaining > Duration::ZERO);
838 }
839
840 #[test]
841 fn an_old_claim_changes_nothing_after_reset() {
842 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
843 let now = Instant::now();
844 let old = only_claim(&scheduler, now);
845 let fresh = Generation::next();
846 scheduler.reset(1, fresh, secs(30));
847 let current = only_claim(&scheduler, Instant::now());
848 let before = scheduler.view(1, now).unwrap();
849
850 let stale = old.complete(
851 now,
852 failed(FailureKind::Network, None),
853 FetchResult::Superseded,
854 );
855 let after = scheduler.view(1, now).unwrap();
856
857 assert!(stale.is_none());
858 assert_eq!(after, before);
859 assert!(after.in_flight);
860 assert_eq!(after.failures, 0);
861 assert_ne!(generations[&1], fresh);
862 assert_eq!(current.generation(), fresh);
863 }
864
865 #[test]
866 fn an_old_claim_dropped_after_reset_leaves_the_new_attempt_alone() {
867 let (scheduler, _) = scheduler_with(&[(1, 30)]);
868 let now = Instant::now();
869 let old = only_claim(&scheduler, now);
870 scheduler.reset(1, Generation::next(), secs(30));
871 let current = only_claim(&scheduler, Instant::now());
872
873 drop(old);
874
875 assert!(scheduler.view(1, now).unwrap().in_flight);
876 drop(current);
877 assert!(!scheduler.view(1, now).unwrap().in_flight);
878 }
879
880 #[test]
881 fn claims_after_remove_are_refused_and_the_entry_is_not_recreated() {
882 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
883 let now = Instant::now();
884 let old = only_claim(&scheduler, now);
885 scheduler.remove(1);
886
887 assert!(matches!(
888 scheduler.claim_forced(1, generations[&1], now),
889 ForcedClaim::Refused
890 ));
891 assert!(scheduler.claim_due(now, &no_floors()).is_empty());
892 old.complete(now, Outcome::Success, FetchResult::Superseded);
893 assert!(scheduler.view(1, now).is_none());
894 assert!(scheduler.claim_due(now, &no_floors()).is_empty());
895 scheduler.set_interval(1, generations[&1], secs(5));
896 assert!(scheduler.view(1, now).is_none());
897 }
898
899 #[test]
900 fn a_reused_id_with_a_new_generation_ignores_old_claims() {
901 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
902 let now = Instant::now();
903 let old = only_claim(&scheduler, now);
904 scheduler.remove(1);
905 let reused = Generation::next();
906 scheduler.register(1, reused, secs(30));
907
908 assert!(matches!(
909 scheduler.claim_forced(1, generations[&1], now),
910 ForcedClaim::Refused
911 ));
912 old.complete(
913 now,
914 failed(FailureKind::Network, None),
915 FetchResult::Superseded,
916 );
917 let view = scheduler.view(1, now).unwrap();
918
919 assert_eq!(view.failures, 0);
920 assert!(!view.in_flight);
921 assert!(matches!(
922 scheduler.claim_forced(1, reused, now),
923 ForcedClaim::Claimed(_)
924 ));
925 }
926
927 #[test]
928 fn register_keeps_an_existing_entry_of_the_same_generation() {
929 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
930 let now = Instant::now();
931 finish_with(
932 only_claim(&scheduler, now),
933 now,
934 failed(FailureKind::Network, None),
935 );
936
937 scheduler.register(1, generations[&1], secs(30));
938
939 assert_eq!(scheduler.view(1, now).unwrap().failures, 1);
940 }
941
942 #[test]
943 fn a_shorter_interval_brings_the_next_attempt_forward() {
944 let (scheduler, generations) = scheduler_with(&[(1, 3600)]);
945 let now = Instant::now();
946 finish_with(only_claim(&scheduler, now), now, Outcome::Success);
947
948 scheduler.set_interval(1, generations[&1], secs(30));
949
950 assert!(scheduler.view(1, Instant::now()).unwrap().next_attempt_in <= secs(30));
951 assert_eq!(due_ids(&scheduler, now + secs(30)), vec![1]);
952 }
953
954 #[test]
955 fn a_quota_pause_postpones_every_outcome_until_the_reset() {
956 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
957 let now = Instant::now();
958 let wall = Utc::now();
959 let reset = wall + chrono::Duration::minutes(20);
960 let claim = only_claim(&scheduler, now);
961
962 let changed = scheduler.set_quota_pause(1, generations[&1], Some(reset), now, wall);
963 let repeated = scheduler.set_quota_pause(1, generations[&1], Some(reset), now, wall);
964 let view = claim.complete(now, Outcome::Cancelled, FetchResult::Cancelled);
965
966 assert!(changed);
967 assert!(!repeated);
968 assert_eq!(view.unwrap().next_attempt_in, secs(20 * 60));
969 assert_eq!(scheduler.paused_until(1, wall), Some(reset));
970 assert!(scheduler
971 .claim_due(now + secs(20 * 60 - 1), &no_floors())
972 .is_empty());
973 assert_eq!(due_ids(&scheduler, now + secs(20 * 60)), vec![1]);
974 }
975
976 #[test]
977 fn a_quota_pause_turns_manual_refreshes_into_a_rate_limit_answer() {
978 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
979 let now = Instant::now();
980 let wall = Utc::now();
981 scheduler.set_quota_pause(
982 1,
983 generations[&1],
984 Some(wall + chrono::Duration::minutes(5)),
985 now,
986 wall,
987 );
988
989 let forced = scheduler.claim_forced(1, generations[&1], now + secs(60));
990
991 let ForcedClaim::Recent(Outcome::Failed { kind, retry_after }) = forced else {
992 panic!("expected a recent rate-limit answer, got {forced:?}");
993 };
994 assert_eq!(kind, FailureKind::RateLimited);
995 assert_eq!(retry_after, Some(secs(240)));
996 }
997
998 #[test]
999 fn a_quota_pause_clears_when_the_quota_is_back_and_ignores_stale_generations() {
1000 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
1001 let now = Instant::now();
1002 let wall = Utc::now();
1003 let reset = wall + chrono::Duration::minutes(5);
1004 let stale = Generation::next();
1005
1006 let ignored = scheduler.set_quota_pause(1, stale, Some(reset), now, wall);
1007 scheduler.set_quota_pause(1, generations[&1], Some(reset), now, wall);
1008 let cleared = scheduler.set_quota_pause(1, generations[&1], None, now, wall);
1009 let already_clear = scheduler.set_quota_pause(1, generations[&1], None, now, wall);
1010
1011 assert!(!ignored);
1012 assert!(cleared);
1013 assert!(!already_clear);
1014 assert_eq!(scheduler.paused_until(1, wall), None);
1015 let past = wall - chrono::Duration::seconds(1);
1016 assert!(!scheduler.set_quota_pause(1, generations[&1], Some(past), now, wall));
1017 assert_eq!(scheduler.paused_until(1, wall), None);
1018 }
1019
1020 #[test]
1021 fn generations_are_never_reused() {
1022 let first = Generation::next();
1023 let second = Generation::next();
1024
1025 assert_ne!(first, second);
1026 assert!(second > first);
1027 }
1028}