refactor: cache

Author: webfansplz

packages/kit/src/client/rpc.ts

@@ -12,12 +12,60 @@ function isNumeric(str: string | number | undefined) {
   return `${+str}` === `${str}`
 }
 
+interface RpcCacheOptions {
+  functions: string[]
+}
+
+// @TODO: should be moved to birpc-x?
+class RpcCacheManager {
+  private cacheMap = new Map<string, Map<string, unknown>>()
+  private options: RpcCacheOptions
+
+  constructor(options: RpcCacheOptions) {
+    this.options = options
+  }
+
+  updateOptions(options: Partial<RpcCacheOptions>) {
+    this.options = {
+      ...this.options,
+      ...options,
+    }
+  }
+
+  cached(m: string, a: unknown[]) {
+    const methodCache = this.cacheMap.get(m)
+    if (methodCache) {
+      return methodCache.get(JSON.stringify(a))
+    }
+    return undefined
+  }
+
+  apply(req: { m: string, a: unknown[] }, res: unknown) {
+    const methodCache = this.cacheMap.get(req.m) || new Map<string, unknown>()
+    methodCache.set(JSON.stringify(req.a), res)
+    this.cacheMap.set(req.m, methodCache)
+  }
+
+  validate(m: string) {
+    return this.options.functions.includes(m)
+  }
+
+  invalidate(key?: string) {
+    if (key) {
+      this.cacheMap.delete(key)
+    }
+    else {
+      this.cacheMap.clear()
+    }
+  }
+}
+
 export interface DevToolsRpcClientOptions {
   connectionMeta?: ConnectionMeta
   baseURL?: string[]
-  cacheResponse?: boolean
   wsOptions?: Partial<WebSocketRpcClientOptions>
   rpcOptions?: Partial<BirpcOptions<DevToolsRpcServerFunctions>>
+  cacheOptions?: boolean | Partial<RpcCacheOptions>
 }
 
 export type DevToolsRpcClient = BirpcReturn<DevToolsRpcServerFunctions, DevToolsRpcClientFunctions>
@@ -26,19 +74,29 @@ export interface ClientRpcReturn {
   connectionMeta: ConnectionMeta
   rpc: DevToolsRpcClient
   clientRpc: DevToolsClientRpcHost
-  invalidateCache: () => void
 }
 
+export async function getDevToolsRpcClient(
+  options: DevToolsRpcClientOptions & { cacheOptions: false },
+): Promise<ClientRpcReturn>
+export async function getDevToolsRpcClient(
+  options: DevToolsRpcClientOptions & { cacheOptions: true },
+): Promise<ClientRpcReturn & { cacheManager: RpcCacheManager }>
+export async function getDevToolsRpcClient(
+  options: DevToolsRpcClientOptions & { cacheOptions: Partial<RpcCacheOptions> },
+): Promise<ClientRpcReturn & { cacheManager: RpcCacheManager }>
+export async function getDevToolsRpcClient(
+  options?: DevToolsRpcClientOptions,
+): Promise<ClientRpcReturn>
 export async function getDevToolsRpcClient(
   options: DevToolsRpcClientOptions = {},
 ): Promise<ClientRpcReturn> {
   const {
     baseURL = '/.devtools/',
     rpcOptions = {},
-    cacheResponse = false,
+    cacheOptions = false,
   } = options
   const urls = Array.isArray(baseURL) ? baseURL : [baseURL]
-  const responseCacheMap = new Map<string, unknown>()
   let connectionMeta: ConnectionMeta | undefined = options.connectionMeta
 
   if (!connectionMeta) {
@@ -64,6 +122,7 @@ export async function getDevToolsRpcClient(
     ? `${location.protocol.replace('http', 'ws')}//${location.hostname}:${connectionMeta.websocket}`
     : connectionMeta.websocket as string
 
+  const cacheManager = cacheOptions ? new RpcCacheManager({ functions: [], ...(typeof options.cacheOptions === 'object' ? options.cacheOptions : {}) }) : null
   const context: DevToolsClientContext = {
     rpc: undefined!,
   }
@@ -79,13 +138,14 @@ export async function getDevToolsRpcClient(
         ...rpcOptions,
         onRequest: async (req, next, resolve) => {
           await rpcOptions.onRequest?.(req, next, resolve)
-          if (cacheResponse) {
-            const cacheKey = `${req.m}-${JSON.stringify(req.a)}`
-            if (responseCacheMap.has(cacheKey)) {
-              resolve(responseCacheMap.get(cacheKey))
+          if (cacheOptions && cacheManager?.validate(req.m)) {
+            const cached = cacheManager.cached(req.m, req.a)
+            if (cached) {
+              return resolve(cached)
             }
             else {
-              responseCacheMap.set(cacheKey, await next(req))
+              const res = await next(req)
+              cacheManager?.apply(req, res)
             }
           }
           else {
@@ -98,19 +158,10 @@ export async function getDevToolsRpcClient(
   // @ts-expect-error assign to readonly property
   context.rpc = rpc
 
-  function invalidateCache(key?: string) {
-    if (key) {
-      responseCacheMap.delete(key)
-    }
-    else {
-      responseCacheMap.clear()
-    }
-  }
-
   return {
     connectionMeta,
     rpc,
     clientRpc,
-    invalidateCache,
+    ...(cacheOptions ? { cacheManager } : {}),
   }
 }

packages/vite/src/app/composables/rpc.ts

@@ -23,7 +23,7 @@ export async function connect() {
         '/.devtools/',
         runtimeConfig.app.baseURL,
       ],
-      cacheResponse: true,
+      cacheOptions: true,
       connectionMeta: runtimeConfig.app.connection,
       wsOptions: {
         onConnected: () => {
@@ -45,6 +45,16 @@ export async function connect() {
     })
 
     rpc.value = result.rpc
+
+    const functions = await rpc.value.$call('vite:core:list-rpc-functions')
+
+    // TODO: add cacheable option to birpc-x and use it here
+    // @ts-expect-error skip type check
+    const cacheableFunctions = Object.keys(functions).filter(name => functions[name]?.cacheable)
+    result.cacheManager.updateOptions({
+      functions: [...cacheableFunctions],
+    })
+
     connectionState.connected = true
   }
   catch (e) {