A new retry policy has been created for region discovery (#5336)

Author: Robbie-Microsoft

src/client/Microsoft.Identity.Client/Http/Retry/HttpRetryCondition.cs

@@ -47,6 +47,21 @@ public static bool Imds(HttpResponse response, Exception exception)
             };
         }
 
+        /// <summary>
+        /// Retry policy specific to Region Discovery.
+        /// Extends Imds retry policy but excludes 404 and 408 status codes.
+        /// </summary>
+        public static bool RegionDiscovery(HttpResponse response, Exception exception)
+        {
+            if (!Imds(response, exception))
+            {
+                return false;
+            }
+
+            // If Imds would retry but the status code is 404 or 408, don't retry
+            return (int)response.StatusCode is not (404 or 408);
+        }
+
         /// <summary>
         /// Retry condition for /token and /authorize endpoints
         /// </summary>

src/client/Microsoft.Identity.Client/Http/Retry/RegionDiscoveryRetryPolicy.cs

@@ -0,0 +1,50 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Threading.Tasks;
+using Microsoft.Identity.Client.Core;
+
+namespace Microsoft.Identity.Client.Http.Retry
+{
+    internal class RegionDiscoveryRetryPolicy : IRetryPolicy
+    {
+        // Referenced in unit tests
+        public const int NumRetries = 3;
+
+        private const int MinExponentialBackoffMs = 1000;
+        private const int MaxExponentialBackoffMs = 4000;
+        private const int ExponentialDeltaBackoffMs = 2000;
+
+        private readonly ExponentialRetryStrategy _exponentialRetryStrategy = new ExponentialRetryStrategy(
+            MinExponentialBackoffMs,
+            MaxExponentialBackoffMs,
+            ExponentialDeltaBackoffMs
+        );
+
+        internal virtual Task DelayAsync(int milliseconds)
+        {
+            return Task.Delay(milliseconds);
+        }
+
+        public async Task<bool> PauseForRetryAsync(HttpResponse response, Exception exception, int retryCount, ILoggerAdapter logger)
+        {
+            // Check if the status code is retriable and if the current retry count is less than max retries
+            if (HttpRetryConditions.RegionDiscovery(response, exception) &&
+                retryCount < NumRetries)
+            {
+                int retryAfterDelay = _exponentialRetryStrategy.CalculateDelay(retryCount);
+
+                logger.Warning($"Retrying request in {retryAfterDelay}ms (retry attempt: {retryCount + 1})");
+
+                // Pause execution for the calculated delay
+                await DelayAsync(retryAfterDelay).ConfigureAwait(false);
+
+                return true;
+            }
+
+            // If the status code is not retriable or max retries have been reached, do not retry
+            return false;
+        }
+    }
+}

src/client/Microsoft.Identity.Client/Http/Retry/RetryPolicyFactory.cs

@@ -16,6 +16,8 @@ public virtual IRetryPolicy GetRetryPolicy(RequestType requestType)
                     return new DefaultRetryPolicy(requestType);
                 case RequestType.Imds:
                     return new ImdsRetryPolicy();
+                case RequestType.RegionDiscovery:
+                    return new RegionDiscoveryRetryPolicy();
                 default:
                     throw new ArgumentOutOfRangeException(nameof(requestType), requestType, "Unknown request type.");
             }

src/client/Microsoft.Identity.Client/Instance/Region/RegionManager.cs

@@ -78,7 +78,7 @@ public async Task<string> GetAzureRegionAsync(RequestContext requestContext)
                 "Do not call GetAzureRegionAsync outside of a request. This can happen if you perform instance discovery outside a request, for example as part of validating input params.");
 
             IRetryPolicyFactory retryPolicyFactory = requestContext.ServiceBundle.Config.RetryPolicyFactory;
-            IRetryPolicy retryPolicy = retryPolicyFactory.GetRetryPolicy(RequestType.STS);
+            IRetryPolicy retryPolicy = retryPolicyFactory.GetRetryPolicy(RequestType.RegionDiscovery);
 
             // MSAL always performs region auto-discovery, even if the user configured an actual region
             // in order to detect inconsistencies and report via telemetry

src/client/Microsoft.Identity.Client/RequestType.cs

@@ -21,6 +21,11 @@ internal enum RequestType
         /// <summary>
         /// Instance Metadata Service (IMDS) request, used for obtaining tokens from the Azure VM metadata endpoint.
         /// </summary>
-        Imds
+        Imds,
+
+        /// <summary>
+        /// Region Discovery request, used for region discovery operations with exponential backoff retry strategy.
+        /// </summary>
+        RegionDiscovery
     }
 }

tests/Microsoft.Identity.Test.Common/Core/Mocks/MockHttpManagerExtensions.cs

