From 46a9e3b4c071a2e94e8f5048c75f981a5b10a252 Mon Sep 17 00:00:00 2001
From: George Fu <kuhe@users.noreply.github.com>
Date: Fri, 30 Aug 2024 14:51:24 +0000
Subject: [PATCH] feat(smithy-client): add handler cache

---
 .../middleware-stack/src/MiddlewareStack.ts   | 22 ++++++++--
 packages/smithy-client/src/client.spec.ts     | 23 +++++++++++
 packages/smithy-client/src/client.ts          | 41 +++++++++++++++----
 3 files changed, 76 insertions(+), 10 deletions(-)

diff --git a/packages/middleware-stack/src/MiddlewareStack.ts b/packages/middleware-stack/src/MiddlewareStack.ts
index 52caaac9ad1..bfc6f18584b 100644
--- a/packages/middleware-stack/src/MiddlewareStack.ts
+++ b/packages/middleware-stack/src/MiddlewareStack.ts
@@ -32,7 +32,14 @@ const getMiddlewareNameWithAliases = (name: string | undefined, aliases: Array<s
   return `${name || "anonymous"}${aliases && aliases.length > 0 ? ` (a.k.a. ${aliases.join(",")})` : ""}`;
 };
 
-export const constructStack = <Input extends object, Output extends object>(): MiddlewareStack<Input, Output> => {
+export const constructStack = <Input extends object, Output extends object>(
+  stackOptions: {
+    /**
+     * Optional change listener, called with stack instance when middleware added/removed.
+     */
+    onChange?: (middlewareStack: MiddlewareStack<any, any>) => void;
+  } = {}
+): MiddlewareStack<Input, Output> => {
   let absoluteEntries: AbsoluteMiddlewareEntry<Input, Output>[] = [];
   let relativeEntries: RelativeMiddlewareEntry<Input, Output>[] = [];
   let identifyOnResolve = false;
@@ -222,6 +229,7 @@ export const constructStack = <Input extends object, Output extends object>(): M
         }
       }
       absoluteEntries.push(entry);
+      stackOptions.onChange?.(stack);
     },
 
     addRelativeTo: (middleware: MiddlewareType<Input, Output>, options: HandlerOptions & RelativeLocation) => {
@@ -258,6 +266,7 @@ export const constructStack = <Input extends object, Output extends object>(): M
         }
       }
       relativeEntries.push(entry);
+      stackOptions.onChange?.(stack);
     },
 
     clone: () => cloneTo(constructStack<Input, Output>()),
@@ -267,8 +276,14 @@ export const constructStack = <Input extends object, Output extends object>(): M
     },
 
     remove: (toRemove: MiddlewareType<Input, Output> | string): boolean => {
-      if (typeof toRemove === "string") return removeByName(toRemove);
-      else return removeByReference(toRemove);
+      let isRemoved: boolean;
+      if (typeof toRemove === "string") {
+        isRemoved = removeByName(toRemove);
+      } else {
+        isRemoved = removeByReference(toRemove);
+      }
+      stackOptions.onChange?.(stack);
+      return isRemoved;
     },
 
     removeByTag: (toRemove: string): boolean => {
@@ -287,6 +302,7 @@ export const constructStack = <Input extends object, Output extends object>(): M
       };
       absoluteEntries = absoluteEntries.filter(filterCb);
       relativeEntries = relativeEntries.filter(filterCb);
+      stackOptions?.onChange?.(stack);
       return isRemoved;
     },
 
diff --git a/packages/smithy-client/src/client.spec.ts b/packages/smithy-client/src/client.spec.ts
index 56ed3944e00..86971e5ecbd 100644
--- a/packages/smithy-client/src/client.spec.ts
+++ b/packages/smithy-client/src/client.spec.ts
@@ -1,4 +1,5 @@
 import { Client } from "./client";
