diff --git a/src/consts/mod.rs b/src/consts/mod.rs index 9d827ac..c9fe174 100644 --- a/src/consts/mod.rs +++ b/src/consts/mod.rs @@ -1 +1,3 @@ -pub mod region; +mod region; + +pub use region::*; diff --git a/src/lib.rs b/src/lib.rs index c74213b..5c06ea9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,7 @@ pub mod consts; mod req; +mod util; mod riot_api_config; #[cfg(test)] diff --git a/src/req/mod.rs b/src/req/mod.rs index 203e475..88c8e46 100644 --- a/src/req/mod.rs +++ b/src/req/mod.rs @@ -1,5 +1,11 @@ -pub mod rate_limit; -pub mod rate_limit_type; -pub mod token_bucket; -pub mod regional_requester; -pub mod requester_manager; +mod rate_limit; +mod rate_limit_type; +mod token_bucket; +mod regional_requester; +mod requester_manager; + +pub use rate_limit::*; +pub use rate_limit_type::*; +pub use token_bucket::*; +pub use regional_requester::*; +pub use requester_manager::*; diff --git a/src/req/rate_limit.rs b/src/req/rate_limit.rs index 1a92601..9fddc97 100644 --- a/src/req/rate_limit.rs +++ b/src/req/rate_limit.rs @@ -8,11 +8,11 @@ use parking_lot::{ RwLock, }; -use super::token_bucket::{ +use super::{ TokenBucket, VectorTokenBucket, }; -use super::rate_limit_type::RateLimitType; +use super::RateLimitType; pub struct RateLimit { rate_limit_type: RateLimitType, diff --git a/src/req/regional_requester.rs b/src/req/regional_requester.rs index 9dcbf1b..5d36c68 100644 --- a/src/req/regional_requester.rs +++ b/src/req/regional_requester.rs @@ -1,4 +1,3 @@ -use std::collections::HashMap; use std::future::Future; use std::sync::Arc; @@ -7,15 +6,14 @@ use reqwest::{ Client, StatusCode, }; -use parking_lot::{ - Mutex, -}; -use serde::de::DeserializeOwned; +use parking_lot::Mutex; -use super::rate_limit::RateLimit; -use super::rate_limit_type::RateLimitType; use crate::riot_api_config::RiotApiConfig; -use crate::consts::region::Region; +use crate::consts::Region; +use crate::util::InsertOnlyCHashMap; + +use super::RateLimit; +use super::RateLimitType; pub struct RegionalRequester<'a> { /// Configuration settings. @@ -26,7 +24,7 @@ pub struct RegionalRequester<'a> { /// Represents the app rate limit. app_rate_limit: RateLimit, /// Represents method rate limits. - method_rate_limits: Mutex>>, + method_rate_limits: InsertOnlyCHashMap<&'a str, RateLimit>, } impl<'a> RegionalRequester<'a> { @@ -42,23 +40,24 @@ impl<'a> RegionalRequester<'a> { riot_api_config: riot_api_config, client: client, app_rate_limit: RateLimit::new(RateLimitType::Application), - method_rate_limits: Mutex::new(HashMap::new()), + method_rate_limits: InsertOnlyCHashMap::new(), } } - pub async fn get( - &mut self, method_id: &'a str, relative_url: &'_ str, - region: &'_ Region<'_>, query: &[(&'_ str, &'_ str)]) -> Result, reqwest::Error> { + pub async fn get( + &self, method_id: &'a str, region: &'_ Region<'_>, relative_url: &'_ str, + query: &[(&'_ str, &'_ str)]) -> Result, reqwest::Error> + { let mut attempts: u8 = 0; for _ in 0..=self.riot_api_config.retries { attempts += 1; + let method_rate_limit: Arc = self.method_rate_limits + .get_or_insert_with(method_id, || RateLimit::new(RateLimitType::Method)); + // Rate limiting. - while let Some(delay) = { - let method_rate_limit = self.get_method_rate_limit(method_id); - RateLimit::get_both_or_delay(&self.app_rate_limit, &*method_rate_limit) - } { + while let Some(delay) = RateLimit::get_both_or_delay(&self.app_rate_limit, &*method_rate_limit) { task::sleep(delay).await; } @@ -75,11 +74,9 @@ impl<'a> RegionalRequester<'a> { Ok(r) => r, }; - // Update rate limits (if needed). - { - self.app_rate_limit.on_response(&response); - self.get_method_rate_limit(method_id).on_response(&response); - } + // Maybe update rate limits (based on response headers). + self.app_rate_limit.on_response(&response); + method_rate_limit.on_response(&response); // Handle response. let status = response.status(); @@ -109,16 +106,10 @@ impl<'a> RegionalRequester<'a> { panic!("FAILED AFTER {} ATTEMPTS!", attempts); } - pub fn get2(&'a mut self, method_id: &'a str, relative_url: &'a str, - region: &'a Region<'_>, query: &'a [(&'a str, &'a str)]) -> impl Future, reqwest::Error>> + 'a { - - self.get(method_id, relative_url, region, query) - } - - fn get_method_rate_limit(&self, method_id: &'a str) -> Arc { - Arc::clone(self.method_rate_limits.lock() - .entry(method_id) - .or_insert_with(|| Arc::new(RateLimit::new(RateLimitType::Method)))) + pub fn get2(&'a self, method_id: &'a str, region: &'a Region<'_>, relative_url: &'a str, + query: &'a [(&'a str, &'a str)]) -> impl Future, reqwest::Error>> + 'a + { + self.get(method_id, region, relative_url, query) } fn is_none_status_code(status: &StatusCode) -> bool { diff --git a/src/req/requester_manager.rs b/src/req/requester_manager.rs index e0dda73..55174fd 100644 --- a/src/req/requester_manager.rs +++ b/src/req/requester_manager.rs @@ -5,9 +5,11 @@ use reqwest::{ Client, }; -use super::regional_requester::RegionalRequester; use crate::riot_api_config::RiotApiConfig; -use crate::consts::region::Region; +use crate::consts::Region; +use crate::util::InsertOnlyCHashMap; + +use super::RegionalRequester; pub struct RequesterManager<'a> { /// Configuration settings. @@ -16,5 +18,24 @@ pub struct RequesterManager<'a> { client: &'a Client, /// Per-region requesters. - regional_requesters: HashMap<&'a Region<'a>, Arc>>, -} \ No newline at end of file + regional_requesters: InsertOnlyCHashMap<&'a Region<'a>, RegionalRequester<'a>>, +} + +impl<'a> RequesterManager<'a> { + pub fn new(riot_api_config: &'a RiotApiConfig<'a>, client: &'a Client) -> Self { + Self { + riot_api_config: riot_api_config, + client: client, + regional_requesters: InsertOnlyCHashMap::new(), + } + } + + pub async fn get( + &mut self, method_id: &'a str, region: &'a Region<'a>, relative_url: &'_ str, + query: &[(&'_ str, &'_ str)]) -> Result, reqwest::Error> + { + let regional_requester = self.regional_requesters + .get_or_insert_with(region, || RegionalRequester::new(self.riot_api_config, self.client)); + regional_requester.get(method_id, region, relative_url, query).await + } +} diff --git a/src/util/insert_only_chashmap.rs b/src/util/insert_only_chashmap.rs new file mode 100644 index 0000000..bc51796 --- /dev/null +++ b/src/util/insert_only_chashmap.rs @@ -0,0 +1,36 @@ +use std::borrow::Borrow; +use std::collections::HashMap; +use std::hash::Hash; +use std::sync::Arc; + +use parking_lot::Mutex; + +pub struct InsertOnlyCHashMap { + base: Mutex>>, +} + +impl InsertOnlyCHashMap { + #[inline] + pub fn new() -> Self { + Self { + base: Mutex::new(HashMap::new()) + } + } + + #[inline] + pub fn get(&self, key: &Q) -> Option> + where + K: Borrow, + Q: Hash + Eq, + { + self.base.lock().get(key).map(|v| Arc::clone(v)) + } + + #[inline] + pub fn get_or_insert_with V>(&self, key: K, default: F) -> Arc + { + Arc::clone(self.base.lock() + .entry(key) + .or_insert_with(|| Arc::new(default()))) + } +} diff --git a/src/util/mod.rs b/src/util/mod.rs new file mode 100644 index 0000000..2cc54eb --- /dev/null +++ b/src/util/mod.rs @@ -0,0 +1,3 @@ +mod insert_only_chashmap; + +pub use insert_only_chashmap::*;