fix: lru cache to support range requests

Author: manzt

src/lru-store.ts

@@ -1,21 +1,55 @@
 import type { Readable } from '@zarrita/storage';
 import QuickLRU from 'quick-lru';
 
-export class LRUCacheStore<S extends Readable> {
-  cache: QuickLRU<string, Promise<Uint8Array | undefined>>;
-  constructor(public store: S, maxSize: number = 100) {
-    this.cache = new QuickLRU({ maxSize });
-  }
-  get(...args: Parameters<S['get']>) {
-    const [key, opts] = args;
-    if (this.cache.has(key)) {
-      return this.cache.get(key)!;
+type RangeQuery =
+  | {
+      offset: number;
+      length: number;
     }
-    const value = Promise.resolve(this.store.get(key, opts)).catch((err) => {
-      this.cache.delete(key);
+  | {
+      suffixLength: number;
+    };
+
+function normalizeKey(key: string, range?: RangeQuery) {
+  if (!range) return key;
+  if ('suffixLength' in range) return `${key}:-${range.suffixLength}`;
+  return `${key}:${range.offset}:${range.offset + range.length - 1}`;
+}
+
+export function lru<S extends Readable>(store: S, maxSize: number = 100) {
+  const cache = new QuickLRU<string, Promise<Uint8Array | undefined>>({ maxSize });
+  let getRange = store.getRange ? store.getRange.bind(store) : undefined;
+  function get(...args: Parameters<S['get']>) {
+    const [key, opts] = args;
+    const cacheKey = normalizeKey(key);
+    const cached = cache.get(cacheKey);
+    if (cached) return cached;
+    const result = Promise.resolve(store.get(key, opts)).catch((err) => {
+      cache.delete(cacheKey);
       throw err;
     });
-    this.cache.set(key, value);
-    return value;
+    cache.set(cacheKey, result);
+    return result;
+  }
+  if (getRange) {
+    getRange = (...args: Parameters<NonNullable<S['getRange']>>) => {
+      const [key, range, opts] = args;
+      const cacheKey = normalizeKey(key, range);
+      const cached = cache.get(cacheKey);
+      if (cached) return cached;
+      const result = Promise.resolve(getRange!(key, range, opts)).catch((err) => {
+        cache.delete(cacheKey);
+        throw err;
+      });
+      cache.set(cacheKey, result);
+      return result;
+    };
   }
+  return new Proxy(store, {
+    get(target, prop, receiver) {
+      if (prop === 'get') return get;
+      if (prop === 'getRange' && getRange) return getRange;
+      return Reflect.get(target, prop, receiver);
+    },
+  });
 }

src/utils.ts

@@ -1,11 +1,12 @@
 import type { ZarrPixelSource } from '@hms-dbmi/viv';
-import { Matrix4 } from 'math.gl';
-import { LRUCacheStore } from './lru-store';
-import type { ViewState } from './state';
-
+import type { Slice } from '@zarrita/indexing';
 import * as zarr from '@zarrita/core';
-import { slice, get, type Slice } from '@zarrita/indexing';
+import { slice, get } from '@zarrita/indexing';
 import { FetchStore, Readable } from '@zarrita/storage';
+import { Matrix4 } from 'math.gl';
+
+import { lru } from './lru-store';
+import type { ViewState } from './state';
 
 export const MAX_CHANNELS = 6;
 
@@ -38,7 +39,7 @@ async function normalizeStore(source: string | Readable): Promise<zarr.Location<
     }
 
     // Wrap remote stores in a cache
-    return zarr.root(new LRUCacheStore(store));
+    return zarr.root(lru(store));
   }
 
   return zarr.root(source);