Skip to content

Commit 0c9a143

Browse files
trivikrtrivikr
authored
feat(middleware-retry): add Adaptive Retry Strategy (#2454)
1 parent fc0a5da commit 0c9a143

8 files changed

+179
-21
lines changed

‎packages/middleware-retry/src/AdaptiveRetryStrategy.spec.ts

+97
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
import { AdaptiveRetryStrategy } from "./AdaptiveRetryStrategy";
2+
import { RETRY_MODES } from "./config";
3+
import { DefaultRateLimiter } from "./DefaultRateLimiter";
4+
import { StandardRetryStrategy } from "./StandardRetryStrategy";
5+
import { RateLimiter, RetryQuota } from "./types";
6+
7+
jest.mock("./StandardRetryStrategy");
8+
jest.mock("./DefaultRateLimiter");
9+
10+
describe(AdaptiveRetryStrategy.name, () => {
11+
const maxAttemptsProvider = jest.fn();
12+
const mockDefaultRateLimiter = {
13+
getSendToken: jest.fn(),
14+
updateClientSendingRate: jest.fn(),
15+
};
16+
17+
beforeEach(() => {
18+
(DefaultRateLimiter as jest.Mock).mockReturnValue(mockDefaultRateLimiter);
19+
});
20+
21+
afterEach(() => {
22+
jest.clearAllMocks();
23+
});
24+
25+
describe("constructor", () => {
26+
it("calls super constructor", () => {
27+
const retryDecider = jest.fn();
28+
const delayDecider = jest.fn();
29+
const retryQuota = {} as RetryQuota;
30+
const rateLimiter = {} as RateLimiter;
31+
32+
new AdaptiveRetryStrategy(maxAttemptsProvider, {
33+
retryDecider,
34+
delayDecider,
35+
retryQuota,
36+
rateLimiter,
37+
});
38+
expect(StandardRetryStrategy).toHaveBeenCalledWith(maxAttemptsProvider, {
39+
retryDecider,
40+
delayDecider,
41+
retryQuota,
42+
});
43+
});
44+
45+
it(`sets mode=${RETRY_MODES.ADAPTIVE}`, () => {
46+
const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider);
47+
expect(retryStrategy.mode).toStrictEqual(RETRY_MODES.ADAPTIVE);
48+
});
49+
50+
describe("rateLimiter init", () => {
51+
it("sets getDefaultrateLimiter if options is undefined", () => {
52+
const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider);
53+
expect(retryStrategy["rateLimiter"]).toBe(mockDefaultRateLimiter);
54+
});
55+
56+
it("sets getDefaultrateLimiter if options.delayDecider undefined", () => {
57+
const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider, {});
58+
expect(retryStrategy["rateLimiter"]).toBe(mockDefaultRateLimiter);
59+
});
60+
61+
it("sets options.rateLimiter if defined", () => {
62+
const rateLimiter = {} as RateLimiter;
63+
const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider, {
64+
rateLimiter,
65+
});
66+
expect(retryStrategy["rateLimiter"]).toBe(rateLimiter);
67+
});
68+
});
69+
});
70+
71+
describe("retry", () => {
72+
const mockedSuperRetry = jest.spyOn(StandardRetryStrategy.prototype, "retry");
73+
74+
beforeEach(async () => {
75+
const next = jest.fn();
76+
const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider);
77+
await retryStrategy.retry(next, { request: { headers: {} } } as any);
78+
expect(mockedSuperRetry).toHaveBeenCalledTimes(1);
79+
});
80+
81+
afterEach(() => {
82+
jest.clearAllMocks();
83+
});
84+
85+
it("calls rateLimiter.getSendToken in beforeRequest", async () => {
86+
expect(mockDefaultRateLimiter.getSendToken).toHaveBeenCalledTimes(0);
87+
await mockedSuperRetry.mock.calls[0][2].beforeRequest();
88+
expect(mockDefaultRateLimiter.getSendToken).toHaveBeenCalledTimes(1);
89+
});
90+
91+
it("calls rateLimiter.updateClientSendingRate in afterRequest", async () => {
92+
expect(mockDefaultRateLimiter.updateClientSendingRate).toHaveBeenCalledTimes(0);
93+
await mockedSuperRetry.mock.calls[0][2].afterRequest();
94+
expect(mockDefaultRateLimiter.updateClientSendingRate).toHaveBeenCalledTimes(1);
95+
});
96+
});
97+
});

