Merge pull request #1036 from polywrap/namesty/cache-improvements

Allow custom cache object to be passed to client config

packages/js/client/src/PolywrapClient.ts

@@ -45,6 +45,7 @@ import {
   PluginPackage,
   RunOptions,
   GetManifestOptions,
+  SimpleCache,
 } from "@polywrap/core-js";
 import { msgpackEncode, msgpackDecode } from "@polywrap/msgpack-js";
 import { WrapManifest } from "@polywrap/wrap-manifest-types-js";
@@ -56,11 +57,7 @@ export interface PolywrapClientConfig<TUri extends Uri | string = string>
 }
 
 export class PolywrapClient implements Client {
-  // TODO: the Wrapper cache needs to be more like a routing table.
-  // It should help us keep track of what URI's map to what Wrappers,
-  // and handle cases where the are multiple jumps. For example, if
-  // A => B => C, then the cache should have A => C, and B => C.
-  private _wrapperCache: WrapperCache = new Map<string, Wrapper>();
+  private _wrapperCache: WrapperCache;
   private _config: PolywrapClientConfig<Uri> = {
     redirects: [],
     plugins: [],
@@ -97,6 +94,14 @@ export class PolywrapClient implements Client {
           uriResolvers: config.uriResolvers ?? [],
           tracingEnabled: !!config.tracingEnabled,
         };
+
+        if (config.wrapperCache) {
+          this._wrapperCache = config.wrapperCache;
+        }
+      }
+
+      if (!this._wrapperCache) {
+        this._wrapperCache = new SimpleCache();
       }
 
       if (!options?.noDefaults) {
@@ -486,9 +491,8 @@ export class PolywrapClient implements Client {
 
     // Update cache for all URIs in the chain
     if (cacheWrite && wrapper) {
-      for (const item of uriHistory.getResolutionPath().stack) {
-        this._wrapperCache.set(item.sourceUri.uri, wrapper);
-      }
+      const uris = uriHistory.getResolutionPath().stack.map((x) => x.sourceUri);
+      this._wrapperCache.set(uris, wrapper);
     }
 
     if (shouldClearContext) {

packages/js/client/src/index.ts

@@ -2,4 +2,5 @@ export * from "./PolywrapClient";
 export * from "./createPolywrapClient";
 export * from "./default-client-config";
 export * from "./wasm";
+export * from "./plugin";
 export * from "@polywrap/core-js";

packages/js/client/src/plugin/index.ts

@@ -0,0 +1 @@
+export * from "./PluginWrapper";

packages/js/core/src/__tests__/resolveUri.spec.ts

@@ -23,6 +23,8 @@ import {
   SubscribeOptions,
   Subscription,
   PluginPackage,
+  SimpleCache,
+  GetManifestOptions,
 } from "..";
 
 import { WrapManifest } from "@polywrap/wrap-manifest-types-js";
@@ -106,7 +108,7 @@ describe("resolveUri", () => {
           encoded: false
         }),
       getFile: (options: GetFileOptions, client: Client) => Promise.resolve(""),
-      getManifest: (client: Client) => Promise.resolve({} as WrapManifest)
+      getManifest: (options: GetManifestOptions, client: Client) => Promise.resolve({} as WrapManifest)
     };
   };
 
@@ -238,7 +240,7 @@ describe("resolveUri", () => {
       new Uri("ens/test.eth"),
       uriResolvers,
       client(wrappers, plugins, interfaces),
-      new Map<string, Wrapper>(),
+      new SimpleCache(),
     );
 
     expect(result.wrapper).toBeTruthy();
@@ -260,7 +262,7 @@ describe("resolveUri", () => {
       new Uri("my/something-different"),
       uriResolvers,
       client(wrappers, plugins, interfaces),
-      new Map<string, Wrapper>(),
+      new SimpleCache(),
     );
 
     expect(result.wrapper).toBeTruthy();
@@ -282,7 +284,7 @@ describe("resolveUri", () => {
       new Uri("ens/ens"),
       uriResolvers,
       client(wrappers, plugins, interfaces),
-      new Map<string, Wrapper>(),
+      new SimpleCache(),
     );
 
     expect(result.wrapper).toBeTruthy();
@@ -304,7 +306,7 @@ describe("resolveUri", () => {
       new Uri("my/something-different"),
       uriResolvers,
       client(wrappers, plugins, interfaces),
-      new Map<string, Wrapper>(),
+      new SimpleCache(),
     );
 
     expect(result.wrapper).toBeTruthy();
@@ -339,7 +341,7 @@ describe("resolveUri", () => {
       new Uri("some/wrapper"),
       uriResolvers,
       client(wrappers, plugins, interfaces, circular),
-      new Map<string, Wrapper>(),
+      new SimpleCache(),
     ).catch((e: Error) =>
       expect(e.message).toMatch(/Infinite loop while resolving URI/)
     );
@@ -363,7 +365,7 @@ describe("resolveUri", () => {
       new Uri("some/wrapper"),
       uriResolvers,
       client(wrappers, plugins, interfaces, missingFromProperty),
-      new Map<string, Wrapper>(),
+      new SimpleCache(),
     ).catch((e: Error) =>
       expect(e.message).toMatch(
         "Redirect missing the from property.\nEncountered while resolving wrap://some/wrapper"
@@ -387,7 +389,7 @@ describe("resolveUri", () => {
       new Uri("some/wrapper"),
       uriResolvers,
       client(wrappers, pluginRegistrations, interfaces),
-      new Map<string, Wrapper>(),
+      new SimpleCache(),
     );
 
     expect(result.wrapper).toBeTruthy();
@@ -425,7 +427,7 @@ describe("resolveUri", () => {
         plugins, 
         interfaces
       ),
-      new Map<string, Wrapper>(),
+      new SimpleCache(),
     );
 
     expect(resolvedUri).toEqual(uri);

packages/js/core/src/cache/SimpleWrapperCache.ts

@@ -0,0 +1,23 @@
+import { Uri, Wrapper, WrapperCache } from "../types";
+
+export class SimpleCache implements WrapperCache {
+  private _map: Map<string, Wrapper> = new Map();
+
+  get(uri: Uri): Wrapper | undefined {
+    return this._map.get(uri.uri);
+  }
+
+  has(uri: Uri): boolean {
+    return this._map.has(uri.uri);
+  }
+
+  set(uris: Uri | Uri[], wrapper: Wrapper): void {
+    if (Array.isArray(uris)) {
+      for (const uri of uris) {
+        this._map.set(uri.uri, wrapper);
+      }
+    } else {
+      this._map.set(uris.uri, wrapper);
+    }
+  }
+}

packages/js/core/src/cache/index.ts

@@ -0,0 +1 @@
+export * from "./SimpleWrapperCache";

packages/js/core/src/index.ts

@@ -1,5 +1,6 @@
 export * from "./types";
 export * from "./algorithms";
+export * from "./cache";
 export * from "./interfaces";
 export * from "./uri-resolution/core";
 export * from "./uri-resolution/resolvers";

packages/js/core/src/types/Client.ts

@@ -11,6 +11,7 @@ import {
 } from "./";
 import { UriResolver } from "../uri-resolution/core";
 import { UriResolverHandler } from "./UriResolver";
+import { WrapperCache } from "./WrapperCache";
 
 import { WrapManifest } from "@polywrap/wrap-manifest-types-js";
 
@@ -20,6 +21,7 @@ export interface ClientConfig<TUri extends Uri | string = string> {
   interfaces: InterfaceImplementations<TUri>[];
   envs: Env<TUri>[];
   uriResolvers: UriResolver[];
+  wrapperCache?: WrapperCache;
 }
 
 export interface Contextualized {

packages/js/core/src/types/Wrapper.ts

@@ -53,5 +53,3 @@ export abstract class Wrapper implements Invocable {
     client: Client
   ): Promise<WrapManifest>;
 }
-
-export type WrapperCache = Map<string, Wrapper>;

packages/js/core/src/types/WrapperCache.ts

@@ -0,0 +1,8 @@
+import { Uri, Wrapper } from ".";
+
+export interface WrapperCache {
+  get(uri: Uri): Wrapper | undefined;
+  has(uri: Uri): boolean;
+  set(uris: Uri[], wrapper: Wrapper): void;
+  set(uri: Uri, wrapper: Wrapper): void;
+}

packages/js/core/src/types/index.ts

@@ -12,3 +12,4 @@ export * from "./InterfaceImplementations";
 export * from "./PluginRegistration";
 export * from "./UriResolver";
 export * from "./Workflow";
+export * from "./WrapperCache";

packages/js/core/src/uri-resolution/resolvers/cache/CacheResolver.ts

@@ -11,7 +11,7 @@ export class CacheResolver implements UriResolver {
     client: Client,
     cache: WrapperCache
   ): Promise<UriResolutionResult> {
-    const wrapper = cache.get(uri.uri);
+    const wrapper = cache.get(uri);
 
     return Promise.resolve({
       uri: uri,

packages/js/core/src/uri-resolution/resolvers/extendable/ExtendableUriResolver.ts

@@ -105,7 +105,7 @@ export class ExtendableUriResolver implements UriResolver {
     const implementationsToLoad = new Queue<Uri>();
 
     for (const implementationUri of implementationUris) {
-      if (!cache.has(implementationUri.uri)) {
+      if (!cache.has(implementationUri)) {
         implementationsToLoad.enqueue(implementationUri);
       }
     }
@@ -148,7 +148,7 @@ export class ExtendableUriResolver implements UriResolver {
           };
         }
       } else {
-        cache.set(implementationUri.uri, wrapper);
+        cache.set(implementationUri, wrapper);
         loadedImplementations.push(implementationUri.uri);
         failedAttempts = 0;
       }