@@ -493,8 +493,9 @@ public static void AddManagedIdentityWSTrustMockHandler(
                     });
         }
 
-        public static void AddRegionDiscoveryMockHandlerNotFound(
-            this MockHttpManager httpManager)
+        public static void AddRegionDiscoveryMockHandlerWithError(
+            this MockHttpManager httpManager,
+            HttpStatusCode statusCode)
         {
             httpManager.AddMockHandler(
                     new MockHttpMessageHandler
@@ -505,7 +506,7 @@ public static void AddRegionDiscoveryMockHandlerNotFound(
                          {
                             {"Metadata", "true"}
                          },
-                        ResponseMessage = MockHelpers.CreateFailureMessage(HttpStatusCode.NotFound, "")
+                        ResponseMessage = MockHelpers.CreateFailureMessage(statusCode, "")
                     });
         }
     }

tests/Microsoft.Identity.Test.Common/TestConstants.cs

@@ -419,6 +419,7 @@ public static MsalTokenResponse CreateMsalTokenResponseWithTokenSource()
         //Region Discovery Failures
         public const string RegionAutoDetectOkFailureMessage = "Call to local IMDS failed with status code OK or an empty response.";
         public const string RegionAutoDetectNotFoundFailureMessage = "Call to local IMDS failed with status code NotFound or an empty response.";
+        public const string RegionAutoDetectInternalServerErrorFailureMessage = "Service is unavailable to process the request";
         public const string RegionDiscoveryNotSupportedErrorMessage = "Region discovery can only be made if the service resides in Azure function or Azure VM";
         public const string RegionDiscoveryIMDSCallFailedMessage = "IMDS call failed";
 

tests/Microsoft.Identity.Test.Unit/CoreTests/RegionDiscoveryProviderTests.cs

@@ -4,15 +4,18 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.Net;
 using System.Net.Http;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client;
+using Microsoft.Identity.Client.Http.Retry;
 using Microsoft.Identity.Client.Instance.Discovery;
 using Microsoft.Identity.Client.Internal;
 using Microsoft.Identity.Client.Region;
 using Microsoft.Identity.Client.TelemetryCore.Internal.Events;
 using Microsoft.Identity.Test.Common.Core.Mocks;
+using Microsoft.Identity.Test.Unit.Helpers;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 
 namespace Microsoft.Identity.Test.Unit.CoreTests
@@ -28,13 +31,15 @@ public class RegionDiscoveryProviderTests : TestBase
         private ApiEvent _apiEvent;
         private CancellationTokenSource _userCancellationTokenSource;
         private IRegionDiscoveryProvider _regionDiscoveryProvider;
+        private readonly TestRetryPolicyFactory _testRetryPolicyFactory = new TestRetryPolicyFactory();
 
         [TestInitialize]
         public override void TestInitialize()
         {
             base.TestInitialize();
 
             _harness = base.CreateTestHarness();
+            _harness.ServiceBundle.Config.RetryPolicyFactory = _testRetryPolicyFactory;
             _httpManager = _harness.HttpManager;
             _userCancellationTokenSource = new CancellationTokenSource();
             _testRequestContext = new RequestContext(
@@ -145,10 +150,18 @@ public async Task FetchRegionFromLocalImdsThenGetMetadataFromCacheAsync()
             ValidateInstanceMetadata(regionalMetadata);
         }
 
-        [TestMethod]
-        public async Task SuccessfulResponseFromUserProvidedRegionAsync()
+        [DataTestMethod]
+        [DataRow(HttpStatusCode.NotFound, 0, TestConstants.RegionAutoDetectNotFoundFailureMessage)]  // No retries for 404 errors
+        [DataRow(HttpStatusCode.InternalServerError, TestRegionDiscoveryRetryPolicy.NumRetries, TestConstants.RegionAutoDetectInternalServerErrorFailureMessage)]
+        public async Task SuccessfulResponseFromUserProvidedRegionAsync(
+            HttpStatusCode statusCode,
+            int expectedRetries,
+            string expectedFailureMessage)
         {
-            AddMockedResponse(MockHelpers.CreateNullMessage(System.Net.HttpStatusCode.NotFound));
+            for (int i = 0; i < (1 + expectedRetries); i++)
+            {
+                AddMockedResponse(MockHelpers.CreateNullMessage(statusCode));
+            }
 
             _testRequestContext.ServiceBundle.Config.AzureRegion = TestConstants.Region;
 
@@ -162,7 +175,10 @@ public async Task SuccessfulResponseFromUserProvidedRegionAsync()
             Assert.AreEqual(TestConstants.Region, _testRequestContext.ApiEvent.RegionUsed);
             Assert.AreEqual(RegionAutodetectionSource.FailedAutoDiscovery, _testRequestContext.ApiEvent.RegionAutodetectionSource);
             Assert.AreEqual(RegionOutcome.UserProvidedAutodetectionFailed, _testRequestContext.ApiEvent.RegionOutcome);
-            Assert.IsTrue(_testRequestContext.ApiEvent.RegionDiscoveryFailureReason.Contains(TestConstants.RegionAutoDetectNotFoundFailureMessage));
+            Assert.IsTrue(_testRequestContext.ApiEvent.RegionDiscoveryFailureReason.Contains(expectedFailureMessage));
+
+            // Verify all mock responses were consumed
+            Assert.AreEqual(0, _httpManager.QueueSize);
         }
 
         [TestMethod]
@@ -303,10 +319,18 @@ public async Task ResponseMissingRegionFromLocalImdsAsync()
             Assert.IsTrue(_testRequestContext.ApiEvent.RegionDiscoveryFailureReason.Contains(TestConstants.RegionAutoDetectOkFailureMessage));
         }
 
