Implement retry logic for IMDS managed identity (#3306)

Author: chlowell

sdk/identity/azure_identity/CHANGELOG.md

@@ -17,6 +17,7 @@
 ### Bugs Fixed
 
 - `ClientCertificateCredential::get_token()` returned an error when given multiple scopes
+- `ManagedIdentityCredential` didn't follow IMDS retry guidance
 
 ### Other Changes
 

sdk/identity/azure_identity/src/app_service_managed_identity_credential.rs

@@ -48,6 +48,7 @@ impl AppServiceManagedIdentityCredential {
                 SECRET_ENV,
                 id,
                 client_options,
+                None,
                 env,
             ),
         }))

sdk/identity/azure_identity/src/imds_managed_identity_credential.rs

@@ -6,7 +6,7 @@ use azure_core::{
     credentials::{AccessToken, Secret, TokenCredential, TokenRequestOptions},
     error::{Error, ErrorKind},
     http::{
-        headers::HeaderName, request::Request, ClientOptions, Method, Pipeline,
+        headers::HeaderName, request::Request, ClientOptions, Method, Pipeline, PipelineOptions,
         PipelineSendOptions, StatusCode, Url,
     },
     json::from_json,
@@ -62,13 +62,15 @@ pub(crate) struct ImdsManagedIdentityCredential {
 }
 
 impl ImdsManagedIdentityCredential {
+    #[allow(clippy::too_many_arguments, reason = "private API")]
     pub fn new(
         endpoint: Url,
         api_version: &str,
         secret_header: HeaderName,
         secret_env: &str,
         id: ImdsId,
         client_options: ClientOptions,
+        pipeline_options: Option<PipelineOptions>,
         env: Env,
     ) -> Self {
         let pipeline = Pipeline::new(
@@ -77,7 +79,7 @@ impl ImdsManagedIdentityCredential {
             client_options,
             Vec::default(),
             Vec::default(),
-            None,
+            pipeline_options,
         );
         Self {
             pipeline,

sdk/identity/azure_identity/src/virtual_machine_managed_identity_credential.rs

@@ -3,10 +3,13 @@
 
 use crate::env::Env;
 use crate::{ImdsId, ImdsManagedIdentityCredential};
-use azure_core::http::ClientOptions;
 use azure_core::{
     credentials::{AccessToken, TokenCredential, TokenRequestOptions},
-    http::{headers::HeaderName, Url},
+    http::{
+        headers::HeaderName, ClientOptions, ExponentialRetryOptions, PipelineOptions, RetryOptions,
+        StatusCode, Url,
+    },
+    time::Duration,
 };
 use std::sync::Arc;
 
@@ -27,6 +30,37 @@ impl VirtualMachineManagedIdentityCredential {
         env: Env,
     ) -> azure_core::Result<Arc<Self>> {
         let endpoint = Url::parse(ENDPOINT).unwrap(); // valid url constant
+        let pipeline_options = Some(PipelineOptions {
+            // https://learn.microsoft.com/entra/identity/managed-identities-azure-resources/how-to-use-vm-token#error-handling
+            retry_status_codes: Vec::from([
+                StatusCode::NotFound,
+                StatusCode::Gone,
+                StatusCode::TooManyRequests,
+                StatusCode::InternalServerError,
+                StatusCode::NotImplemented,
+                StatusCode::BadGateway,
+                StatusCode::ServiceUnavailable,
+                StatusCode::GatewayTimeout,
+                StatusCode::HttpVersionNotSupported,
+                StatusCode::VariantAlsoNegotiates,
+                StatusCode::InsufficientStorage,
+                StatusCode::LoopDetected,
+                StatusCode::NotExtended,
+                StatusCode::NetworkAuthenticationRequired,
+            ]),
+            ..Default::default()
+        });
+        // these settings approximate the recommendations at
+        // https://learn.microsoft.com/entra/identity/managed-identities-azure-resources/how-to-use-vm-token#retry-guidance
+        let client_options = ClientOptions {
+            retry: RetryOptions::exponential(ExponentialRetryOptions {
+                initial_delay: Duration::milliseconds(1340),
+                max_retries: 6,
+                max_total_elapsed: Duration::seconds(72),
+                ..Default::default()
+            }),
+            ..client_options
+        };
         Ok(Arc::new(Self {
             credential: ImdsManagedIdentityCredential::new(
                 endpoint,
@@ -35,6 +69,7 @@ impl VirtualMachineManagedIdentityCredential {
                 SECRET_ENV,
                 id,
                 client_options,
+                pipeline_options,
                 env,
             ),
         }))