diff --git a/src/lib.rs b/src/lib.rs index b8d02d4..4a38c25 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,4 +4,5 @@ pub mod error; pub mod integrations; pub mod services; pub mod state; +pub mod utils; pub mod web; diff --git a/src/services/discord.rs b/src/services/discord.rs index 7ccb9c6..54cadf9 100644 --- a/src/services/discord.rs +++ b/src/services/discord.rs @@ -11,6 +11,7 @@ use crate::{ integrations::discord::{ DiscordClient, DiscordUserModel, GuildMemberModel, PartialGuildModel, Snowflake, }, + utils::tokens::RefreshLock, web::WebError, }; @@ -48,11 +49,17 @@ pub struct DiscordService { db: MantleDb, client: DiscordClient, config: SharedConfig, + refresh_lock: RefreshLock, } impl DiscordService { pub fn new(db: MantleDb, client: DiscordClient, config: SharedConfig) -> Self { - Self { db, client, config } + Self { + db, + client, + config, + refresh_lock: Default::default(), + } } pub async fn get_current_user( @@ -212,19 +219,25 @@ impl DiscordService { } async fn access_token(&self, identity: &Identity) -> Result { - // TODO: Current implementation can allow making refresh races - // where hundreds of requests can try to get access_token and see expired one -> - // they all will try to update it and may end up being rate limited. - // This stuff need some per-identity lock or smth - let tokens = db::oauth_tokens::find_by_id(&self.db, identity.id) - .await - .map_err(InternalError::from)?; + loop { + let tokens = db::oauth_tokens::find_by_id(&self.db, identity.id) + .await + .map_err(InternalError::from)? + .ok_or(DiscordError::IdentityHasNoTokens(identity.id))?; - if let Some(tokens) = tokens { if tokens.expires_at > Utc::now() + Duration::minutes(1) { return Ok(tokens.access_token); } + // try acquire a lock before refresh + let guard = match self.refresh_lock.try_lock(identity) { + Ok(guard) => guard, + Err(notify) => { + notify.notified().await; + continue; + } + }; + let credentials = self.client_credentials(); let new_tokens = self .client @@ -243,9 +256,11 @@ impl DiscordService { .await .map_err(InternalError::from)?; + // we will explicitly drop it here so rustc won't do its magic + // (probably it won't do it anyway because of Drop impl, but still) + drop(guard); + return Ok(new_tokens.access_token); } - - Err(DiscordError::IdentityHasNoTokens(identity.id)) } } diff --git a/src/utils/mod.rs b/src/utils/mod.rs new file mode 100644 index 0000000..5c76635 --- /dev/null +++ b/src/utils/mod.rs @@ -0,0 +1 @@ +pub mod tokens; diff --git a/src/utils/tokens.rs b/src/utils/tokens.rs new file mode 100644 index 0000000..15b983d --- /dev/null +++ b/src/utils/tokens.rs @@ -0,0 +1,49 @@ +use std::{ + collections::{HashMap, hash_map::Entry}, + sync::{Arc, Mutex}, +}; + +use tokio::sync::Notify; +use uuid::Uuid; + +use crate::db::identities::Identity; + +// TODO: it should be tested properly, but currently im too lazy to do this +#[derive(Debug, Default, Clone)] +pub struct RefreshLock { + refreshes: Arc>>>, +} + +pub struct RefreshGuard { + identity_id: Uuid, + refreshes: Arc>>>, +} + +impl Drop for RefreshGuard { + fn drop(&mut self) { + let identity_id = self.identity_id; + let refreshes = self.refreshes.clone(); + + let mut refreshes = refreshes.lock().unwrap(); + if let Some(current) = refreshes.remove(&identity_id) { + current.notify_waiters(); + } + } +} + +impl RefreshLock { + pub fn try_lock(&self, identity: &Identity) -> Result> { + let mut refreshes = self.refreshes.lock().unwrap(); + + match refreshes.entry(identity.id) { + Entry::Occupied(entry) => Err(entry.get().clone()), + Entry::Vacant(entry) => { + entry.insert(Arc::new(Notify::new())); + Ok(RefreshGuard { + identity_id: identity.id, + refreshes: self.refreshes.clone(), + }) + } + } + } +}