@milaboratories/pl-drivers 1.5.63 → 1.5.65

package/src/drivers/download_blob/sparse_cache/cache.test.ts

@@ -0,0 +1,371 @@
+import { SparseCache, type SparseFileStorage, type SparseCacheRanges } from './cache';
+import { ConsoleLoggerAdapter } from '@milaboratories/ts-helpers';
+import type { Ranges } from './ranges';
+import { describe, it, expect } from 'vitest';
+
+/** gen helpers */
+
+/** Generates Uint8Array from string. */
+function genData(string: string) {
+  const enc = new TextEncoder();
+  return enc.encode(string);
+}
+
+/** Converts uint 8 array to string. */
+function toString(arr: Uint8Array) {
+  const dec = new TextDecoder();
+  return dec.decode(arr);
+}
+
+/** In-memory implementation for cache ranges. */
+class InMemoryRanges implements SparseCacheRanges {
+  constructor(
+    public readonly keyToRanges: Record<string, Ranges>,
+  ) {}
+
+  async get(key: string) {
+    const result = this.keyToRanges[key];
+    return result ?? { ranges: [] };
+  }
+
+  async set(key: string, ranges: Ranges) {
+    this.keyToRanges[key] = ranges;
+  }
+
+  async delete(key: string) {
+    delete this.keyToRanges[key];
+  }
+}
+
+/** In-memory implementation for a sparse file. */
+class InMemoryFile implements SparseFileStorage {
+  constructor(
+    public readonly keyToFromToData: Record<string, Record<number, string>>,
+  ) {}
+
+  async all() {
+    return Object.entries(this.keyToFromToData).map(([k, _]) => k);
+  }
+
+  async exists(key: string) {
+    return key in this.keyToFromToData;
+  }
+
+  path(key: string) {
+    return key;
+  }
+
+  async write(key: string, data: Uint8Array, from: number) {
+    this.keyToFromToData[key] = this.keyToFromToData[key] ?? {};
+    this.keyToFromToData[key][from] = toString(data);
+  }
+
+  async delete(key: string) {
+    delete this.keyToFromToData[key];
+  }
+}
+
+function genCache(
+  keyToLastAccessTime: Record<string, Date>,
+  size: number,
+  maxSize: number,
+  ranges: Record<string, Ranges>,
+  storage: Record<string, Record<number, string>>,
+): {
+  cache: SparseCache,
+  ranges: InMemoryRanges,
+  storage: InMemoryFile,
+} {
+  const inMemRanges = new InMemoryRanges(ranges);
+  const inMemStorage = new InMemoryFile(storage);
+
+  const c = new SparseCache(
+    new ConsoleLoggerAdapter(),
+    maxSize,
+    inMemRanges,
+    inMemStorage,
+  );
+  c.keyToLastAccessTime = new Map(Object.entries(keyToLastAccessTime));
+  c.size = size;
+
+  return {
+    cache: c,
+    ranges: inMemRanges,
+    storage: inMemStorage,
+  };
+}
+
+/** Tests */
+
+describe('SparseCache.get', () => {
+  const cases: {
+    name: string;
+    initialCache: SparseCache;
+    key: string;
+    range: { from: number; to: number };
+    expectedPath?: string;
+  }[] = [
+      {
+        name: 'key does not exist in storage',
+        initialCache: genCache({}, 0, 100, {}, {}).cache,
+        key: 'key1',
+        range: { from: 0, to: 10 },
+        expectedPath: undefined,
+      },
+      {
+        name: 'key exists, but range does not exist in cache',
+        initialCache: genCache(
+          { 'key1': new Date() },
+          10,
+          100,
+          { 'key1': { ranges: [{ from: 0, to: 10 }] } },
+          { 'key1': { 0: 'data' } },
+        ).cache,
+        key: 'key1',
+        range: { from: 5, to: 15 },
+        expectedPath: undefined,
+      },
+      {
+        name: 'key exists, and exact range exists in cache',
+        initialCache: genCache(
+          { 'key1': new Date() },
+          10,
+          100,
+          { 'key1': { ranges: [{ from: 0, to: 10 }] } },
+          { 'key1': { 0: 'data' } },
+        ).cache,
+        key: 'key1',
+        range: { from: 0, to: 10 },
+        expectedPath: 'key1',
+      },
+      {
+        name: 'key exists, and sub-range exists in cache',
+        initialCache: genCache(
+          { 'key1': new Date() },
+          20,
+          100,
+          { 'key1': { ranges: [{ from: 0, to: 20 }] } },
+          { 'key1': { 0: 'data' } },
+        ).cache,
+        key: 'key1',
+        range: { from: 5, to: 15 },
+        expectedPath: 'key1',
+      },
+    ];
+
+  for (const tc of cases) {
+    it(tc.name, async () => {
+      const result = await tc.initialCache.get(tc.key, tc.range);
+
+      expect(result).toEqual(tc.expectedPath);
+    });
+  }
+});
+
+describe('SparseCache.setWithoutEviction', () => {
+  it('should add a new key and range', async () => {
+    const c = genCache(
+      {},
+      0,
+      100,
+      {},
+      {},
+    );
+
+    await c.cache.setWithoutEviction('key1', { from: 0, to: 10 }, genData('abc123abc1'));
+
+    expect(c.ranges.keyToRanges).toMatchObject({
+      'key1': { ranges: [{ from: 0, to: 10 }] },
+    });
+    expect(c.storage.keyToFromToData).toMatchObject({
+      'key1': { 0: 'abc123abc1' },
+    });
+    expect(c.cache.size).toEqual(10);
+  });
+
+  it('should overwrite data when adding the same key and range', async () => {
+    const c = genCache(
+      { 'key1': new Date(Date.now() - 1000) },
+      10,
+      100,
+      { 'key1': { ranges: [{ from: 0, to: 10 }] } },
+      { 'key1': { 0: 'data' } },
+    );
+
+    await c.cache.setWithoutEviction('key1', { from: 0, to: 10 }, genData('abc123abc1'));
+
+    expect(c.ranges.keyToRanges).toMatchObject({
+      'key1': { ranges: [{ from: 0, to: 10 }] },
+    });
+    expect(c.storage.keyToFromToData).toMatchObject({
+      'key1': { 0: 'abc123abc1' },
+    });
+    expect(c.cache.size).toEqual(10);
+  });
+
+  it('should add a new non-overlapping range to an existing key', async () => {
+    const c = genCache(
+      { 'key1': new Date(Date.now() - 10000) },
+      10,
+      100,
+      { 'key1': { ranges: [{ from: 0, to: 10 }] } },
+      { 'key1': { 0: 'original  ' } } // 10 chars
+    );
+
+    await c.cache.setWithoutEviction('key1', { from: 20, to: 30 }, genData('new       ')); // 10 chars
+
+    expect(c.ranges.keyToRanges['key1']).toEqual({ ranges: [{ from: 0, to: 10 }, { from: 20, to: 30 }] });
+    expect(c.storage.keyToFromToData['key1']).toMatchObject({
+      0: 'original  ',
+      20: 'new       '
+    });
+    expect(c.cache.size).toEqual(20); // 10 for original + 10 for new
+  });
+
+  it('should add an overlapping range and merge with existing range', async () => {
+    const c = genCache(
+      { 'key1': new Date(Date.now() - 10000) },
+      10,
+      100,
+      { 'key1': { ranges: [{ from: 0, to: 10 }] } }, // 0-10
+      { 'key1': { 0: 'original  ' } }      // data for 0-10
+    );
+
+    // Add range 5-15, overlaps with 0-10
+    await c.cache.setWithoutEviction('key1', { from: 5, to: 15 }, genData('nal1231231')); // data for 5-15
+
+    expect(c.ranges.keyToRanges['key1']).toEqual({ ranges: [{ from: 0, to: 15 }] });
+    expect(c.storage.keyToFromToData['key1']).toMatchObject({
+      0: 'original  ',
+      5: 'nal1231231'
+    });
+    expect(c.cache.size).toEqual(15); // Size of merged range 0-15
+  });
+
+  it('should update size correctly when ranges are modified (subsumed range)', async () => {
+    const c = genCache(
+      { 'key1': new Date() },
+      20,
+      100,
+      { 'key1': { ranges: [{ from: 0, to: 20 }] } },
+      { 'key1': { 0: 'longoriginaldata    ' } }
+    );
+
+    // Add range 5-15, which is a sub-range of 0-20. The total range 0-20 should remain.
+    // The size should remain 20 as normalizeRanges will keep the outer [0, 20] range.
+    await c.cache.setWithoutEviction('key1', { from: 5, to: 15 }, genData('subrange  '));
+
+    expect(c.cache.size).toEqual(20);
+    expect(c.ranges.keyToRanges['key1']).toEqual({ ranges: [{ from: 0, to: 20 }] });
+    expect(c.storage.keyToFromToData['key1']).toMatchObject({
+      0: 'longoriginaldata    ',
+      5: 'subrange  '
+    });
+  });
+});
+
+describe('SparseCache.ensureEvicted', () => {
+  it('should do nothing if cache size is below or equal to maxSize', async () => {
+    const c = genCache(
+      {
+        'key1': new Date(2023, 0, 1, 10, 0, 0), // oldest
+        'key2': new Date(2023, 0, 1, 11, 0, 0), // newest
+      },
+      20,
+      30,
+      {
+        'key1': { ranges: [{ from: 0, to: 10 }] },
+        'key2': { ranges: [{ from: 0, to: 10 }] },
+      },
+      {
+        'key1': { 0: 'data1' },
+        'key2': { 0: 'data2' },
+      }
+    );
+
+    await c.cache.ensureEvicted();
+
+    expect(c.cache.size).toEqual(20);
+    expect(c.ranges.keyToRanges).toMatchObject({
+      'key1': { ranges: [{ from: 0, to: 10 }] },
+      'key2': { ranges: [{ from: 0, to: 10 }] },
+    });
+    expect(c.storage.keyToFromToData).toMatchObject({
+      'key1': { 0: 'data1' },
+      'key2': { 0: 'data2' },
+    });
+  });
+
+  it('should evict the oldest item(s) if cache size is above maxSize', async () => {
+    const c = genCache(
+      {
+        'key1': new Date(2023, 0, 1, 10, 0, 0),
+        'key2': new Date(2023, 0, 1, 11, 0, 0),
+        'key3': new Date(2023, 0, 1, 12, 0, 0),
+      },
+      41,
+      25,
+      {
+        'key1': { ranges: [{ from: 0, to: 11 }] },
+        'key2': { ranges: [{ from: 0, to: 20 }] },
+        'key3': { ranges: [{ from: 0, to: 10 }] },
+      },
+      {
+        'key1': { 0: 'data1' },
+        'key2': { 0: 'data2' },
+        'key3': { 0: 'data3' },
+      }
+    );
+
+    await c.cache.ensureEvicted();
+
+    expect(c.cache.size).toEqual(10);
+    expect(c.ranges.keyToRanges).toMatchObject({
+      'key3': { ranges: [{ from: 0, to: 10 }] },
+    });
+    expect(c.storage.keyToFromToData).toMatchObject({
+      'key3': { 0: 'data3' },
+    });
+  });
+
+  it('should evict all items if necessary to meet maxSize (e.g., maxSize is 0)', async () => {
+    const c = genCache(
+      {
+        'key1': new Date(2023, 0, 1, 10, 0, 0),
+        'key2': new Date(2023, 0, 1, 11, 0, 0),
+      },
+      30, // size = 10 (key1) + 20 (key2)
+      0,  // maxSize
+      {
+        'key1': { ranges: [{ from: 0, to: 10 }] },
+        'key2': { ranges: [{ from: 0, to: 20 }] },
+      },
+      {
+        'key1': { 0: 'data1' },
+        'key2': { 0: 'data2' },
+      }
+    );
+
+    await c.cache.ensureEvicted();
+
+    expect(c.cache.size).toEqual(0);
+    expect(Object.keys(c.ranges.keyToRanges).length).toEqual(0);
+    expect(Object.keys(c.storage.keyToFromToData).length).toEqual(0);
+  });
+
+  it('should handle an empty cache initially', async () => {
+    const c = genCache(
+      {},
+      0,
+      10,
+      {},
+      {}
+    );
+
+    await c.cache.ensureEvicted();
+
+    expect(c.cache.size).toEqual(0);
+    expect(Object.keys(c.ranges.keyToRanges).length).toEqual(0);
+    expect(Object.keys(c.storage.keyToFromToData).length).toEqual(0);
+  });
+});