‎packages/middleware-retry/src/AdaptiveRetryStrategy.ts

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
import { FinalizeHandler, FinalizeHandlerArguments, MetadataBearer, Provider } from "@aws-sdk/types";
2+
3+
import { RETRY_MODES } from "./config";
4+
import { DefaultRateLimiter } from "./DefaultRateLimiter";
5+
import { StandardRetryStrategy, StandardRetryStrategyOptions } from "./StandardRetryStrategy";
6+
import { RateLimiter } from "./types";
7+
8+
/**
9+
* Strategy options to be passed to AdaptiveRetryStrategy
10+
*/
11+
export interface AdaptiveRetryStrategyOptions extends StandardRetryStrategyOptions {
12+
rateLimiter?: RateLimiter;
13+
}
14+
15+
export class AdaptiveRetryStrategy extends StandardRetryStrategy {
16+
private rateLimiter: RateLimiter;
17+
18+
constructor(maxAttemptsProvider: Provider<number>, options?: AdaptiveRetryStrategyOptions) {
19+
const { rateLimiter, ...superOptions } = options ?? {};
20+
super(maxAttemptsProvider, superOptions);
21+
this.rateLimiter = rateLimiter ?? new DefaultRateLimiter();
22+
this.mode = RETRY_MODES.ADAPTIVE;
23+
}
24+
25+
async retry<Input extends object, Ouput extends MetadataBearer>(
26+
next: FinalizeHandler<Input, Ouput>,
27+
args: FinalizeHandlerArguments<Input>
28+
) {
29+
return super.retry(next, args, {
30+
beforeRequest: async () => {
31+
return this.rateLimiter.getSendToken();
32+
},
33+
afterRequest: (response: any) => {
34+
this.rateLimiter.updateClientSendingRate(response);
35+
},
36+
});
37+
}
38+
}

‎packages/middleware-retry/src/defaultStrategy.spec.ts renamed to ‎packages/middleware-retry/src/StandardRetryStrategy.spec.ts

+7-2
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,12 @@ import { HttpRequest } from "@aws-sdk/protocol-http";
22
import { isThrottlingError } from "@aws-sdk/service-error-classification";
33
import { v4 } from "uuid";
44

5-
import { DEFAULT_MAX_ATTEMPTS } from "./configurations";
5+
import { DEFAULT_MAX_ATTEMPTS, RETRY_MODES } from "./config";
66
import { DEFAULT_RETRY_DELAY_BASE, INITIAL_RETRY_TOKENS, THROTTLING_RETRY_DELAY_BASE } from "./constants";
77
import { getDefaultRetryQuota } from "./defaultRetryQuota";
8-
import { StandardRetryStrategy } from "./defaultStrategy";
98
import { defaultDelayDecider } from "./delayDecider";
109
import { defaultRetryDecider } from "./retryDecider";
10+
import { StandardRetryStrategy } from "./StandardRetryStrategy";
1111
import { RetryQuota } from "./types";
1212

1313
jest.mock("@aws-sdk/service-error-classification");
@@ -102,6 +102,11 @@ describe("defaultStrategy", () => {
102102
});
103103
});
104104

