Merge pull request #963 from AzureAD/avdunn/retry-behavior-disable

Add API to disable internal retry behavior

Author: Avery-Dunn

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

@@ -32,6 +32,7 @@ public abstract class AbstractApplicationBase implements IApplicationBase {
     private IHttpClient httpClient;
     private Integer connectTimeoutForDefaultHttpClient;
     private Integer readTimeoutForDefaultHttpClient;
+    private boolean retryDisabled;
     String tenant;
 
     //The following fields are set in only some applications and/or set internally by the library. To avoid excessive
@@ -150,6 +151,10 @@ public Integer readTimeoutForDefaultHttpClient() {
         return this.readTimeoutForDefaultHttpClient;
     }
 
+    boolean isRetryDisabled() {
+        return this.retryDisabled;
+    }
+
     String tenant() {
         return this.tenant;
     }
@@ -190,6 +195,7 @@ public abstract static class Builder<T extends Builder<T>> {
         Boolean onlySendFailureTelemetry = false;
         Integer connectTimeoutForDefaultHttpClient;
         Integer readTimeoutForDefaultHttpClient;
+        boolean disableInternalRetries;
         private String clientId;
         private Authority authenticationAuthority = createDefaultAADAuthority();
 
@@ -319,6 +325,18 @@ public T readTimeoutForDefaultHttpClient(Integer val) {
             return self();
         }
 
+        /**
+         * The library has a number of policies for retrying HTTP calls in different scenarios.
+         * <p>
+         * This will disable all internal retry behavior, allowing customized retry behavior via your own implementation of {@link IHttpClient}
+         *
+         * @return instance of the Builder on which method was called
+         */
+        public T disableInternalRetries() {
+            disableInternalRetries = true;
+            return self();
+        }
+
         T telemetryConsumer(Consumer<List<HashMap<String, String>>> val) {
             validateNotNull("telemetryConsumer", val);
 
@@ -356,5 +374,16 @@ private static Authority createDefaultAADAuthority() {
         readTimeoutForDefaultHttpClient = builder.readTimeoutForDefaultHttpClient;
         authenticationAuthority = builder.authenticationAuthority;
         clientId = builder.clientId;
+        retryDisabled = builder.disableInternalRetries;
+
+        if (builder.httpClient == null) {
+            httpClient = new DefaultHttpClient(
+                    builder.proxy,
+                    builder.sslSocketFactory,
+                    builder.connectTimeoutForDefaultHttpClient,
+                    builder.readTimeoutForDefaultHttpClient);
+        } else {
+            httpClient = builder.httpClient;
+        }
     }
 }

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

@@ -566,10 +566,7 @@ public T correlationId(String val) {
         super.serviceBundle = new ServiceBundle(
                 builder.executorService,
                 new TelemetryManager(telemetryConsumer, builder.onlySendFailureTelemetry),
-                new HttpHelper(builder.httpClient == null ?
-                        new DefaultHttpClient(builder.proxy, builder.sslSocketFactory, builder.connectTimeoutForDefaultHttpClient, builder.readTimeoutForDefaultHttpClient) :
-                        builder.httpClient,
-                        new DefaultRetryPolicy())
+                new HttpHelper(this, new DefaultRetryPolicy())
         );
 
         if (aadAadInstanceDiscoveryResponse != null) {

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

@@ -21,12 +21,19 @@ class HttpHelper implements IHttpHelper {
 
     private IHttpClient httpClient;
     private IRetryPolicy retryPolicy;
+    private boolean retryDisabled;
 
     HttpHelper(IHttpClient httpClient, IRetryPolicy retryPolicy) {
         this.httpClient = httpClient;
         this.retryPolicy = retryPolicy != null ? retryPolicy : new DefaultRetryPolicy();
     }
 
+    HttpHelper(AbstractApplicationBase application, IRetryPolicy retryPolicy) {
+        this.httpClient = application.httpClient();
+        this.retryDisabled = application.isRetryDisabled();
+        this.retryPolicy = retryPolicy != null ? retryPolicy : new DefaultRetryPolicy();
+    }
+
     public IHttpResponse executeHttpRequest(HttpRequest httpRequest,
                                             RequestContext requestContext,
                                             ServiceBundle serviceBundle) {
@@ -140,6 +147,10 @@ IHttpResponse executeHttpRequestWithRetries(HttpRequest httpRequest, IHttpClient
             throws Exception {
         IHttpResponse httpResponse = httpClient.send(httpRequest);
 
+        if (retryDisabled) {
+            return httpResponse;
+        }
+
         int retryCount = 0;
         int maxRetries = retryPolicy.getMaxRetryCount(httpResponse);
 

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

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

msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/ManagedIdentityTests.java

@@ -586,6 +586,40 @@ void managedIdentityTest_Retry(ManagedIdentitySourceType source, String endpoint
             fail("MsalServiceException is expected but not thrown.");
         }
 
+        @ParameterizedTest
+        @MethodSource("com.microsoft.aad.msal4j.ManagedIdentityTestDataProvider#createDataWrongScope")
+        void managedIdentityTest_RetriesDisabled(ManagedIdentitySourceType source, String endpoint, String resource) throws Exception {
+            IEnvironmentVariables environmentVariables = new EnvironmentVariablesHelper(source, endpoint);
+            ManagedIdentityApplication.setEnvironmentVariables(environmentVariables);
+
+            DefaultHttpClientManagedIdentity httpClientMock = mock(DefaultHttpClientManagedIdentity.class);
+            if (source == SERVICE_FABRIC) {
+                ServiceFabricManagedIdentitySource.setHttpClient(httpClientMock);
+            }
+
+            miApp = ManagedIdentityApplication
+                    .builder(ManagedIdentityId.systemAssigned())
+                    .httpClient(httpClientMock)
+                    .disableInternalRetries()
+                    .build();
+
+            //Several specific 4xx and 5xx errors, such as 500, should trigger MSAL's retry logic
+            when(httpClientMock.send(expectedRequest(source, resource))).thenReturn(expectedResponse(500, ManagedIdentityTestConstants.MSI_ERROR_RESPONSE_500));
+
+            try {
+                acquireTokenCommon(resource).get();
+            } catch (Exception exception) {
+                assert(exception.getCause() instanceof MsalServiceException);
+
+                //But because retries are disabled at the app level, there should be no retries for any source except Service Fabric, which uses its own HttpClient
+                if (source == SERVICE_FABRIC) {
+                    verify(httpClientMock, times(4)).send(any());
+                } else {
+                    verify(httpClientMock, times(1)).send(any());
+                }
+            }
+        }
+
         @Test
         void managedIdentityTest_IMDSRetry() throws Exception {
             IEnvironmentVariables environmentVariables = new EnvironmentVariablesHelper(IMDS, ManagedIdentityTestConstants.IMDS_ENDPOINT);