use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use tokio::sync::watch;
use tokio::time::Instant;
const SERVER_OVERLOADED_SIGNATURES: &[&str] = &[
"too many simultaneous connections",
"too many connections",
"too many requests",
"try again later",
"temporarily unavailable",
"server unavailable",
"over quota",
"rate limit",
"throttl",
];
pub(crate) fn is_server_overloaded_error(err: &str) -> bool {
let lowered = err.to_ascii_lowercase();
SERVER_OVERLOADED_SIGNATURES
.iter()
.any(|signature| lowered.contains(signature))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum GateDecision {
Retry,
GiveUp,
}
pub(crate) struct RecoveryGate {
open: watch::Sender<bool>,
prober: tokio::sync::Mutex<()>,
rounds: AtomicUsize,
max_rounds: usize,
backoff: Duration,
leg: &'static str,
}
impl RecoveryGate {
pub(crate) fn new(leg: &'static str, max_rounds: usize, backoff: Duration) -> Self {
Self {
open: watch::Sender::new(true),
prober: tokio::sync::Mutex::new(()),
rounds: AtomicUsize::new(0),
max_rounds,
backoff,
leg,
}
}
pub(crate) async fn wait_for_recovery(&self, subject: &str, err: &str) -> GateDecision {
match self.prober.try_lock() {
Ok(_election) => {
let _ = self.open.send(false);
let round = self.rounds.fetch_add(1, Ordering::SeqCst) + 1;
if round > self.max_rounds {
let _ = self.open.send(true);
return GateDecision::GiveUp;
}
tracing::warn!(
pigeon_job = "email-import",
leg = self.leg,
subject,
round,
rounds = self.max_rounds,
error = %err,
"email-import: endpoint refusing work; pausing"
);
let paused = Instant::now();
tokio::time::sleep(self.backoff * round as u32).await;
let _ = self.open.send(true);
tracing::info!(
pigeon_job = "email-import",
leg = self.leg,
round,
paused_ms = paused.elapsed().as_millis(),
"email-import: resuming"
);
GateDecision::Retry
}
Err(_) => {
let mut open = self.open.subscribe();
while !*open.borrow_and_update() {
if open.changed().await.is_err() {
break;
}
}
if self.rounds.load(Ordering::SeqCst) > self.max_rounds {
GateDecision::GiveUp
} else {
GateDecision::Retry
}
}
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::*;
const BACKOFF: Duration = Duration::from_secs(60);
#[test]
fn is_server_overloaded_error_matches_the_signature_adr_0071_actually_observed() {
assert!(is_server_overloaded_error(
"failed to connect: Too many simultaneous connections"
));
}
#[test]
fn is_server_overloaded_error_matches_every_named_signature_case_insensitively() {
for signature in SERVER_OVERLOADED_SIGNATURES {
let shouted = signature.to_ascii_uppercase();
assert!(
is_server_overloaded_error(&format!("server said: {shouted}")),
"{signature} must match regardless of case"
);
}
}
#[test]
fn is_server_overloaded_error_rejects_an_unrelated_failure() {
assert!(!is_server_overloaded_error(
"LOGIN failed: invalid credentials"
));
assert!(!is_server_overloaded_error("failed to parse BODYSTRUCTURE"));
}
#[test]
fn is_server_overloaded_error_does_not_claim_an_unreachable_endpoint_error() {
let unreachable = "dial tcp: lookup s3.fr-par.scw.cloud: no such host";
assert!(!is_server_overloaded_error(unreachable));
assert!(crate::core::rclone::is_endpoint_unreachable_error(
unreachable
));
}
#[tokio::test(start_paused = true)]
async fn first_round_sleeps_the_backoff_and_then_retries() {
let gate = RecoveryGate::new("imap", 6, BACKOFF);
let start = Instant::now();
let decision = gate.wait_for_recovery("INBOX", "rate limit exceeded").await;
assert_eq!(decision, GateDecision::Retry);
assert_eq!(Instant::now(), start + BACKOFF, "round 1 sleeps 1x backoff");
}
#[tokio::test(start_paused = true)]
async fn a_task_losing_the_election_parks_until_the_prober_reopens() {
let gate = Arc::new(RecoveryGate::new("imap", 6, BACKOFF));
let start = Instant::now();
let prober = {
let gate = Arc::clone(&gate);
tokio::spawn(async move { gate.wait_for_recovery("INBOX", "rate limit").await })
};
tokio::task::yield_now().await;
let parked = gate.wait_for_recovery("Archive", "rate limit").await;
assert_eq!(parked, GateDecision::Retry);
assert_eq!(prober.await.unwrap(), GateDecision::Retry);
assert_eq!(
Instant::now(),
start + BACKOFF,
"one shared wait, not two rounds"
);
assert_eq!(
gate.rounds.load(Ordering::SeqCst),
1,
"the parked task must not have started its own round"
);
}
#[tokio::test(start_paused = true)]
async fn backoff_scales_linearly_with_the_round_number() {
let gate = RecoveryGate::new("imap", 6, BACKOFF);
let start = Instant::now();
gate.wait_for_recovery("INBOX", "rate limit").await;
gate.wait_for_recovery("INBOX", "rate limit").await;
gate.wait_for_recovery("INBOX", "rate limit").await;
assert_eq!(Instant::now(), start + BACKOFF * 6);
}
#[tokio::test(start_paused = true)]
async fn gives_up_past_the_budget_and_stays_latched() {
let gate = RecoveryGate::new("imap", 2, BACKOFF);
assert_eq!(
gate.wait_for_recovery("INBOX", "rate limit").await,
GateDecision::Retry
);
assert_eq!(
gate.wait_for_recovery("INBOX", "rate limit").await,
GateDecision::Retry
);
let latched_at = Instant::now();
for _ in 0..3 {
assert_eq!(
gate.wait_for_recovery("INBOX", "rate limit").await,
GateDecision::GiveUp
);
}
assert_eq!(
Instant::now(),
latched_at,
"a latched gate must not sleep again"
);
}
#[tokio::test(start_paused = true)]
async fn a_parked_task_also_gives_up_once_the_budget_is_spent() {
let gate = Arc::new(RecoveryGate::new("imap", 1, BACKOFF));
gate.wait_for_recovery("INBOX", "rate limit").await;
assert_eq!(
gate.wait_for_recovery("INBOX", "rate limit").await,
GateDecision::GiveUp
);
let prober = {
let gate = Arc::clone(&gate);
tokio::spawn(async move { gate.wait_for_recovery("INBOX", "rate limit").await })
};
tokio::task::yield_now().await;
assert_eq!(
gate.wait_for_recovery("Archive", "rate limit").await,
GateDecision::GiveUp
);
assert_eq!(prober.await.unwrap(), GateDecision::GiveUp);
}
}