Skip to content
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

Jump to bottom

[core-rest-pipeline] Idea for improving the tokenCycler #15120

Closed
wants to merge 3 commits into from
Closed

[core-rest-pipeline] Idea for improving the tokenCycler #15120

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8f253b6
[core-rest-pipeline] Idea for improving the tokenCycler
sadasant May 3, 2021
8217ca6
recent feedback by Will
sadasant May 19, 2021
7539956
Merge remote-tracking branch 'Azure/master' into core-rest-pipeline/t…
sadasant May 19, 2021
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
4 changes: 2 additions & 2 deletions sdk/core/core-rest-pipeline/src/policies/bearerTokenAuthenticationPolicy.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,12 +38,12 @@ export function bearerTokenAuthenticationPolicy(
// The options are left out of the public API until there's demand to configure this.
// Remember to extend `BearerTokenAuthenticationPolicyOptions` with `TokenCyclerOptions`
// in order to pass through the `options` object.
const cycler = createTokenCycler(credential /* , options */);
const getToken = createTokenCycler(credential /* , options */);

return {
name: bearerTokenAuthenticationPolicyName,
async sendRequest(request: PipelineRequest, next: SendRequest): Promise<PipelineResponse> {
const { token } = await cycler.getToken(scopes, {
sadasant marked this conversation as resolved.
Show resolved Hide resolved
const { token } = await getToken(scopes, {
abortSignal: request.abortSignal,
tracingOptions: request.tracingOptions
});
Expand Down
2 changes: 1 addition & 1 deletion sdk/core/core-rest-pipeline/src/policies/bearerTokenChallengeAuthenticationPolicy.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ export function bearerTokenChallengeAuthenticationPolicy(
// Remember to extend `BearerTokenChallengeAuthenticationPolicyOptions` with `TokenCyclerOptions`
// in order to pass through the `options` object.
const getAccessToken = credential
? createTokenCycler(credential /* , options */).getToken
? createTokenCycler(credential /* , options */)
: () => Promise.resolve(null);

return {
Expand Down
146 changes: 71 additions & 75 deletions sdk/core/core-rest-pipeline/src/util/tokenCycler.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,6 @@ export type AccessTokenGetter = (
options: GetTokenOptions
) => Promise<AccessToken>;

/**
* The response of the
*/
export interface AccessTokenRefresher {
cachedToken: AccessToken | null;
getToken: AccessTokenGetter;
}

export interface TokenCyclerOptions {
/**
* The window of time before token expiration during which the token will be
Expand Down Expand Up @@ -52,50 +44,59 @@ export const DEFAULT_CYCLER_OPTIONS: TokenCyclerOptions = {
};

/**
* Converts an an unreliable access token getter (which may resolve with null)
* into an AccessTokenGetter by retrying the unreliable getter in a regular
* interval.
* Calls to producer until the timeout is reached.
* Prevents errors from bubbling up until the refresh timeout is reached.
* Once the timeout is reached, errors will bubble up, and null return values will throw with the received "errorMessage".
*
* @param getAccessToken - A function that produces a promise of an access token that may fail by returning null.
* @param retryIntervalInMs - The time (in milliseconds) to wait between retry attempts.
* @param refreshTimeout - The timestamp after which the refresh attempt will fail, throwing an exception.
* @returns - A promise that, if it resolves, will resolve with an access token.
* @param producer - The asynchronous function to call.
* @param timeout - Unix time when to stop ignoring errors or null values.
* @param errorMessage - Message to be used if the final value after the timeout is reached is still "null".
* @returns - A promise that can be resolved with T or null until the time is reached, after which it will throw if it can't resolve with a value of the type T.
*/
async function beginRefresh(
getAccessToken: () => Promise<AccessToken | null>,
retryIntervalInMs: number,
refreshTimeout: number
): Promise<AccessToken> {
// This wrapper handles exceptions gracefully as long as we haven't exceeded
// the timeout.
async function tryGetAccessToken(): Promise<AccessToken | null> {
if (Date.now() < refreshTimeout) {
try {
return await getAccessToken();
} catch {
return null;
}
} else {
const finalToken = await getAccessToken();

// Timeout is up, so throw if it's still null
if (finalToken === null) {
throw new Error("Failed to refresh access token.");
}

return finalToken;
async function tryUntilTimeout<T>(
Copy link
Member

Choose a reason for hiding this comment

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

To me the names tryUntilTimeout and retryUntilTimeout mean the same thing. This function is specifically made to be part of the inner loop of retryUntilTimeout, and I think it has limited utility on its own.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Inside of retryUntilTimeout I think it adds too much complexity for one function. The name could be better though. This could become something like "run silently unless a condition is met", and receive a condition instead of a timeout, which would be generic.

Copy link
Member

Choose a reason for hiding this comment

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

I just think it's a very particular function. "tryOnceAndThrowIfTimeoutExpired"... It only makes sense to be used with retryUntilTimeout, which still parses in my brain as meaning the exact same thing as tryUntilTimeout. Maybe this function could simply be called tryOnce.

I generally like the idea of accepting a predicate representing the error condition instead of checking the date, but like I mentioned in the other comment, I'm skeptical that this makes much sense to use with other "refreshable things" as long as | null is baked so deeply in. That feels like a leaky abstraction where things would start wanting to return null as a failure mode just for the sake of fitting into that abstraction, even if it's more natural to reject/call a handler, etc.

It might also be possible to use some clever promise chains to eliminate this secondary function entirely.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Oh I see! I like that. Alright, let's come back to this later when we have more time 👍

producer: () => Promise<T | null>,
timeout: number
): Promise<T | null> {
if (Date.now() < timeout) {
try {
return await producer();
} catch {
return null;
}
} else {
const result = await producer();

// Timeout is up, so throw if it's still null
if (result === null) {
throw new Error("operation timed out");
}
return result;
}
}

let token: AccessToken | null = await tryGetAccessToken();
/**
* Retries an asynchronous function as often as the provided interval milliseconds,
* until the provided Unix date timeout is reached.
* Once the timeout is reached, errors will bubble up, and null return values will throw with the received "errorMessage".
*
* @param producer - The asynchronous function to call.
* @param retryIntervalInMs - The time (in milliseconds) to wait between retry attempts.
* @param refreshTimeout - The timestamp after which the refresh attempt will fail, throwing an exception.
* @returns - A promise that, if it resolves, will resolve with the T type.
*/
async function retryUntilTimeout<T>(
producer: () => Promise<T | null>,
retryIntervalInMs: number,
retryTimeout: number
): Promise<T> {
let result = await tryUntilTimeout<T>(producer, retryTimeout);

while (token === null) {
while (result === null) {
Copy link
Member

Choose a reason for hiding this comment

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

The issue I have with this is that it's still particular to the contract of getToken in core-auth. You wouldn't be able to use this abstraction with something that doesn't return null as a failure mode. The way this reads to me is "retry every X milliseconds as long as this function is returning null," and then in the inner function we catch errors and return them as null to retry, but that might not always be safe to do in the general case.

If we want to do this refactor with the intention of making this cycler abstraction easier to use with other types of functions, then I think we have to re-think the failure modes to come up with something truly generic. For example, what if the retryUntilTimeout function accepts as a parameter a predicate that can be used to determine whether an Error can be retried or should be treated as fatal?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I agree that the failure modes are an issue on themselves and that null as a failure mode is not generic, but null as a failure mode is easy for two reasons:

  • Has way less overhead.
  • It's as clear as possible on the type that uses null as the failure mode. It has been clear enough for us to use it in the TokenCredential, for example 🤔

As an alternative to null as a failure mode, what would be better while still not make it much more complicated?

await delay(retryIntervalInMs);

token = await tryGetAccessToken();
result = await tryUntilTimeout(producer, retryTimeout);
}

return token;
return result;
}

/**
Expand All @@ -115,7 +116,7 @@ async function beginRefresh(
export function createTokenCycler(
credential: TokenCredential,
tokenCyclerOptions?: Partial<TokenCyclerOptions>
): AccessTokenRefresher {
): AccessTokenGetter {
let refreshWorker: Promise<AccessToken> | null = null;
let token: AccessToken | null = null;

Expand All @@ -128,7 +129,7 @@ export function createTokenCycler(
* This little holder defines several predicates that we use to construct
* the rules of refreshing the token.
*/
const cycler = {
const self = {
/**
* Produces true if a refresh job is currently in progress.
*/
Expand All @@ -141,7 +142,7 @@ export function createTokenCycler(
*/
get shouldRefresh(): boolean {
return (
!cycler.isRefreshing &&
!self.isRefreshing &&
(token?.expiresOnTimestamp ?? 0) - options.refreshWindowInMs < Date.now()
);
},
Expand All @@ -164,14 +165,14 @@ export function createTokenCycler(
scopes: string | string[],
getTokenOptions: GetTokenOptions
): Promise<AccessToken> {
if (!cycler.isRefreshing) {
if (!self.isRefreshing) {
// We bind `scopes` here to avoid passing it around a lot
const tryGetAccessToken = (): Promise<AccessToken | null> =>
credential.getToken(scopes, getTokenOptions);

// Take advantage of promise chaining to insert an assignment to `token`
// before the refresh can be considered done.
refreshWorker = beginRefresh(
refreshWorker = retryUntilTimeout<AccessToken>(
tryGetAccessToken,
options.retryIntervalInMs,
// If we don't have a token, then we should timeout immediately
Expand All @@ -182,44 +183,39 @@ export function createTokenCycler(
token = _token;
return token;
})
.catch((reason) => {
.catch((error) => {
// We also should reset the refresher if we enter a failed state. All
// existing awaiters will throw, but subsequent requests will start a
// new retry chain.
refreshWorker = null;
token = null;
throw reason;
throw new Error(`Failed to refresh access token: ${error.message}`);
});
}

return refreshWorker as Promise<AccessToken>;
}

return {
get cachedToken(): AccessToken | null {
return token;
},
getToken: async (
scopes: string | string[],
tokenOptions: GetTokenOptions
): Promise<AccessToken> => {
//
// Simple rules:
// - If we MUST refresh, then return the refresh task, blocking
// the pipeline until a token is available.
// - If we SHOULD refresh, then run refresh but don't return it
// (we can still use the cached token).
// - Return the token, since it's fine if we didn't return in
// step 1.
//

if (cycler.mustRefresh) return refresh(scopes, tokenOptions);

if (cycler.shouldRefresh) {
refresh(scopes, tokenOptions);
}

return token as AccessToken;
return async function(
scopes: string | string[],
tokenOptions: GetTokenOptions
): Promise<AccessToken> {
//
// Simple rules:
// - If we MUST refresh, then return the refresh task, blocking
// the pipeline until a token is available.
// - If we SHOULD refresh, then run refresh but don't return it
// (we can still use the cached token).
// - Return the token, since it's fine if we didn't return in
// step 1.
//

if (self.mustRefresh) return refresh(scopes, tokenOptions);

if (self.shouldRefresh) {
refresh(scopes, tokenOptions);
}

return token as AccessToken;
};
}
4 changes: 2 additions & 2 deletions sdk/core/core-rest-pipeline/test/bearerTokenAuthenticationPolicy.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,7 +172,7 @@ describe("BearerTokenAuthenticationPolicy", function() {
} catch (e) {
error = e;
}
assert.equal(error?.message, "Failed to retrieve the token");
assert.equal(error?.message, "Failed to refresh access token: Test Error");

assert.strictEqual(
credential.authCount,
Expand Down Expand Up @@ -248,7 +248,7 @@ class MockRefreshAzureCredential implements TokenCredential {
this.authCount++;

if (this.shouldThrow) {
throw new Error("Failed to retrieve the token");
throw new Error("Test Error");
}

// Allowing getToken to take a while
Expand Down