-        [TestMethod]
-        public async Task ErrorResponseFromLocalImdsAsync()
+        [DataTestMethod]
+        [DataRow(HttpStatusCode.NotFound, 0, TestConstants.RegionAutoDetectNotFoundFailureMessage)]  // No retries for 404 errors
+        [DataRow(HttpStatusCode.InternalServerError, TestRegionDiscoveryRetryPolicy.NumRetries, TestConstants.RegionAutoDetectInternalServerErrorFailureMessage)]
+        public async Task ErrorResponseFromLocalImdsAsync(
+            HttpStatusCode statusCode,
+            int expectedRetries,
+            string expectedFailureMessage)
         {
-            AddMockedResponse(MockHelpers.CreateNullMessage(System.Net.HttpStatusCode.NotFound));
+            for (int i = 0; i < (1 + expectedRetries); i++)
+            {
+                AddMockedResponse(MockHelpers.CreateNullMessage(statusCode));
+            }
             _testRequestContext.ServiceBundle.Config.AzureRegion = ConfidentialClientApplication.AttemptRegionDiscovery;
 
             InstanceDiscoveryMetadataEntry regionalMetadata = await _regionDiscoveryProvider.
@@ -318,7 +342,10 @@ public async Task ErrorResponseFromLocalImdsAsync()
             Assert.AreEqual(null, _testRequestContext.ApiEvent.RegionUsed);
             Assert.AreEqual(RegionAutodetectionSource.FailedAutoDiscovery, _testRequestContext.ApiEvent.RegionAutodetectionSource);
             Assert.AreEqual(RegionOutcome.FallbackToGlobal, _testRequestContext.ApiEvent.RegionOutcome);
-            Assert.IsTrue(_testRequestContext.ApiEvent.RegionDiscoveryFailureReason.Contains(TestConstants.RegionAutoDetectNotFoundFailureMessage));
+            Assert.IsTrue(_testRequestContext.ApiEvent.RegionDiscoveryFailureReason.Contains(expectedFailureMessage));
+
+            // Verify all mock responses were consumed
+            Assert.AreEqual(0, _httpManager.QueueSize);
         }
 
         [TestMethod]
@@ -382,6 +409,75 @@ public async Task UpdateApiversionFailsWithNoNewestVersionsAsync()
             Assert.IsTrue(_testRequestContext.ApiEvent.RegionDiscoveryFailureReason.Contains(TestConstants.RegionDiscoveryNotSupportedErrorMessage));
         }
 