package/src/drivers/download_blob/sparse_cache/cache.ts

@@ -0,0 +1,240 @@
+import { RangeBytes } from '@milaboratories/pl-model-common';
+import { ensureDirExists, fileExists, mapEntries, MiLogger } from '@milaboratories/ts-helpers';
+import { promises as fs } from 'node:fs';
+import path from 'node:path';
+import { addRange, doesRangeExist, Ranges, rangesFileName, rangesFilePostfix, rangesSize, readRangesFile, writeRangesFile } from './ranges';
+import { writeToSparseFile } from './file';
+import { functions } from '@milaboratories/helpers';
+
+/** The implementer of SparseCacheRanges could throw it if ranges were corrupted. */
+export class CorruptedRangesError extends Error {}
+
+/** Extracted ranges methods to be able to store ranges somewhere else (e.g. in memory for tests). */
+export interface SparseCacheRanges {
+  get(key: string): Promise<Ranges>;
+  set(key: string, ranges: Ranges): Promise<void>;
+  delete(key: string): Promise<void>;
+}
+
+/** Stores ranges in a directory as JSON files (the default implementation). */
+export class SparseCacheFsRanges implements SparseCacheRanges {
+  constructor(
+    private readonly logger: MiLogger,
+    private readonly cacheDir: string,
+  ) {}
+
+  private fPath(key: string): string {
+    return path.join(this.cacheDir, rangesFileName(key));
+  }
+
+  async get(key: string): Promise<Ranges> {
+    return await readRangesFile(this.logger, this.fPath(key));
+  }
+
+  async set(key: string, ranges: Ranges) {
+    return await writeRangesFile(this.logger, this.fPath(key), ranges);
+  }
+
+  async delete(key: string) {
+    await fs.rm(this.fPath(key));
+  }
+}
+
+/** Extracted interface for storing sparse files. */
+export interface SparseFileStorage {
+  all(): Promise<string[]>;
+  exists(key: string): Promise<boolean>;
+  path(key: string): string;
+  write(key: string, data: Uint8Array, from: number): Promise<void>;
+  delete(key: string): Promise<void>;
+}
+
+/** Stores sparse files in a directory (the default implementation). */
+export class SparseCacheFsFile implements SparseFileStorage {
+  private readonly suffix = '.sparse.bin';
+
+  constructor(
+    private readonly logger: MiLogger,
+    private readonly cacheDir: string,
+  ) {}
+
+  async all(): Promise<string[]> {
+    await ensureDirExists(this.cacheDir);
+    const files = await fs.readdir(this.cacheDir);
+    return files.filter((f) => f.endsWith(this.suffix));
+  }
+
+  async exists(key: string): Promise<boolean> {
+    return await fileExists(this.path(key));
+  }
+
+  path(key: string): string {
+    return path.join(this.cacheDir, key + this.suffix);
+  }
+
+  async write(key: string, data: Uint8Array, from: number): Promise<void> {
+    await ensureDirExists(this.cacheDir);
+    await writeToSparseFile(this.logger, process.platform, this.path(key), data, from);
+  }
+
+  async delete(key: string): Promise<void> {
+    await fs.rm(this.path(key));
+  }
+}
+
+/** LRU cache for ranges of sparse files. */
+export class SparseCache {
+  /** Fields are public for tests. */
+
+  /** The lock to make sure cache requests are done one by one. */
+  private lock = new functions.AwaitLock()
+
+  public keyToLastAccessTime = new Map<string, Date>();
+  public size = 0;
+
+  constructor(
+    public readonly logger: MiLogger,
+    /** The hard limit in bytes. */
+    public readonly maxSize: number,
+    public readonly ranges: SparseCacheRanges,
+    public readonly storage: SparseFileStorage,
+  ) {}
+
+  /** Resets a cache's size by rereading everything we already store.
+   * Safe for concurrent use. */
+  async reset() {
+    await withLock(this.lock, async () => {
+      await this.resetUnsafe();
+    })
+  }
+
+  /** Returns a path to the key if the range exists in a cache, otherwise returns undefined.
+   * Safe for concurrent use. */
+  async get(key: string, range: RangeBytes): Promise<string | undefined> {
+    return await withLock(this.lock, async () => {
+      return await this.getUnsafe(key, range);
+    });
+  }
+
+  /** Sets data to the cache's file and clear the cache if it's needed.
+   * Safe for concurrent use. */
+  async set(key: string, range: RangeBytes, data: Uint8Array): Promise<void> {
+    await withLock(this.lock, async () => {
+      await this.setUnsafe(key, range, data);
+    })
+  }
+
+  private async resetUnsafe() {
+    this.size = 0;
+    this.keyToLastAccessTime = new Map<string, Date>();
+
+    const now = new Date();
+    // In rmKey method we first deletes the key from a storage and only then from ranges,
+    // so if something goes wrong between 2 operations,
+    // on reset the logic will be correct.
+    for (const key of await this.storage.all()) {
+      const ranges = await this.ranges.get(key);
+      this.size += rangesSize(ranges);
+      this.keyToLastAccessTime.set(key, now);
+    }
+  }
+
+  private async getUnsafe(key: string, range: RangeBytes): Promise<string | undefined> {
+    // It first checks the storage, and then the ranges.
+    // In another method, when we remove a key, it first deletes a key from a storage and then from ranges,
+    // so if we don't have a key in storage but have it in ranges, the logic here is correct.
+    // We probably could reverse the operations here and there, and everywhere we work with both storage and ranges.
+    if (await this.storage.exists(key)) {
+      this.keyToLastAccessTime.set(key, new Date());
+
+      const ranges = await this.getRanges(key);
+      if (doesRangeExist(ranges, range)) {
+        return this.storage.path(key);
+      }
+
+      return undefined;
+    }
+
+    return undefined;
+  }
+
+  private async setUnsafe(key: string, range: { from: number; to: number; }, data: Uint8Array) {
+    await this.setWithoutEviction(key, range, data);
+    await this.ensureEvicted();
+  }
+
+  /** Sets a key and recalculates a size, but doesn't ensures that the size is less than the hard limit. */
+  async setWithoutEviction(key: string, range: RangeBytes, data: Uint8Array): Promise<void> {
+    if (range.to - range.from !== data.length) {
+      throw new Error(
+        `SparseCache.set: trying to set ${key} with wrong range length: `
+        + `range: ${JSON.stringify(range)}, data: ${data.length}`
+      );
+    }
+
+    this.keyToLastAccessTime.set(key, new Date());
+
+    const ranges = await this.getRanges(key);
+    this.size -= rangesSize(ranges);
+
+    await this.storage.write(key, data, range.from);
+
+    const newRanges = addRange(ranges, range);
+    this.size += rangesSize(newRanges);
+
+    await this.ranges.set(key, newRanges);
+  }
+
+  /** Ensures the size is less than hard limit by deleting the oldest keys. */
+  async ensureEvicted(): Promise<void> {
+    const byTime = mapEntries(this.keyToLastAccessTime);
+    byTime.sort(([_, aDate], [__, bDate]) => bDate.getTime() - aDate.getTime());
+
+    while (this.size > this.maxSize) {
+      const keyAndDate = byTime.pop(); // removes the oldest
+      if (!keyAndDate) {
+        break;
+      }
+      const [key, _] = keyAndDate;
+
+      const ranges = await this.getRanges(key);
+      this.size -= rangesSize(ranges);
+      this.rmKey(key);
+    }
+  }
+
+  /** Gets ranges and if they were corrupted, then remove the file from the cache and reset the cache's size. */
+  private async getRanges(key: string) {
+    try {
+      return await this.ranges.get(key);
+    } catch (e: unknown) {
+      if (e instanceof CorruptedRangesError) {
+        // We need to reset a state of the cache and update current size,
+        // it's the only way to calculate the real size when one of the ranges were corrupted.
+        await this.rmKey(key);
+        await this.resetUnsafe();
+
+        return await this.ranges.get(key);
+      }
+
+      throw e;
+    }
+  }
+
+  /** Removes a key the state of the cache. The size should be updated. */
+  private async rmKey(key: string) {
+    await this.storage.delete(key);
+    await this.ranges.delete(key);
+    this.keyToLastAccessTime.delete(key);
+  }
+}
+
+/** Acquires the lock and executes a callback. */
+async function withLock<T>(lock: functions.AwaitLock, cb: () => Promise<T>): Promise<T> {
+  try {
+    await lock.acquireAsync();
+    return await cb();
+  } finally {
+    lock.release();
+  }
+}