use std::time::{Duration, SystemTime, SystemTimeError};
use crate::headers::{HeaderName, HeaderValue, Headers, RETRY_AFTER};
use crate::utils::{fmt_http_date, parse_http_date};
#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct RetryAfter {
inner: RetryDirective,
}
#[allow(clippy::len_without_is_empty)]
impl RetryAfter {
pub fn new(dur: Duration) -> Self {
Self {
inner: RetryDirective::Duration(dur),
}
}
pub fn new_at(at: SystemTime) -> Self {
Self {
inner: RetryDirective::SystemTime(at),
}
}
pub fn from_headers(headers: impl AsRef<Headers>) -> crate::Result<Option<Self>> {
let header = match headers.as_ref().get(RETRY_AFTER) {
Some(headers) => headers.last(),
None => return Ok(None),
};
let inner = match header.as_str().parse::<u64>() {
Ok(dur) => RetryDirective::Duration(Duration::from_secs(dur)),
Err(_) => {
let at = parse_http_date(header.as_str())?;
RetryDirective::SystemTime(at)
}
};
Ok(Some(Self { inner }))
}
pub fn duration_since(&self, earlier: SystemTime) -> Result<Duration, SystemTimeError> {
let at = match self.inner {
RetryDirective::Duration(dur) => SystemTime::now() + dur,
RetryDirective::SystemTime(at) => at,
};
at.duration_since(earlier)
}
pub fn apply(&self, mut headers: impl AsMut<Headers>) {
headers.as_mut().insert(self.name(), self.value());
}
pub fn name(&self) -> HeaderName {
RETRY_AFTER
}
pub fn value(&self) -> HeaderValue {
let output = match self.inner {
RetryDirective::Duration(dur) => format!("{}", dur.as_secs()),
RetryDirective::SystemTime(at) => fmt_http_date(at),
};
unsafe { HeaderValue::from_bytes_unchecked(output.into()) }
}
}
impl From<RetryAfter> for SystemTime {
fn from(retry_after: RetryAfter) -> Self {
match retry_after.inner {
RetryDirective::Duration(dur) => SystemTime::now() + dur,
RetryDirective::SystemTime(at) => at,
}
}
}
#[derive(Clone, Debug, Eq, PartialEq, Hash, PartialOrd, Ord)]
enum RetryDirective {
Duration(Duration),
SystemTime(SystemTime),
}
#[cfg(test)]
mod test {
use super::*;
use crate::headers::Headers;
#[test]
fn smoke() -> crate::Result<()> {
let retry = RetryAfter::new(Duration::from_secs(10));
let mut headers = Headers::new();
retry.apply(&mut headers);
let now = SystemTime::now();
let retry = RetryAfter::from_headers(headers)?.unwrap();
assert_eq!(
retry.duration_since(now)?.as_secs(),
Duration::from_secs(10).as_secs()
);
Ok(())
}
#[test]
fn new_at() -> crate::Result<()> {
let now = SystemTime::now();
let retry = RetryAfter::new_at(now + Duration::from_secs(10));
let mut headers = Headers::new();
retry.apply(&mut headers);
let retry = RetryAfter::from_headers(headers)?.unwrap();
let delta = retry.duration_since(now)?;
assert!(delta >= Duration::from_secs(9));
assert!(delta <= Duration::from_secs(10));
Ok(())
}
}