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,27 @@ 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 res = 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 res = client.request(request).await;
+            match res {
+                Ok(res) => break Ok(res),
+                Err(e) => {
+                    retries += 1;
+                    if retries >= RETRY_COUNT {
+                        break Err(Error::OAuthConnectionError(e));
+                    }
+                }
+            }
+        }?;
+
+        let token = res.deserialize::<Token>().await?;
 
         let key = scopes.iter().map(|x| (*x).to_string()).collect();
         self.tokens.write().unwrap().insert(key, token.clone());
@@ -154,3 +163,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,33 @@ 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 res = 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 res = client.request(req).await;
+
+            match res {
+                Ok(res) => break Ok(res),
+                Err(e) => {
+                    tracing::warn!(
+                        "Failed to get token from GCP instance metadata server: {}, trying again...",
+                        e
+                    );
+
+                    retries += 1;
+                    if retries >= RETRY_COUNT {
+                        break Err(Error::OAuthConnectionError(e));
+                    }
+                }
+            }
+        }?;
+
+        res.deserialize().await.map_err(Into::into)
     }
 }
 
@@ -106,3 +119,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,29 @@ 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 res = loop {
+            let req = Self::build_token_request(Self::DEFAULT_TOKEN_GCP_URI);
+            let res = client.request(req).await;
+
+            match res {
+                Ok(res) => break Ok(res),
+                Err(e) => {
+                    tracing::warn!(
+                        "Failed to get token from GCP instance metadata server: {}, trying again...",
+                        e
+                    );
+
+                    retries += 1;
+                    if retries >= RETRY_COUNT {
+                        break Err(Error::ConnectionError(e));
+                    }
+                }
+            }
+        }?;
+
+        res.deserialize().await.map_err(Into::into)
     }
 }
 
@@ -75,3 +89,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;