feat(middleware-user-agent): add retry mode metrics

Author: AllanZhengYP

packages/middleware-retry/src/defaultStrategy.ts

@@ -74,6 +74,7 @@ export class StandardRetryStrategy implements RetryStrategy {
   private retryDecider: RetryDecider;
   private delayDecider: DelayDecider;
   private retryQuota: RetryQuota;
+  public readonly mode = DEFAULT_RETRY_MODE;
 
   constructor(private readonly maxAttemptsProvider: Provider<number>, options?: StandardRetryStrategyOptions) {
     this.retryDecider = options?.retryDecider ?? defaultRetryDecider;

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

@@ -1,4 +1,4 @@
-import { FinalizeHandlerArguments, MiddlewareStack, RetryStrategy } from "@aws-sdk/types";
+import { FinalizeHandlerArguments, HandlerExecutionContext, MiddlewareStack, RetryStrategy } from "@aws-sdk/types";
 
 import { getRetryPlugin, retryMiddleware, retryMiddlewareOptions } from "./retryMiddleware";
 
@@ -37,15 +37,21 @@ describe("retryMiddleware", () => {
       request: {},
     };
     const mockRetryStrategy = {
+      mode: "mock",
       maxAttempts,
       retry: jest.fn(),
     };
+    const context: HandlerExecutionContext = {};
 
     await retryMiddleware({
       maxAttempts: () => Promise.resolve(maxAttempts),
       retryStrategy: mockRetryStrategy,
-    })(next)(args as FinalizeHandlerArguments<any>);
+    })(
+      next,
+      context
+    )(args as FinalizeHandlerArguments<any>);
     expect(mockRetryStrategy.retry).toHaveBeenCalledTimes(1);
     expect(mockRetryStrategy.retry).toHaveBeenCalledWith(next, args);
+    expect(context.userAgent).toContainEqual(["cfg/retry-mode", "mock"]);
   });
 });

packages/middleware-retry/src/retryMiddleware.ts

@@ -4,17 +4,23 @@ import {
   FinalizeHandlerArguments,
   FinalizeHandlerOutput,
   FinalizeRequestHandlerOptions,
+  HandlerExecutionContext,
   MetadataBearer,
   Pluggable,
 } from "@aws-sdk/types";
 
 import { RetryResolvedConfig } from "./configurations";
 
 export const retryMiddleware = (options: RetryResolvedConfig) => <Output extends MetadataBearer = MetadataBearer>(
-  next: FinalizeHandler<any, Output>
+  next: FinalizeHandler<any, Output>,
+  context: HandlerExecutionContext
 ): FinalizeHandler<any, Output> => async (
   args: FinalizeHandlerArguments<any>
-): Promise<FinalizeHandlerOutput<Output>> => options.retryStrategy.retry(next, args);
+): Promise<FinalizeHandlerOutput<Output>> => {
+  if (options?.retryStrategy?.mode)
+    context.userAgent = [...(context.userAgent || []), ["cfg/retry-mode", options.retryStrategy.mode]];
+  return options.retryStrategy.retry(next, args);
+};
 
 export const retryMiddlewareOptions: FinalizeRequestHandlerOptions & AbsoluteLocation = {
   name: "retryMiddleware",

packages/middleware-user-agent/src/user-agent-middleware.ts

@@ -1,12 +1,12 @@
 import { HttpRequest } from "@aws-sdk/protocol-http";
 import {
-  BuildHandler,
-  BuildHandlerArguments,
-  BuildHandlerOptions,
-  BuildHandlerOutput,
+  FinalizeHandler,
+  FinalizeHandlerArguments,
+  FinalizeHandlerOutput,
   HandlerExecutionContext,
   MetadataBearer,
   Pluggable,
+  RelativeMiddlewareOptions,
   UserAgentPair,
 } from "@aws-sdk/types";
 
@@ -26,9 +26,9 @@ import { SPACE, UA_ESCAPE_REGEX, USER_AGENT, X_AMZ_USER_AGENT } from "./constant
  * agent.
  */
 export const userAgentMiddleware = (options: UserAgentResolvedConfig) => <Output extends MetadataBearer>(
-  next: BuildHandler<any, any>,
+  next: FinalizeHandler<any, any>,
   context: HandlerExecutionContext
-): BuildHandler<any, any> => async (args: BuildHandlerArguments<any>): Promise<BuildHandlerOutput<Output>> => {
+): FinalizeHandler<any, any> => async (args: FinalizeHandlerArguments<any>): Promise<FinalizeHandlerOutput<Output>> => {
   const { request } = args;
   if (!HttpRequest.isInstance(request)) return next(args);
   const { headers } = request;
@@ -57,7 +57,7 @@ export const userAgentMiddleware = (options: UserAgentResolvedConfig) => <Output
  * User agent name may include prefix like `md/`, `api/`, `os/` etc., we should not escape the `/` after the prefix.
  * @private
  */
-export const escapeUserAgent = ([name, version]: UserAgentPair): string => {
+const escapeUserAgent = ([name, version]: UserAgentPair): string => {
   const prefixSeparatorIndex = name.indexOf("/");
   const prefix = name.substring(0, prefixSeparatorIndex); // If no prefix, prefix is just ""
   const uaName = name.substring(prefixSeparatorIndex + 1);
@@ -67,14 +67,15 @@ export const escapeUserAgent = ([name, version]: UserAgentPair): string => {
     .join("/");
 };
 
-export const getUserAgentMiddlewareOptions: BuildHandlerOptions = {
+export const getUserAgentMiddlewareOptions: RelativeMiddlewareOptions = {
   name: "getUserAgentMiddleware",
-  step: "build",
+  relation: "before",
+  toMiddleware: "awsAuthMiddleware",
   tags: ["SET_USER_AGENT", "USER_AGENT"],
 };
 
 export const getUserAgentPlugin = (config: UserAgentResolvedConfig): Pluggable<any, any> => ({
   applyToStack: (clientStack) => {
-    clientStack.add(userAgentMiddleware(config), getUserAgentMiddlewareOptions);
+    clientStack.addRelativeTo(userAgentMiddleware(config), getUserAgentMiddlewareOptions);
   },
 });

packages/types/src/util.ts

@@ -1,5 +1,10 @@
 import { Endpoint } from "./http";
-import { FinalizeHandler, FinalizeHandlerArguments, FinalizeHandlerOutput } from "./middleware";
+import {
+  FinalizeHandler,
+  FinalizeHandlerArguments,
+  FinalizeHandlerOutput,
+  HandlerExecutionContext,
+} from "./middleware";
 import { MetadataBearer } from "./response";
 
 /**
@@ -53,6 +58,10 @@ export interface BodyLengthCalculator {
  * Interface that specifies the retry behavior
  */
 export interface RetryStrategy {
+  /**
+   * The retry mode describing how the retry strategy control the traffic flow.
+   */
+  mode?: string;
   /**
    * the retry behavior the will invoke the next handler and handle the retry accordingly.
    * This function should also update the $metadata from the response accordingly.