Skip to content

Refactor instance/region discovery #763

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Dec 20, 2023
Merged

Refactor instance/region discovery #763

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a7f6b08
Refactor instance/region discovery logic
Avery-Dunn Dec 19, 2023
6df147f
Address PR feedback
Avery-Dunn Dec 19, 2023
f37b2b7
Address PR feedback
Avery-Dunn Dec 19, 2023
fb06253
Return earlier if instance discovery is disabled
Avery-Dunn Dec 19, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 50 additions & 37 deletions msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AadInstanceDiscoveryProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All this top level logic needs logging statements, which will help investiagte issues.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the latest commit I added debug logs for all the major if blocks: when instance discovery is disabled, when there's no cached metadata, when region API is used, etc.

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);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if the request provides an azureRegion, do you still need to do the discovery? Might be better to use the if(masalRequest.application().azureRegion() != null) logic before making this call.

Suggested change
String detectedRegion = discoverRegion(msalRequest, serviceBundle);
if (msalRequest.application().azureRegion() != null) {
host = getRegionalizedHost(authorityUrl.getHost(), msalRequest.application().azureRegion());
}
else if (msalRequest.application().autoDetectRegion() ) {
//Server side telemetry requires the result from region discovery when any part of the region API is used
String detectedRegion = discoverRegion(msalRequest, serviceBundle);
msalRequest.application().azureRegion = detectedRegion;
host = getRegionalizedHost(authorityUrl.getHost(), msalRequest.application().azureRegion());
}

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm having trouble finding the design doc, but there's a list of region telemetry values for different scenarios, including:

  • A developer set the region, and MSAL is able to detect the same one
  • A developer set the region, but MSAL found a different one
  • A developer set the region, but MSAL's autodetect failed for some reason

This means that we need to do autodetection whenever any part of the region API is used, unless we decide to just not collect this telemetry. I believe this behavior is the same in .NET and other MSALs too (feel free to correct me on that @bgavrilMS).

Our RegionTelemetry enum gives a good overview of all the different scenarios we look for: https://github.com/AzureAD/microsoft-authentication-library-for-java/blob/dev/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/RegionTelemetry.java

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, we will try to auto-detect the region even if the app developer tells us "use region X". Auto-detection should only happen once per process (even if it fails) and I believe is set to timeout after 2 sec.

If this is too slow, we can discuss about removing this logic.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Prior to these changes we were incorrectly doing autodetection every time a region was used, since the code to do it was happening before the code to check if instance metadata was cached.

With these changes autodetection happens only if no instance is cached, so it'll only happen once for a given authority/region and thus the 2 second timeout will also only happen once.


//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
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Avery-Dunn don't you need to record the telemetry before this block? As is if the developer didn't provide the region and only did auto-detect it is going to look like the developer provided the region and the auto-detect got the same result as provided.

&& 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;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should you also update the host in this case?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the latest commit I refactored some of the logic here:

  1. The part in the if block has been moved up closer to the start of the method, so if the application's azureRegion value is set it will do the host = getRegionalizedHost... and the regional host will be used when looking through the cache
  2. The part in the else if block has been adjusted so the host = getRegionalizedHost... is done if autodetection succeeds and the azureRegion value was not already set

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);
Expand Down Expand Up @@ -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) {
Expand All @@ -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).
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -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);

Expand All @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -351,7 +364,7 @@ private static void doInstanceDiscoveryAndCache(URL authorityUrl,
}
}

cacheInstanceDiscoveryMetadata(authorityUrl.getHost(), aadInstanceDiscoveryResponse);
cacheInstanceDiscoveryResponse(authorityUrl.getHost(), aadInstanceDiscoveryResponse);
}

private static void validate(AadInstanceDiscoveryResponse aadInstanceDiscoveryResponse) {
Expand Down
2 changes: 1 addition & 1 deletion msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AbstractClientApplicationBase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -708,7 +708,7 @@ public T instanceDiscovery(boolean val) {
instanceDiscovery = builder.instanceDiscovery;

if (aadAadInstanceDiscoveryResponse != null) {
AadInstanceDiscoveryProvider.cacheInstanceDiscoveryMetadata(
AadInstanceDiscoveryProvider.cacheInstanceDiscoveryResponse(
authenticationAuthority.host,
aadAadInstanceDiscoveryResponse);
}
Expand Down