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 forced_delay(&self, policy: &BackoffPolicy, now: Instant) -> Duration {
143 let mut resume = self
144 .quota
145 .as_ref()
146 .map_or(now, |quota| quota.resume.max(now));
147 if matches!(
148 self.last_outcome,
149 Some(Outcome::Failed {
150 kind: FailureKind::RateLimited,
151 ..
152 })
153 ) {
154 resume = resume.max(self.next_attempt);
155 }
156 if let Some(last) = self.last_attempt {
157 resume = resume.max(after(last, policy.forced_min_gap));
158 }
159 resume.saturating_duration_since(now)
160 }
161
162 fn start_attempt(&mut self, now: Instant) -> (AttemptSender, Option<Instant>) {
163 let (sender, receiver) = watch::channel(None);
164 self.attempt = Some(receiver);
165 (sender, self.last_attempt.replace(now))
166 }
167
168 fn release(&mut self, previous_attempt: Option<Instant>) {
169 self.attempt = None;
170 self.last_attempt = previous_attempt;
171 }
172
173 fn finish(&mut self, policy: &BackoffPolicy, now: Instant, outcome: Outcome) {
174 let interval = self.interval;
175 let wait = match &outcome {
176 Outcome::Success => {
177 self.failures = 0;
178 self.last_failure = None;
179 interval
180 }
181 Outcome::Cancelled => interval,
182 Outcome::Failed { kind, retry_after } => {
183 self.last_failure = Some(*kind);
184 if *kind == FailureKind::RateLimited {
185 self.failures = self.failures.saturating_add(1);
186 let wait = (*retry_after).unwrap_or_else(|| doubled(interval, self.failures));
187 interval.max(wait.min(policy.max_rate_limit))
188 } else if kind.is_transient() {
189 self.failures = self.failures.saturating_add(1);
190 let hint = (*retry_after)
191 .unwrap_or_default()
192 .min(policy.max_rate_limit);
193 interval
194 .max(doubled(interval, self.failures).min(policy.max_backoff))
195 .max(hint)
196 } else {
197 interval
198 }
199 }
200 };
201 self.attempt = None;
202 self.last_outcome = Some(outcome);
203 self.next_attempt = after(now, wait);
204 match self.quota.as_ref().map(|quota| quota.resume) {
205 Some(resume) if resume > now => self.next_attempt = self.next_attempt.max(resume),
206 Some(_) => self.quota = None,
207 None => {}
208 }
209 }
210}
211
212fn doubled(interval: Duration, failures: u32) -> Duration {
213 interval.saturating_mul(1u32.checked_shl(failures).unwrap_or(u32::MAX))
214}
215
216fn after(now: Instant, wait: Duration) -> Instant {
217 now.checked_add(wait).unwrap_or_else(|| now + FAR_FUTURE)
218}
219
220pub struct FetchScheduler {
221 entries: Mutex<HashMap<i64, Entry>>,
222 policy: BackoffPolicy,
223}
224
225impl Default for FetchScheduler {
226 fn default() -> Self {
227 Self::new(BackoffPolicy::default())
228 }
229}
230
231impl FetchScheduler {
232 pub fn new(policy: BackoffPolicy) -> Self {
233 Self {
234 entries: Mutex::new(HashMap::new()),
235 policy,
236 }
237 }
238
239 pub fn policy(&self) -> &BackoffPolicy {
240 &self.policy
241 }
242
243 fn entries(&self) -> MutexGuard<'_, HashMap<i64, Entry>> {
244 self.entries.lock().unwrap_or_else(PoisonError::into_inner)
245 }
246
247 pub fn register(&self, id: i64, generation: Generation, interval: Duration) {
248 let mut entries = self.entries();
249 if entries
250 .get(&id)
251 .is_some_and(|entry| entry.generation == generation)
252 {
253 return;
254 }
255 entries.insert(id, Entry::new(generation, interval, Instant::now()));
256 }
257
258 pub fn set_interval(&self, id: i64, generation: Generation, interval: Duration) {
259 let mut entries = self.entries();
260 let Some(entry) = entries
261 .get_mut(&id)
262 .filter(|entry| entry.generation == generation)
263 else {
264 return;
265 };
266 entry.interval = interval.min(MAX_INTERVAL);
267 if entry.failures == 0 && entry.attempt.is_none() {
268 if let Some(last) = entry.last_attempt {
269 entry.next_attempt = after(last, entry.interval);
270 }
271 }
272 }
273
274 pub fn reset(&self, id: i64, generation: Generation, interval: Duration) {
275 self.entries()
276 .insert(id, Entry::new(generation, interval, Instant::now()));
277 }
278
279 pub fn remove(&self, id: i64) {
280 self.entries().remove(&id);
281 }
282
283 pub fn claim_due(
284 self: &Arc<Self>, now: Instant, floors: &HashMap<i64, Duration>,
285 ) -> Vec<FetchClaim> {
286 let mut entries = self.entries();
287 let mut claims: Vec<FetchClaim> = entries
288 .iter_mut()
289 .filter(|(id, entry)| {
290 !entry.interval.is_zero()
291 && entry.attempt.is_none()
292 && now >= entry.due_at(floors.get(*id).copied())
293 })
294 .map(|(id, entry)| {
295 let (result, previous_attempt) = entry.start_attempt(now);
296 FetchClaim {
297 scheduler: Arc::clone(self),
298 id: *id,
299 generation: entry.generation,
300 result,
301 previous_attempt,
302 done: false,
303 }
304 })
305 .collect();
306 claims.sort_by_key(FetchClaim::id);
307
308 claims
309 }
310
311 pub fn claim_forced(
312 self: &Arc<Self>, id: i64, generation: Generation, now: Instant,
313 ) -> ForcedClaim {
314 let mut entries = self.entries();
315 let Some(entry) = entries
316 .get_mut(&id)
317 .filter(|entry| entry.generation == generation)
318 else {
319 return ForcedClaim::Refused;
320 };
321 if let Some(attempt) = &entry.attempt {
322 return ForcedClaim::Joined(AttemptHandle(attempt.clone()));
323 }
324 if let Some(resume) = entry
325 .quota
326 .as_ref()
327 .map(|quota| quota.resume)
328 .filter(|resume| *resume > now)
329 {
330 return ForcedClaim::Recent(Outcome::Failed {
331 kind: FailureKind::RateLimited,
332 retry_after: Some(resume.saturating_duration_since(now)),
333 });
334 }
335 if !entry.forced_delay(&self.policy, now).is_zero() {
336 return ForcedClaim::Recent(entry.last_outcome.clone().unwrap_or(Outcome::Cancelled));
337 }
338
339 let (result, previous_attempt) = entry.start_attempt(now);
340 ForcedClaim::Claimed(FetchClaim {
341 scheduler: Arc::clone(self),
342 id,
343 generation,
344 result,
345 previous_attempt,
346 done: false,
347 })
348 }
349
350 pub(crate) fn forced_retry_in(
351 &self, id: i64, generation: Generation, now: Instant,
352 ) -> Option<Duration> {
353 self.entries()
354 .get(&id)
355 .filter(|entry| entry.generation == generation)
356 .map(|entry| entry.forced_delay(&self.policy, now))
357 }
358
359 pub fn view(&self, id: i64, now: Instant) -> Option<ScheduleView> {
360 self.entries().get(&id).map(|entry| entry.view(now))
361 }
362
363 pub fn next_due_in(&self, id: i64, now: Instant, floor: Option<Duration>) -> Option<Duration> {
364 self.entries()
365 .get(&id)
366 .map(|entry| entry.due_at(floor).saturating_duration_since(now))
367 }
368
369 pub fn set_quota_pause(
370 &self, id: i64, generation: Generation, until: Option<DateTime<Utc>>, now: Instant,
371 wall: DateTime<Utc>,
372 ) -> bool {
373 let mut entries = self.entries();
374 let Some(entry) = entries
375 .get_mut(&id)
376 .filter(|entry| entry.generation == generation)
377 else {
378 return false;
379 };
380 let previous = entry.quota.as_ref().map(|quota| quota.until);
381 let pause = until.and_then(|until| {
382 (until - wall)
383 .to_std()
384 .ok()
385 .filter(|wait| !wait.is_zero())
386 .map(|wait| QuotaPause {
387 until,
388 resume: after(now, wait),
389 })
390 });
391 if let Some(pause) = &pause {
392 entry.next_attempt = entry.next_attempt.max(pause.resume);
393 }
394 let current = pause.as_ref().map(|pause| pause.until);
395 entry.quota = pause;
396
397 previous != current
398 }
399
400 pub fn paused_until(&self, id: i64, wall: DateTime<Utc>) -> Option<DateTime<Utc>> {
401 self.entries()
402 .get(&id)
403 .and_then(|entry| entry.quota.as_ref())
404 .map(|quota| quota.until)
405 .filter(|until| *until > wall)
406 }
407
408 fn finish(
409 &self, id: i64, generation: Generation, now: Instant, outcome: Outcome,
410 ) -> Option<ScheduleView> {
411 let mut entries = self.entries();
412 let entry = entries
413 .get_mut(&id)
414 .filter(|entry| entry.generation == generation)?;
415 entry.finish(&self.policy, now, outcome);
416
417 Some(entry.view(now))
418 }
419
420 fn release(&self, id: i64, generation: Generation, previous_attempt: Option<Instant>) {
421 if let Some(entry) = self
422 .entries()
423 .get_mut(&id)
424 .filter(|entry| entry.generation == generation)
425 {
426 entry.release(previous_attempt);
427 }
428 }
429}
430
431#[derive(Debug)]
432pub enum ForcedClaim {
433 Claimed(FetchClaim),
434 Joined(AttemptHandle),
435 Recent(Outcome),
436 Refused,
437}
438
439pub struct FetchClaim {
440 scheduler: Arc<FetchScheduler>,
441 id: i64,
442 generation: Generation,
443 result: AttemptSender,
444 previous_attempt: Option<Instant>,
445 done: bool,
446}
447
448impl fmt::Debug for FetchClaim {
449 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
450 f.debug_struct("FetchClaim")
451 .field("id", &self.id)
452 .field("generation", &self.generation)
453 .finish()
454 }
455}
456
457impl FetchClaim {
458 pub fn id(&self) -> i64 {
459 self.id
460 }
461
462 pub fn generation(&self) -> Generation {
463 self.generation
464 }
465
466 pub fn handle(&self) -> AttemptHandle {
467 AttemptHandle(self.result.subscribe())
468 }
469
470 pub fn complete(
471 mut self, now: Instant, outcome: Outcome, result: FetchResult,
472 ) -> Option<ScheduleView> {
473 self.done = true;
474 let view = self
475 .scheduler
476 .finish(self.id, self.generation, now, outcome);
477 self.result.send_replace(Some(result));
478 view
479 }
480
481 pub fn release(mut self) {
482 self.done = true;
483 self.scheduler
484 .release(self.id, self.generation, self.previous_attempt);
485 self.result.send_replace(Some(FetchResult::Cancelled));
486 }
487}
488
489impl Drop for FetchClaim {
490 fn drop(&mut self) {
491 if !self.done {
492 self.scheduler
493 .finish(self.id, self.generation, Instant::now(), Outcome::Cancelled);
494 }
495 }
496}
497
498pub struct AttemptHandle(AttemptReceiver);
499
500impl fmt::Debug for AttemptHandle {
501 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
502 f.write_str("AttemptHandle")
503 }
504}
505
506impl AttemptHandle {
507 pub async fn wait(mut self) -> FetchResult {
508 match self.0.wait_for(Option::is_some).await {
509 Ok(result) => result.clone().unwrap_or(FetchResult::Cancelled),
510 Err(_) => FetchResult::Cancelled,
511 }
512 }
513}
514
515#[cfg(test)]
516mod tests {
517 use super::*;
518 use crate::application::test_support::{
519 pipeline,
520 Harness,
521 StubProvider,
522 };
523 use crate::application::RefreshMode;
524 use crate::domain::{
525 PipelineStatus,
526 RefreshState,
527 };
528
529 async fn wait_for_joiners(claim: &FetchClaim, receivers: usize) {
530 tokio::time::timeout(Duration::from_secs(10), async {
531 while claim.result.receiver_count() < receivers {
532 tokio::task::yield_now().await;
533 }
534 })
535 .await
536 .expect("refresh did not join the scheduler claim");
537 }
538
539 #[tokio::test]
540 async fn joined_refreshes_fetch_real_pipeline_data_after_a_run_claim_is_released() {
541 let h = Harness::new().await;
542 let (id, stub) = h.add_stub(StubProvider::new).await;
543 *stub.pipelines.lock().unwrap() = Ok(vec![pipeline(id, "fresh", PipelineStatus::Success)]);
544 h.repository
545 .update_pipelines_cache(id, &[pipeline(id, "stale", PipelineStatus::Running)])
546 .await
547 .unwrap();
548 let lease = h.provider_service.acquire(id).await.unwrap();
549 let ForcedClaim::Claimed(claim) =
550 h.scheduler
551 .claim_forced(id, lease.generation, Instant::now())
552 else {
553 panic!("expected a run-detail claim");
554 };
555 let receivers = claim.result.receiver_count();
556 h.refresh_manager.set_mode(RefreshMode::Idle);
557 h.start().await;
558 let service = Arc::clone(&h.pipeline_service);
559 let first = tokio::spawn(async move { service.refresh_provider(id).await });
560 let service = Arc::clone(&h.pipeline_service);
561 let second = tokio::spawn(async move { service.refresh_provider(id).await });
562 wait_for_joiners(&claim, receivers + 2).await;
563 claim.release();
564
565 let states = [first.await.unwrap().state, second.await.unwrap().state];
566 assert!(states.contains(&RefreshState::Ok));
567 assert!(states
568 .iter()
569 .all(|state| matches!(state, RefreshState::Ok | RefreshState::Recent)));
570 assert_eq!(stub.fetch_count(), 1);
571 assert_eq!(h.cached_names(id).await, ["fresh"]);
572 }
573
574 #[tokio::test]
575 async fn joined_refresh_does_not_retry_through_a_new_quota_pause() {
576 let h = Harness::new().await;
577 let (id, stub) = h.add_stub(StubProvider::new).await;
578 let lease = h.provider_service.acquire(id).await.unwrap();
579 let ForcedClaim::Claimed(claim) =
580 h.scheduler
581 .claim_forced(id, lease.generation, Instant::now())
582 else {
583 panic!("expected a run-detail claim");
584 };
585 let receivers = claim.result.receiver_count();
586 h.refresh_manager.set_mode(RefreshMode::Idle);
587 h.start().await;
588 let service = Arc::clone(&h.pipeline_service);
589 let refresh = tokio::spawn(async move { service.refresh_provider(id).await });
590 wait_for_joiners(&claim, receivers + 1).await;
591 h.scheduler.set_quota_pause(
592 id,
593 lease.generation,
594 Some(Utc::now() + chrono::Duration::minutes(10)),
595 Instant::now(),
596 Utc::now(),
597 );
598 claim.release();
599
600 let result = refresh.await.unwrap();
601 assert_eq!(result.state, RefreshState::Recent);
602 assert_eq!(result.kind, Some(FailureKind::RateLimited));
603 assert_eq!(stub.fetch_count(), 0);
604 }
605
606 fn secs(seconds: u64) -> Duration {
607 Duration::from_secs(seconds)
608 }
609
610 fn scheduler_with(entries: &[(i64, u64)]) -> (Arc<FetchScheduler>, HashMap<i64, Generation>) {
611 let scheduler = Arc::new(FetchScheduler::default());
612 let generations = entries
613 .iter()
614 .map(|(id, interval)| {
615 let generation = Generation::next();
616 scheduler.register(*id, generation, secs(*interval));
617 (*id, generation)
618 })
619 .collect();
620
621 (scheduler, generations)
622 }
623
624 fn no_floors() -> HashMap<i64, Duration> {
625 HashMap::new()
626 }
627
628 fn due_ids(scheduler: &Arc<FetchScheduler>, now: Instant) -> Vec<i64> {
629 scheduler
630 .claim_due(now, &no_floors())
631 .iter()
632 .map(FetchClaim::id)
633 .collect()
634 }
635
636 fn only_claim(scheduler: &Arc<FetchScheduler>, now: Instant) -> FetchClaim {
637 let mut claims = scheduler.claim_due(now, &no_floors());
638 assert_eq!(claims.len(), 1);
639 claims.remove(0)
640 }
641
642 fn failed(kind: FailureKind, retry_after: Option<Duration>) -> Outcome {
643 Outcome::Failed { kind, retry_after }
644 }
645
646 fn finish_with(claim: FetchClaim, now: Instant, outcome: Outcome) -> ScheduleView {
647 let result = match &outcome {
648 Outcome::Failed { .. } => {
649 FetchResult::Committed(Err(DomainError::NetworkError("unreachable".to_string())))
650 }
651 _ => FetchResult::Committed(Ok(0)),
652 };
653 claim.complete(now, outcome, result).unwrap()
654 }
655
656 fn due_after(seconds: u64) -> Vec<i64> {
657 let (scheduler, generations) = scheduler_with(&[(1, 10), (2, 600), (3, 0), (4, 30)]);
658 let start = Instant::now();
659 for id in [1, 2] {
660 let ForcedClaim::Claimed(claim) = scheduler.claim_forced(id, generations[&id], start)
661 else {
662 panic!("expected a claim for {id}");
663 };
664 finish_with(claim, start, Outcome::Success);
665 }
666
667 due_ids(&scheduler, start + secs(seconds))
668 }
669
670 #[test]
671 fn each_provider_waits_for_its_own_interval() {
672 assert_eq!(due_after(5), vec![4]);
673 assert_eq!(due_after(60), vec![1, 4]);
674 assert_eq!(due_after(599), vec![1, 4]);
675 assert_eq!(due_after(600), vec![1, 2, 4]);
676 }
677
678 #[test]
679 fn counters_do_not_change_from_view_or_from_claiming() {
680 let (scheduler, _) = scheduler_with(&[(1, 30)]);
681 let now = Instant::now();
682
683 let before = scheduler.view(1, now).unwrap();
684 assert_eq!(scheduler.view(1, now).unwrap(), before);
685 let claim = only_claim(&scheduler, now);
686 let claimed = scheduler.view(1, now).unwrap();
687
688 assert_eq!(claimed.failures, before.failures);
689 assert_eq!(claimed.next_attempt_in, before.next_attempt_in);
690 assert!(claimed.in_flight);
691 drop(claim);
692 }
693
694 #[test]
695 fn transient_failures_double_up_to_the_cap() {
696 let (scheduler, _) = scheduler_with(&[(1, 30)]);
697 let mut now = Instant::now();
698 let mut waits = Vec::new();
699 for _ in 0..7 {
700 let view = finish_with(
701 only_claim(&scheduler, now),
702 now,
703 failed(FailureKind::Network, None),
704 );
705 waits.push(view.next_attempt_in.as_secs());
706 now += view.next_attempt_in;
707 }
708
709 assert_eq!(waits, vec![60, 120, 240, 480, 600, 600, 600]);
710 assert_eq!(scheduler.view(1, now).unwrap().failures, 7);
711 }
712
713 #[test]
714 fn timeouts_and_unavailable_count_like_network_failures() {
715 for kind in [FailureKind::Timeout, FailureKind::Unavailable] {
716 let (scheduler, _) = scheduler_with(&[(1, 30)]);
717 let now = Instant::now();
718
719 let view = finish_with(only_claim(&scheduler, now), now, failed(kind, None));
720
721 assert_eq!(view.failures, 1, "{kind:?}");
722 assert_eq!(view.next_attempt_in, secs(60), "{kind:?}");
723 }
724 }
725
726 #[test]
727 fn a_server_hint_extends_the_wait_but_never_shortens_it() {
728 let (scheduler, _) = scheduler_with(&[(1, 30)]);
729 let now = Instant::now();
730
731 let long = finish_with(
732 only_claim(&scheduler, now),
733 now,
734 failed(FailureKind::Unavailable, Some(secs(300))),
735 );
736 let later = now + long.next_attempt_in;
737 let short = finish_with(
738 only_claim(&scheduler, later),
739 later,
740 failed(FailureKind::Unavailable, Some(secs(5))),
741 );
742
743 assert_eq!(long.next_attempt_in, secs(300));
744 assert_eq!(short.next_attempt_in, secs(120));
745 }
746
747 #[test]
748 fn rate_limits_use_retry_after_capped_at_an_hour() {
749 let (scheduler, _) = scheduler_with(&[(1, 30)]);
750 let mut now = Instant::now();
751 let mut waits = Vec::new();
752 for retry_after in [Some(secs(120)), Some(secs(2 * 3600)), None, Some(secs(5))] {
753 let view = finish_with(
754 only_claim(&scheduler, now),
755 now,
756 failed(FailureKind::RateLimited, retry_after),
757 );
758 waits.push(view.next_attempt_in.as_secs());
759 now += view.next_attempt_in;
760 }
761
762 assert_eq!(waits, vec![120, 3600, 30 * 8, 30]);
763 }
764
765 #[test]
766 fn an_interval_longer_than_the_wait_caps_is_kept() {
767 let (scheduler, _) = scheduler_with(&[(1, 2 * 3600), (2, 2 * 3600)]);
768 let now = Instant::now();
769 let claims = scheduler.claim_due(now, &no_floors());
770 let mut views = HashMap::new();
771 for claim in claims {
772 let id = claim.id();
773 let outcome = if id == 1 {
774 failed(FailureKind::RateLimited, Some(secs(3600)))
775 } else {
776 failed(FailureKind::Network, None)
777 };
778 views.insert(id, finish_with(claim, now, outcome));
779 }
780
781 assert_eq!(views[&1].next_attempt_in, secs(2 * 3600));
782 assert_eq!(views[&2].next_attempt_in, secs(2 * 3600));
783 }
784
785 #[test]
786 fn auth_failures_do_not_count_and_success_resets() {
787 let (scheduler, _) = scheduler_with(&[(1, 30)]);
788 let mut now = Instant::now();
789
790 let network = finish_with(
791 only_claim(&scheduler, now),
792 now,
793 failed(FailureKind::Network, None),
794 );
795 now += network.next_attempt_in;
796 let auth = finish_with(
797 only_claim(&scheduler, now),
798 now,
799 failed(FailureKind::Auth, None),
800 );
801 now += auth.next_attempt_in;
802 let success = finish_with(only_claim(&scheduler, now), now, Outcome::Success);
803
804 assert_eq!(network.failures, 1);
805 assert_eq!(auth.failures, 1);
806 assert_eq!(auth.last_failure, Some(FailureKind::Auth));
807 assert_eq!(auth.next_attempt_in, secs(30));
808 assert_eq!(success.failures, 0);
809 assert_eq!(success.last_failure, None);
810 assert_eq!(success.next_attempt_in, secs(30));
811 }
812
813 #[test]
814 fn the_outage_floor_delays_but_still_lets_a_probe_through() {
815 let (scheduler, _) = scheduler_with(&[(1, 10)]);
816 let floors = HashMap::from([(1, secs(120))]);
817 let start = Instant::now();
818 finish_with(only_claim(&scheduler, start), start, Outcome::Success);
819
820 assert!(scheduler.claim_due(start + secs(10), &floors).is_empty());
821 assert!(scheduler.claim_due(start + secs(119), &floors).is_empty());
822 assert_eq!(
823 scheduler.next_due_in(1, start + secs(10), Some(secs(120))),
824 Some(secs(110))
825 );
826 assert_eq!(
827 scheduler.next_due_in(1, start + secs(10), None),
828 Some(Duration::ZERO)
829 );
830 let probe = scheduler.claim_due(start + secs(120), &floors);
831
832 assert_eq!(probe.len(), 1);
833 let view = finish_with(
834 probe.into_iter().next().unwrap(),
835 start + secs(120),
836 Outcome::Success,
837 );
838 assert_eq!(view.failures, 0);
839 assert_eq!(view.next_attempt_in, secs(10));
840 }
841
842 #[test]
843 fn interval_zero_is_never_scheduled_but_can_be_forced() {
844 let (scheduler, generations) = scheduler_with(&[(1, 0)]);
845 let now = Instant::now();
846
847 assert!(scheduler
848 .claim_due(now + secs(86_400), &no_floors())
849 .is_empty());
850 assert!(matches!(
851 scheduler.claim_forced(1, generations[&1], now),
852 ForcedClaim::Claimed(_)
853 ));
854 }
855
856 #[test]
857 fn forced_claims_bypass_next_attempt_and_the_outage_floor() {
858 let (scheduler, generations) = scheduler_with(&[(1, 600)]);
859 let start = Instant::now();
860 finish_with(only_claim(&scheduler, start), start, Outcome::Success);
861 let later = start + secs(20);
862
863 assert!(scheduler.claim_due(later, &no_floors()).is_empty());
864 assert!(matches!(
865 scheduler.claim_forced(1, generations[&1], later),
866 ForcedClaim::Claimed(_)
867 ));
868 }
869
870 #[test]
871 fn forced_claims_wait_out_a_rate_limit_and_the_minimum_gap() {
872 let (scheduler, generations) = scheduler_with(&[(1, 30), (2, 30)]);
873 let start = Instant::now();
874 let mut claims = scheduler.claim_due(start, &no_floors());
875 let second = claims.pop().unwrap();
876 let first = claims.pop().unwrap();
877 finish_with(
878 first,
879 start,
880 failed(FailureKind::RateLimited, Some(secs(120))),
881 );
882 finish_with(second, start, Outcome::Success);
883
884 let during_rate_limit = scheduler.claim_forced(1, generations[&1], start + secs(60));
885 let within_gap = scheduler.claim_forced(2, generations[&2], start + secs(9));
886 let after_gap = scheduler.claim_forced(2, generations[&2], start + secs(10));
887 let after_rate_limit = scheduler.claim_forced(1, generations[&1], start + secs(120));
888
889 assert!(matches!(
890 during_rate_limit,
891 ForcedClaim::Recent(Outcome::Failed {
892 kind: FailureKind::RateLimited,
893 ..
894 })
895 ));
896 assert!(matches!(within_gap, ForcedClaim::Recent(Outcome::Success)));
897 assert!(matches!(after_gap, ForcedClaim::Claimed(_)));
898 assert!(matches!(after_rate_limit, ForcedClaim::Claimed(_)));
899 }
900
901 #[tokio::test]
902 async fn forced_claims_join_an_attempt_from_claim_time_until_it_completes() {
903 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
904 let now = Instant::now();
905 let claim = only_claim(&scheduler, now);
906
907 let ForcedClaim::Joined(joiner) = scheduler.claim_forced(1, generations[&1], now) else {
908 panic!("expected to join the in-flight attempt");
909 };
910 let waiting = tokio::spawn(joiner.wait());
911 tokio::task::yield_now().await;
912 assert!(!waiting.is_finished());
913 let ForcedClaim::Joined(late) = scheduler.claim_forced(1, generations[&1], now) else {
914 panic!("expected to join the in-flight attempt");
915 };
916 claim.complete(now, Outcome::Success, FetchResult::Committed(Ok(3)));
917
918 assert!(matches!(
919 waiting.await.unwrap(),
920 FetchResult::Committed(Ok(3))
921 ));
922 assert!(matches!(late.wait().await, FetchResult::Committed(Ok(3))));
923 assert!(!scheduler.view(1, now).unwrap().in_flight);
924 }
925
926 #[tokio::test]
927 async fn a_dropped_claim_cancels_its_joiners_without_counting_a_failure() {
928 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
929 let now = Instant::now();
930 let claim = only_claim(&scheduler, now);
931 let ForcedClaim::Joined(joiner) = scheduler.claim_forced(1, generations[&1], now) else {
932 panic!("expected to join the in-flight attempt");
933 };
934
935 drop(claim);
936
937 assert!(matches!(joiner.wait().await, FetchResult::Cancelled));
938 let view = scheduler.view(1, now).unwrap();
939 assert_eq!(view.failures, 0);
940 assert_eq!(view.last_failure, None);
941 assert!(!view.in_flight);
942 let remaining = scheduler.view(1, Instant::now()).unwrap().next_attempt_in;
943 assert!(remaining <= secs(30));
944 assert!(remaining > Duration::ZERO);
945 }
946
947 #[test]
948 fn an_old_claim_changes_nothing_after_reset() {
949 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
950 let now = Instant::now();
951 let old = only_claim(&scheduler, now);
952 let fresh = Generation::next();
953 scheduler.reset(1, fresh, secs(30));
954 let current = only_claim(&scheduler, Instant::now());
955 let before = scheduler.view(1, now).unwrap();
956
957 let stale = old.complete(
958 now,
959 failed(FailureKind::Network, None),
960 FetchResult::Superseded,
961 );
962 let after = scheduler.view(1, now).unwrap();
963
964 assert!(stale.is_none());
965 assert_eq!(after, before);
966 assert!(after.in_flight);
967 assert_eq!(after.failures, 0);
968 assert_ne!(generations[&1], fresh);
969 assert_eq!(current.generation(), fresh);
970 }
971
972 #[test]
973 fn an_old_claim_dropped_after_reset_leaves_the_new_attempt_alone() {
974 let (scheduler, _) = scheduler_with(&[(1, 30)]);
975 let now = Instant::now();
976 let old = only_claim(&scheduler, now);
977 scheduler.reset(1, Generation::next(), secs(30));
978 let current = only_claim(&scheduler, Instant::now());
979
980 drop(old);
981
982 assert!(scheduler.view(1, now).unwrap().in_flight);
983 drop(current);
984 assert!(!scheduler.view(1, now).unwrap().in_flight);
985 }
986
987 #[test]
988 fn claims_after_remove_are_refused_and_the_entry_is_not_recreated() {
989 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
990 let now = Instant::now();
991 let old = only_claim(&scheduler, now);
992 scheduler.remove(1);
993
994 assert!(matches!(
995 scheduler.claim_forced(1, generations[&1], now),
996 ForcedClaim::Refused
997 ));
998 assert!(scheduler.claim_due(now, &no_floors()).is_empty());
999 old.complete(now, Outcome::Success, FetchResult::Superseded);
1000 assert!(scheduler.view(1, now).is_none());
1001 assert!(scheduler.claim_due(now, &no_floors()).is_empty());
1002 scheduler.set_interval(1, generations[&1], secs(5));
1003 assert!(scheduler.view(1, now).is_none());
1004 }
1005
1006 #[test]
1007 fn a_reused_id_with_a_new_generation_ignores_old_claims() {
1008 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
1009 let now = Instant::now();
1010 let old = only_claim(&scheduler, now);
1011 scheduler.remove(1);
1012 let reused = Generation::next();
1013 scheduler.register(1, reused, secs(30));
1014
1015 assert!(matches!(
1016 scheduler.claim_forced(1, generations[&1], now),
1017 ForcedClaim::Refused
1018 ));
1019 old.complete(
1020 now,
1021 failed(FailureKind::Network, None),
1022 FetchResult::Superseded,
1023 );
1024 let view = scheduler.view(1, now).unwrap();
1025
1026 assert_eq!(view.failures, 0);
1027 assert!(!view.in_flight);
1028 assert!(matches!(
1029 scheduler.claim_forced(1, reused, now),
1030 ForcedClaim::Claimed(_)
1031 ));
1032 }
1033
1034 #[test]
1035 fn register_keeps_an_existing_entry_of_the_same_generation() {
1036 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
1037 let now = Instant::now();
1038 finish_with(
1039 only_claim(&scheduler, now),
1040 now,
1041 failed(FailureKind::Network, None),
1042 );
1043
1044 scheduler.register(1, generations[&1], secs(30));
1045
1046 assert_eq!(scheduler.view(1, now).unwrap().failures, 1);
1047 }
1048
1049 #[test]
1050 fn a_shorter_interval_brings_the_next_attempt_forward() {
1051 let (scheduler, generations) = scheduler_with(&[(1, 3600)]);
1052 let now = Instant::now();
1053 finish_with(only_claim(&scheduler, now), now, Outcome::Success);
1054
1055 scheduler.set_interval(1, generations[&1], secs(30));
1056
1057 assert!(scheduler.view(1, Instant::now()).unwrap().next_attempt_in <= secs(30));
1058 assert_eq!(due_ids(&scheduler, now + secs(30)), vec![1]);
1059 }
1060
1061 #[test]
1062 fn a_quota_pause_postpones_every_outcome_until_the_reset() {
1063 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
1064 let now = Instant::now();
1065 let wall = Utc::now();
1066 let reset = wall + chrono::Duration::minutes(20);
1067 let claim = only_claim(&scheduler, now);
1068
1069 let changed = scheduler.set_quota_pause(1, generations[&1], Some(reset), now, wall);
1070 let repeated = scheduler.set_quota_pause(1, generations[&1], Some(reset), now, wall);
1071 let view = claim.complete(now, Outcome::Cancelled, FetchResult::Cancelled);
1072
1073 assert!(changed);
1074 assert!(!repeated);
1075 assert_eq!(view.unwrap().next_attempt_in, secs(20 * 60));
1076 assert_eq!(scheduler.paused_until(1, wall), Some(reset));
1077 assert!(scheduler
1078 .claim_due(now + secs(20 * 60 - 1), &no_floors())
1079 .is_empty());
1080 assert_eq!(due_ids(&scheduler, now + secs(20 * 60)), vec![1]);
1081 }
1082
1083 #[test]
1084 fn a_quota_pause_turns_manual_refreshes_into_a_rate_limit_answer() {
1085 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
1086 let now = Instant::now();
1087 let wall = Utc::now();
1088 scheduler.set_quota_pause(
1089 1,
1090 generations[&1],
1091 Some(wall + chrono::Duration::minutes(5)),
1092 now,
1093 wall,
1094 );
1095
1096 let forced = scheduler.claim_forced(1, generations[&1], now + secs(60));
1097
1098 let ForcedClaim::Recent(Outcome::Failed { kind, retry_after }) = forced else {
1099 panic!("expected a recent rate-limit answer, got {forced:?}");
1100 };
1101 assert_eq!(kind, FailureKind::RateLimited);
1102 assert_eq!(retry_after, Some(secs(240)));
1103 }
1104
1105 #[test]
1106 fn a_quota_pause_clears_when_the_quota_is_back_and_ignores_stale_generations() {
1107 let (scheduler, generations) = scheduler_with(&[(1, 30)]);
1108 let now = Instant::now();
1109 let wall = Utc::now();
1110 let reset = wall + chrono::Duration::minutes(5);
1111 let stale = Generation::next();
1112
1113 let ignored = scheduler.set_quota_pause(1, stale, Some(reset), now, wall);
1114 scheduler.set_quota_pause(1, generations[&1], Some(reset), now, wall);
1115 let cleared = scheduler.set_quota_pause(1, generations[&1], None, now, wall);
1116 let already_clear = scheduler.set_quota_pause(1, generations[&1], None, now, wall);
1117
1118 assert!(!ignored);
1119 assert!(cleared);
1120 assert!(!already_clear);
1121 assert_eq!(scheduler.paused_until(1, wall), None);
1122 let past = wall - chrono::Duration::seconds(1);
1123 assert!(!scheduler.set_quota_pause(1, generations[&1], Some(past), now, wall));
1124 assert_eq!(scheduler.paused_until(1, wall), None);
1125 }
1126
1127 #[test]
1128 fn generations_are_never_reused() {
1129 let first = Generation::next();
1130 let second = Generation::next();
1131
1132 assert_ne!(first, second);
1133 assert!(second > first);
1134 }
1135}