+import { clientHandlers } from "./client-handlers";
 
 describe("SmithyClient", () => {
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
@@ -50,4 +51,26 @@ describe("SmithyClient", () => {
     };
     client.send(getCommandWithOutput("foo") as any, options, callback);
   });
+
+  describe("handler caching", () => {
+    beforeEach(() => {
+      clientHandlers.delete(client);
+    });
+
+    it("should cache the resolved handler", async () => {
+      await expect(client.send(getCommandWithOutput("foo") as any)).resolves.toEqual("foo");
+      expect(clientHandlers.get(client)!.get({}.constructor)).toBeDefined();
+    });
+
+    it("should not cache the resolved handler if called with request options", async () => {
+      await expect(client.send(getCommandWithOutput("foo") as any, {})).resolves.toEqual("foo");
+      expect(clientHandlers.get(client)).toBeUndefined();
+    });
+
+    it("should clear the handler cache when the middleareStack is modified", async () => {
+      await expect(client.send(getCommandWithOutput("foo") as any)).resolves.toEqual("foo");
+      client.middlewareStack.add((n) => (a) => n(a));
+      expect(clientHandlers.get(client)).toBeUndefined();
+    });
+  });
 });
diff --git a/packages/smithy-client/src/client.ts b/packages/smithy-client/src/client.ts
index c0fd856b53c..a2768de40cd 100644
--- a/packages/smithy-client/src/client.ts
+++ b/packages/smithy-client/src/client.ts
@@ -3,12 +3,15 @@ import {
   Client as IClient,
   Command,
   FetchHttpHandlerOptions,
+  Handler,
   MetadataBearer,
   MiddlewareStack,
   NodeHttpHandlerOptions,
   RequestHandler,
 } from "@smithy/types";
 
+import { clientHandlers } from "./client-handlers";
+
 /**
  * @public
  */
@@ -44,11 +47,14 @@ export class Client<
   ResolvedClientConfiguration extends SmithyResolvedConfiguration<HandlerOptions>,
 > implements IClient<ClientInput, ClientOutput, ResolvedClientConfiguration>
 {
-  public middlewareStack: MiddlewareStack<ClientInput, ClientOutput> = constructStack<ClientInput, ClientOutput>();
-  readonly config: ResolvedClientConfiguration;
-  constructor(config: ResolvedClientConfiguration) {
-    this.config = config;
-  }
+  public middlewareStack: MiddlewareStack<ClientInput, ClientOutput> = constructStack<ClientInput, ClientOutput>({
+    onChange: () => {
+      clientHandlers.delete(this);
+    },
+  });
+
+  constructor(public readonly config: ResolvedClientConfiguration) {}
+
   send<InputType extends ClientInput, OutputType extends ClientOutput>(
     command: Command<ClientInput, InputType, ClientOutput, OutputType, SmithyResolvedConfiguration<HandlerOptions>>,
     options?: HandlerOptions
@@ -69,7 +75,27 @@ export class Client<
   ): Promise<OutputType> | void {
     const options = typeof optionsOrCb !== "function" ? optionsOrCb : undefined;
     const callback = typeof optionsOrCb === "function" ? (optionsOrCb as (err: any, data?: OutputType) => void) : cb;
-    const handler = command.resolveMiddleware(this.middlewareStack as any, this.config, options);
+
+    const useHandlerCache = options === undefined;
+
+    let handler: Handler<any, any>;
+
+    if (useHandlerCache) {
+      if (!clientHandlers.has(this)) {
+        clientHandlers.set(this, new WeakMap());
+      }
+      const handlers = clientHandlers.get(this)!;
+
+      if (handlers.has(command.constructor)) {
+        handler = handlers.get(command.constructor)!;
+      } else {
+        handler = command.resolveMiddleware(this.middlewareStack as any, this.config, options);
+        handlers.set(command.constructor, handler);
+      }
+    } else {
+      handler = command.resolveMiddleware(this.middlewareStack as any, this.config, options);
+    }
+
     if (callback) {
       handler(command)
         .then(
@@ -87,6 +113,7 @@ export class Client<
   }
 
   destroy() {
-    if (this.config.requestHandler.destroy) this.config.requestHandler.destroy();
+    this.config.requestHandler.destroy?.();
+    clientHandlers.delete(this);
   }
 }