+        [TestMethod]
+        public async Task RegionDiscoveryFails500OnceThenSucceeds200Async()
+        {
+            AddMockedResponse(MockHelpers.CreateNullMessage(HttpStatusCode.InternalServerError));
+            AddMockedResponse(MockHelpers.CreateSuccessResponseMessage(TestConstants.Region));
+
+            _testRequestContext.ServiceBundle.Config.AzureRegion = ConfidentialClientApplication.AttemptRegionDiscovery;
+
+            InstanceDiscoveryMetadataEntry regionalMetadata = await _regionDiscoveryProvider.GetMetadataAsync(
+                new Uri("https://login.microsoftonline.com/common/"), _testRequestContext).ConfigureAwait(false);
+
+            ValidateInstanceMetadata(regionalMetadata);
+            Assert.AreEqual(TestConstants.Region, _testRequestContext.ApiEvent.RegionUsed);
+            Assert.AreEqual(RegionAutodetectionSource.Imds, _testRequestContext.ApiEvent.RegionAutodetectionSource);
+            Assert.AreEqual(RegionOutcome.AutodetectSuccess, _testRequestContext.ApiEvent.RegionOutcome);
+            Assert.IsNull(_testRequestContext.ApiEvent.RegionDiscoveryFailureReason);
+
+            const int NumRequests = 2; // initial request + one retry
+            int requestsMade = NumRequests - _httpManager.QueueSize;
+            Assert.AreEqual(NumRequests, requestsMade);
+        }
+
+        [TestMethod]
+        public async Task RegionDiscoveryFails500PermanentlyAsync()
+        {
+            // Simulate permanent 500s (to trigger the maximum number of retries)
+            const int Num500Errors = 1 + RegionDiscoveryRetryPolicy.NumRetries; // initial request + maximum number of retries
+            for (int i = 0; i < Num500Errors; i++)
+            {
+                AddMockedResponse(MockHelpers.CreateNullMessage(HttpStatusCode.InternalServerError));
+            }
+
+            _testRequestContext.ServiceBundle.Config.AzureRegion = ConfidentialClientApplication.AttemptRegionDiscovery;
+
+            InstanceDiscoveryMetadataEntry regionalMetadata = await _regionDiscoveryProvider.GetMetadataAsync(
+                new Uri("https://login.microsoftonline.com/common/"), _testRequestContext).ConfigureAwait(false);
+
+            Assert.IsNull(regionalMetadata, "Discovery should fail after max retries");
+            Assert.AreEqual(null, _testRequestContext.ApiEvent.RegionUsed);
+            Assert.AreEqual(RegionAutodetectionSource.FailedAutoDiscovery, _testRequestContext.ApiEvent.RegionAutodetectionSource);
+            Assert.AreEqual(RegionOutcome.FallbackToGlobal, _testRequestContext.ApiEvent.RegionOutcome);
+
+            const int NumRequests = Num500Errors; // initial request + three retries
+            int requestsMade = NumRequests - _httpManager.QueueSize;
+            Assert.AreEqual(NumRequests, requestsMade);
+        }
+
+        [DataTestMethod]
+        [DataRow(HttpStatusCode.NotFound)]
+        [DataRow(HttpStatusCode.RequestTimeout)]
+        public async Task RegionDiscoveryDoesNotRetryOnNonRetryableStatusCodesAsync(HttpStatusCode statusCode)
+        {
+            AddMockedResponse(MockHelpers.CreateNullMessage(statusCode));
+
+            _testRequestContext.ServiceBundle.Config.AzureRegion = ConfidentialClientApplication.AttemptRegionDiscovery;
+
+            InstanceDiscoveryMetadataEntry regionalMetadata = await _regionDiscoveryProvider.GetMetadataAsync(
+                new Uri("https://login.microsoftonline.com/common/"), _testRequestContext).ConfigureAwait(false);
+
+            Assert.IsNull(regionalMetadata, "Discovery should fail and not retry");
+            Assert.AreEqual(null, _testRequestContext.ApiEvent.RegionUsed);
+            Assert.AreEqual(RegionAutodetectionSource.FailedAutoDiscovery, _testRequestContext.ApiEvent.RegionAutodetectionSource);
+            Assert.AreEqual(RegionOutcome.FallbackToGlobal, _testRequestContext.ApiEvent.RegionOutcome);
+
+            const int NumRequests = 1; // initial request + 0 retries (non-retryable status codes should not trigger retry)
+            int requestsMade = NumRequests - _httpManager.QueueSize;
+            Assert.AreEqual(NumRequests, requestsMade);
+        }
+
         private void AddMockedResponse(HttpResponseMessage responseMessage, string apiVersion = "2020-06-01", bool expectedParams = true)
         {
             var queryParams = new Dictionary<string, string>();

tests/Microsoft.Identity.Test.Unit/Helpers/TestRetryPolicies.cs

@@ -28,4 +28,15 @@ internal override Task DelayAsync(int milliseconds)
             return Task.CompletedTask;
         }
     }
+
+    internal class TestRegionDiscoveryRetryPolicy : RegionDiscoveryRetryPolicy
+    {
+        public TestRegionDiscoveryRetryPolicy() : base() { }
+
+        internal override Task DelayAsync(int milliseconds)
+        {
+            // No delay for tests
+            return Task.CompletedTask;
+        }
+    }
 }

tests/Microsoft.Identity.Test.Unit/Helpers/TestRetryPolicyFactory.cs

@@ -18,6 +18,8 @@ public virtual IRetryPolicy GetRetryPolicy(RequestType requestType)
                     return new TestDefaultRetryPolicy(requestType);
                 case RequestType.Imds:
                     return new TestImdsRetryPolicy();
+                case RequestType.RegionDiscovery:
+                    return new TestRegionDiscoveryRetryPolicy();
                 default:
                     throw new ArgumentOutOfRangeException(nameof(requestType), requestType, "Unknown request type.");
             }