Skip to content

Commit fb44543

Browse files
trivikrtrivikr
committed
chore: modified getDefaultRetryQuota to accept initialRetryTokens
1 parent 01ad868 commit fb44543

File tree

4 files changed

+92
-34
lines changed

4 files changed

+92
-34
lines changed

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

Lines changed: 54 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -14,27 +14,63 @@ describe("defaultRetryQuota", () => {
1414
name: "TimeoutError"
1515
}) as SdkError;
1616

17-
const getDrainedRetryQuota = (targetCapacity: number, error: SdkError) => {
18-
const retryQuota = getDefaultRetryQuota();
19-
let availableCapacity = INITIAL_RETRY_TOKENS;
17+
const getDrainedRetryQuota = (
18+
targetCapacity: number,
19+
error: SdkError,
20+
initialRetryTokens: number = INITIAL_RETRY_TOKENS
21+
) => {
22+
const retryQuota = getDefaultRetryQuota(initialRetryTokens);
23+
let availableCapacity = initialRetryTokens;
2024
while (availableCapacity >= targetCapacity) {
2125
retryQuota.retrieveRetryTokens(error);
2226
availableCapacity -= targetCapacity;
2327
}
2428
return retryQuota;
2529
};
2630

31+
describe("custom initial retry tokens", () => {
32+
it("hasRetryTokens returns false if capacity is not available", () => {
33+
const customRetryTokens = 100;
34+
const error = getMockError();
35+
const retryQuota = getDrainedRetryQuota(
36+
RETRY_COST,
37+
error,
38+
customRetryTokens
39+
);
40+
expect(retryQuota.hasRetryTokens(error)).toBe(false);
41+
});
42+
43+
it("retrieveRetryToken throws error if retry tokens not available", () => {
44+
const customRetryTokens = 100;
45+
const error = getMockError();
46+
const retryQuota = getDrainedRetryQuota(
47+
RETRY_COST,
48+
error,
49+
customRetryTokens
50+
);
51+
expect(() => {
52+
retryQuota.retrieveRetryTokens(error);
53+
}).toThrowError(new Error("No retry token available"));
54+
});
55+
});
56+
2757
describe("hasRetryTokens", () => {
2858
describe("returns true if capacity is available", () => {
2959
it("when it's TimeoutError", () => {
3060
const timeoutError = getMockTimeoutError();
31-
expect(getDefaultRetryQuota().hasRetryTokens(timeoutError)).toBe(true);
61+
expect(
62+
getDefaultRetryQuota(INITIAL_RETRY_TOKENS).hasRetryTokens(
63+
timeoutError
64+
)
65+
).toBe(true);
3266
});
3367

3468
it("when it's not TimeoutError", () => {
35-
expect(getDefaultRetryQuota().hasRetryTokens(getMockError())).toBe(
36-
true
37-
);
69+
expect(
70+
getDefaultRetryQuota(INITIAL_RETRY_TOKENS).hasRetryTokens(
71+
getMockError()
72+
)
73+
).toBe(true);
3874
});
3975
});
4076

@@ -60,15 +96,19 @@ describe("defaultRetryQuota", () => {
6096
describe("returns retry tokens amount if available", () => {
6197
it("when it's TimeoutError", () => {
6298
const timeoutError = getMockTimeoutError();
63-
expect(getDefaultRetryQuota().retrieveRetryTokens(timeoutError)).toBe(
64-
TIMEOUT_RETRY_COST
65-
);
99+
expect(
100+
getDefaultRetryQuota(INITIAL_RETRY_TOKENS).retrieveRetryTokens(
101+
timeoutError
102+
)
103+
).toBe(TIMEOUT_RETRY_COST);
66104
});
67105

68106
it("when it's not TimeoutError", () => {
69-
expect(getDefaultRetryQuota().retrieveRetryTokens(getMockError())).toBe(
70-
RETRY_COST
71-
);
107+
expect(
108+
getDefaultRetryQuota(INITIAL_RETRY_TOKENS).retrieveRetryTokens(
109+
getMockError()
110+
)
111+
).toBe(RETRY_COST);
72112
});
73113
});
74114

@@ -128,7 +168,7 @@ describe("defaultRetryQuota", () => {
128168

129169
it("ensures availableCapacity is maxed at INITIAL_RETRY_TOKENS", () => {
130170
const error = getMockError();
131-
const retryQuota = getDefaultRetryQuota();
171+
const retryQuota = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
132172

133173
// release 100 tokens.
134174
[...Array(100).keys()].forEach(key => {

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,16 @@
11
import { RetryQuota } from "./defaultStrategy";
22
import { SdkError } from "@aws-sdk/smithy-client";
33
import {
4-
INITIAL_RETRY_TOKENS,
54
RETRY_COST,
65
TIMEOUT_RETRY_COST,
76
NO_RETRY_INCREMENT
87
} from "./constants";
98

10-
export const getDefaultRetryQuota = (): RetryQuota => {
11-
const MAX_CAPACITY = INITIAL_RETRY_TOKENS;
12-
let availableCapacity = INITIAL_RETRY_TOKENS;
9+
export const getDefaultRetryQuota = (
10+
initialRetryTokens: number
11+
): RetryQuota => {
12+
const MAX_CAPACITY = initialRetryTokens;
13+
let availableCapacity = initialRetryTokens;
1314

1415
const getCapacityAmount = (error: SdkError) =>
1516
error.name === "TimeoutError" ? TIMEOUT_RETRY_COST : RETRY_COST;

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

Lines changed: 29 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import {
22
DEFAULT_RETRY_DELAY_BASE,
3-
THROTTLING_RETRY_DELAY_BASE
3+
THROTTLING_RETRY_DELAY_BASE,
4+
INITIAL_RETRY_TOKENS
45
} from "./constants";
56
import { isThrottlingError } from "@aws-sdk/service-error-classification";
67
import { defaultDelayDecider } from "./delayDecider";
@@ -154,12 +155,16 @@ describe("defaultStrategy", () => {
154155
describe("retryQuota init", () => {
155156
it("sets getDefaultRetryQuota if options is undefined", () => {
156157
const retryStrategy = new StandardRetryStrategy(maxAttempts);
157-
expect(retryStrategy["retryQuota"]).toBe(getDefaultRetryQuota());
158+
expect(retryStrategy["retryQuota"]).toBe(
159+
getDefaultRetryQuota(INITIAL_RETRY_TOKENS)
160+
);
158161
});
159162

160163
it("sets getDefaultRetryQuota if options.delayDecider undefined", () => {
161164
const retryStrategy = new StandardRetryStrategy(maxAttempts, {});
162-
expect(retryStrategy["retryQuota"]).toBe(getDefaultRetryQuota());
165+
expect(retryStrategy["retryQuota"]).toBe(
166+
getDefaultRetryQuota(INITIAL_RETRY_TOKENS)
167+
);
163168
});
164169

165170
it("sets options.retryQuota if defined", () => {
@@ -253,27 +258,29 @@ describe("defaultStrategy", () => {
253258
describe("retryQuota", () => {
254259
describe("hasRetryTokens", () => {
255260
it("not called on successful operation", async () => {
256-
const { hasRetryTokens } = getDefaultRetryQuota();
261+
const { hasRetryTokens } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
257262
await mockSuccessfulOperation(maxAttempts);
258263
expect(hasRetryTokens).not.toHaveBeenCalled();
259264
});
260265

261266
it("called once in case of single failure", async () => {
262-
const { hasRetryTokens } = getDefaultRetryQuota();
267+
const { hasRetryTokens } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
263268
await mockSuccessAfterOneFail(maxAttempts);
264269
expect(hasRetryTokens).toHaveBeenCalledTimes(1);
265270
});
266271

267272
it("called once on each retry request", async () => {
268-
const { hasRetryTokens } = getDefaultRetryQuota();
273+
const { hasRetryTokens } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
269274
await mockFailedOperation(maxAttempts);
270275
expect(hasRetryTokens).toHaveBeenCalledTimes(maxAttempts - 1);
271276
});
272277
});
273278

274279
describe("releaseRetryTokens", () => {
275280
it("called once without param on successful operation", async () => {
276-
const { releaseRetryTokens } = getDefaultRetryQuota();
281+
const { releaseRetryTokens } = getDefaultRetryQuota(
282+
INITIAL_RETRY_TOKENS
283+
);
277284
await mockSuccessfulOperation(maxAttempts);
278285
expect(releaseRetryTokens).toHaveBeenCalledTimes(1);
279286
expect(releaseRetryTokens).toHaveBeenCalledWith(undefined);
@@ -284,7 +291,7 @@ describe("defaultStrategy", () => {
284291
const {
285292
releaseRetryTokens,
286293
retrieveRetryTokens
287-
} = getDefaultRetryQuota();
294+
} = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
288295
(retrieveRetryTokens as jest.Mock).mockReturnValueOnce(retryTokens);
289296

290297
await mockSuccessAfterOneFail(maxAttempts);
@@ -299,7 +306,7 @@ describe("defaultStrategy", () => {
299306
const {
300307
releaseRetryTokens,
301308
retrieveRetryTokens
302-
} = getDefaultRetryQuota();
309+
} = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
303310

304311
(retrieveRetryTokens as jest.Mock)
305312
.mockReturnValueOnce(retryTokensFirst)
@@ -311,27 +318,35 @@ describe("defaultStrategy", () => {
311318
});
312319

313320
it("not called on unsuccessful operation", async () => {
314-
const { releaseRetryTokens } = getDefaultRetryQuota();
321+
const { releaseRetryTokens } = getDefaultRetryQuota(
322+
INITIAL_RETRY_TOKENS
323+
);
315324
await mockFailedOperation(maxAttempts);
316325
expect(releaseRetryTokens).not.toHaveBeenCalled();
317326
});
318327
});
319328

320329
describe("retrieveRetryTokens", () => {
321330
it("not called on successful operation", async () => {
322-
const { retrieveRetryTokens } = getDefaultRetryQuota();
331+
const { retrieveRetryTokens } = getDefaultRetryQuota(
332+
INITIAL_RETRY_TOKENS
333+
);
323334
await mockSuccessfulOperation(maxAttempts);
324335
expect(retrieveRetryTokens).not.toHaveBeenCalled();
325336
});
326337

327338
it("called once in case of single failure", async () => {
328-
const { retrieveRetryTokens } = getDefaultRetryQuota();
339+
const { retrieveRetryTokens } = getDefaultRetryQuota(
340+
INITIAL_RETRY_TOKENS
341+
);
329342
await mockSuccessAfterOneFail(maxAttempts);
330343
expect(retrieveRetryTokens).toHaveBeenCalledTimes(1);
331344
});
332345

333346
it("called once on each retry request", async () => {
334-
const { retrieveRetryTokens } = getDefaultRetryQuota();
347+
const { retrieveRetryTokens } = getDefaultRetryQuota(
348+
INITIAL_RETRY_TOKENS
349+
);
335350
await mockFailedOperation(maxAttempts);
336351
expect(retrieveRetryTokens).toHaveBeenCalledTimes(maxAttempts - 1);
337352
});
@@ -372,7 +387,7 @@ describe("defaultStrategy", () => {
372387
hasRetryTokens,
373388
retrieveRetryTokens,
374389
releaseRetryTokens
375-
} = getDefaultRetryQuota();
390+
} = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
376391
(hasRetryTokens as jest.Mock).mockReturnValueOnce(false);
377392

378393
const mockError = new Error();

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import {
22
DEFAULT_RETRY_DELAY_BASE,
3-
THROTTLING_RETRY_DELAY_BASE
3+
THROTTLING_RETRY_DELAY_BASE,
4+
INITIAL_RETRY_TOKENS
45
} from "./constants";
56
import { defaultDelayDecider } from "./delayDecider";
67
import { defaultRetryDecider } from "./retryDecider";
@@ -75,7 +76,8 @@ export class StandardRetryStrategy implements RetryStrategy {
7576
) {
7677
this.retryDecider = options?.retryDecider ?? defaultRetryDecider;
7778
this.delayDecider = options?.delayDecider ?? defaultDelayDecider;
78-
this.retryQuota = options?.retryQuota ?? getDefaultRetryQuota();
79+
this.retryQuota =
80+
options?.retryQuota ?? getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
7981
}
8082

8183
private shouldRetry(error: SdkError, attempts: number) {

0 commit comments

Comments
 (0)