Skip to content

Commit

Permalink
feat: enable pre-emptive async oauth token refreshes (#646)
Browse files Browse the repository at this point in the history
* feat: add pre-emptive async oauth token refreshes

This is currently a rough sketch and should not be merged. I just wanted to get some feedback here.

The current implementation of oauth refresh offloads the IO to a separate executor, however when a token expires, all calls to getRequestMetadata would hang until the token is refreshed.
This PR is a rough sketch to improve the situation by adding a stale state to token. If a token is within a minute of expiration, the first request to notice this, will spawn a refresh on the executor and immediately return with the existing token. This avoids hourly latency spikes for grpc.

The implementation uses guava's ListenableFutures to manage the async refresh state. Although the apis are marked BetaApi in guava 19, they are GA in guava 30

* The class introduces 3 distinct states:
  * Expired -  the token is not usable
  * Stale - the token is useable, but its time to refresh
  * Fresh - token can be used without any extra effort
* All of the funtionality of getRequestMetadata has been extracted to asyncFetch. The new function will check if the token is unfresh and schedule a refresh using the executor
* asyncFetch uses ListenableFutures to wrap state: if the token is not expired, an immediate future is returned. If the token is expired the future of the refresh task is returned
* A helper refreshAsync 	& finishRefreshAsync are also introduced. They schedule the actual refresh and propagate the side effects
* To handle blocking invocation: the async functionality is re-used but a DirectExecutor is used. All ExecutionErrors are unwrapped. In most cases the stack trace will be preserved because of the DirectExecutor. However if the async & sync methods are interleaved, it's possible that a sync call will await an async refresh task. In this case the callers stacktrace will not be present.

* minor doc

* update broken test

* prep for merging with master: The initial async refresh feature was implemented on top of 0.8, so now I'm backporting features to minimize the merge conflicts

* in blocking mode, when a token is stale, only block the first caller and allow subsequent callers to use the stale token

* use private monitor to minimize change noise

* minor tweaks and test

* format

* fix ExternalAccountCredentials

* fix boundary checks and add a few more tests

* add another test for making sure that blocking stacktraces include the caller

* address feedback

* optimize for the common case

* codestyle

* use Date to calculate cache state to fix tests that mock access token

* remove accidental double call

Co-authored-by: Les Vogel <lesv@users.noreply.github.com>
  • Loading branch information
igorbernstein2 and lesv committed May 20, 2021
1 parent 469b160 commit e3f4c7e
Show file tree
Hide file tree
Showing 5 changed files with 684 additions and 102 deletions.
Expand Up @@ -35,6 +35,7 @@

import com.google.api.client.json.GenericJson;
import com.google.api.client.json.JsonObjectParser;
import com.google.auth.RequestMetadataCallback;
import com.google.auth.http.HttpTransportFactory;
import com.google.auth.oauth2.AwsCredentials.AwsCredentialSource;
import com.google.auth.oauth2.IdentityPoolCredentials.IdentityPoolCredentialSource;
Expand All @@ -48,6 +49,7 @@
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Executor;
import javax.annotation.Nullable;

/**
Expand Down Expand Up @@ -208,6 +210,26 @@ private ImpersonatedCredentials initializeImpersonatedCredentials() {
.build();
}

@Override
public void getRequestMetadata(
URI uri, Executor executor, final RequestMetadataCallback callback) {
super.getRequestMetadata(
uri,
executor,
new RequestMetadataCallback() {
@Override
public void onSuccess(Map<String, List<String>> metadata) {
metadata = addQuotaProjectIdToRequestMetadata(quotaProjectId, metadata);
callback.onSuccess(metadata);
}

@Override
public void onFailure(Throwable exception) {
callback.onFailure(exception);
}
});
}

@Override
public Map<String, List<String>> getRequestMetadata(URI uri) throws IOException {
Map<String, List<String>> requestMetadata = super.getRequestMetadata(uri);
Expand Down

0 comments on commit e3f4c7e

Please sign in to comment.