diff --git a/packages/middleware-retry/src/AdaptiveRetryStrategy.spec.ts b/packages/middleware-retry/src/AdaptiveRetryStrategy.spec.ts new file mode 100644 index 000000000000..27eee8021ffa --- /dev/null +++ b/packages/middleware-retry/src/AdaptiveRetryStrategy.spec.ts @@ -0,0 +1,97 @@ +import { AdaptiveRetryStrategy } from "./AdaptiveRetryStrategy"; +import { RETRY_MODES } from "./config"; +import { DefaultRateLimiter } from "./DefaultRateLimiter"; +import { StandardRetryStrategy } from "./StandardRetryStrategy"; +import { RateLimiter, RetryQuota } from "./types"; + +jest.mock("./StandardRetryStrategy"); +jest.mock("./DefaultRateLimiter"); + +describe(AdaptiveRetryStrategy.name, () => { + const maxAttemptsProvider = jest.fn(); + const mockDefaultRateLimiter = { + getSendToken: jest.fn(), + updateClientSendingRate: jest.fn(), + }; + + beforeEach(() => { + (DefaultRateLimiter as jest.Mock).mockReturnValue(mockDefaultRateLimiter); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + describe("constructor", () => { + it("calls super constructor", () => { + const retryDecider = jest.fn(); + const delayDecider = jest.fn(); + const retryQuota = {} as RetryQuota; + const rateLimiter = {} as RateLimiter; + + new AdaptiveRetryStrategy(maxAttemptsProvider, { + retryDecider, + delayDecider, + retryQuota, + rateLimiter, + }); + expect(StandardRetryStrategy).toHaveBeenCalledWith(maxAttemptsProvider, { + retryDecider, + delayDecider, + retryQuota, + }); + }); + + it(`sets mode=${RETRY_MODES.ADAPTIVE}`, () => { + const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider); + expect(retryStrategy.mode).toStrictEqual(RETRY_MODES.ADAPTIVE); + }); + + describe("rateLimiter init", () => { + it("sets getDefaultrateLimiter if options is undefined", () => { + const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider); + expect(retryStrategy["rateLimiter"]).toBe(mockDefaultRateLimiter); + }); + + it("sets getDefaultrateLimiter if options.delayDecider undefined", () => { + const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider, {}); + expect(retryStrategy["rateLimiter"]).toBe(mockDefaultRateLimiter); + }); + + it("sets options.rateLimiter if defined", () => { + const rateLimiter = {} as RateLimiter; + const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider, { + rateLimiter, + }); + expect(retryStrategy["rateLimiter"]).toBe(rateLimiter); + }); + }); + }); + + describe("retry", () => { + const mockedSuperRetry = jest.spyOn(StandardRetryStrategy.prototype, "retry"); + + beforeEach(async () => { + const next = jest.fn(); + const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider); + await retryStrategy.retry(next, { request: { headers: {} } } as any); + expect(mockedSuperRetry).toHaveBeenCalledTimes(1); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + it("calls rateLimiter.getSendToken in beforeRequest", async () => { + expect(mockDefaultRateLimiter.getSendToken).toHaveBeenCalledTimes(0); + await mockedSuperRetry.mock.calls[0][2].beforeRequest(); + expect(mockDefaultRateLimiter.getSendToken).toHaveBeenCalledTimes(1); + }); + + it("calls rateLimiter.updateClientSendingRate in afterRequest", async () => { + expect(mockDefaultRateLimiter.updateClientSendingRate).toHaveBeenCalledTimes(0); + await mockedSuperRetry.mock.calls[0][2].afterRequest(); + expect(mockDefaultRateLimiter.updateClientSendingRate).toHaveBeenCalledTimes(1); + }); + }); +}); diff --git a/packages/middleware-retry/src/AdaptiveRetryStrategy.ts b/packages/middleware-retry/src/AdaptiveRetryStrategy.ts new file mode 100644 index 000000000000..31423fdce1c3 --- /dev/null +++ b/packages/middleware-retry/src/AdaptiveRetryStrategy.ts @@ -0,0 +1,38 @@ +import { FinalizeHandler, FinalizeHandlerArguments, MetadataBearer, Provider } from "@aws-sdk/types"; + +import { RETRY_MODES } from "./config"; +import { DefaultRateLimiter } from "./DefaultRateLimiter"; +import { StandardRetryStrategy, StandardRetryStrategyOptions } from "./StandardRetryStrategy"; +import { RateLimiter } from "./types"; + +/** + * Strategy options to be passed to AdaptiveRetryStrategy + */ +export interface AdaptiveRetryStrategyOptions extends StandardRetryStrategyOptions { + rateLimiter?: RateLimiter; +} + +export class AdaptiveRetryStrategy extends StandardRetryStrategy { + private rateLimiter: RateLimiter; + + constructor(maxAttemptsProvider: Provider, options?: AdaptiveRetryStrategyOptions) { + const { rateLimiter, ...superOptions } = options ?? {}; + super(maxAttemptsProvider, superOptions); + this.rateLimiter = rateLimiter ?? new DefaultRateLimiter(); + this.mode = RETRY_MODES.ADAPTIVE; + } + + async retry( + next: FinalizeHandler, + args: FinalizeHandlerArguments + ) { + return super.retry(next, args, { + beforeRequest: async () => { + return this.rateLimiter.getSendToken(); + }, + afterRequest: (response: any) => { + this.rateLimiter.updateClientSendingRate(response); + }, + }); + } +} diff --git a/packages/middleware-retry/src/defaultStrategy.spec.ts b/packages/middleware-retry/src/StandardRetryStrategy.spec.ts similarity index 98% rename from packages/middleware-retry/src/defaultStrategy.spec.ts rename to packages/middleware-retry/src/StandardRetryStrategy.spec.ts index 8093b5e7d41a..f5089137dc0c 100644 --- a/packages/middleware-retry/src/defaultStrategy.spec.ts +++ b/packages/middleware-retry/src/StandardRetryStrategy.spec.ts @@ -2,12 +2,12 @@ import { HttpRequest } from "@aws-sdk/protocol-http"; import { isThrottlingError } from "@aws-sdk/service-error-classification"; import { v4 } from "uuid"; -import { DEFAULT_MAX_ATTEMPTS } from "./configurations"; +import { DEFAULT_MAX_ATTEMPTS, RETRY_MODES } from "./config"; import { DEFAULT_RETRY_DELAY_BASE, INITIAL_RETRY_TOKENS, THROTTLING_RETRY_DELAY_BASE } from "./constants"; import { getDefaultRetryQuota } from "./defaultRetryQuota"; -import { StandardRetryStrategy } from "./defaultStrategy"; import { defaultDelayDecider } from "./delayDecider"; import { defaultRetryDecider } from "./retryDecider"; +import { StandardRetryStrategy } from "./StandardRetryStrategy"; import { RetryQuota } from "./types"; jest.mock("@aws-sdk/service-error-classification"); @@ -102,6 +102,11 @@ describe("defaultStrategy", () => { }); }); + it(`sets mode=${RETRY_MODES.STANDARD}`, () => { + const retryStrategy = new StandardRetryStrategy(() => Promise.resolve(maxAttempts)); + expect(retryStrategy.mode).toStrictEqual(RETRY_MODES.STANDARD); + }); + it("handles non-standard errors", () => { const nonStandardErrors = [undefined, "foo", { foo: "bar" }, 123, false, null]; const maxAttempts = 1; diff --git a/packages/middleware-retry/src/defaultStrategy.ts b/packages/middleware-retry/src/StandardRetryStrategy.ts similarity index 89% rename from packages/middleware-retry/src/defaultStrategy.ts rename to packages/middleware-retry/src/StandardRetryStrategy.ts index ab621b0687d7..6d13bafde436 100644 --- a/packages/middleware-retry/src/defaultStrategy.ts +++ b/packages/middleware-retry/src/StandardRetryStrategy.ts @@ -4,7 +4,7 @@ import { SdkError } from "@aws-sdk/smithy-client"; import { FinalizeHandler, FinalizeHandlerArguments, MetadataBearer, Provider, RetryStrategy } from "@aws-sdk/types"; import { v4 } from "uuid"; -import { DEFAULT_MAX_ATTEMPTS, DEFAULT_RETRY_MODE } from "./configurations"; +import { DEFAULT_MAX_ATTEMPTS, RETRY_MODES } from "./config"; import { DEFAULT_RETRY_DELAY_BASE, INITIAL_RETRY_TOKENS, @@ -30,7 +30,7 @@ export class StandardRetryStrategy implements RetryStrategy { private retryDecider: RetryDecider; private delayDecider: DelayDecider; private retryQuota: RetryQuota; - public readonly mode = DEFAULT_RETRY_MODE; + public mode: string = RETRY_MODES.STANDARD; constructor(private readonly maxAttemptsProvider: Provider, options?: StandardRetryStrategyOptions) { this.retryDecider = options?.retryDecider ?? defaultRetryDecider; @@ -54,7 +54,11 @@ export class StandardRetryStrategy implements RetryStrategy { async retry( next: FinalizeHandler, - args: FinalizeHandlerArguments + args: FinalizeHandlerArguments, + options?: { + beforeRequest: Function; + afterRequest: Function; + } ) { let retryTokenAmount; let attempts = 0; @@ -72,7 +76,14 @@ export class StandardRetryStrategy implements RetryStrategy { if (HttpRequest.isInstance(request)) { request.headers[REQUEST_HEADER] = `attempt=${attempts + 1}; max=${maxAttempts}`; } + + if (options?.beforeRequest) { + await options.beforeRequest(); + } const { response, output } = await next(args); + if (options?.afterRequest) { + options.afterRequest(response); + } this.retryQuota.releaseRetryTokens(retryTokenAmount); output.$metadata.attempts = attempts + 1; diff --git a/packages/middleware-retry/src/config.ts b/packages/middleware-retry/src/config.ts new file mode 100644 index 000000000000..96e4f9169cfb --- /dev/null +++ b/packages/middleware-retry/src/config.ts @@ -0,0 +1,15 @@ +export enum RETRY_MODES { + STANDARD = "standard", + ADAPTIVE = "adaptive", +} + +/** + * The default value for how many HTTP requests an SDK should make for a + * single SDK operation invocation before giving up + */ +export const DEFAULT_MAX_ATTEMPTS = 3; + +/** + * The default retry algorithm to use. + */ +export const DEFAULT_RETRY_MODE = RETRY_MODES.STANDARD; diff --git a/packages/middleware-retry/src/configurations.spec.ts b/packages/middleware-retry/src/configurations.spec.ts index cb79a0491375..0f0a7a31d0d2 100644 --- a/packages/middleware-retry/src/configurations.spec.ts +++ b/packages/middleware-retry/src/configurations.spec.ts @@ -1,13 +1,13 @@ +import { DEFAULT_MAX_ATTEMPTS } from "./config"; import { CONFIG_MAX_ATTEMPTS, - DEFAULT_MAX_ATTEMPTS, ENV_MAX_ATTEMPTS, NODE_MAX_ATTEMPT_CONFIG_OPTIONS, resolveRetryConfig, } from "./configurations"; -import { StandardRetryStrategy } from "./defaultStrategy"; +import { StandardRetryStrategy } from "./StandardRetryStrategy"; -jest.mock("./defaultStrategy"); +jest.mock("./StandardRetryStrategy"); describe("resolveRetryConfig", () => { afterEach(() => { diff --git a/packages/middleware-retry/src/configurations.ts b/packages/middleware-retry/src/configurations.ts index ff9b97cbf8e4..3fb9a8f5f9a5 100644 --- a/packages/middleware-retry/src/configurations.ts +++ b/packages/middleware-retry/src/configurations.ts @@ -1,22 +1,12 @@ import { LoadedConfigSelectors } from "@aws-sdk/node-config-provider"; import { Provider, RetryStrategy } from "@aws-sdk/types"; -import { StandardRetryStrategy } from "./defaultStrategy"; +import { DEFAULT_MAX_ATTEMPTS, DEFAULT_RETRY_MODE } from "./config"; +import { StandardRetryStrategy } from "./StandardRetryStrategy"; export const ENV_MAX_ATTEMPTS = "AWS_MAX_ATTEMPTS"; export const CONFIG_MAX_ATTEMPTS = "max_attempts"; -/** - * The default value for how many HTTP requests an SDK should make for a - * single SDK operation invocation before giving up - */ -export const DEFAULT_MAX_ATTEMPTS = 3; - -/** - * The default retry algorithm to use. - */ -export const DEFAULT_RETRY_MODE = "standard"; - export const NODE_MAX_ATTEMPT_CONFIG_OPTIONS: LoadedConfigSelectors = { environmentVariableSelector: (env) => { const value = env[ENV_MAX_ATTEMPTS]; diff --git a/packages/middleware-retry/src/index.ts b/packages/middleware-retry/src/index.ts index 6013d7d0e74c..45562797a295 100644 --- a/packages/middleware-retry/src/index.ts +++ b/packages/middleware-retry/src/index.ts @@ -1,6 +1,8 @@ export * from "./retryMiddleware"; export * from "./omitRetryHeadersMiddleware"; -export * from "./defaultStrategy"; +export * from "./StandardRetryStrategy"; +export * from "./AdaptiveRetryStrategy"; +export * from "./config"; export * from "./configurations"; export * from "./delayDecider"; export * from "./retryDecider";