Skip to content
Permalink
Browse files
feat: enable pre-emptive async oauth token refreshes (#646)
* 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 e3f4c7eac0417705553ef8259599ec29fc8ad9b4
@@ -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;
@@ -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;

/**
@@ -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);

0 comments on commit e3f4c7e

Please sign in to comment.