Skip to content

Commit 33364cb

Browse files
trivikrtrivikr
authored
feat: add retry quotas to StandardRetryStrategy (#1258)
1 parent be8e048 commit 33364cb

File tree

5 files changed

+611
-101
lines changed

5 files changed

+611
-101
lines changed

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,3 +15,25 @@ export const MAXIMUM_RETRY_DELAY = 20 * 1000;
1515
* encountered.
1616
*/
1717
export const THROTTLING_RETRY_DELAY_BASE = 500;
18+
19+
/**
20+
* Initial number of retry tokens in Retry Quota
21+
*/
22+
export const INITIAL_RETRY_TOKENS = 500;
23+
24+
/**
25+
* The total amount of retry tokens to be decremented from retry token balance.
26+
*/
27+
export const RETRY_COST = 5;
28+
29+
/**
30+
* The total amount of retry tokens to be decremented from retry token balance
31+
* when a throttling error is encountered.
32+
*/
33+
export const TIMEOUT_RETRY_COST = 10;
34+
35+
/**
36+
* The total amount of retry token to be incremented from retry token balance
37+
* if an SDK operation invocation succeeds without requiring a retry request.
38+
*/
39+
export const NO_RETRY_INCREMENT = 1;

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

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
import { getDefaultRetryQuota } from "./defaultRetryQuota";
2+
import { SdkError } from "@aws-sdk/smithy-client";
3+
import {
4+
INITIAL_RETRY_TOKENS,
5+
TIMEOUT_RETRY_COST,
6+
RETRY_COST,
7+
NO_RETRY_INCREMENT
8+
} from "./constants";
9+
10+
describe("defaultRetryQuota", () => {
11+
const getMockError = () => new Error() as SdkError;
12+
const getMockTimeoutError = () =>
13+
Object.assign(new Error(), {
14+
name: "TimeoutError"
15+
}) as SdkError;
16+
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;
24+
while (availableCapacity >= targetCapacity) {
25+
retryQuota.retrieveRetryTokens(error);
26+
availableCapacity -= targetCapacity;
27+
}
28+
return retryQuota;
29+
};
30+
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+
57+
describe("hasRetryTokens", () => {
58+
describe("returns true if capacity is available", () => {
59+
it("when it's TimeoutError", () => {
60+
const timeoutError = getMockTimeoutError();
61+
expect(
62+
getDefaultRetryQuota(INITIAL_RETRY_TOKENS).hasRetryTokens(
63+
timeoutError
64+
)
65+
).toBe(true);
66+
});
67+
68+
it("when it's not TimeoutError", () => {
69+
expect(
70+
getDefaultRetryQuota(INITIAL_RETRY_TOKENS).hasRetryTokens(
71+
getMockError()
72+
)
73+
).toBe(true);
74+
});
75+
});
76+
77+
describe("returns false if capacity is not available", () => {
78+
it("when it's TimeoutError", () => {
79+
const timeoutError = getMockTimeoutError();
80+
const retryQuota = getDrainedRetryQuota(
81+
TIMEOUT_RETRY_COST,
82+
timeoutError
83+
);
84+
expect(retryQuota.hasRetryTokens(timeoutError)).toBe(false);
85+
});
86+
87+
it("when it's not TimeoutError", () => {
88+
const error = getMockError();
89+
const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
90+
expect(retryQuota.hasRetryTokens(error)).toBe(false);
91+
});
92+
});
93+
});
94+
95+
describe("retrieveRetryToken", () => {
96+
describe("returns retry tokens amount if available", () => {
97+
it("when it's TimeoutError", () => {
98+
const timeoutError = getMockTimeoutError();
99+
expect(
100+
getDefaultRetryQuota(INITIAL_RETRY_TOKENS).retrieveRetryTokens(
101+
timeoutError
102+
)
103+
).toBe(TIMEOUT_RETRY_COST);
104+
});
105+
106+
it("when it's not TimeoutError", () => {
107+
expect(
108+
getDefaultRetryQuota(INITIAL_RETRY_TOKENS).retrieveRetryTokens(
109+
getMockError()
110+
)
111+
).toBe(RETRY_COST);
112+
});
113+
});
114+
115+
describe("throws error if retry tokens not available", () => {
116+
it("when it's TimeoutError", () => {
117+
const timeoutError = getMockTimeoutError();
118+
const retryQuota = getDrainedRetryQuota(
119+
TIMEOUT_RETRY_COST,
120+
timeoutError
121+
);
122+
expect(() => {
123+
retryQuota.retrieveRetryTokens(timeoutError);
124+
}).toThrowError(new Error("No retry token available"));
125+
});
126+
127+
it("when it's not TimeoutError", () => {
128+
const error = getMockError();
129+
const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
130+
expect(() => {
131+
retryQuota.retrieveRetryTokens(error);
132+
}).toThrowError(new Error("No retry token available"));
133+
});
134+
});
135+
});
136+
137+
describe("releaseRetryToken", () => {
138+
it("adds capacityReleaseAmount if passed", () => {
139+
const error = getMockError();
140+
const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
141+
142+
// Ensure that retry tokens are not available.
143+
expect(retryQuota.hasRetryTokens(error)).toBe(false);
144+
145+
// Release RETRY_COST tokens.
146+
retryQuota.releaseRetryTokens(RETRY_COST);
147+
expect(retryQuota.hasRetryTokens(error)).toBe(true);
148+
expect(retryQuota.retrieveRetryTokens(error)).toBe(RETRY_COST);
149+
expect(retryQuota.hasRetryTokens(error)).toBe(false);
150+
});
151+
152+
it("adds NO_RETRY_INCREMENT if capacityReleaseAmount not passed", () => {
153+
const error = getMockError();
154+
const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
155+
156+
// retry tokens will not be available till NO_RETRY_INCREMENT is added
157+
// till it's equal to RETRY_COST - (INITIAL_RETRY_TOKENS % RETRY_COST)
158+
let tokensReleased = 0;
159+
const tokensToBeReleased =
160+
RETRY_COST - (INITIAL_RETRY_TOKENS % RETRY_COST);
161+
while (tokensReleased < tokensToBeReleased) {
162+
expect(retryQuota.hasRetryTokens(error)).toBe(false);
163+
retryQuota.releaseRetryTokens();
164+
tokensReleased += NO_RETRY_INCREMENT;
165+
}
166+
expect(retryQuota.hasRetryTokens(error)).toBe(true);
167+
});
168+
169+
it("ensures availableCapacity is maxed at INITIAL_RETRY_TOKENS", () => {
170+
const error = getMockError();
171+
const retryQuota = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
172+
173+
// release 100 tokens.
174+
[...Array(100).keys()].forEach(key => {
175+
retryQuota.releaseRetryTokens();
176+
});
177+
178+
// availableCapacity is still maxed at INITIAL_RETRY_TOKENS
179+
// hasRetryTokens would be true only till INITIAL_RETRY_TOKENS/RETRY_COST times
180+
[...Array(Math.floor(INITIAL_RETRY_TOKENS / RETRY_COST)).keys()].forEach(
181+
key => {
182+
expect(retryQuota.hasRetryTokens(error)).toBe(true);
183+
retryQuota.retrieveRetryTokens(error);
184+
}
185+
);
186+
expect(retryQuota.hasRetryTokens(error)).toBe(false);
187+
});
188+
});
189+
});

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

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
import { RetryQuota } from "./defaultStrategy";
2+
import { SdkError } from "@aws-sdk/smithy-client";
3+
import {
4+
RETRY_COST,
5+
TIMEOUT_RETRY_COST,
6+
NO_RETRY_INCREMENT
7+
} from "./constants";
8+
9+
export const getDefaultRetryQuota = (
10+
initialRetryTokens: number
11+
): RetryQuota => {
12+
const MAX_CAPACITY = initialRetryTokens;
13+
let availableCapacity = initialRetryTokens;
14+
15+
const getCapacityAmount = (error: SdkError) =>
16+
error.name === "TimeoutError" ? TIMEOUT_RETRY_COST : RETRY_COST;
17+
18+
const hasRetryTokens = (error: SdkError) =>
19+
getCapacityAmount(error) <= availableCapacity;
20+
21+
const retrieveRetryTokens = (error: SdkError) => {
22+
if (!hasRetryTokens(error)) {
23+
// retryStrategy should stop retrying, and return last error
24+
throw new Error("No retry token available");
25+
}
26+
const capacityAmount = getCapacityAmount(error);
27+
availableCapacity -= capacityAmount;
28+
return capacityAmount;
29+
};
30+
31+
const releaseRetryTokens = (capacityReleaseAmount?: number) => {
32+
availableCapacity += capacityReleaseAmount ?? NO_RETRY_INCREMENT;
33+
availableCapacity = Math.min(availableCapacity, MAX_CAPACITY);
34+
};
35+
36+
return Object.freeze({
37+
hasRetryTokens,
38+
retrieveRetryTokens,
39+
releaseRetryTokens
40+
});
41+
};

0 commit comments

Comments
 (0)