Refactor retry logic and add special behavior for IMDS

Author: Avery-Dunn

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AbstractClientApplicationBase.java

@@ -568,7 +568,8 @@ public T correlationId(String val) {
                 new TelemetryManager(telemetryConsumer, builder.onlySendFailureTelemetry),
                 new HttpHelper(builder.httpClient == null ?
                         new DefaultHttpClient(builder.proxy, builder.sslSocketFactory, builder.connectTimeoutForDefaultHttpClient, builder.readTimeoutForDefaultHttpClient) :
-                        builder.httpClient)
+                        builder.httpClient,
+                        new DefaultRetryPolicy())
         );
 
         if (aadAadInstanceDiscoveryResponse != null) {

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/DefaultRetryPolicy.java

@@ -0,0 +1,38 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.aad.msal4j;
+
+/**
+ * Default retry policy for most MSAL Java flows
+ */
+class DefaultRetryPolicy implements IRetryPolicy {
+    private static final int RETRY_NUM = 1;
+    private static int RETRY_DELAY_MS = 1000;
+
+    @Override
+    public boolean isRetryable(IHttpResponse httpResponse) {
+        return httpResponse.statusCode() >= 500 &&
+                httpResponse.statusCode() < 600 &&
+                HttpHelper.getRetryAfterHeader(httpResponse) == null;
+    }
+
+    @Override
+    public int getMaxRetryCount(IHttpResponse httpResponse) {
+        return RETRY_NUM;
+    }
+
+    @Override
+    public int getRetryDelayMs(IHttpResponse httpResponse) {
+        return RETRY_DELAY_MS;
+    }
+
+    //Package-private methods to allow much quicker testing. The delay values should be treated as constants in any non-test scenario.
+    static void setRetryDelayMs(int retryDelayMs) {
+        RETRY_DELAY_MS = retryDelayMs;
+    }
+
+    static void resetToDefaults() {
+        RETRY_DELAY_MS = 1000;
+    }
+}

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/HttpHelper.java

@@ -18,18 +18,18 @@ class HttpHelper implements IHttpHelper {
 
     private static final Logger log = LoggerFactory.getLogger(HttpHelper.class);
     public static final String RETRY_AFTER_HEADER = "Retry-After";
-    public static final int RETRY_NUM = 2;
-    public static final int RETRY_DELAY_MS = 1000;
 
     public static final int HTTP_STATUS_200 = 200;
     public static final int HTTP_STATUS_400 = 400;
     public static final int HTTP_STATUS_429 = 429;
     public static final int HTTP_STATUS_500 = 500;
 
     private IHttpClient httpClient;
+    private IRetryPolicy retryPolicy;
 
-    HttpHelper(IHttpClient httpClient) {
+    HttpHelper(IHttpClient httpClient, IRetryPolicy retryPolicy) {
         this.httpClient = httpClient;
+        this.retryPolicy = retryPolicy != null ? retryPolicy : new DefaultRetryPolicy();
     }
 
     public IHttpResponse executeHttpRequest(HttpRequest httpRequest,
@@ -141,20 +141,19 @@ private String getRequestThumbprint(RequestContext requestContext) {
         return StringHelper.createSha256Hash(sb.toString());
     }
 
-    boolean isRetryable(IHttpResponse httpResponse) {
-        return httpResponse.statusCode() >= HTTP_STATUS_500 &&
-                getRetryAfterHeader(httpResponse) == null;
-    }
-
     IHttpResponse executeHttpRequestWithRetries(HttpRequest httpRequest, IHttpClient httpClient)
             throws Exception {
-        IHttpResponse httpResponse = null;
-        for (int i = 0; i < RETRY_NUM; i++) {
+        IHttpResponse httpResponse = httpClient.send(httpRequest);
+
+        int retryCount = 0;
+        int maxRetries = retryPolicy.getMaxRetryCount(httpResponse);
+
+        while (retryPolicy.isRetryable(httpResponse) && retryCount < maxRetries) {
+            Thread.sleep(retryPolicy.getRetryDelayMs(httpResponse));
+
+            retryCount++;
+
             httpResponse = httpClient.send(httpRequest);
-            if (!isRetryable(httpResponse)) {
-                break;
-            }
-            Thread.sleep(RETRY_DELAY_MS);
         }
 
         return httpResponse;
@@ -191,7 +190,7 @@ private void processThrottlingInstructions(IHttpResponse httpResponse, RequestCo
         }
     }
 
-    private Integer getRetryAfterHeader(IHttpResponse httpResponse) {
+    static Integer getRetryAfterHeader(IHttpResponse httpResponse) {
 
         if (httpResponse.headers() != null) {
             TreeMap<String, List<String>> headers = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
@@ -279,4 +278,8 @@ private static void verifyReturnedCorrelationId(final HttpRequest httpRequest,
             log.info(msg);
         }
     }
+
+    void setRetryPolicy(IRetryPolicy retryPolicy) {
+        this.retryPolicy = retryPolicy;
+    }
 }

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/HttpHelperManagedIdentity.java

@@ -36,6 +36,12 @@ public IMDSManagedIdentitySource(MsalRequest msalRequest,
         super(msalRequest, serviceBundle, ManagedIdentitySourceType.IMDS);
         IEnvironmentVariables environmentVariables = getEnvironmentVariables();
 
+        //IMDS uses a different retry policy than the default used in other MI flows
+        IHttpHelper httpHelper = serviceBundle.getHttpHelper();
+        if (httpHelper instanceof HttpHelper) {
+            ((HttpHelper) httpHelper).setRetryPolicy(new IMDSRetryPolicy());
+        }
+
         if (!StringHelper.isNullOrBlank(environmentVariables.getEnvironmentVariable(Constants.AZURE_POD_IDENTITY_AUTHORITY_HOST))){
             LOG.info(String.format("[Managed Identity] Environment variable AZURE_POD_IDENTITY_AUTHORITY_HOST for IMDS returned endpoint: %s", environmentVariables.getEnvironmentVariable(Constants.AZURE_POD_IDENTITY_AUTHORITY_HOST)));
             try {

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/IMDSManagedIdentitySource.java

@@ -0,0 +1,53 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.aad.msal4j;
+
+//IMDS uses a different try policy than other MI flows, see https://github.com/AzureAD/microsoft-authentication-library-for-dotnet/blob/main/docs/imds_retry_based_on_errors.md
+class IMDSRetryPolicy extends ManagedIdentityRetryPolicy {
+    private static final int LINEAR_RETRY_NUM = 7;
+    private static int LINEAR_RETRY_DELAY_MS = 10000; // 10 seconds
+    private static final int EXPONENTIAL_RETRY_NUM = 3;
+    private static int EXPONENTIAL_RETRY_DELAY_MS = 1000; // 1 second
+
+    private int currentRetryCount;
+    private int lastStatusCode;
+    
+    @Override
+    public boolean isRetryable(IHttpResponse httpResponse) {
+        currentRetryCount++;
+        lastStatusCode = httpResponse.statusCode();
+
+        return (lastStatusCode >= 500 && lastStatusCode < 600) ||
+                lastStatusCode == 404 || // Not Found
+                lastStatusCode == 408 || // Request Timeout
+                lastStatusCode == 410 || // Gone
+                lastStatusCode == 429;   // Too Many Requests
+    }
+
+    @Override
+    public int getMaxRetryCount(IHttpResponse httpResponse) {
+        return (httpResponse.statusCode() == 410) ? LINEAR_RETRY_NUM : EXPONENTIAL_RETRY_NUM;
+    }
+
+    @Override
+    public int getRetryDelayMs(IHttpResponse httpResponse) {
+        // Use exponential backoff for non-410 status codes
+        if (lastStatusCode == 410) {
+            return LINEAR_RETRY_DELAY_MS;
+        } else {
+            return (int) (Math.pow(2, currentRetryCount) * EXPONENTIAL_RETRY_DELAY_MS);
+        }
+    }
+
+    //Package-private methods to allow much quicker testing. The delay values should be treated as constants in any non-test scenario.
+    static void setRetryDelayMs(int retryDelayMs) {
+        LINEAR_RETRY_DELAY_MS = retryDelayMs;
+        EXPONENTIAL_RETRY_DELAY_MS = retryDelayMs;
+    }
+
+    static void resetToDefaults() {
+        LINEAR_RETRY_DELAY_MS = 10000;
+        EXPONENTIAL_RETRY_DELAY_MS = 1000;
+    }
+}

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/IMDSRetryPolicy.java

@@ -0,0 +1,25 @@
+package com.microsoft.aad.msal4j;
+
+/**
+ * Interface for HTTP request retry policies
+ */
+ interface IRetryPolicy {
+    /**
+     * Determines whether a request should be retried based on the HTTP response
+     * @param httpResponse The HTTP response to evaluate
+     * @return true if retry should be attempted, false otherwise
+     */
+    boolean isRetryable(IHttpResponse httpResponse);
+
+    /**
+     * Gets the number of retries to attempt based on the HTTP response
+     * @return maximum retry count
+     */
+    int getMaxRetryCount(IHttpResponse httpResponse);
+
+    /**
+     * Gets the delay in milliseconds between retry attempts
+     * @return delay in milliseconds
+     */
+    int getRetryDelayMs(IHttpResponse httpResponse);
+}

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/IRetryPolicy.java

@@ -37,9 +37,10 @@ private ManagedIdentityApplication(Builder builder) {
         super.serviceBundle = new ServiceBundle(
                 builder.executorService,
                 new TelemetryManager(telemetryConsumer, builder.onlySendFailureTelemetry),
-                new HttpHelperManagedIdentity(builder.httpClient == null ?
+                new HttpHelper(builder.httpClient == null ?
                         new DefaultHttpClient(builder.proxy, builder.sslSocketFactory, builder.connectTimeoutForDefaultHttpClient, builder.readTimeoutForDefaultHttpClient) :
-                        builder.httpClient)
+                        builder.httpClient,
+                        new ManagedIdentityRetryPolicy())
         );
         log = LoggerFactory.getLogger(ManagedIdentityApplication.class);
 

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ManagedIdentityApplication.java

@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.aad.msal4j;
+
+/**
+ * Retry policy for most Managed Identity scenarios
+ */
+class ManagedIdentityRetryPolicy implements IRetryPolicy {
+    private static final int RETRY_NUM = 3;
+    private static int RETRY_DELAY_MS = 1000;
+
+    @Override
+    public boolean isRetryable(IHttpResponse httpResponse) {
+        int statusCode = httpResponse.statusCode();
+
+        return statusCode == 404 || // Not Found
+                statusCode == 408 || // Request Timeout
+                statusCode == 429 || // Too Many Requests
+                statusCode == 500 || // Internal Server Error
+                statusCode == 503 || // Service Unavailable
+                statusCode == 504;   // Gateway Timeout
+    }
+
+    @Override
+    public int getMaxRetryCount(IHttpResponse httpResponse) {
+        return RETRY_NUM;
+    }
+
+    @Override
+    public int getRetryDelayMs(IHttpResponse httpResponse) {
+        return RETRY_DELAY_MS;
+    }
+
+    //Package-private methods to allow much quicker testing. The delay values should be treated as constants in any non-test scenario.
+    static void setRetryDelayMs(int retryDelayMs) {
+        RETRY_DELAY_MS = retryDelayMs;
+    }
+
+    static void resetToDefaults() {
+        RETRY_DELAY_MS = 1000;
+    }
+}

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ManagedIdentityRetryPolicy.java

@@ -25,7 +25,7 @@ class ServiceFabricManagedIdentitySource extends AbstractManagedIdentitySource {
     //No other flow need this and an app developer may not be aware of it, so it was decided that for the Service Fabric flow we will simply override
     // any HttpClient that may have been set by the app developer with our own client which performs the validation logic.
     private static IHttpClient httpClient = new DefaultHttpClientManagedIdentity(null, null, null, null);
-    private static HttpHelper httpHelper = new HttpHelperManagedIdentity(httpClient);
+    private static HttpHelper httpHelper = new HttpHelper(httpClient, new ManagedIdentityRetryPolicy());
 
     @Override
     public void createManagedIdentityRequest(String resource) {
@@ -122,6 +122,6 @@ private static URI validateAndGetUri(String msiEndpoint)
     //However, unit tests often need to mock HttpClient and need a way to inject the mocked object into this class.
     static void setHttpClient(IHttpClient client) {
         httpClient = client;
-        httpHelper = new HttpHelperManagedIdentity(httpClient);
+        httpHelper = new HttpHelper(httpClient, new ManagedIdentityRetryPolicy());
     }
 }