use std::marker::PhantomData;
use polyoxide_core::{retry_after_header, truncate_for_log, ApiError, HttpClient};
use reqwest::{Response, StatusCode};
use serde::de::DeserializeOwned;
use crate::{
error::{retry_after_secs, BinanceError},
weight::{Cost, Route, WeightBudget, DEFAULT_BAN},
};
pub const USED_WEIGHT_HEADER: &str = "x-mbx-used-weight-1m";
#[must_use = "a request does nothing until it is sent"]
pub struct WeightedRequest<T> {
http: HttpClient,
budget: WeightBudget,
route: Route,
query: Vec<(&'static str, String)>,
_marker: PhantomData<fn() -> T>,
}
impl<T> WeightedRequest<T> {
pub(crate) fn new(http: &HttpClient, budget: &WeightBudget, route: Route) -> Self {
Self {
http: http.clone(),
budget: budget.clone(),
route,
query: Vec::new(),
_marker: PhantomData,
}
}
pub(crate) fn query(mut self, key: &'static str, value: impl ToString) -> Self {
let value = value.to_string();
match self.query.iter_mut().find(|(k, _)| *k == key) {
Some(slot) => slot.1 = value,
None => self.query.push((key, value)),
}
self
}
pub(crate) fn route(mut self, route: Route) -> Self {
self.route = route;
self
}
pub fn cost(&self) -> Cost {
self.route.cost()
}
}
impl<T: DeserializeOwned> WeightedRequest<T> {
pub async fn send(self) -> Result<T, BinanceError> {
let text = self.send_text().await?;
serde_json::from_str(&text).map_err(|err| {
tracing::error!(
"Failed to decode {}: {err}: {}",
self.route.path(),
truncate_for_log(&text)
);
BinanceError::from(ApiError::from(err))
})
}
async fn send_text(&self) -> Result<String, BinanceError> {
let path = self.route.path();
let url = self.http.base_url.join(path)?;
let cost = self.route.cost();
let mut attempt = 0u32;
loop {
let permit = self.http.acquire_concurrency().await;
let charge = self.budget.acquire(cost).await;
let mut request = self.http.client.get(url.clone());
if !self.query.is_empty() {
request = request.query(&self.query);
}
let response = request.send().await?;
let status = response.status();
let retry_after = retry_after_header(&response);
if let Some(used) = used_weight(&response) {
self.budget.record_used(charge, used);
}
if status == StatusCode::IM_A_TEAPOT {
let ban = retry_after_secs(retry_after.as_deref()).unwrap_or(DEFAULT_BAN);
self.budget.begin_cooldown(ban);
let body = response.text().await.unwrap_or_default();
tracing::warn!(
"418 on {path}: IP banned, every request held {ban:?}: {}",
truncate_for_log(&body)
);
return Err(BinanceError::from_response_parts(
status.as_u16(),
retry_after.as_deref(),
&body,
));
}
if status == StatusCode::TOO_MANY_REQUESTS {
let retry = self
.http
.should_retry(status, attempt, retry_after.as_deref());
let asked = retry_after_secs(retry_after.as_deref());
let cooldown = match (retry, asked) {
(None, None) => self.budget.hold_until_next_minute(),
_ => {
let cooldown = retry.unwrap_or_default().max(asked.unwrap_or_default());
self.budget.begin_cooldown(cooldown);
cooldown
}
};
if retry.is_some() {
attempt += 1;
tracing::warn!("429 on {path}, retry {attempt} after {cooldown:?}");
drop(permit);
continue;
}
tracing::warn!("429 on {path}, no retry left: every request held {cooldown:?}");
}
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
return Err(BinanceError::from_response_parts(
status.as_u16(),
retry_after.as_deref(),
&body,
));
}
return Ok(response.text().await?);
}
}
}
fn used_weight(response: &Response) -> Option<u32> {
response
.headers()
.get(USED_WEIGHT_HEADER)?
.to_str()
.ok()?
.trim()
.parse()
.ok()
}