#[cfg(proxima_alloc)]
use alloc::collections::BTreeSet;
use crate::pipe::capabilities::Retryable;
#[cfg(not(proxima_alloc))]
use crate::pipe::sized::RETRY_STATUS_CAP;
#[derive(Debug, Clone)]
pub struct RetryRules {
pub retry_on_status: StatusSet,
pub retry_on_error: bool,
pub idempotent_only: bool,
}
#[cfg(proxima_alloc)]
pub type StatusSet = BTreeSet<u16>;
#[cfg(not(proxima_alloc))]
#[derive(Debug, Clone, Default)]
pub struct StatusSet {
inner: heapless::Vec<u16, RETRY_STATUS_CAP>,
}
#[cfg(not(proxima_alloc))]
impl StatusSet {
#[must_use]
pub fn new() -> Self {
Self {
inner: heapless::Vec::new(),
}
}
pub fn insert(&mut self, status: u16) -> bool {
match self.inner.binary_search(&status) {
Ok(_) => false,
Err(position) => self.inner.insert(position, status).is_ok(),
}
}
#[must_use]
pub fn contains(&self, status: &u16) -> bool {
self.inner.binary_search(status).is_ok()
}
}
#[cfg(not(proxima_alloc))]
impl Extend<u16> for StatusSet {
fn extend<I: IntoIterator<Item = u16>>(&mut self, statuses: I) {
for status in statuses {
self.insert(status);
}
}
}
#[cfg(not(proxima_alloc))]
impl IntoIterator for StatusSet {
type Item = u16;
type IntoIter = <heapless::Vec<u16, RETRY_STATUS_CAP> as IntoIterator>::IntoIter;
fn into_iter(self) -> Self::IntoIter {
self.inner.into_iter()
}
}
impl Default for RetryRules {
fn default() -> Self {
let mut retry_on_status = StatusSet::default();
for status in [502u16, 503, 504] {
retry_on_status.insert(status);
}
Self {
retry_on_status,
retry_on_error: true,
idempotent_only: false,
}
}
}
impl RetryRules {
pub fn should_retry<Out: Retryable, Err>(&self, outcome: &Result<Out, Err>) -> bool {
match outcome {
Ok(out) => out
.retry_status()
.is_some_and(|status| self.retry_on_status.contains(&status)),
Err(_) => self.retry_on_error,
}
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
struct Outcome(Option<u16>);
impl Retryable for Outcome {
fn retry_status(&self) -> Option<u16> {
self.0
}
fn is_success(&self) -> bool {
self.0.is_none()
}
}
#[test]
fn default_retries_the_gateway_statuses() {
let rules = RetryRules::default();
for status in [502u16, 503, 504] {
assert!(
rules.retry_on_status.contains(&status),
"{status} should retry"
);
}
assert!(
!rules.retry_on_status.contains(&200),
"200 should not retry"
);
}
#[test]
fn should_retry_matches_status_membership() {
let rules = RetryRules::default();
let ok: Result<Outcome, ()> = Ok(Outcome(Some(503)));
let not_listed: Result<Outcome, ()> = Ok(Outcome(Some(418)));
let statusless: Result<Outcome, ()> = Ok(Outcome(None));
assert!(rules.should_retry(&ok), "503 is in the default set");
assert!(!rules.should_retry(¬_listed), "418 is not");
assert!(
!rules.should_retry(&statusless),
"no status => no status retry"
);
}
#[test]
fn should_retry_on_error_follows_the_flag() {
let err: Result<Outcome, ()> = Err(());
assert!(
RetryRules::default().should_retry(&err),
"default retries on error"
);
let no_error_retry = RetryRules {
retry_on_error: false,
..RetryRules::default()
};
assert!(
!no_error_retry.should_retry(&err),
"flag off => no error retry"
);
}
#[test]
fn insert_is_idempotent_and_sorted() {
let mut set = StatusSet::default();
set.insert(503);
set.insert(502);
set.insert(503);
assert!(set.contains(&502) && set.contains(&503));
assert!(!set.contains(&504));
}
}