forked from mirror/Riven
1
0
Fork 0

using parking_lot

users/mingwei/surf
Mingwei Samuel 2019-10-14 01:00:20 -07:00
parent 5700547c05
commit 346ef13ed4
4 changed files with 41 additions and 15 deletions

View File

@ -13,5 +13,6 @@ exclude = [
[dependencies]
async-std = "0.99"
parking_lot = { version = "0.9", features = [ "nightly" ] }
reqwest = { version = "0.10.0-alpha.1", features = [ "gzip", "json" ] }
serde = "^1.0"

View File

@ -3,7 +3,8 @@ use std::time::{
Duration,
Instant,
};
use std::sync::{
use parking_lot::{
RwLock,
};
@ -48,8 +49,8 @@ impl RateLimit {
return retry_after_delay
}
// Check buckets.
let app_buckets = app_rate_limit.buckets.read().unwrap();
let method_buckets = app_rate_limit.buckets.read().unwrap();
let app_buckets = app_rate_limit.buckets.read();
let method_buckets = method_rate_limit.buckets.read();
for bucket in app_buckets.iter().chain(method_buckets.iter()) {
let delay = bucket.get_delay();
if delay.is_some() {

View File

@ -1,6 +1,12 @@
use std::collections::HashMap;
use std::future::Future;
use std::sync::RwLock;
use parking_lot::{
RwLock,
RwLockReadGuard,
RwLockWriteGuard,
MappedRwLockReadGuard,
MappedRwLockWriteGuard,
};
use async_std::task;
use reqwest::{
@ -23,7 +29,7 @@ pub struct RegionalRequester<'a> {
/// Represents the app rate limit.
app_rate_limit: RateLimit,
/// Represents method rate limits.
method_rate_limits: HashMap<&'a str, RateLimit>,
method_rate_limits: RwLock<HashMap<&'a str, RateLimit>>,
}
impl <'a> RegionalRequester<'a> {
@ -39,7 +45,7 @@ impl <'a> RegionalRequester<'a> {
riot_api_config: riot_api_config,
client: client,
app_rate_limit: RateLimit::new(RateLimitType::Application),
method_rate_limits: HashMap::new(),
method_rate_limits: RwLock::new(HashMap::new()),
}
}
@ -52,11 +58,10 @@ impl <'a> RegionalRequester<'a> {
attempts += 1;
// Rate limiting.
let app_rate_limit = &self.app_rate_limit;
let method_rate_limit = self.method_rate_limits.entry(method_id)
.or_insert_with(|| RateLimit::new(RateLimitType::Method));
while let Some(delay) = RateLimit::get_both_or_delay(app_rate_limit, method_rate_limit) {
while let Some(delay) = {
let method_rate_limit = &self.get_insert_rate_limit(method_id);
RateLimit::get_both_or_delay(&self.app_rate_limit, method_rate_limit)
} {
task::sleep(delay).await;
}
@ -74,8 +79,10 @@ impl <'a> RegionalRequester<'a> {
};
// Update rate limits (if needed).
app_rate_limit.on_response(&response);
method_rate_limit.on_response(&response);
{
self.app_rate_limit.on_response(&response);
self.method_rate_limits.read().get(method_id).unwrap().on_response(&response);
}
// Handle response.
let status = response.status();
@ -111,6 +118,22 @@ impl <'a> RegionalRequester<'a> {
self.get(method_id, relative_url, region, query)
}
fn get_insert_rate_limit(&self, method_id: &'a str) -> MappedRwLockReadGuard<RateLimit> {
// This is really stupid?
{
let map_guard = self.method_rate_limits.read();
if map_guard.contains_key(method_id) {
return RwLockReadGuard::map(map_guard, |mrl| mrl.get(method_id).unwrap());
}
}
let map_guard = self.method_rate_limits.write();
let val_write = RwLockWriteGuard::map(
map_guard, |mrl| mrl.entry(method_id)
.or_insert(RateLimit::new(RateLimitType::Method))
);
MappedRwLockWriteGuard::downgrade(val_write)
}
fn is_none_status_code(status: &StatusCode) -> bool {
Self::NONE_STATUS_CODES.contains(&status.as_u16())
}

View File

@ -1,6 +1,7 @@
use std::collections::VecDeque;
use std::time::{Duration, Instant};
use std::sync::{Mutex, MutexGuard};
use parking_lot::{Mutex, MutexGuard};
pub trait TokenBucket {
/// Get the duration til the next available token, or 0 duration if a token is available.
@ -45,7 +46,7 @@ impl VectorTokenBucket {
}
fn update_get_timestamps(&self) -> MutexGuard<VecDeque<Instant>> {
let mut timestamps = self.timestamps.lock().unwrap();
let mut timestamps = self.timestamps.lock();
let cutoff = Instant::now() - self.duration;
while timestamps.back().map_or(false, |ts| ts < &cutoff) {
timestamps.pop_back();