feat(util-endpoint): add endpoint ruleset cache (#1385)

* feat(util-endpoint): add endpoint ruleset cache

* handle shorthand argv

* handle delimiter in cache key segment

* no caching for missing param list

.changeset/moody-fishes-buy.md

@@ -0,0 +1,5 @@
+---
+"@smithy/util-endpoints": minor
+---
+
+add endpoint ruleset cache

packages/util-endpoints/src/cache/EndpointCache.spec.ts

@@ -0,0 +1,83 @@
+import { EndpointCache } from "./EndpointCache";
+
+describe(EndpointCache.name, () => {
+  const endpoint1: any = {};
+  const endpoint2: any = {};
+
+  it("should store and retrieve items", () => {
+    const cache = new EndpointCache({
+      size: 50,
+      params: ["A", "B", "C"],
+    });
+
+    expect(cache.get({ A: "b", B: "b" }, () => endpoint1)).toBe(endpoint1);
+    expect(cache.get({ A: "b", B: "b" }, () => endpoint1)).toBe(endpoint1);
+    expect(cache.get({ A: "b", B: "b", C: "c" }, () => endpoint1)).toBe(endpoint1);
+    expect(cache.get({ A: "b", B: "b", C: "c" }, () => endpoint1)).toBe(endpoint1);
+    expect(cache.get({ A: "b", B: "b", C: "cc" }, () => endpoint2)).toBe(endpoint2);
+    expect(cache.get({ A: "b", B: "b", C: "cc" }, () => endpoint2)).toBe(endpoint2);
+
+    expect(cache.size()).toEqual(3);
+  });
+
+  it("should accept a custom parameter list", () => {
+    const cache = new EndpointCache({
+      size: 50,
+      params: ["A", "B"],
+    });
+
+    expect(cache.get({ A: "b", B: "b" }, () => endpoint1)).toBe(endpoint1);
+    expect(cache.get({ A: "b", B: "b", C: "c" }, () => endpoint1)).toBe(endpoint1);
+    expect(cache.get({ A: "b", B: "b", C: "cc" }, () => endpoint2)).toBe(endpoint1);
+
+    expect(cache.size()).toEqual(1);
+  });
+
+  it("bypasses caching if param values include the cache key delimiter", () => {
+    const cache = new EndpointCache({
+      size: 50,
+      params: ["A", "B"],
+    });
+
+    expect(cache.get({ A: "b", B: "aaa|;aaa" }, () => endpoint1)).toBe(endpoint1);
+    expect(cache.size()).toEqual(0);
+  });
+
+  it("bypasses caching if param list is empty", () => {
+    const cache = new EndpointCache({
+      size: 50,
+      params: [],
+    });
+
+    expect(cache.get({ A: "b", B: "b" }, () => endpoint1)).toBe(endpoint1);
+    expect(cache.size()).toEqual(0);
+  });
+
+  it("bypasses caching if no param list is supplied", () => {
+    const cache = new EndpointCache({
+      size: 50,
+    });
+
+    expect(cache.get({ A: "b", B: "b" }, () => endpoint1)).toBe(endpoint1);
+    expect(cache.size()).toEqual(0);
+  });
+
+  it("should be an LRU cache", () => {
+    const cache = new EndpointCache({
+      size: 5,
+      params: ["A", "B"],
+    });
+
+    for (let i = 0; i < 50; ++i) {
+      cache.get({ A: "b", B: "b" + i }, () => endpoint1);
+    }
+
+    const size = cache.size();
+    expect(size).toBeLessThan(16);
+    expect(cache.get({ A: "b", B: "b49" }, () => endpoint2)).toBe(endpoint1);
+    expect(cache.size()).toEqual(size);
+
+    expect(cache.get({ A: "b", B: "b1" }, () => endpoint2)).toBe(endpoint2);
+    expect(cache.size()).toEqual(size + 1);
+  });
+});

packages/util-endpoints/src/cache/EndpointCache.ts

@@ -0,0 +1,78 @@
+import type { EndpointParams, EndpointV2 } from "@smithy/types";
+
+/**
+ * @internal
+ *
+ * Cache for endpoint ruleSet resolution.
+ */
+export class EndpointCache {
+  private capacity: number;
+  private data = new Map<string, EndpointV2>();
+  private parameters: string[] = [];
+
+  /**
+   * @param [size] - desired average maximum capacity. A buffer of 10 additional keys will be allowed
+   *                 before keys are dropped.
+   * @param [params] - list of params to consider as part of the cache key.
+   *
+   * If the params list is not populated, no caching will happen.
+   * This may be out of order depending on how the object is created and arrives to this class.
+   */
+  public constructor({ size, params }: { size?: number; params?: string[] }) {
+    this.capacity = size ?? 50;
+    if (params) {
+      this.parameters = params;
+    }
+  }
+
+  /**
+   * @param endpointParams - query for endpoint.
+   * @param resolver - provider of the value if not present.
+   * @returns endpoint corresponding to the query.
+   */
+  public get(endpointParams: EndpointParams, resolver: () => EndpointV2): EndpointV2 {
+    const key = this.hash(endpointParams);
+    if (key === false) {
+      return resolver();
+    }
+
+    if (!this.data.has(key)) {
+      if (this.data.size > this.capacity + 10) {
+        const keys = this.data.keys();
+        let i = 0;
+        while (true) {
+          const { value, done } = keys.next();
+          this.data.delete(value);
+          if (done || ++i > 10) {
+            break;
+          }
+        }
+      }
+      this.data.set(key, resolver());
+    }
+    return this.data.get(key)!;
+  }
+
+  public size() {
+    return this.data.size;
+  }
+
+  /**
+   * @returns cache key or false if not cachable.
+   */
+  private hash(endpointParams: EndpointParams): string | false {
+    let buffer = "";
+    const { parameters } = this;
+    if (parameters.length === 0) {
+      return false;
+    }
+    for (const param of parameters) {
+      const val = String(endpointParams[param] ?? "");
+      if (val.includes("|;")) {
+        return false;
+      }
+      buffer += val + "|;";
+    }
+    return buffer;
+  }
+}

packages/util-endpoints/src/index.ts

@@ -1,3 +1,4 @@
+export * from "./cache/EndpointCache";
 export * from "./lib/isIpAddress";
 export * from "./lib/isValidHostLabel";
 export * from "./utils/customEndpointFunctions";

scripts/build-generated-test-packages.js

@@ -5,6 +5,8 @@
  */
 
 const path = require("node:path");
+const fs = require("node:fs");
+
 const { spawnProcess } = require("./utils/spawn-process");
 
 const root = path.join(__dirname, "..");
@@ -15,16 +17,6 @@ const codegenTestDir = path.join(testProjectDir, "build", "smithyprojections", "
 
 const weatherClientDir = path.join(codegenTestDir, "source", "typescript-client-codegen");
 
-const releasedClientDir = path.join(
-  testProjectDir,
-  "released-version-test",
-  "build",
-  "smithyprojections",
-  "released-version-test",
-  "source",
-  "typescript-codegen"
-);
-
 // Build generic legacy auth client for integration tests
 const weatherLegacyAuthClientDir = path.join(codegenTestDir, "client-legacy-auth", "typescript-client-codegen");
 
@@ -54,7 +46,27 @@ const buildAndCopyToNodeModules = async (packageName, codegenDir, nodeModulesDir
     // as its own package.
     await spawnProcess("touch", ["yarn.lock"], { cwd: codegenDir });
     await spawnProcess("yarn", { cwd: codegenDir });
+    const smithyPackages = path.join(__dirname, "..", "packages");
+    const node_modules = path.join(codegenDir, "node_modules");
+    const localSmithyPkgs = fs.readdirSync(smithyPackages);
+
+    for (const smithyPkg of localSmithyPkgs) {
+      if (!fs.existsSync(path.join(smithyPackages, smithyPkg, "dist-cjs"))) {
+        continue;
+      }
+      await Promise.all(
+        ["dist-cjs", "dist-types", "dist-es", "package.json"].map((folder) =>
+          spawnProcess("cp", [
+            "-r",
+            path.join(smithyPackages, smithyPkg, folder),
+            path.join(node_modules, "@smithy", smithyPkg),
+          ])
+        )
+      );
+    }
+
     await spawnProcess("yarn", ["build"], { cwd: codegenDir });
+
     // Optionally, after building the package, it's packed and copied to node_modules so that
     // it can be used in integration tests by other packages within the monorepo.
     if (nodeModulesDir != undefined) {
@@ -89,6 +101,18 @@ const buildAndCopyToNodeModules = async (packageName, codegenDir, nodeModulesDir
     httpBearerAuthClientDir,
     nodeModulesDir
   );
+
   // TODO(released-version-test): Test released version of smithy-typescript codegenerators, but currently is not working
+  /*
+  const releasedClientDir = path.join(
+    testProjectDir,
+    "released-version-test",
+    "build",
+    "smithyprojections",
+    "released-version-test",
+    "source",
+    "typescript-codegen"
+  );
+  */
   // await buildAndCopyToNodeModules("released", releasedClientDir, undefined);
 })();

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/endpointsV2/EndpointsV2Generator.java

@@ -21,6 +21,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 import software.amazon.smithy.codegen.core.SymbolDependency;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.node.ObjectNode;
@@ -79,6 +80,7 @@ public List<SymbolDependency> getDependencies() {
     private final EndpointRuleSetTrait endpointRuleSetTrait;
     private final ServiceShape service;
     private final TypeScriptSettings settings;
+    private final RuleSetParameterFinder ruleSetParameterFinder;
 
     public EndpointsV2Generator(
             TypeScriptDelegator delegator,
@@ -90,6 +92,7 @@ public EndpointsV2Generator(
         this.settings = settings;
         endpointRuleSetTrait = service.getTrait(EndpointRuleSetTrait.class)
             .orElseThrow(() -> new RuntimeException("service missing EndpointRuleSetTrait"));
+        ruleSetParameterFinder = new RuleSetParameterFinder(service);
     }
 
     @Override
@@ -114,8 +117,6 @@ private void generateEndpointParameters() {
                     "export interface ClientInputEndpointParameters {",
                     "}",
                     () -> {
-                        RuleSetParameterFinder ruleSetParameterFinder = new RuleSetParameterFinder(service);
-
                         Map<String, String> clientInputParams = ruleSetParameterFinder.getClientContextParams();
                         //Omit Endpoint params that should not be a part of the ClientInputEndpointParameters interface
                         Map<String, String> builtInParams = ruleSetParameterFinder.getBuiltInParams();
@@ -164,10 +165,9 @@ private void generateEndpointParameters() {
                 writer.openBlock(
                     "export const commonParams = {", "} as const",
                     () -> {
-                        RuleSetParameterFinder parameterFinder = new RuleSetParameterFinder(service);
                         Set<String> paramNames = new HashSet<>();
 
-                        parameterFinder.getClientContextParams().forEach((name, type) -> {
+                        ruleSetParameterFinder.getClientContextParams().forEach((name, type) -> {
                             if (!paramNames.contains(name)) {
                                 writer.write(
                                     "$L: { type: \"clientContextParams\", name: \"$L\" },",
@@ -176,7 +176,7 @@ private void generateEndpointParameters() {
                             paramNames.add(name);
                         });
 
-                        parameterFinder.getBuiltInParams().forEach((name, type) -> {
+                        ruleSetParameterFinder.getBuiltInParams().forEach((name, type) -> {
                             if (!paramNames.contains(name)) {
                                 writer.write(
                                     "$L: { type: \"builtInParams\", name: \"$L\" },",
@@ -222,25 +222,29 @@ private void generateEndpointResolver() {
                     Paths.get(".", CodegenUtils.SOURCE_FOLDER, ENDPOINT_FOLDER,
                         ENDPOINT_RULESET_FILE.replace(".ts", "")));
 
-                writer.openBlock(
-                    "export const defaultEndpointResolver = ",
-                    "",
-                    () -> {
-                        writer.openBlock(
-                            "(endpointParams: EndpointParameters, context: { logger?: Logger } = {}): EndpointV2 => {",
-                            "};",
-                            () -> {
-                                writer.openBlock(
-                                    "return resolveEndpoint(ruleSet, {",
-                                    "});",
-                                    () -> {
-                                        writer.write("endpointParams: endpointParams as EndpointParams,");
-                                        writer.write("logger: context.logger,");
-                                    }
-                                );
-                            }
-                        );
-                    }
+                writer.addImport("EndpointCache", null, TypeScriptDependency.UTIL_ENDPOINTS);
+                writer.write("""
+                    const cache = new EndpointCache({
+                        size: 50,
+                        params: [$L]
+                    });
+                    """,
+                    ruleSetParameterFinder.getEffectiveParams()
+                        .stream().collect(Collectors.joining("\",\n \"", "\"", "\""))
+                );
+
+                writer.write(
+                """
+                    export const defaultEndpointResolver = (
+                        endpointParams: EndpointParameters,
+                        context: { logger?: Logger } = {}
+                    ): EndpointV2 => {
+                        return cache.get(endpointParams as EndpointParams, () => resolveEndpoint(ruleSet, {
+                            endpointParams: endpointParams as EndpointParams,
+                            logger: context.logger,
+                        }));
+                    };
+                    """
                 );
             }
         );