use std::future::Future;
use std::time::Duration;
use http::header::{HeaderMap, RETRY_AFTER};
use http::{Method, StatusCode};
use polyoxide_venue::{Class, Classify};
use crate::error::ApiError;
use crate::rate_limit::RetryConfig;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct LayerId(pub &'static str);
impl std::fmt::Display for LayerId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.0)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Cost {
pub layer: LayerId,
pub units: u32,
pub exact: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct LayerCharge {
pub layer: LayerId,
pub units: u32,
pub window: Option<u64>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Charge {
layers: Vec<LayerCharge>,
}
impl Charge {
pub fn none() -> Self {
Self::default()
}
pub fn with(mut self, charge: LayerCharge) -> Self {
self.layers.push(charge);
self
}
pub fn layers(&self) -> &[LayerCharge] {
&self.layers
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[error("the {layer} layer holds at most {capacity} units, and this request costs {units}")]
pub struct Refused {
pub layer: LayerId,
pub units: u32,
pub capacity: u32,
}
impl Classify for Refused {
fn class(&self) -> Class {
Class::InvalidRequest
}
}
#[derive(Debug, Clone, Copy)]
pub struct RequestMeta<'a> {
pub method: &'a Method,
pub path: &'a str,
pub query: &'a [(String, String)],
pub costs: &'a [Cost],
}
#[derive(Debug, Clone, Copy)]
pub struct ResponseMeta<'a> {
pub status: StatusCode,
pub headers: &'a HeaderMap,
}
impl ResponseMeta<'_> {
pub fn retry_after(&self) -> Option<&str> {
self.headers.get(RETRY_AFTER)?.to_str().ok()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AttemptInfo {
pub attempt: u32,
pub retries_left: u32,
}
#[derive(Debug, Clone)]
pub struct RequestParts {
pub method: Method,
pub path: String,
pub query: Vec<(String, String)>,
pub headers: HeaderMap,
pub body: Option<bytes::Bytes>,
pub timeout: Option<Duration>,
}
impl RequestParts {
pub fn new(method: Method, path: impl Into<String>) -> Self {
Self {
method,
path: path.into(),
query: Vec::new(),
headers: HeaderMap::new(),
body: None,
timeout: None,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Outcome {
Done,
Retry(Duration),
Fail,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Decision {
pub outcome: Outcome,
pub hold: Option<Duration>,
}
#[dynosaur::dynosaur(pub DynThrottle = dyn(box) Throttle)]
pub trait Throttle: Send + Sync {
fn acquire(
&self,
meta: &RequestMeta<'_>,
) -> impl Future<Output = Result<Charge, Refused>> + Send;
fn observe(&self, charge: &Charge, response: &ResponseMeta<'_>, attempt: &AttemptInfo);
fn hold(&self, delay: Duration);
}
#[dynosaur::dynosaur(pub DynRetryPolicy = dyn(box) RetryPolicy)]
pub trait RetryPolicy: Send + Sync {
fn decide(
&self,
response: &ResponseMeta<'_>,
attempt: &AttemptInfo,
schedule: &RetryConfig,
) -> Decision;
}
#[dynosaur::dynosaur(pub DynAuthenticator = dyn(box) Authenticator)]
pub trait Authenticator: Send + Sync {
fn sign(
&self,
parts: &mut RequestParts,
attempt: u32,
) -> impl Future<Output = Result<(), ApiError>> + Send;
}
#[derive(Debug, Clone, Copy, Default)]
pub struct NoThrottle;
impl Throttle for NoThrottle {
async fn acquire(&self, _meta: &RequestMeta<'_>) -> Result<Charge, Refused> {
Ok(Charge::none())
}
fn observe(&self, _charge: &Charge, _response: &ResponseMeta<'_>, _attempt: &AttemptInfo) {}
fn hold(&self, _delay: Duration) {}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct DefaultRetryPolicy;
impl RetryPolicy for DefaultRetryPolicy {
fn decide(
&self,
response: &ResponseMeta<'_>,
_attempt: &AttemptInfo,
schedule: &RetryConfig,
) -> Decision {
if response.status == StatusCode::TOO_MANY_REQUESTS {
Decision {
outcome: Outcome::Retry(Duration::ZERO),
hold: Some(schedule.retry_delay(0, response.retry_after())),
}
} else if response.status.is_success() {
Decision {
outcome: Outcome::Done,
hold: None,
}
} else {
Decision {
outcome: Outcome::Fail,
hold: None,
}
}
}
}