2023-04-12 22:59:08 +02:00
|
|
|
use crate::api::limits::{Limit, LimitType, Limits};
|
2023-04-10 17:20:02 +02:00
|
|
|
|
2023-04-12 18:33:16 +02:00
|
|
|
use reqwest::{Client, RequestBuilder, Response};
|
2023-04-12 22:59:08 +02:00
|
|
|
use std::collections::{HashMap, VecDeque};
|
2023-04-07 21:01:48 +02:00
|
|
|
|
2023-04-08 14:51:36 +02:00
|
|
|
// Note: There seem to be some overlapping request limiters. We need to make sure that sending a
|
|
|
|
// request checks for all the request limiters that apply, and blocks if any of the limiters are 0
|
|
|
|
|
2023-04-12 17:10:12 +02:00
|
|
|
#[allow(dead_code)]
|
2023-04-07 21:51:50 +02:00
|
|
|
pub struct LimitedRequester {
|
2023-04-07 21:01:48 +02:00
|
|
|
http: Client,
|
2023-04-12 18:33:16 +02:00
|
|
|
requests: VecDeque<RequestBuilder>,
|
2023-04-10 20:58:07 +02:00
|
|
|
last_reset_epoch: i64,
|
2023-04-12 22:59:08 +02:00
|
|
|
limits_rate: HashMap<LimitType, Limit>,
|
2023-04-07 21:01:48 +02:00
|
|
|
}
|
|
|
|
|
2023-04-07 21:51:50 +02:00
|
|
|
impl LimitedRequester {
|
|
|
|
/// Create a new `LimitedRequester`. `LimitedRequester`s use a `VecDeque` to store requests and
|
|
|
|
/// send them to the server using a `Client`. It keeps track of the remaining requests that can
|
|
|
|
/// be send within the `Limit` of an external API Ratelimiter, and looks at the returned request
|
|
|
|
/// headers to see if it can find Ratelimit info to update itself.
|
2023-04-12 17:10:12 +02:00
|
|
|
#[allow(dead_code)]
|
2023-04-08 14:51:36 +02:00
|
|
|
pub async fn new(api_url: String) -> Self {
|
2023-04-07 21:51:50 +02:00
|
|
|
LimitedRequester {
|
2023-04-07 21:01:48 +02:00
|
|
|
http: Client::new(),
|
2023-04-07 21:51:50 +02:00
|
|
|
requests: VecDeque::new(),
|
2023-04-10 20:58:07 +02:00
|
|
|
last_reset_epoch: chrono::Utc::now().timestamp(),
|
2023-04-12 16:49:18 +02:00
|
|
|
limits_rate: Limits::check_limits(api_url).await,
|
2023-04-07 21:01:48 +02:00
|
|
|
}
|
|
|
|
}
|
2023-04-12 18:33:16 +02:00
|
|
|
|
|
|
|
fn add_to_queue(request: RequestBuilder, queue: &mut VecDeque<RequestBuilder>) {
|
|
|
|
queue.push_back(request);
|
|
|
|
}
|
|
|
|
|
2023-04-14 17:09:50 +02:00
|
|
|
fn update_limit_entry(entry: &mut Limit, reset: u64, remaining: u64, limit: u64) {
|
|
|
|
if reset != entry.reset {
|
|
|
|
entry.reset = reset;
|
|
|
|
entry.remaining = limit;
|
|
|
|
entry.limit = limit;
|
|
|
|
} else {
|
|
|
|
entry.remaining = remaining;
|
|
|
|
entry.limit = limit;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
fn decrement_limit_entry(entry: &mut Limit) {
|
|
|
|
entry.remaining -= 1;
|
|
|
|
}
|
|
|
|
|
2023-04-12 22:59:08 +02:00
|
|
|
fn update_limits(&mut self, response: Response, limit_type: LimitType) {
|
|
|
|
// TODO: Make this work
|
2023-04-12 18:33:16 +02:00
|
|
|
let remaining = match response.headers().get("X-RateLimit-Remaining") {
|
2023-04-12 22:59:08 +02:00
|
|
|
Some(remaining) => remaining.to_str().unwrap().parse::<u64>().unwrap(),
|
2023-04-12 18:34:20 +02:00
|
|
|
None => return, //false,
|
2023-04-12 18:33:16 +02:00
|
|
|
};
|
|
|
|
let limit = match response.headers().get("X-RateLimit-Limit") {
|
2023-04-12 22:59:08 +02:00
|
|
|
Some(limit) => limit.to_str().unwrap().parse::<u64>().unwrap(),
|
2023-04-12 18:34:20 +02:00
|
|
|
None => return, //false,
|
2023-04-12 18:33:16 +02:00
|
|
|
};
|
|
|
|
let reset = match response.headers().get("X-RateLimit-Reset") {
|
2023-04-12 22:59:08 +02:00
|
|
|
Some(reset) => reset.to_str().unwrap().parse::<u64>().unwrap(),
|
2023-04-12 18:34:20 +02:00
|
|
|
None => return, //false,
|
2023-04-12 18:33:16 +02:00
|
|
|
};
|
2023-04-12 22:59:08 +02:00
|
|
|
|
|
|
|
let mut limits_copy = self.limits_rate.clone();
|
|
|
|
let status = response.status();
|
|
|
|
let status_str = status.as_str();
|
|
|
|
|
|
|
|
if status_str.chars().next().unwrap() == '4' {
|
|
|
|
limits_copy.get_mut(&LimitType::Error).unwrap().remaining -= 1;
|
|
|
|
}
|
|
|
|
|
|
|
|
limits_copy.get_mut(&LimitType::Global).unwrap().remaining -= 1;
|
|
|
|
limits_copy.get_mut(&LimitType::Ip).unwrap().remaining -= 1;
|
|
|
|
|
|
|
|
match limit_type {
|
|
|
|
// Error, Global and Ip get handled seperately.
|
|
|
|
LimitType::Error => {}
|
|
|
|
LimitType::Global => {}
|
|
|
|
LimitType::Ip => {}
|
|
|
|
LimitType::AuthLogin => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::AuthLogin).unwrap();
|
2023-04-14 17:11:13 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, limit, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
}
|
|
|
|
LimitType::AbsoluteRegister => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::AbsoluteRegister).unwrap();
|
2023-04-14 17:11:13 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, limit, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
// AbsoluteRegister and AuthRegister both need to be updated, if a Register event
|
|
|
|
// happens.
|
|
|
|
limits_copy
|
|
|
|
.get_mut(&LimitType::AuthRegister)
|
|
|
|
.unwrap()
|
|
|
|
.remaining -= 1;
|
|
|
|
}
|
|
|
|
LimitType::AuthRegister => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::AuthRegister).unwrap();
|
2023-04-14 17:11:13 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, limit, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
// AbsoluteRegister and AuthRegister both need to be updated, if a Register event
|
|
|
|
// happens.
|
|
|
|
limits_copy
|
|
|
|
.get_mut(&LimitType::AbsoluteRegister)
|
|
|
|
.unwrap()
|
|
|
|
.remaining -= 1;
|
|
|
|
}
|
|
|
|
LimitType::AbsoluteMessage => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::AbsoluteMessage).unwrap();
|
2023-04-14 17:11:13 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, limit, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
}
|
|
|
|
LimitType::Channel => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::Channel).unwrap();
|
2023-04-14 17:11:13 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, limit, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
}
|
|
|
|
LimitType::Guild => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::Guild).unwrap();
|
2023-04-14 17:11:13 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, limit, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
}
|
|
|
|
LimitType::Webhook => {
|
|
|
|
let entry = limits_copy.get_mut(&LimitType::Webhook).unwrap();
|
2023-04-14 17:11:13 +02:00
|
|
|
LimitedRequester::update_limit_entry(entry, reset, limit, limit);
|
2023-04-12 22:59:08 +02:00
|
|
|
}
|
|
|
|
}
|
2023-04-12 18:33:16 +02:00
|
|
|
}
|
2023-04-07 21:01:48 +02:00
|
|
|
}
|