implement retries for fetching tokens from the various endpoints

The official SDKs implement a retry mechanism for fetching the tokens
from the metadata server in the case of I/O errors, etc.
This adds a similar mechanism to the provided ServiceAccount
implementations

Author: valkum

src/custom_service_account.rs

@@ -104,18 +104,30 @@ impl ServiceAccount for CustomServiceAccount {
             .extend_pairs(&[("grant_type", GRANT_TYPE), ("assertion", jwt.as_str())])
             .finish();
 
-        let request = hyper::Request::post(&self.credentials.token_uri)
-            .header(header::CONTENT_TYPE, "application/x-www-form-urlencoded")
-            .body(hyper::Body::from(rqbody))
-            .unwrap();
-
-        tracing::debug!("requesting token from service account: {:?}", request);
-        let token = client
-            .request(request)
-            .await
-            .map_err(Error::OAuthConnectionError)?
-            .deserialize::<Token>()
-            .await?;
+        let mut retries = 0;
+        let response = loop {
+            let request = hyper::Request::post(&self.credentials.token_uri)
+                .header(header::CONTENT_TYPE, "application/x-www-form-urlencoded")
+                .body(hyper::Body::from(rqbody.clone()))
+                .unwrap();
+
+            tracing::debug!("requesting token from service account: {request:?}");
+            let err = match client.request(request).await {
+                // Early return when the request succeeds
+                Ok(response) => break response,
+                Err(err) => err,
+            };
+
+            tracing::warn!(
+                "Failed to refresh token with GCP oauth2 token endpoint: {err}, trying again..."
+            );
+            retries += 1;
+            if retries >= RETRY_COUNT {
+                return Err(Error::OAuthConnectionError(err));
+            }
+        };
+
+        let token = response.deserialize::<Token>().await?;
 
         let key = scopes.iter().map(|x| (*x).to_string()).collect();
         self.tokens.write().unwrap().insert(key, token.clone());
@@ -154,3 +166,6 @@ impl fmt::Debug for ApplicationCredentials {
             .finish()
     }
 }
+
+/// How many times to attempt to fetch a token from the set credentials token endpoint.
+const RETRY_COUNT: u8 = 5;

src/default_authorized_user.rs

@@ -48,20 +48,31 @@ impl DefaultAuthorizedUser {
 
     #[tracing::instrument]
     async fn get_token(cred: &UserCredentials, client: &HyperClient) -> Result<Token, Error> {
-        let req = Self::build_token_request(&RefreshRequest {
-            client_id: &cred.client_id,
-            client_secret: &cred.client_secret,
-            grant_type: "refresh_token",
-            refresh_token: &cred.refresh_token,
-        });
-
-        let token = client
-            .request(req)
-            .await
-            .map_err(Error::OAuthConnectionError)?
-            .deserialize()
-            .await?;
-        Ok(token)
+        let mut retries = 0;
+        let response = loop {
+            let req = Self::build_token_request(&RefreshRequest {
+                client_id: &cred.client_id,
+                client_secret: &cred.client_secret,
+                grant_type: "refresh_token",
+                refresh_token: &cred.refresh_token,
+            });
+
+            let err = match client.request(req).await {
+                // Early return when the request succeeds
+                Ok(response) => break response,
+                Err(err) => err,
+            };
+
+            tracing::warn!(
+                "Failed to get token from GCP oauth2 token endpoint: {err}, trying again..."
+            );
+            retries += 1;
+            if retries >= RETRY_COUNT {
+                return Err(Error::OAuthConnectionError(err));
+            }
+        };
+
+        response.deserialize().await.map_err(Into::into)
     }
 }
 
@@ -106,3 +117,6 @@ struct UserCredentials {
     /// Type
     pub(crate) r#type: String,
 }
+
+/// How many times to attempt to fetch a token from the GCP token endpoint.
+const RETRY_COUNT: u8 = 5;

src/default_service_account.rs

@@ -36,15 +36,27 @@ impl DefaultServiceAccount {
 
     #[tracing::instrument]
     async fn get_token(client: &HyperClient) -> Result<Token, Error> {
+        let mut retries = 0;
         tracing::debug!("Getting token from GCP instance metadata server");
-        let req = Self::build_token_request(Self::DEFAULT_TOKEN_GCP_URI);
-        let token = client
-            .request(req)
-            .await
-            .map_err(Error::ConnectionError)?
-            .deserialize()
-            .await?;
-        Ok(token)
+        let response = loop {
+            let req = Self::build_token_request(Self::DEFAULT_TOKEN_GCP_URI);
+
+            let err = match client.request(req).await {
+                // Early return when the request succeeds
+                Ok(response) => break response,
+                Err(err) => err,
+            };
+
+            tracing::warn!(
+                "Failed to get token from GCP instance metadata server: {err}, trying again..."
+            );
+            retries += 1;
+            if retries >= RETRY_COUNT {
+                return Err(Error::ConnectionError(err));
+            }
+        };
+
+        response.deserialize().await.map_err(Into::into)
     }
 }
 
@@ -75,3 +87,6 @@ impl ServiceAccount for DefaultServiceAccount {
         Ok(token)
     }
 }
+
+/// How many times to attempt to fetch a token from the GCP metadata server.
+const RETRY_COUNT: u8 = 5;