Refactor instance/region discovery (#763)

* Refactor instance/region discovery logic

* Address PR feedback

* Address PR feedback

* Return earlier if instance discovery is disabled

Author: Avery-Dunn

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

@@ -62,40 +62,50 @@ static InstanceDiscoveryMetadataEntry getMetadataEntry(URL authorityUrl,
                                                            boolean validateAuthority,
                                                            MsalRequest msalRequest,
                                                            ServiceBundle serviceBundle) {
-        String host = authorityUrl.getHost();
 
-        if (shouldUseRegionalEndpoint(msalRequest)) {
-            //Server side telemetry requires the result from region discovery when any part of the region API is used
-            String detectedRegion = discoverRegion(msalRequest, serviceBundle);
+        String host = authorityUrl.getHost();
 
-            if (msalRequest.application().azureRegion() != null) {
-                host = getRegionalizedHost(authorityUrl.getHost(), msalRequest.application().azureRegion());
+        //If instanceDiscovery flag set to false, cache a basic instance metadata entry to skip future lookups
+        if (!msalRequest.application().instanceDiscovery()) {
+            if (cache.get(host) == null) {
+                log.debug("Instance discovery set to false, caching a default entry.");
+                cacheInstanceDiscoveryMetadata(host);
             }
+            return cache.get(host);
+        }
 
-            //If region autodetection is enabled and a specific region not already set,
-            // set the application's region to the discovered region so that future requests can skip the IMDS endpoint call
-            if (null == msalRequest.application().azureRegion() && msalRequest.application().autoDetectRegion()
-                    && null != detectedRegion) {
-                msalRequest.application().azureRegion = detectedRegion;
-            }
-            cacheRegionInstanceMetadata(authorityUrl.getHost(), msalRequest.application().azureRegion());
-            serviceBundle.getServerSideTelemetry().getCurrentRequest().regionOutcome(
-                    determineRegionOutcome(detectedRegion, msalRequest.application().azureRegion(), msalRequest.application().autoDetectRegion()));
+        //If a region was set by an app developer or previously found through autodetection, adjust the authority host to use it
+        if (shouldUseRegionalEndpoint(msalRequest) && msalRequest.application().azureRegion() != null) {
+            host = getRegionalizedHost(authorityUrl.getHost(), msalRequest.application().azureRegion());
         }
 
-        InstanceDiscoveryMetadataEntry result = cache.get(host);
+        //If there is no cached instance metadata, do instance discovery cache the result
+        if (cache.get(host) == null) {
+            log.debug("No cached instance metadata, will attempt instance discovery.");
 
-        if (result == null) {
-            if(msalRequest.application().instanceDiscovery() && !instanceDiscoveryFailed){
-                doInstanceDiscoveryAndCache(authorityUrl, validateAuthority, msalRequest, serviceBundle);
-            } else {
-                // instanceDiscovery flag is set to False. Do not perform instanceDiscovery.
-                return InstanceDiscoveryMetadataEntry.builder().
-                        preferredCache(host).
-                        preferredNetwork(host).
-                        aliases(Collections.singleton(host)).
-                        build();
+            if (shouldUseRegionalEndpoint(msalRequest)) {
+                log.debug("Region API used, will attempt to discover Azure region.");
+
+                //Server side telemetry requires the result from region discovery when any part of the region API is used
+                String detectedRegion = discoverRegion(msalRequest, serviceBundle);
+
+                //If region autodetection is enabled and a specific region was not already set, set the application's
+                // region to the discovered region so that future requests can skip the IMDS endpoint call
+                if (msalRequest.application().azureRegion() == null
+                        && msalRequest.application().autoDetectRegion()
+                        && detectedRegion != null) {
+                    log.debug(String.format("Region autodetection found %s, this region will be used for future calls.", detectedRegion));
+
+                    msalRequest.application().azureRegion = detectedRegion;
+                    host = getRegionalizedHost(authorityUrl.getHost(), msalRequest.application().azureRegion());
+                }
+
+                cacheRegionInstanceMetadata(authorityUrl.getHost(), host);
+                serviceBundle.getServerSideTelemetry().getCurrentRequest().regionOutcome(
+                        determineRegionOutcome(detectedRegion, msalRequest.application().azureRegion(), msalRequest.application().autoDetectRegion()));
             }
+
+            doInstanceDiscoveryAndCache(authorityUrl, validateAuthority, msalRequest, serviceBundle);
         }
 
         return cache.get(host);
@@ -126,7 +136,7 @@ static AadInstanceDiscoveryResponse parseInstanceDiscoveryMetadata(String instan
         return aadInstanceDiscoveryResponse;
     }
 
-    static void cacheInstanceDiscoveryMetadata(String host,
+    static void cacheInstanceDiscoveryResponse(String host,
                                                AadInstanceDiscoveryResponse aadInstanceDiscoveryResponse) {
 
         if (aadInstanceDiscoveryResponse != null && aadInstanceDiscoveryResponse.metadata() != null) {
@@ -136,6 +146,11 @@ static void cacheInstanceDiscoveryMetadata(String host,
                 }
             }
         }
+
+        cacheInstanceDiscoveryMetadata(host);
+    }
+
+    static void cacheInstanceDiscoveryMetadata(String host) {
         cache.putIfAbsent(host, InstanceDiscoveryMetadataEntry.builder().
                 preferredCache(host).
                 preferredNetwork(host).
@@ -164,14 +179,13 @@ private static boolean shouldUseRegionalEndpoint(MsalRequest msalRequest){
         return false;
     }
 
-    static void cacheRegionInstanceMetadata(String host, String region) {
+    static void cacheRegionInstanceMetadata(String originalHost, String regionalHost) {
 
         Set<String> aliases = new HashSet<>();
-        aliases.add(host);
-        String regionalHost = getRegionalizedHost(host, region);
+        aliases.add(originalHost);
 
         cache.putIfAbsent(regionalHost, InstanceDiscoveryMetadataEntry.builder().
-                preferredCache(host).
+                preferredCache(originalHost).
                 preferredNetwork(regionalHost).
                 aliases(aliases).
                 build());
@@ -229,12 +243,10 @@ static AadInstanceDiscoveryResponse sendInstanceDiscoveryRequest(URL authorityUr
                                                                              MsalRequest msalRequest,
                                                                              ServiceBundle serviceBundle) {
 
-        IHttpResponse httpResponse = null;
-
         String instanceDiscoveryRequestUrl = getInstanceDiscoveryEndpoint(authorityUrl) +
                 formInstanceDiscoveryParameters(authorityUrl);
 
-        httpResponse = executeRequest(instanceDiscoveryRequestUrl, msalRequest.headers().getReadonlyHeaderMap(), msalRequest, serviceBundle);
+        IHttpResponse httpResponse = executeRequest(instanceDiscoveryRequestUrl, msalRequest.headers().getReadonlyHeaderMap(), msalRequest, serviceBundle);
 
         AadInstanceDiscoveryResponse response = JsonHelper.convertJsonToObject(httpResponse.body(), AadInstanceDiscoveryResponse.class);
 
@@ -244,7 +256,8 @@ static AadInstanceDiscoveryResponse sendInstanceDiscoveryRequest(URL authorityUr
                 throw MsalServiceExceptionFactory.fromHttpResponse(httpResponse);
             }
             // instance discovery failed due to reasons other than an invalid authority, do not perform instance discovery again in this environment.
-            instanceDiscoveryFailed = true;
+            log.debug("Instance discovery failed due to an unknown error, no more instance discovery attempts will be made.");
+            cacheInstanceDiscoveryMetadata(authorityUrl.getHost());
         }
 
         return response;
@@ -295,7 +308,7 @@ static String discoverRegion(MsalRequest msalRequest, ServiceBundle serviceBundl
 
         //Check if the REGION_NAME environment variable has a value for the region
         if (System.getenv(REGION_NAME) != null) {
-            log.info("Region found in environment variable: " + System.getenv(REGION_NAME));
+            log.info(String.format("Region found in environment variable: %s",System.getenv(REGION_NAME)));
             currentRequest.regionSource(RegionTelemetry.REGION_SOURCE_ENV_VARIABLE.telemetryValue);
 
             return System.getenv(REGION_NAME);
@@ -351,7 +364,7 @@ private static void doInstanceDiscoveryAndCache(URL authorityUrl,
             }
         }
 
-        cacheInstanceDiscoveryMetadata(authorityUrl.getHost(), aadInstanceDiscoveryResponse);
+        cacheInstanceDiscoveryResponse(authorityUrl.getHost(), aadInstanceDiscoveryResponse);
     }
 
     private static void validate(AadInstanceDiscoveryResponse aadInstanceDiscoveryResponse) {

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

@@ -708,7 +708,7 @@ public T instanceDiscovery(boolean val) {
         instanceDiscovery = builder.instanceDiscovery;
 
         if (aadAadInstanceDiscoveryResponse != null) {
-            AadInstanceDiscoveryProvider.cacheInstanceDiscoveryMetadata(
+            AadInstanceDiscoveryProvider.cacheInstanceDiscoveryResponse(
                     authenticationAuthority.host,
                     aadAadInstanceDiscoveryResponse);
         }