Hold a mutex while refreshing the token (#41)

There is no point in asking the provider for multiple tokens at the same
time.  With a high enough request rate this results in a thundering
herd of refresh requests and eventually the metadata service rejecting
the requests with HTTP 429 status code.

Could be improved by having a per-scope lock for `CustomServiceAccount`.
Smarter approaches are possible, but this is a good first approach to
avoid a thundering herd on the refresh.

Co-authored-by: Sebastian Zivota <[email protected]>

Author: Floris Bruynooghe

Cargo.toml

@@ -24,7 +24,7 @@ rustls = "0.20.2"
 rustls-pemfile = "0.2.1"
 serde = {version = "1.0", features = ["derive", "rc"]}
 serde_json = "1.0"
-tokio = { version = "1.1", features = ["fs"] }
+tokio = { version = "1.1", features = ["fs", "sync"] }
 url = "2"
 which = "4.2"
 async-trait = "0.1"

src/authentication_manager.rs

@@ -1,4 +1,5 @@
 use async_trait::async_trait;
+use tokio::sync::Mutex;
 
 use crate::error::Error;
 use crate::types::{HyperClient, Token};
@@ -16,9 +17,18 @@ pub(crate) trait ServiceAccount: Send + Sync {
 pub struct AuthenticationManager {
     pub(crate) client: HyperClient,
     pub(crate) service_account: Box<dyn ServiceAccount>,
+    refresh_mutex: Mutex<()>,
 }
 
 impl AuthenticationManager {
+    pub(crate) fn new(client: HyperClient, service_account: Box<dyn ServiceAccount>) -> Self {
+        Self {
+            client,
+            service_account,
+            refresh_mutex: Mutex::new(()),
+        }
+    }
+
     /// Requests Bearer token for the provided scope
     ///
     /// Token can be used in the request authorization header in format "Bearer {token}"
@@ -27,6 +37,15 @@ impl AuthenticationManager {
         if let Some(token) = token.filter(|token| !token.has_expired()) {
             return Ok(token);
         }
+
+        let _guard = self.refresh_mutex.lock().await;
+
+        // Check if refresh happened while we were waiting.
+        let token = self.service_account.get_token(scopes);
+        if let Some(token) = token.filter(|token| !token.has_expired()) {
+            return Ok(token);
+        }
+
         self.service_account
             .refresh_token(&self.client, scopes)
             .await

src/lib.rs

@@ -111,31 +111,28 @@ async fn get_authentication_manager(
     };
 
     if let Ok(service_account) = custom {
-        return Ok(AuthenticationManager {
+        return Ok(AuthenticationManager::new(
             client,
-            service_account: Box::new(service_account),
-        });
+            Box::new(service_account),
+        ));
     }
     let gcloud = gcloud_authorized_user::GCloudAuthorizedUser::new().await;
     if let Ok(service_account) = gcloud {
-        return Ok(AuthenticationManager {
-            client: client.clone(),
-            service_account: Box::new(service_account),
-        });
+        return Ok(AuthenticationManager::new(
+            client.clone(),
+            Box::new(service_account),
+        ));
     }
     let default = default_service_account::DefaultServiceAccount::new(&client).await;
     if let Ok(service_account) = default {
-        return Ok(AuthenticationManager {
-            client: client.clone(),
-            service_account: Box::new(service_account),
-        });
+        return Ok(AuthenticationManager::new(
+            client.clone(),
+            Box::new(service_account),
+        ));
     }
     let user = default_authorized_user::DefaultAuthorizedUser::new(&client).await;
     if let Ok(user_account) = user {
-        return Ok(AuthenticationManager {
-            client,
-            service_account: Box::new(user_account),
-        });
+        return Ok(AuthenticationManager::new(client, Box::new(user_account)));
     }
     Err(Error::NoAuthMethod(
         Box::new(custom.unwrap_err()),