Skip to content

Commit 2817cad

Browse files
srchasesrchase
committed
feat(middleware-retry): set attempts and total delay on response metadata
1 parent ee2ef78 commit 2817cad

File tree

4 files changed

+88
-54
lines changed

4 files changed

+88
-54
lines changed

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

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import { getDefaultRetryQuota } from "./defaultRetryQuota";
1717
import { defaultDelayDecider } from "./delayDecider";
1818
import { defaultRetryDecider } from "./retryDecider";
1919
import { DelayDecider, RetryDecider, RetryQuota } from "./types";
20+
import { asSdkError } from "./util";
2021

2122
/**
2223
* Strategy options to be passed to StandardRetryStrategy
@@ -133,7 +134,6 @@ const getDelayFromRetryAfterHeader = (response: unknown): number | undefined =>
133134

134135
const retryAfterHeaderName = Object.keys(response.headers).find((key) => key.toLowerCase() === "retry-after");
135136
if (!retryAfterHeaderName) return;
136-
137137
const retryAfter = response.headers[retryAfterHeaderName];
138138

139139
const retryAfterSeconds = Number(retryAfter);
@@ -142,10 +142,3 @@ const getDelayFromRetryAfterHeader = (response: unknown): number | undefined =>
142142
const retryAfterDate = new Date(retryAfter);
143143
return retryAfterDate.getTime() - Date.now();
144144
};
145-
146-
const asSdkError = (error: unknown): SdkError => {
147-
if (error instanceof Error) return error;
148-
if (error instanceof Object) return Object.assign(new Error(), error);
149-
if (typeof error === "string") return new Error(error);
150-
return new Error(`AWS SDK error wrapper for ${error}`);
151-
};

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

Lines changed: 49 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { HttpRequest, HttpResponse } from "@aws-sdk/protocol-http";
22
import { isServerError, isThrottlingError, isTransientError } from "@aws-sdk/service-error-classification";
3-
import { FinalizeHandlerArguments, HandlerExecutionContext, MiddlewareStack, RetryErrorType } from "@aws-sdk/types";
3+
import { FinalizeHandlerArguments, HandlerExecutionContext, MiddlewareStack } from "@aws-sdk/types";
44
import { INVOCATION_ID_HEADER, REQUEST_HEADER } from "@aws-sdk/util-retry";
55
import { v4 } from "uuid";
66

@@ -99,6 +99,13 @@ describe(retryMiddleware.name, () => {
9999
refreshRetryTokenForRetry: jest.fn().mockResolvedValue(mockRetryToken),
100100
recordSuccess: jest.fn(),
101101
};
102+
const mockResponse = "mockResponse";
103+
const mockSuccess = {
104+
response: mockResponse,
105+
output: {
106+
$metadata: {},
107+
},
108+
};
102109
const getErrorWithValues = (retryAfter: number | string, retryAfterHeaderName?: string) => {
103110
const error = new Error("mockError");
104111
Object.defineProperty(error, "$response", {
@@ -110,8 +117,8 @@ describe(retryMiddleware.name, () => {
110117
};
111118

112119
it("calls acquireInitialRetryToken and records success when next succeeds", async () => {
113-
const next = jest.fn();
114-
await retryMiddleware({
120+
const next = jest.fn().mockResolvedValueOnce(mockSuccess);
121+
const { response, output } = await retryMiddleware({
115122
maxAttempts: () => Promise.resolve(maxAttempts),
116123
retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),
117124
})(
@@ -122,6 +129,7 @@ describe(retryMiddleware.name, () => {
122129
expect(mockRetryStrategy.acquireInitialRetryToken).toHaveBeenCalledWith(partitionId);
123130
expect(mockRetryStrategy.recordSuccess).toHaveBeenCalledTimes(1);
124131
expect(mockRetryStrategy.recordSuccess).toHaveBeenCalledWith(mockRetryToken);
132+
expect(output.$metadata.attempts).toBe(1);
125133
});
126134

127135
describe("throws when token cannot be refreshed", () => {
@@ -138,8 +146,6 @@ describe(retryMiddleware.name, () => {
138146
refreshRetryTokenForRetry: jest.fn().mockRejectedValue(new Error("Cannot refresh token")),
139147
recordSuccess: jest.fn(),
140148
};
141-
// const maxAttempts = 3;
142-
// const retryCount = maxAttempts - 1;
143149
try {
144150
await retryMiddleware({
145151
maxAttempts: () => Promise.resolve(maxAttempts),
@@ -154,6 +160,8 @@ describe(retryMiddleware.name, () => {
154160
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledTimes(1);
155161
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledWith(mockRetryToken, errorInfo);
156162
expect(error).toStrictEqual(requestError);
163+
expect(error.$metadata.attempts).toBe(1);
164+
expect(error.$metadata.totalRetryDelay).toBeDefined();
157165
}
158166
});
159167
});
@@ -162,11 +170,11 @@ describe(retryMiddleware.name, () => {
162170
const mockError = new Error("mockError");
163171
it("sets throttling error type", async () => {
164172
(isThrottlingError as jest.Mock).mockReturnValue(true);
165-
const next = jest.fn().mockRejectedValueOnce(mockError).mockResolvedValueOnce({});
173+
const next = jest.fn().mockRejectedValueOnce(mockError).mockResolvedValueOnce(mockSuccess);
166174
const errorInfo = {
167175
errorType: "THROTTLING",
168176
};
169-
await retryMiddleware({
177+
const { response, output } = await retryMiddleware({
170178
maxAttempts: () => Promise.resolve(maxAttempts),
171179
retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),
172180
})(
@@ -177,15 +185,17 @@ describe(retryMiddleware.name, () => {
177185
expect(mockRetryStrategy.acquireInitialRetryToken).toHaveBeenCalledWith(partitionId);
178186
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledTimes(1);
179187
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledWith(mockRetryToken, errorInfo);
188+
expect(output.$metadata.attempts).toBe(2);
189+
expect(output.$metadata.totalRetryDelay).toBeDefined();
180190
});
181191
it("sets transient error type", async () => {
182192
(isTransientError as jest.Mock).mockReturnValue(true);
183193
(isThrottlingError as jest.Mock).mockReturnValue(false);
184-
const next = jest.fn().mockRejectedValueOnce(mockError).mockResolvedValueOnce({});
194+
const next = jest.fn().mockRejectedValueOnce(mockError).mockResolvedValueOnce(mockSuccess);
185195
const errorInfo = {
186196
errorType: "TRANSIENT",
187197
};
188-
await retryMiddleware({
198+
const { response, output } = await retryMiddleware({
189199
maxAttempts: () => Promise.resolve(maxAttempts),
190200
retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),
191201
})(
@@ -196,16 +206,18 @@ describe(retryMiddleware.name, () => {
196206
expect(mockRetryStrategy.acquireInitialRetryToken).toHaveBeenCalledWith(partitionId);
197207
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledTimes(1);
198208
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledWith(mockRetryToken, errorInfo);
209+
expect(output.$metadata.attempts).toBe(2);
210+
expect(output.$metadata.totalRetryDelay).toBeDefined();
199211
});
200212
it("sets server error type", async () => {
201213
(isServerError as jest.Mock).mockReturnValue(true);
202214
(isTransientError as jest.Mock).mockReturnValue(false);
203215
(isThrottlingError as jest.Mock).mockReturnValue(false);
204-
const next = jest.fn().mockRejectedValueOnce(mockError).mockResolvedValueOnce({});
216+
const next = jest.fn().mockRejectedValueOnce(mockError).mockResolvedValueOnce(mockSuccess);
205217
const errorInfo = {
206218
errorType: "SERVER_ERROR",
207219
};
208-
await retryMiddleware({
220+
const { response, output } = await retryMiddleware({
209221
maxAttempts: () => Promise.resolve(maxAttempts),
210222
retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),
211223
})(
@@ -216,16 +228,18 @@ describe(retryMiddleware.name, () => {
216228
expect(mockRetryStrategy.acquireInitialRetryToken).toHaveBeenCalledWith(partitionId);
217229
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledTimes(1);
218230
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledWith(mockRetryToken, errorInfo);
231+
expect(output.$metadata.attempts).toBe(2);
232+
expect(output.$metadata.totalRetryDelay).toBeDefined();
219233
});
220234
it("sets client error type", async () => {
221235
(isServerError as jest.Mock).mockReturnValue(false);
222236
(isTransientError as jest.Mock).mockReturnValue(false);
223237
(isThrottlingError as jest.Mock).mockReturnValue(false);
224-
const next = jest.fn().mockRejectedValueOnce(mockError).mockResolvedValueOnce({});
238+
const next = jest.fn().mockRejectedValueOnce(mockError).mockResolvedValueOnce(mockSuccess);
225239
const errorInfo = {
226240
errorType: "CLIENT_ERROR",
227241
};
228-
await retryMiddleware({
242+
const { response, output } = await retryMiddleware({
229243
maxAttempts: () => Promise.resolve(maxAttempts),
230244
retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),
231245
})(
@@ -236,6 +250,8 @@ describe(retryMiddleware.name, () => {
236250
expect(mockRetryStrategy.acquireInitialRetryToken).toHaveBeenCalledWith(partitionId);
237251
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledTimes(1);
238252
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledWith(mockRetryToken, errorInfo);
253+
expect(output.$metadata.attempts).toBe(2);
254+
expect(output.$metadata.totalRetryDelay).toBeDefined();
239255
});
240256

241257
describe("when retry-after is not set", () => {
@@ -247,11 +263,11 @@ describe(retryMiddleware.name, () => {
247263
headers: { ["other-header"]: "foo" },
248264
},
249265
});
250-
const next = jest.fn().mockRejectedValueOnce(mockError).mockResolvedValueOnce({});
266+
const next = jest.fn().mockRejectedValueOnce(mockError).mockResolvedValueOnce(mockSuccess);
251267
const errorInfo = {
252268
errorType: "CLIENT_ERROR",
253269
};
254-
await retryMiddleware({
270+
const { response, output } = await retryMiddleware({
255271
maxAttempts: () => Promise.resolve(maxAttempts),
256272
retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),
257273
})(
@@ -262,6 +278,8 @@ describe(retryMiddleware.name, () => {
262278
expect(mockRetryStrategy.acquireInitialRetryToken).toHaveBeenCalledWith(partitionId);
263279
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledTimes(1);
264280
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledWith(mockRetryToken, errorInfo);
281+
expect(output.$metadata.attempts).toBe(2);
282+
expect(output.$metadata.totalRetryDelay).toBeDefined();
265283
});
266284
});
267285

@@ -276,8 +294,8 @@ describe(retryMiddleware.name, () => {
276294
};
277295
it("parses retry-after from date string", async () => {
278296
const error = getErrorWithValues(retryAfterDate.toISOString());
279-
const next = jest.fn().mockRejectedValueOnce(error).mockResolvedValueOnce({});
280-
await retryMiddleware({
297+
const next = jest.fn().mockRejectedValueOnce(error).mockResolvedValueOnce(mockSuccess);
298+
const { response, output } = await retryMiddleware({
281299
maxAttempts: () => Promise.resolve(maxAttempts),
282300
retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),
283301
})(
@@ -288,11 +306,13 @@ describe(retryMiddleware.name, () => {
288306
expect(mockRetryStrategy.acquireInitialRetryToken).toHaveBeenCalledWith(partitionId);
289307
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledTimes(1);
290308
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledWith(mockRetryToken, errorInfo);
309+
expect(output.$metadata.attempts).toBe(2);
310+
expect(output.$metadata.totalRetryDelay).toBeDefined();
291311
});
292312
it("parses retry-after from seconds", async () => {
293313
const error = getErrorWithValues(retryAfterDate.getTime() / 1000);
294-
const next = jest.fn().mockRejectedValueOnce(error).mockResolvedValueOnce({});
295-
await retryMiddleware({
314+
const next = jest.fn().mockRejectedValueOnce(error).mockResolvedValueOnce(mockSuccess);
315+
const { response, output } = await retryMiddleware({
296316
maxAttempts: () => Promise.resolve(maxAttempts),
297317
retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),
298318
})(
@@ -303,11 +323,13 @@ describe(retryMiddleware.name, () => {
303323
expect(mockRetryStrategy.acquireInitialRetryToken).toHaveBeenCalledWith(partitionId);
304324
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledTimes(1);
305325
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledWith(mockRetryToken, errorInfo);
326+
expect(output.$metadata.attempts).toBe(2);
327+
expect(output.$metadata.totalRetryDelay).toBeDefined();
306328
});
307329
it("parses retry-after from Retry-After header name", async () => {
308330
const error = getErrorWithValues(retryAfterDate.toISOString(), "Retry-After");
309-
const next = jest.fn().mockRejectedValueOnce(error).mockResolvedValueOnce({});
310-
await retryMiddleware({
331+
const next = jest.fn().mockRejectedValueOnce(error).mockResolvedValueOnce(mockSuccess);
332+
const { response, output } = await retryMiddleware({
311333
maxAttempts: () => Promise.resolve(maxAttempts),
312334
retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),
313335
})(
@@ -318,7 +340,8 @@ describe(retryMiddleware.name, () => {
318340
expect(mockRetryStrategy.acquireInitialRetryToken).toHaveBeenCalledWith(partitionId);
319341
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledTimes(1);
320342
expect(mockRetryStrategy.refreshRetryTokenForRetry).toHaveBeenCalledWith(mockRetryToken, errorInfo);
321-
// (isInstance as unknown as jest.Mock).mockReturnValue(false);
343+
expect(output.$metadata.attempts).toBe(2);
344+
expect(output.$metadata.totalRetryDelay).toBeDefined();
322345
});
323346
(isInstance as unknown as jest.Mock).mockReturnValue(false);
324347
});
@@ -327,7 +350,7 @@ describe(retryMiddleware.name, () => {
327350
describe("retry headers", () => {
328351
describe("not added if HttpRequest.isInstance returns false", () => {
329352
it(`retry informational header: ${INVOCATION_ID_HEADER}`, async () => {
330-
const next = jest.fn();
353+
const next = jest.fn().mockResolvedValueOnce(mockSuccess);
331354
await retryMiddleware({
332355
maxAttempts: () => Promise.resolve(maxAttempts),
333356
retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),
@@ -340,7 +363,7 @@ describe(retryMiddleware.name, () => {
340363
});
341364
});
342365
it(`header for each attempt as ${REQUEST_HEADER}`, async () => {
343-
const next = jest.fn();
366+
const next = jest.fn().mockResolvedValueOnce(mockSuccess);
344367
await retryMiddleware({
345368
maxAttempts: () => Promise.resolve(maxAttempts),
346369
retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),
@@ -359,7 +382,7 @@ describe(retryMiddleware.name, () => {
359382
const { isInstance } = HttpRequest;
360383
(isInstance as unknown as jest.Mock).mockReturnValue(true);
361384
(isThrottlingError as jest.Mock).mockReturnValue(true);
362-
const next = jest.fn().mockRejectedValueOnce(error).mockResolvedValueOnce({});
385+
const next = jest.fn().mockRejectedValueOnce(error).mockResolvedValueOnce(mockSuccess);
363386
await retryMiddleware({
364387
maxAttempts: () => Promise.resolve(maxAttempts),
365388
retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),
@@ -377,7 +400,7 @@ describe(retryMiddleware.name, () => {
377400
const { isInstance } = HttpRequest;
378401
(isInstance as unknown as jest.Mock).mockReturnValue(true);
379402
(isThrottlingError as jest.Mock).mockReturnValue(true);
380-
const next = jest.fn().mockRejectedValueOnce(error).mockResolvedValueOnce({});
403+
const next = jest.fn().mockRejectedValueOnce(error).mockResolvedValueOnce(mockSuccess);
381404
await retryMiddleware({
382405
maxAttempts: () => Promise.resolve(maxAttempts),
383406
retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),

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

Lines changed: 30 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import { INVOCATION_ID_HEADER, REQUEST_HEADER } from "@aws-sdk/util-retry";
2020
import { v4 } from "uuid";
2121

2222
import { RetryResolvedConfig } from "./configurations";
23+
import { asSdkError } from "./util";
2324

2425
export const retryMiddleware =
2526
(options: RetryResolvedConfig) =>
@@ -35,29 +36,38 @@ export const retryMiddleware =
3536
retryStrategy = retryStrategy as RetryStrategyV2;
3637
let retryToken: RetryToken = await retryStrategy.acquireInitialRetryToken(context["partition_id"]);
3738
let lastError: SdkError = new Error();
38-
let retryCount = retryToken.getRetryCount();
39+
let attempts = 0;
40+
let totalRetryDelay = 0;
3941
const { request } = args;
4042
if (HttpRequest.isInstance(request)) {
4143
request.headers[INVOCATION_ID_HEADER] = v4();
4244
}
4345
while (true) {
4446
try {
4547
if (HttpRequest.isInstance(request)) {
46-
request.headers[REQUEST_HEADER] = `attempt=${retryCount + 1}; max=${maxAttempts}`;
48+
request.headers[REQUEST_HEADER] = `attempt=${attempts + 1}; max=${maxAttempts}`;
4749
}
48-
const response = await next(args);
50+
const { response, output } = await next(args);
4951
retryStrategy.recordSuccess(retryToken);
50-
return response;
52+
output.$metadata.attempts = attempts + 1;
53+
output.$metadata.totalRetryDelay = totalRetryDelay;
54+
return { response, output };
5155
} catch (e) {
5256
const retryErrorInfo = getRetyErrorInto(e);
53-
lastError = e;
57+
lastError = asSdkError(e);
5458
try {
5559
retryToken = await retryStrategy.refreshRetryTokenForRetry(retryToken, retryErrorInfo);
5660
} catch (refreshError) {
61+
if (!lastError.$metadata) {
62+
lastError.$metadata = {};
63+
}
64+
lastError.$metadata.attempts = attempts + 1;
65+
lastError.$metadata.totalRetryDelay = totalRetryDelay;
5766
throw lastError;
5867
}
59-
retryCount = retryToken.getRetryCount();
68+
attempts = retryToken.getRetryCount();
6069
const delay = retryToken.getRetryDelay();
70+
totalRetryDelay += delay;
6171
await new Promise((resolve) => setTimeout(resolve, delay));
6272
}
6373
}
@@ -86,20 +96,6 @@ const getRetyErrorInto = (error: SdkError): RetryErrorInfo => {
8696
return errorInfo;
8797
};
8898

89-
const getRetryAfterHint = (response: unknown): Date | undefined => {
90-
if (!HttpResponse.isInstance(response)) return;
91-
92-
const retryAfterHeaderName = Object.keys(response.headers).find((key) => key.toLowerCase() === "retry-after");
93-
if (!retryAfterHeaderName) return;
94-
const retryAfter = response.headers[retryAfterHeaderName];
95-
96-
const retryAfterSeconds = Number(retryAfter);
97-
const retryAfterDate = new Date(retryAfter);
98-
if (!Number.isNaN(retryAfterSeconds)) return new Date(retryAfterSeconds * 1000);
99-
100-
return retryAfterDate;
101-
};
102-
10399
const getRetryErrorType = (error: SdkError): RetryErrorType => {
104100
if (isThrottlingError(error)) return "THROTTLING";
105101
if (isTransientError(error)) return "TRANSIENT";
@@ -120,3 +116,17 @@ export const getRetryPlugin = (options: RetryResolvedConfig): Pluggable<any, any
120116
clientStack.add(retryMiddleware(options), retryMiddlewareOptions);
121117
},
122118
});
119+
120+
export const getRetryAfterHint = (response: unknown): Date | undefined => {
121+
if (!HttpResponse.isInstance(response)) return;
122+
123+
const retryAfterHeaderName = Object.keys(response.headers).find((key) => key.toLowerCase() === "retry-after");
124+
if (!retryAfterHeaderName) return;
125+
const retryAfter = response.headers[retryAfterHeaderName];
126+
127+
const retryAfterSeconds = Number(retryAfter);
128+
if (!Number.isNaN(retryAfterSeconds)) return new Date(retryAfterSeconds * 1000);
129+
130+
const retryAfterDate = new Date(retryAfter);
131+
return retryAfterDate;
132+
};

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import { SdkError } from "@aws-sdk/types";
2+
3+
export const asSdkError = (error: unknown): SdkError => {
4+
if (error instanceof Error) return error;
5+
if (error instanceof Object) return Object.assign(new Error(), error);
6+
if (typeof error === "string") return new Error(error);
7+
return new Error(`AWS SDK error wrapper for ${error}`);
8+
};

0 commit comments

Comments
 (0)