diff --git a/packages/middleware-retry/src/constants.ts b/packages/middleware-retry/src/constants.ts index 5f976928e514..c14baa7b87c3 100644 --- a/packages/middleware-retry/src/constants.ts +++ b/packages/middleware-retry/src/constants.ts @@ -15,3 +15,25 @@ export const MAXIMUM_RETRY_DELAY = 20 * 1000; * encountered. */ export const THROTTLING_RETRY_DELAY_BASE = 500; + +/** + * Initial number of retry tokens in Retry Quota + */ +export const INITIAL_RETRY_TOKENS = 500; + +/** + * The total amount of retry tokens to be decremented from retry token balance. + */ +export const RETRY_COST = 5; + +/** + * The total amount of retry tokens to be decremented from retry token balance + * when a throttling error is encountered. + */ +export const TIMEOUT_RETRY_COST = 10; + +/** + * The total amount of retry token to be incremented from retry token balance + * if an SDK operation invocation succeeds without requiring a retry request. + */ +export const NO_RETRY_INCREMENT = 1; diff --git a/packages/middleware-retry/src/defaultRetryQuota.spec.ts b/packages/middleware-retry/src/defaultRetryQuota.spec.ts new file mode 100644 index 000000000000..985c591771ac --- /dev/null +++ b/packages/middleware-retry/src/defaultRetryQuota.spec.ts @@ -0,0 +1,189 @@ +import { getDefaultRetryQuota } from "./defaultRetryQuota"; +import { SdkError } from "@aws-sdk/smithy-client"; +import { + INITIAL_RETRY_TOKENS, + TIMEOUT_RETRY_COST, + RETRY_COST, + NO_RETRY_INCREMENT +} from "./constants"; + +describe("defaultRetryQuota", () => { + const getMockError = () => new Error() as SdkError; + const getMockTimeoutError = () => + Object.assign(new Error(), { + name: "TimeoutError" + }) as SdkError; + + const getDrainedRetryQuota = ( + targetCapacity: number, + error: SdkError, + initialRetryTokens: number = INITIAL_RETRY_TOKENS + ) => { + const retryQuota = getDefaultRetryQuota(initialRetryTokens); + let availableCapacity = initialRetryTokens; + while (availableCapacity >= targetCapacity) { + retryQuota.retrieveRetryTokens(error); + availableCapacity -= targetCapacity; + } + return retryQuota; + }; + + describe("custom initial retry tokens", () => { + it("hasRetryTokens returns false if capacity is not available", () => { + const customRetryTokens = 100; + const error = getMockError(); + const retryQuota = getDrainedRetryQuota( + RETRY_COST, + error, + customRetryTokens + ); + expect(retryQuota.hasRetryTokens(error)).toBe(false); + }); + + it("retrieveRetryToken throws error if retry tokens not available", () => { + const customRetryTokens = 100; + const error = getMockError(); + const retryQuota = getDrainedRetryQuota( + RETRY_COST, + error, + customRetryTokens + ); + expect(() => { + retryQuota.retrieveRetryTokens(error); + }).toThrowError(new Error("No retry token available")); + }); + }); + + describe("hasRetryTokens", () => { + describe("returns true if capacity is available", () => { + it("when it's TimeoutError", () => { + const timeoutError = getMockTimeoutError(); + expect( + getDefaultRetryQuota(INITIAL_RETRY_TOKENS).hasRetryTokens( + timeoutError + ) + ).toBe(true); + }); + + it("when it's not TimeoutError", () => { + expect( + getDefaultRetryQuota(INITIAL_RETRY_TOKENS).hasRetryTokens( + getMockError() + ) + ).toBe(true); + }); + }); + + describe("returns false if capacity is not available", () => { + it("when it's TimeoutError", () => { + const timeoutError = getMockTimeoutError(); + const retryQuota = getDrainedRetryQuota( + TIMEOUT_RETRY_COST, + timeoutError + ); + expect(retryQuota.hasRetryTokens(timeoutError)).toBe(false); + }); + + it("when it's not TimeoutError", () => { + const error = getMockError(); + const retryQuota = getDrainedRetryQuota(RETRY_COST, error); + expect(retryQuota.hasRetryTokens(error)).toBe(false); + }); + }); + }); + + describe("retrieveRetryToken", () => { + describe("returns retry tokens amount if available", () => { + it("when it's TimeoutError", () => { + const timeoutError = getMockTimeoutError(); + expect( + getDefaultRetryQuota(INITIAL_RETRY_TOKENS).retrieveRetryTokens( + timeoutError + ) + ).toBe(TIMEOUT_RETRY_COST); + }); + + it("when it's not TimeoutError", () => { + expect( + getDefaultRetryQuota(INITIAL_RETRY_TOKENS).retrieveRetryTokens( + getMockError() + ) + ).toBe(RETRY_COST); + }); + }); + + describe("throws error if retry tokens not available", () => { + it("when it's TimeoutError", () => { + const timeoutError = getMockTimeoutError(); + const retryQuota = getDrainedRetryQuota( + TIMEOUT_RETRY_COST, + timeoutError + ); + expect(() => { + retryQuota.retrieveRetryTokens(timeoutError); + }).toThrowError(new Error("No retry token available")); + }); + + it("when it's not TimeoutError", () => { + const error = getMockError(); + const retryQuota = getDrainedRetryQuota(RETRY_COST, error); + expect(() => { + retryQuota.retrieveRetryTokens(error); + }).toThrowError(new Error("No retry token available")); + }); + }); + }); + + describe("releaseRetryToken", () => { + it("adds capacityReleaseAmount if passed", () => { + const error = getMockError(); + const retryQuota = getDrainedRetryQuota(RETRY_COST, error); + + // Ensure that retry tokens are not available. + expect(retryQuota.hasRetryTokens(error)).toBe(false); + + // Release RETRY_COST tokens. + retryQuota.releaseRetryTokens(RETRY_COST); + expect(retryQuota.hasRetryTokens(error)).toBe(true); + expect(retryQuota.retrieveRetryTokens(error)).toBe(RETRY_COST); + expect(retryQuota.hasRetryTokens(error)).toBe(false); + }); + + it("adds NO_RETRY_INCREMENT if capacityReleaseAmount not passed", () => { + const error = getMockError(); + const retryQuota = getDrainedRetryQuota(RETRY_COST, error); + + // retry tokens will not be available till NO_RETRY_INCREMENT is added + // till it's equal to RETRY_COST - (INITIAL_RETRY_TOKENS % RETRY_COST) + let tokensReleased = 0; + const tokensToBeReleased = + RETRY_COST - (INITIAL_RETRY_TOKENS % RETRY_COST); + while (tokensReleased < tokensToBeReleased) { + expect(retryQuota.hasRetryTokens(error)).toBe(false); + retryQuota.releaseRetryTokens(); + tokensReleased += NO_RETRY_INCREMENT; + } + expect(retryQuota.hasRetryTokens(error)).toBe(true); + }); + + it("ensures availableCapacity is maxed at INITIAL_RETRY_TOKENS", () => { + const error = getMockError(); + const retryQuota = getDefaultRetryQuota(INITIAL_RETRY_TOKENS); + + // release 100 tokens. + [...Array(100).keys()].forEach(key => { + retryQuota.releaseRetryTokens(); + }); + + // availableCapacity is still maxed at INITIAL_RETRY_TOKENS + // hasRetryTokens would be true only till INITIAL_RETRY_TOKENS/RETRY_COST times + [...Array(Math.floor(INITIAL_RETRY_TOKENS / RETRY_COST)).keys()].forEach( + key => { + expect(retryQuota.hasRetryTokens(error)).toBe(true); + retryQuota.retrieveRetryTokens(error); + } + ); + expect(retryQuota.hasRetryTokens(error)).toBe(false); + }); + }); +}); diff --git a/packages/middleware-retry/src/defaultRetryQuota.ts b/packages/middleware-retry/src/defaultRetryQuota.ts new file mode 100644 index 000000000000..3a78abf2a3ae --- /dev/null +++ b/packages/middleware-retry/src/defaultRetryQuota.ts @@ -0,0 +1,41 @@ +import { RetryQuota } from "./defaultStrategy"; +import { SdkError } from "@aws-sdk/smithy-client"; +import { + RETRY_COST, + TIMEOUT_RETRY_COST, + NO_RETRY_INCREMENT +} from "./constants"; + +export const getDefaultRetryQuota = ( + initialRetryTokens: number +): RetryQuota => { + const MAX_CAPACITY = initialRetryTokens; + let availableCapacity = initialRetryTokens; + + const getCapacityAmount = (error: SdkError) => + error.name === "TimeoutError" ? TIMEOUT_RETRY_COST : RETRY_COST; + + const hasRetryTokens = (error: SdkError) => + getCapacityAmount(error) <= availableCapacity; + + const retrieveRetryTokens = (error: SdkError) => { + if (!hasRetryTokens(error)) { + // retryStrategy should stop retrying, and return last error + throw new Error("No retry token available"); + } + const capacityAmount = getCapacityAmount(error); + availableCapacity -= capacityAmount; + return capacityAmount; + }; + + const releaseRetryTokens = (capacityReleaseAmount?: number) => { + availableCapacity += capacityReleaseAmount ?? NO_RETRY_INCREMENT; + availableCapacity = Math.min(availableCapacity, MAX_CAPACITY); + }; + + return Object.freeze({ + hasRetryTokens, + retrieveRetryTokens, + releaseRetryTokens + }); +}; diff --git a/packages/middleware-retry/src/defaultStrategy.spec.ts b/packages/middleware-retry/src/defaultStrategy.spec.ts index 0b9207a34eeb..8809f55b6dce 100644 --- a/packages/middleware-retry/src/defaultStrategy.spec.ts +++ b/packages/middleware-retry/src/defaultStrategy.spec.ts @@ -1,11 +1,13 @@ import { DEFAULT_RETRY_DELAY_BASE, - THROTTLING_RETRY_DELAY_BASE + THROTTLING_RETRY_DELAY_BASE, + INITIAL_RETRY_TOKENS } from "./constants"; import { isThrottlingError } from "@aws-sdk/service-error-classification"; import { defaultDelayDecider } from "./delayDecider"; import { defaultRetryDecider } from "./retryDecider"; -import { StandardRetryStrategy } from "./defaultStrategy"; +import { StandardRetryStrategy, RetryQuota } from "./defaultStrategy"; +import { getDefaultRetryQuota } from "./defaultRetryQuota"; jest.mock("@aws-sdk/service-error-classification", () => ({ isThrottlingError: jest.fn().mockReturnValue(true) @@ -19,8 +21,85 @@ jest.mock("./retryDecider", () => ({ defaultRetryDecider: jest.fn().mockReturnValue(true) })); +jest.mock("./defaultRetryQuota", () => { + const mockDefaultRetryQuota = { + hasRetryTokens: jest.fn().mockReturnValue(true), + retrieveRetryTokens: jest.fn().mockReturnValue(1), + releaseRetryTokens: jest.fn() + }; + return { getDefaultRetryQuota: () => mockDefaultRetryQuota }; +}); + describe("defaultStrategy", () => { - const maxAttempts = 2; + const maxAttempts = 3; + + const mockSuccessfulOperation = ( + maxAttempts: number, + options?: { mockResponse?: string } + ) => { + const next = jest.fn().mockResolvedValueOnce({ + response: options?.mockResponse, + output: { $metadata: {} } + }); + + const retryStrategy = new StandardRetryStrategy(maxAttempts); + return retryStrategy.retry(next, {} as any); + }; + + const mockFailedOperation = async ( + maxAttempts: number, + options?: { mockError?: Error } + ) => { + const mockError = options?.mockError ?? new Error("mockError"); + const next = jest.fn().mockRejectedValue(mockError); + + const retryStrategy = new StandardRetryStrategy(maxAttempts); + try { + await retryStrategy.retry(next, {} as any); + } catch (error) { + expect(error).toStrictEqual(mockError); + return error; + } + }; + + const mockSuccessAfterOneFail = ( + maxAttempts: number, + options?: { mockError?: Error; mockResponse?: string } + ) => { + const mockError = options?.mockError ?? new Error("mockError"); + const mockResponse = { + response: options?.mockResponse, + output: { $metadata: {} } + }; + + const next = jest + .fn() + .mockRejectedValueOnce(mockError) + .mockResolvedValueOnce(mockResponse); + + const retryStrategy = new StandardRetryStrategy(maxAttempts); + return retryStrategy.retry(next, {} as any); + }; + + const mockSuccessAfterTwoFails = ( + maxAttempts: number, + options?: { mockError?: Error; mockResponse?: string } + ) => { + const mockError = options?.mockError ?? new Error("mockError"); + const mockResponse = { + response: options?.mockResponse, + output: { $metadata: {} } + }; + + const next = jest + .fn() + .mockRejectedValueOnce(mockError) + .mockRejectedValueOnce(mockError) + .mockResolvedValueOnce(mockResponse); + + const retryStrategy = new StandardRetryStrategy(maxAttempts); + return retryStrategy.retry(next, {} as any); + }; afterEach(() => { jest.clearAllMocks(); @@ -33,7 +112,7 @@ describe("defaultStrategy", () => { }); }); - describe("retryDecider", () => { + describe("retryDecider init", () => { it("sets defaultRetryDecider if options is undefined", () => { const retryStrategy = new StandardRetryStrategy(maxAttempts); expect(retryStrategy["retryDecider"]).toBe(defaultRetryDecider); @@ -53,7 +132,7 @@ describe("defaultStrategy", () => { }); }); - describe("delayDecider", () => { + describe("delayDecider init", () => { it("sets defaultDelayDecider if options is undefined", () => { const retryStrategy = new StandardRetryStrategy(maxAttempts); expect(retryStrategy["delayDecider"]).toBe(defaultDelayDecider); @@ -73,58 +152,215 @@ describe("defaultStrategy", () => { }); }); - describe("delayBase passed to delayDecider", () => { - const testDelayBasePassed = async ( - delayBaseToTest: number, - mockThrottlingError: boolean - ) => { - (isThrottlingError as jest.Mock).mockReturnValueOnce(mockThrottlingError); - - const mockError = new Error("mockError"); - const mockResponse = { - response: "mockResponse", - output: { $metadata: {} } + describe("retryQuota init", () => { + it("sets getDefaultRetryQuota if options is undefined", () => { + const retryStrategy = new StandardRetryStrategy(maxAttempts); + expect(retryStrategy["retryQuota"]).toBe( + getDefaultRetryQuota(INITIAL_RETRY_TOKENS) + ); + }); + + it("sets getDefaultRetryQuota if options.delayDecider undefined", () => { + const retryStrategy = new StandardRetryStrategy(maxAttempts, {}); + expect(retryStrategy["retryQuota"]).toBe( + getDefaultRetryQuota(INITIAL_RETRY_TOKENS) + ); + }); + + it("sets options.retryQuota if defined", () => { + const retryQuota = {} as RetryQuota; + const retryStrategy = new StandardRetryStrategy(maxAttempts, { + retryQuota + }); + expect(retryStrategy["retryQuota"]).toBe(retryQuota); + }); + }); + + describe("delayDecider", () => { + describe("delayBase value passed", () => { + const testDelayBasePassed = async ( + delayBaseToTest: number, + mockThrottlingError: boolean + ) => { + (isThrottlingError as jest.Mock).mockReturnValueOnce( + mockThrottlingError + ); + + const mockError = new Error(); + await mockSuccessAfterOneFail(maxAttempts, { mockError }); + + expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(1); + expect(isThrottlingError as jest.Mock).toHaveBeenCalledWith(mockError); + expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(1); + expect((defaultDelayDecider as jest.Mock).mock.calls[0][0]).toBe( + delayBaseToTest + ); }; - const next = jest - .fn() - .mockRejectedValueOnce(mockError) - .mockResolvedValueOnce(mockResponse); + it("should be equal to THROTTLING_RETRY_DELAY_BASE if error is throttling error", async () => { + return testDelayBasePassed(THROTTLING_RETRY_DELAY_BASE, true); + }); - const retryStrategy = new StandardRetryStrategy(maxAttempts); - await retryStrategy.retry(next, {} as any); + it("should be equal to DEFAULT_RETRY_DELAY_BASE in error is not a throttling error", async () => { + return testDelayBasePassed(DEFAULT_RETRY_DELAY_BASE, false); + }); + }); + + describe("attempts value passed", () => { + it("on successful operation", async () => { + await mockSuccessfulOperation(maxAttempts); + expect(defaultDelayDecider as jest.Mock).not.toHaveBeenCalled(); + }); + + it("in case of single failure", async () => { + await mockSuccessAfterOneFail(maxAttempts); + expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(1); + expect((defaultDelayDecider as jest.Mock).mock.calls[0][1]).toBe(1); + }); - expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(1); - expect(isThrottlingError as jest.Mock).toHaveBeenCalledWith(mockError); - expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(1); - expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledWith( - delayBaseToTest, - 1 + it("on all fails", async () => { + await mockFailedOperation(maxAttempts); + expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(2); + expect((defaultDelayDecider as jest.Mock).mock.calls[0][1]).toBe(1); + expect((defaultDelayDecider as jest.Mock).mock.calls[1][1]).toBe(2); + }); + }); + + it("delay value returned", async () => { + jest.spyOn(global, "setTimeout"); + + const FIRST_DELAY = 100; + const SECOND_DELAY = 200; + + (defaultDelayDecider as jest.Mock) + .mockReturnValueOnce(FIRST_DELAY) + .mockReturnValueOnce(SECOND_DELAY); + + const maxAttempts = 3; + const error = await mockFailedOperation(maxAttempts); + expect(error.$metadata.totalRetryDelay).toEqual( + FIRST_DELAY + SECOND_DELAY ); - }; - it("should be equal to THROTTLING_RETRY_DELAY_BASE if error is throttling error", async () => { - return testDelayBasePassed(THROTTLING_RETRY_DELAY_BASE, true); + expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes( + maxAttempts - 1 + ); + expect(setTimeout).toHaveBeenCalledTimes(maxAttempts - 1); + expect(((setTimeout as unknown) as jest.Mock).mock.calls[0][1]).toBe( + FIRST_DELAY + ); + expect(((setTimeout as unknown) as jest.Mock).mock.calls[1][1]).toBe( + SECOND_DELAY + ); }); + }); + + describe("retryQuota", () => { + describe("hasRetryTokens", () => { + it("not called on successful operation", async () => { + const { hasRetryTokens } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS); + await mockSuccessfulOperation(maxAttempts); + expect(hasRetryTokens).not.toHaveBeenCalled(); + }); - it("should be equal to DEFAULT_RETRY_DELAY_BASE in error is not a throttling error", async () => { - return testDelayBasePassed(DEFAULT_RETRY_DELAY_BASE, false); + it("called once in case of single failure", async () => { + const { hasRetryTokens } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS); + await mockSuccessAfterOneFail(maxAttempts); + expect(hasRetryTokens).toHaveBeenCalledTimes(1); + }); + + it("called once on each retry request", async () => { + const { hasRetryTokens } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS); + await mockFailedOperation(maxAttempts); + expect(hasRetryTokens).toHaveBeenCalledTimes(maxAttempts - 1); + }); + }); + + describe("releaseRetryTokens", () => { + it("called once without param on successful operation", async () => { + const { releaseRetryTokens } = getDefaultRetryQuota( + INITIAL_RETRY_TOKENS + ); + await mockSuccessfulOperation(maxAttempts); + expect(releaseRetryTokens).toHaveBeenCalledTimes(1); + expect(releaseRetryTokens).toHaveBeenCalledWith(undefined); + }); + + it("called once with retryTokenAmount in case of single failure", async () => { + const retryTokens = 15; + const { + releaseRetryTokens, + retrieveRetryTokens + } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS); + (retrieveRetryTokens as jest.Mock).mockReturnValueOnce(retryTokens); + + await mockSuccessAfterOneFail(maxAttempts); + expect(releaseRetryTokens).toHaveBeenCalledTimes(1); + expect(releaseRetryTokens).toHaveBeenCalledWith(retryTokens); + }); + + it("called once with second retryTokenAmount in case of two failures", async () => { + const retryTokensFirst = 15; + const retryTokensSecond = 30; + + const { + releaseRetryTokens, + retrieveRetryTokens + } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS); + + (retrieveRetryTokens as jest.Mock) + .mockReturnValueOnce(retryTokensFirst) + .mockReturnValueOnce(retryTokensSecond); + + await mockSuccessAfterTwoFails(maxAttempts); + expect(releaseRetryTokens).toHaveBeenCalledTimes(1); + expect(releaseRetryTokens).toHaveBeenCalledWith(retryTokensSecond); + }); + + it("not called on unsuccessful operation", async () => { + const { releaseRetryTokens } = getDefaultRetryQuota( + INITIAL_RETRY_TOKENS + ); + await mockFailedOperation(maxAttempts); + expect(releaseRetryTokens).not.toHaveBeenCalled(); + }); + }); + + describe("retrieveRetryTokens", () => { + it("not called on successful operation", async () => { + const { retrieveRetryTokens } = getDefaultRetryQuota( + INITIAL_RETRY_TOKENS + ); + await mockSuccessfulOperation(maxAttempts); + expect(retrieveRetryTokens).not.toHaveBeenCalled(); + }); + + it("called once in case of single failure", async () => { + const { retrieveRetryTokens } = getDefaultRetryQuota( + INITIAL_RETRY_TOKENS + ); + await mockSuccessAfterOneFail(maxAttempts); + expect(retrieveRetryTokens).toHaveBeenCalledTimes(1); + }); + + it("called once on each retry request", async () => { + const { retrieveRetryTokens } = getDefaultRetryQuota( + INITIAL_RETRY_TOKENS + ); + await mockFailedOperation(maxAttempts); + expect(retrieveRetryTokens).toHaveBeenCalledTimes(maxAttempts - 1); + }); }); }); describe("should not retry", () => { it("when the handler completes successfully", async () => { - const mockResponse = { - response: "mockResponse", - output: { $metadata: {} } - }; - - const next = jest.fn().mockResolvedValueOnce(mockResponse); - - const retryStrategy = new StandardRetryStrategy(maxAttempts); - const { response, output } = await retryStrategy.retry(next, {} as any); + const mockResponse = "mockResponse"; + const { response, output } = await mockSuccessfulOperation(maxAttempts, { + mockResponse + }); - expect(response).toStrictEqual(mockResponse.response); + expect(response).toStrictEqual(mockResponse); expect(output.$metadata.attempts).toBe(1); expect(output.$metadata.totalRetryDelay).toBe(0); expect(defaultRetryDecider as jest.Mock).not.toHaveBeenCalled(); @@ -133,71 +369,58 @@ describe("defaultStrategy", () => { it("when retryDecider returns false", async () => { (defaultRetryDecider as jest.Mock).mockReturnValueOnce(false); - - const mockError = new Error("mockError"); - const next = jest.fn().mockRejectedValueOnce(mockError); - - const retryStrategy = new StandardRetryStrategy(maxAttempts); - try { - await retryStrategy.retry(next, {} as any); - } catch (error) { - expect(error).toStrictEqual(mockError); - } - + const mockError = new Error(); + await mockFailedOperation(maxAttempts, { mockError }); expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledTimes(1); expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledWith(mockError); }); - it("when the the maximum number of attempts is reached", async () => { - const mockError = new Error("mockError"); - const next = jest.fn().mockRejectedValue(mockError); - - const retryStrategy = new StandardRetryStrategy(maxAttempts); - try { - await retryStrategy.retry(next, {} as any); - } catch (error) { - expect(error).toStrictEqual(mockError); - } + it("when the maximum number of attempts is reached", async () => { + await mockFailedOperation(maxAttempts); expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledTimes( maxAttempts - 1 ); }); - }); - - it("should delay equal to the value returned by delayDecider", async () => { - jest.spyOn(global, "setTimeout"); - - const FIRST_DELAY = 100; - const SECOND_DELAY = 200; - - (defaultDelayDecider as jest.Mock) - .mockReturnValueOnce(FIRST_DELAY) - .mockReturnValueOnce(SECOND_DELAY); - const mockError = new Error("mockError"); - const next = jest.fn().mockRejectedValue(mockError); - - const retryStrategy = new StandardRetryStrategy(3); - try { - await retryStrategy.retry(next, {} as any); - } catch (error) { - expect(error).toStrictEqual(mockError); - expect(error.$metadata.totalRetryDelay).toEqual( - FIRST_DELAY + SECOND_DELAY - ); - } + describe("when retryQuota.hasRetryTokens returns false", () => { + it("in the first request", async () => { + const { + hasRetryTokens, + retrieveRetryTokens, + releaseRetryTokens + } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS); + (hasRetryTokens as jest.Mock).mockReturnValueOnce(false); + + const mockError = new Error(); + await mockFailedOperation(maxAttempts, { mockError }); + + expect(hasRetryTokens).toHaveBeenCalledTimes(1); + expect(hasRetryTokens).toHaveBeenCalledWith(mockError); + expect(retrieveRetryTokens).not.toHaveBeenCalled(); + expect(releaseRetryTokens).not.toHaveBeenCalled(); + }); - expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(2); - expect(setTimeout).toHaveBeenCalledTimes(2); - expect(setTimeout).toHaveBeenNthCalledWith( - 1, - expect.any(Function), - FIRST_DELAY - ); - expect(setTimeout).toHaveBeenNthCalledWith( - 2, - expect.any(Function), - SECOND_DELAY - ); + it("after the first retry", async () => { + const { + hasRetryTokens, + retrieveRetryTokens, + releaseRetryTokens + } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS); + (hasRetryTokens as jest.Mock) + .mockReturnValueOnce(true) + .mockReturnValueOnce(false); + + const mockError = new Error(); + await mockFailedOperation(maxAttempts, { mockError }); + + expect(hasRetryTokens).toHaveBeenCalledTimes(2); + [1, 2].forEach(n => { + expect(hasRetryTokens).toHaveBeenNthCalledWith(n, mockError); + }); + expect(retrieveRetryTokens).toHaveBeenCalledTimes(1); + expect(retrieveRetryTokens).toHaveBeenCalledWith(mockError); + expect(releaseRetryTokens).not.toHaveBeenCalled(); + }); + }); }); }); diff --git a/packages/middleware-retry/src/defaultStrategy.ts b/packages/middleware-retry/src/defaultStrategy.ts index c2b5d798ed8c..cb8c242d13a9 100644 --- a/packages/middleware-retry/src/defaultStrategy.ts +++ b/packages/middleware-retry/src/defaultStrategy.ts @@ -1,6 +1,7 @@ import { DEFAULT_RETRY_DELAY_BASE, - THROTTLING_RETRY_DELAY_BASE + THROTTLING_RETRY_DELAY_BASE, + INITIAL_RETRY_TOKENS } from "./constants"; import { defaultDelayDecider } from "./delayDecider"; import { defaultRetryDecider } from "./retryDecider"; @@ -12,6 +13,7 @@ import { FinalizeHandlerArguments, RetryStrategy } from "@aws-sdk/types"; +import { getDefaultRetryQuota } from "./defaultRetryQuota"; /** * Determines whether an error is retryable based on the number of retries @@ -33,17 +35,40 @@ export interface DelayDecider { (delayBase: number, attempts: number): number; } +/** + * Interface that specifies the retry quota behavior. + */ +export interface RetryQuota { + /** + * returns true if retry tokens are available from the retry quota bucket. + */ + hasRetryTokens: (error: SdkError) => boolean; + + /** + * returns token amount from the retry quota bucket. + * throws error is retry tokens are not available. + */ + retrieveRetryTokens: (error: SdkError) => number; + + /** + * releases tokens back to the retry quota. + */ + releaseRetryTokens: (releaseCapacityAmount?: number) => void; +} + /** * Strategy options to be passed to StandardRetryStrategy */ export interface StandardRetryStrategyOptions { retryDecider?: RetryDecider; delayDecider?: DelayDecider; + retryQuota?: RetryQuota; } export class StandardRetryStrategy implements RetryStrategy { private retryDecider: RetryDecider; private delayDecider: DelayDecider; + private retryQuota: RetryQuota; constructor( public readonly maxAttempts: number, @@ -51,21 +76,30 @@ export class StandardRetryStrategy implements RetryStrategy { ) { this.retryDecider = options?.retryDecider ?? defaultRetryDecider; this.delayDecider = options?.delayDecider ?? defaultDelayDecider; + this.retryQuota = + options?.retryQuota ?? getDefaultRetryQuota(INITIAL_RETRY_TOKENS); } private shouldRetry(error: SdkError, attempts: number) { - return attempts < this.maxAttempts && this.retryDecider(error); + return ( + attempts < this.maxAttempts && + this.retryDecider(error) && + this.retryQuota.hasRetryTokens(error) + ); } async retry( next: FinalizeHandler, args: FinalizeHandlerArguments ) { + let retryTokenAmount; let attempts = 0; let totalDelay = 0; while (true) { try { const { response, output } = await next(args); + + this.retryQuota.releaseRetryTokens(retryTokenAmount); output.$metadata.attempts = attempts + 1; output.$metadata.totalRetryDelay = totalDelay; @@ -73,6 +107,7 @@ export class StandardRetryStrategy implements RetryStrategy { } catch (err) { attempts++; if (this.shouldRetry(err as SdkError, attempts)) { + retryTokenAmount = this.retryQuota.retrieveRetryTokens(err); const delay = this.delayDecider( isThrottlingError(err) ? THROTTLING_RETRY_DELAY_BASE