105+
it(`sets mode=${RETRY_MODES.STANDARD}`, () => {
106+
const retryStrategy = new StandardRetryStrategy(() => Promise.resolve(maxAttempts));
107+
expect(retryStrategy.mode).toStrictEqual(RETRY_MODES.STANDARD);
108+
});
109+
105110
it("handles non-standard errors", () => {
106111
const nonStandardErrors = [undefined, "foo", { foo: "bar" }, 123, false, null];
107112
const maxAttempts = 1;

‎packages/middleware-retry/src/defaultStrategy.ts renamed to ‎packages/middleware-retry/src/StandardRetryStrategy.ts

+14-3
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import { SdkError } from "@aws-sdk/smithy-client";
44
import { FinalizeHandler, FinalizeHandlerArguments, MetadataBearer, Provider, RetryStrategy } from "@aws-sdk/types";
55
import { v4 } from "uuid";
66

7-
import { DEFAULT_MAX_ATTEMPTS, DEFAULT_RETRY_MODE } from "./configurations";
7+
import { DEFAULT_MAX_ATTEMPTS, RETRY_MODES } from "./config";
88
import {
99
DEFAULT_RETRY_DELAY_BASE,
1010
INITIAL_RETRY_TOKENS,
@@ -30,7 +30,7 @@ export class StandardRetryStrategy implements RetryStrategy {
3030
private retryDecider: RetryDecider;
3131
private delayDecider: DelayDecider;
3232
private retryQuota: RetryQuota;
33-
public readonly mode = DEFAULT_RETRY_MODE;
33+
public mode: string = RETRY_MODES.STANDARD;
3434

3535
constructor(private readonly maxAttemptsProvider: Provider<number>, options?: StandardRetryStrategyOptions) {
3636
this.retryDecider = options?.retryDecider ?? defaultRetryDecider;
@@ -54,7 +54,11 @@ export class StandardRetryStrategy implements RetryStrategy {
5454

5555
async retry<Input extends object, Ouput extends MetadataBearer>(
5656
next: FinalizeHandler<Input, Ouput>,
57-
args: FinalizeHandlerArguments<Input>
57+
args: FinalizeHandlerArguments<Input>,
58+
options?: {
59+
beforeRequest: Function;
60+
afterRequest: Function;
61+
}
5862
) {
5963
let retryTokenAmount;
6064
let attempts = 0;
@@ -72,7 +76,14 @@ export class StandardRetryStrategy implements RetryStrategy {
7276
if (HttpRequest.isInstance(request)) {
7377
request.headers[REQUEST_HEADER] = `attempt=${attempts + 1}; max=${maxAttempts}`;
7478
}
79+
80+
if (options?.beforeRequest) {
81+
await options.beforeRequest();
82+
}
7583
const { response, output } = await next(args);
84+
if (options?.afterRequest) {
85+
options.afterRequest(response);
86+
}
7687

7788
this.retryQuota.releaseRetryTokens(retryTokenAmount);
7889
output.$metadata.attempts = attempts + 1;

‎packages/middleware-retry/src/config.ts

+15
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
export enum RETRY_MODES {
2+
STANDARD = "standard",
3+
ADAPTIVE = "adaptive",
4+
}
5+
6+
/**
7+
* The default value for how many HTTP requests an SDK should make for a
8+
* single SDK operation invocation before giving up
9+
*/
10+
export const DEFAULT_MAX_ATTEMPTS = 3;
11+
12+
/**
13+
* The default retry algorithm to use.
14+
*/
15+
export const DEFAULT_RETRY_MODE = RETRY_MODES.STANDARD;

‎packages/middleware-retry/src/configurations.spec.ts

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
1+
import { DEFAULT_MAX_ATTEMPTS } from "./config";
12
import {
23
CONFIG_MAX_ATTEMPTS,
3-
DEFAULT_MAX_ATTEMPTS,
44
ENV_MAX_ATTEMPTS,
55
NODE_MAX_ATTEMPT_CONFIG_OPTIONS,
66
resolveRetryConfig,
77
} from "./configurations";
8-
import { StandardRetryStrategy } from "./defaultStrategy";
8+
import { StandardRetryStrategy } from "./StandardRetryStrategy";
99

10-
jest.mock("./defaultStrategy");
10+
jest.mock("./StandardRetryStrategy");
1111

1212
describe("resolveRetryConfig", () => {
1313
afterEach(() => {

‎packages/middleware-retry/src/configurations.ts

+2-12
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,12 @@
11
import { LoadedConfigSelectors } from "@aws-sdk/node-config-provider";
22
import { Provider, RetryStrategy } from "@aws-sdk/types";
33

4-
import { StandardRetryStrategy } from "./defaultStrategy";
4+
import { DEFAULT_MAX_ATTEMPTS, DEFAULT_RETRY_MODE } from "./config";
5+
import { StandardRetryStrategy } from "./StandardRetryStrategy";
56

67
export const ENV_MAX_ATTEMPTS = "AWS_MAX_ATTEMPTS";
78
export const CONFIG_MAX_ATTEMPTS = "max_attempts";
89

9-
/**
10-
* The default value for how many HTTP requests an SDK should make for a
11-
* single SDK operation invocation before giving up
12-
*/
13-
export const DEFAULT_MAX_ATTEMPTS = 3;
14-
15-
/**
16-
* The default retry algorithm to use.
17-
*/
18-
export const DEFAULT_RETRY_MODE = "standard";
19-
2010
export const NODE_MAX_ATTEMPT_CONFIG_OPTIONS: LoadedConfigSelectors<number> = {
2111
environmentVariableSelector: (env) => {
2212
const value = env[ENV_MAX_ATTEMPTS];

‎packages/middleware-retry/src/index.ts

+3-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
export * from "./retryMiddleware";
22
export * from "./omitRetryHeadersMiddleware";
3-
export * from "./defaultStrategy";
3+
export * from "./StandardRetryStrategy";
4+
export * from "./AdaptiveRetryStrategy";
5+
export * from "./config";
46
export * from "./configurations";
57
export * from "./delayDecider";
68
export * from "./retryDecider";

0 commit comments

Comments
 (0)