@aztec/kv-store 0.0.0-test.0

package/src/lmdb-v2/store.ts

@@ -0,0 +1,237 @@
+import { type Logger, createLogger } from '@aztec/foundation/log';
+import { Semaphore, SerialQueue } from '@aztec/foundation/queue';
+import { MsgpackChannel, NativeLMDBStore } from '@aztec/native';
+
+import { AsyncLocalStorage } from 'async_hooks';
+import { rm } from 'fs/promises';
+
+import type { AztecAsyncArray } from '../interfaces/array.js';
+import type { Key, StoreSize } from '../interfaces/common.js';
+import type { AztecAsyncCounter } from '../interfaces/counter.js';
+import type { AztecAsyncMap, AztecAsyncMultiMap } from '../interfaces/map.js';
+import type { AztecAsyncSet } from '../interfaces/set.js';
+import type { AztecAsyncSingleton } from '../interfaces/singleton.js';
+import type { AztecAsyncKVStore } from '../interfaces/store.js';
+import { LMDBArray } from './array.js';
+import { LMDBMap, LMDBMultiMap } from './map.js';
+import {
+  Database,
+  type LMDBMessageChannel,
+  LMDBMessageType,
+  type LMDBRequestBody,
+  type LMDBResponseBody,
+} from './message.js';
+import { ReadTransaction } from './read_transaction.js';
+import { LMDBSingleValue } from './singleton.js';
+import { WriteTransaction } from './write_transaction.js';
+
+export class AztecLMDBStoreV2 implements AztecAsyncKVStore, LMDBMessageChannel {
+  private open = false;
+  private channel: MsgpackChannel<LMDBMessageType, LMDBRequestBody, LMDBResponseBody>;
+  private writerCtx = new AsyncLocalStorage<WriteTransaction>();
+  private writerQueue = new SerialQueue();
+  private availableCursors: Semaphore;
+
+  private constructor(
+    private dataDir: string,
+    mapSize: number,
+    maxReaders: number,
+    private log: Logger,
+    private cleanup?: () => Promise<void>,
+  ) {
+    this.log.info(`Starting data store with maxReaders ${maxReaders}`);
+    this.channel = new MsgpackChannel(new NativeLMDBStore(dataDir, mapSize, maxReaders));
+    // leave one reader to always be available for regular, atomic, reads
+    this.availableCursors = new Semaphore(maxReaders - 1);
+  }
+
+  public get dataDirectory(): string {
+    return this.dataDir;
+  }
+
+  private async start() {
+    this.writerQueue.start();
+
+    await this.channel.sendMessage(LMDBMessageType.OPEN_DATABASE, {
+      db: Database.DATA,
+      uniqueKeys: true,
+    });
+
+    await this.channel.sendMessage(LMDBMessageType.OPEN_DATABASE, {
+      db: Database.INDEX,
+      uniqueKeys: false,
+    });
+
+    this.open = true;
+  }
+
+  public static async new(
+    dataDir: string,
+    dbMapSizeKb: number = 10 * 1024 * 1024,
+    maxReaders: number = 16,
+    cleanup?: () => Promise<void>,
+    log = createLogger('kv-store:lmdb-v2'),
+  ) {
+    const db = new AztecLMDBStoreV2(dataDir, dbMapSizeKb, maxReaders, log, cleanup);
+    await db.start();
+    return db;
+  }
+
+  public getReadTx(): ReadTransaction {
+    if (!this.open) {
+      throw new Error('Store is closed');
+    }
+    return new ReadTransaction(this);
+  }
+
+  public getCurrentWriteTx(): WriteTransaction | undefined {
+    if (!this.open) {
+      throw new Error('Store is closed');
+    }
+    const currentWrite = this.writerCtx.getStore();
+    return currentWrite;
+  }
+
+  openMap<K extends Key, V>(name: string): AztecAsyncMap<K, V> {
+    return new LMDBMap(this, name);
+  }
+
+  openMultiMap<K extends Key, V>(name: string): AztecAsyncMultiMap<K, V> {
+    return new LMDBMultiMap(this, name);
+  }
+
+  openSingleton<T>(name: string): AztecAsyncSingleton<T> {
+    return new LMDBSingleValue(this, name);
+  }
+
+  openArray<T>(name: string): AztecAsyncArray<T> {
+    return new LMDBArray(this, name);
+  }
+
+  openSet<K extends Key>(_name: string): AztecAsyncSet<K> {
+    throw new Error('Not implemented');
+  }
+
+  openCounter<K extends Key>(_name: string): AztecAsyncCounter<K> {
+    throw new Error('Not implemented');
+  }
+
+  async transactionAsync<T extends Exclude<any, Promise<any>>>(
+    callback: (tx: WriteTransaction) => Promise<T>,
+  ): Promise<T> {
+    if (!this.open) {
+      throw new Error('Store is closed');
+    }
+
+    // transactionAsync might be called recursively
+    // send any writes to the parent tx, but don't close it
+    // if the callback throws then the parent tx will rollback automatically
+    const currentTx = this.getCurrentWriteTx();
+    if (currentTx) {
+      return await callback(currentTx);
+    }
+
+    return this.writerQueue.put(async () => {
+      const tx = new WriteTransaction(this);
+      try {
+        const res = await this.writerCtx.run(tx, callback, tx);
+        await tx.commit();
+        return res;
+      } catch (err) {
+        this.log.error(`Failed to commit transaction`, err);
+        throw err;
+      } finally {
+        tx.close();
+      }
+    });
+  }
+
+  clear(): Promise<void> {
+    return Promise.resolve();
+  }
+
+  fork(): Promise<AztecAsyncKVStore> {
+    throw new Error('Not implemented');
+  }
+
+  async delete(): Promise<void> {
+    await this.close();
+    await rm(this.dataDir, { recursive: true, force: true, maxRetries: 3 });
+    this.log.verbose(`Deleted database files at ${this.dataDir}`);
+    await this.cleanup?.();
+  }
+
+  async close() {
+    if (!this.open) {
+      // already closed
+      return;
+    }
+    this.open = false;
+    await this.writerQueue.cancel();
+    await this.channel.sendMessage(LMDBMessageType.CLOSE, undefined);
+  }
+
+  public async sendMessage<T extends LMDBMessageType>(
+    msgType: T,
+    body: LMDBRequestBody[T],
+  ): Promise<LMDBResponseBody[T]> {
+    if (!this.open) {
+      throw new Error('Store is closed');
+    }
+
+    if (msgType === LMDBMessageType.START_CURSOR) {
+      await this.availableCursors.acquire();
+    }
+
+    let response: LMDBResponseBody[T] | undefined = undefined;
+    try {
+      ({ response } = await this.channel.sendMessage(msgType, body));
+      return response;
+    } finally {
+      if (
+        (msgType === LMDBMessageType.START_CURSOR && response === undefined) ||
+        msgType === LMDBMessageType.CLOSE_CURSOR ||
+        // it's possible for a START_CURSOR command to not return a cursor (e.g. db is empty)
+        (msgType === LMDBMessageType.START_CURSOR &&
+          typeof (response as LMDBResponseBody[LMDBMessageType.START_CURSOR]).cursor !== 'number')
+      ) {
+        this.availableCursors.release();
+      }
+    }
+  }
+
+  public async estimateSize(): Promise<StoreSize> {
+    const resp = await this.sendMessage(LMDBMessageType.STATS, undefined);
+    return {
+      mappingSize: Number(resp.dbMapSizeBytes),
+      actualSize: resp.stats.reduce((s, db) => Number(db.totalUsedSize) + s, 0),
+      numItems: resp.stats.reduce((s, db) => Number(db.numDataItems) + s, 0),
+    };
+  }
+}
+
+export function execInWriteTx<T>(store: AztecLMDBStoreV2, fn: (tx: WriteTransaction) => Promise<T>): Promise<T> {
+  const currentWrite = store.getCurrentWriteTx();
+  if (currentWrite) {
+    return fn(currentWrite);
+  } else {
+    return store.transactionAsync(fn);
+  }
+}
+
+export async function execInReadTx<T>(
+  store: AztecLMDBStoreV2,
+  fn: (tx: ReadTransaction) => T | Promise<T>,
+): Promise<T> {
+  const currentWrite = store.getCurrentWriteTx();
+  if (currentWrite) {
+    return await fn(currentWrite);
+  } else {
+    const tx = store.getReadTx();
+    try {
+      return await fn(tx);
+    } finally {
+      tx.close();
+    }
+  }
+}

package/src/lmdb-v2/utils.ts

@@ -0,0 +1,150 @@
+import { MAXIMUM_KEY, fromBufferKey, toBufferKey } from 'ordered-binary';
+
+import type { Key } from '../interfaces/common.js';
+
+type Cmp<T> = (a: T, b: T) => -1 | 0 | 1;
+
+export function dedupeSortedArray<T>(arr: T[], cmp: Cmp<T>): void {
+  for (let i = 0; i < arr.length; i++) {
+    let j = i + 1;
+    for (; j < arr.length; j++) {
+      const res = cmp(arr[i], arr[j]);
+      if (res === 0) {
+        continue;
+      } else if (res < 0) {
+        break;
+      } else {
+        throw new Error('Array not sorted');
+      }
+    }
+
+    if (j - i > 1) {
+      arr.splice(i + 1, j - i - 1);
+    }
+  }
+}
+
+export function insertIntoSortedArray<T>(arr: T[], item: T, cmp: (a: T, b: T) => number): void {
+  let left = 0;
+  let right = arr.length;
+
+  while (left < right) {
+    const mid = (left + right) >> 1;
+    const comparison = cmp(arr[mid], item);
+
+    if (comparison < 0) {
+      left = mid + 1;
+    } else {
+      right = mid;
+    }
+  }
+
+  arr.splice(left, 0, item);
+}
+
+export function findIndexInSortedArray<T, N>(values: T[], needle: N, cmp: (a: T, b: N) => number): number {
+  let start = 0;
+  let end = values.length - 1;
+
+  while (start <= end) {
+    const mid = start + (((end - start) / 2) | 0);
+    const res = cmp(values[mid], needle);
+    if (res === 0) {
+      return mid;
+    } else if (res > 0) {
+      end = mid - 1;
+    } else {
+      start = mid + 1;
+    }
+  }
+
+  return -1;
+}
+
+export function findInSortedArray<T, N>(values: T[], needle: N, cmp: (a: T, b: N) => number): T | undefined {
+  const idx = findIndexInSortedArray(values, needle, cmp);
+  return idx > -1 ? values[idx] : undefined;
+}
+
+export function removeAnyOf<T, N>(arr: T[], vals: N[], cmp: (a: T, b: N) => -1 | 0 | 1): void {
+  let writeIdx = 0;
+  let readIdx = 0;
+  let valIdx = 0;
+
+  while (readIdx < arr.length && valIdx < vals.length) {
+    const comparison = cmp(arr[readIdx], vals[valIdx]);
+
+    if (comparison < 0) {
+      arr[writeIdx++] = arr[readIdx++];
+    } else if (comparison > 0) {
+      valIdx++;
+    } else {
+      readIdx++;
+    }
+  }
+
+  while (readIdx < arr.length) {
+    arr[writeIdx++] = arr[readIdx++];
+  }
+
+  arr.length = writeIdx;
+}
+
+export function removeFromSortedArray<T, N>(arr: T[], val: N, cmp: (a: T, b: N) => -1 | 0 | 1) {
+  const idx = findIndexInSortedArray(arr, val, cmp);
+  if (idx > -1) {
+    arr.splice(idx, 1);
+  }
+}
+
+export function merge<T>(arr: T[], toInsert: T[], cmp: (a: T, b: T) => -1 | 0 | 1): void {
+  const result = new Array<T>(arr.length + toInsert.length);
+  let i = 0,
+    j = 0,
+    k = 0;
+
+  while (i < arr.length && j < toInsert.length) {
+    result[k++] = cmp(arr[i], toInsert[j]) <= 0 ? arr[i++] : toInsert[j++];
+  }
+
+  while (i < arr.length) {
+    result[k++] = arr[i++];
+  }
+  while (j < toInsert.length) {
+    result[k++] = toInsert[j++];
+  }
+
+  for (i = 0; i < result.length; i++) {
+    arr[i] = result[i];
+  }
+  arr.length = result.length;
+}
+
+export function keyCmp(a: [Uint8Array, Uint8Array[] | null], b: [Uint8Array, Uint8Array[] | null]): -1 | 0 | 1 {
+  return Buffer.compare(a[0], b[0]);
+}
+
+export function singleKeyCmp(a: [Uint8Array, Uint8Array[] | null], b: Uint8Array): -1 | 0 | 1 {
+  return Buffer.compare(a[0], b);
+}
+
+export function minKey(prefix: string) {
+  return toBufferKey([prefix]);
+}
+
+export function maxKey(prefix: string) {
+  return toBufferKey([prefix, MAXIMUM_KEY]);
+}
+
+export function serializeKey(prefix: string, key: Key): Buffer {
+  return toBufferKey([prefix, key]);
+}
+
+export function deserializeKey<K extends Key>(prefix: string, key: Uint8Array): K | false {
+  const buf = Buffer.from(key);
+  const parsed = fromBufferKey(buf);
+  if (!Array.isArray(parsed) || parsed[0] !== prefix) {
+    return false;
+  }
+  return parsed[1] as K;
+}

package/src/lmdb-v2/write_transaction.ts

@@ -0,0 +1,314 @@
+import { type Batch, Database, LMDBMessageType } from './message.js';
+import { ReadTransaction } from './read_transaction.js';
+import {
+  dedupeSortedArray,
+  findInSortedArray,
+  findIndexInSortedArray,
+  insertIntoSortedArray,
+  keyCmp,
+  merge,
+  removeAnyOf,
+  removeFromSortedArray,
+  singleKeyCmp,
+} from './utils.js';
+
+export class WriteTransaction extends ReadTransaction {
+  // exposed for tests
+  public readonly dataBatch: Batch = {
+    addEntries: [],
+    removeEntries: [],
+  };
+  public readonly indexBatch: Batch = {
+    addEntries: [],
+    removeEntries: [],
+  };
+
+  set(key: Uint8Array, value: Uint8Array): Promise<void> {
+    this.assertIsOpen();
+
+    const addEntry = findInSortedArray(this.dataBatch.addEntries, key, singleKeyCmp);
+    if (!addEntry) {
+      insertIntoSortedArray(this.dataBatch.addEntries, [key, [value]], keyCmp);
+    } else {
+      addEntry[1] = [value];
+    }
+
+    const removeEntryIndex = findIndexInSortedArray(this.dataBatch.removeEntries, key, singleKeyCmp);
+    if (removeEntryIndex > -1) {
+      this.dataBatch.removeEntries.splice(removeEntryIndex, 1);
+    }
+
+    return Promise.resolve();
+  }
+
+  remove(key: Uint8Array): Promise<void> {
+    const removeEntryIndex = findIndexInSortedArray(this.dataBatch.removeEntries, key, singleKeyCmp);
+    if (removeEntryIndex === -1) {
+      this.dataBatch.removeEntries.push([key, null]);
+    }
+
+    const addEntryIndex = findIndexInSortedArray(this.dataBatch.addEntries, key, singleKeyCmp);
+    if (addEntryIndex > -1) {
+      this.dataBatch.addEntries.splice(addEntryIndex, 1);
+    }
+
+    return Promise.resolve();
+  }
+
+  public override async get(key: Buffer): Promise<Uint8Array | undefined> {
+    this.assertIsOpen();
+
+    const addEntry = findInSortedArray(this.dataBatch.addEntries, key, singleKeyCmp);
+    if (addEntry) {
+      return addEntry[1][0];
+    }
+    const removeEntryIdx = findIndexInSortedArray(this.dataBatch.removeEntries, key, singleKeyCmp);
+    if (removeEntryIdx > -1) {
+      return undefined;
+    }
+
+    return await super.get(key);
+  }
+
+  setIndex(key: Buffer, ...values: Buffer[]): Promise<void> {
+    this.assertIsOpen();
+
+    const addEntries = findInSortedArray(this.indexBatch.addEntries, key, singleKeyCmp);
+    const removeEntries = findInSortedArray(this.indexBatch.removeEntries, key, singleKeyCmp);
+
+    if (removeEntries) {
+      if (removeEntries[1]) {
+        // check if we were deleting these values and update
+        removeAnyOf(removeEntries[1], values, Buffer.compare);
+      }
+
+      if (!removeEntries[1] || removeEntries[1].length === 0) {
+        // either we were deleting the entire key previously
+        // or after cleaning up duplicates, we don't have anything else to delete
+        removeFromSortedArray(this.indexBatch.removeEntries, removeEntries, keyCmp);
+      }
+    }
+
+    if (addEntries) {
+      merge(addEntries[1], values, Buffer.compare);
+      dedupeSortedArray(addEntries[1], Buffer.compare);
+    } else {
+      insertIntoSortedArray(this.indexBatch.addEntries, [key, values], keyCmp);
+    }
+
+    return Promise.resolve();
+  }
+
+  removeIndex(key: Buffer, ...values: Buffer[]): Promise<void> {
+    this.assertIsOpen();
+
+    const addEntries = findInSortedArray(this.indexBatch.addEntries, key, singleKeyCmp);
+    const removeEntries = findInSortedArray(this.indexBatch.removeEntries, key, singleKeyCmp);
+
+    if (values.length === 0) {
+      // special case, we're deleting the entire key
+      if (addEntries) {
+        removeFromSortedArray(this.indexBatch.addEntries, addEntries, keyCmp);
+      }
+
+      if (removeEntries) {
+        removeEntries[1] = null;
+      } else {
+        insertIntoSortedArray(this.indexBatch.removeEntries, [key, null], keyCmp);
+      }
+
+      return Promise.resolve();
+    }
+
+    if (addEntries) {
+      removeAnyOf(addEntries[1], values, Buffer.compare);
+      if (addEntries[1].length === 0) {
+        removeFromSortedArray(this.indexBatch.addEntries, addEntries, keyCmp);
+      }
+    }
+
+    if (removeEntries) {
+      removeEntries[1] ??= [];
+      merge(removeEntries[1], values, Buffer.compare);
+      dedupeSortedArray(removeEntries[1], Buffer.compare);
+    } else {
+      insertIntoSortedArray(this.indexBatch.removeEntries, [key, values], keyCmp);
+    }
+
+    return Promise.resolve();
+  }
+
+  public override async getIndex(key: Buffer): Promise<Uint8Array[]> {
+    this.assertIsOpen();
+
+    const removeEntries = findInSortedArray(this.indexBatch.removeEntries, key, singleKeyCmp);
+    if (removeEntries && removeEntries[1] === null) {
+      return [];
+    }
+
+    const addEntries = findInSortedArray(this.indexBatch.addEntries, key, singleKeyCmp);
+    const results = await super.getIndex(key);
+
+    if (addEntries) {
+      merge(results, addEntries[1], Buffer.compare);
+      dedupeSortedArray(results, Buffer.compare);
+    }
+
+    if (removeEntries && Array.isArray(removeEntries[1])) {
+      removeAnyOf(results, removeEntries[1], Buffer.compare);
+    }
+
+    return results;
+  }
+
+  public override async *iterate(
+    startKey: Uint8Array,
+    endKey?: Uint8Array | undefined,
+    reverse?: boolean,
+    limit?: number,
+  ): AsyncIterable<[Uint8Array, Uint8Array]> {
+    yield* this.#iterate(
+      super.iterate(startKey, endKey, reverse),
+      this.dataBatch,
+      startKey,
+      endKey,
+      reverse,
+      limit,
+      (committed, toAdd) => (toAdd.length > 0 ? toAdd[0] : committed),
+      vals => vals[0],
+    );
+  }
+
+  public override async *iterateIndex(
+    startKey: Uint8Array,
+    endKey?: Uint8Array | undefined,
+    reverse?: boolean,
+    limit?: number,
+  ): AsyncIterable<[Uint8Array, Uint8Array[]]> {
+    yield* this.#iterate(
+      super.iterateIndex(startKey, endKey, reverse),
+      this.indexBatch,
+      startKey,
+      endKey,
+      reverse,
+      limit,
+      (committed, toAdd, toRemove) => {
+        if (toAdd.length > 0) {
+          merge(committed, toAdd, Buffer.compare);
+          dedupeSortedArray(committed, Buffer.compare);
+        }
+        if (toRemove.length > 0) {
+          removeAnyOf(committed, toRemove, Buffer.compare);
+        }
+        return committed;
+      },
+      vals => vals,
+    );
+  }
+
+  async *#iterate<T>(
+    iterator: AsyncIterable<[Uint8Array, T]>,
+    batch: Batch,
+    startKey: Uint8Array,
+    endKey: Uint8Array | undefined,
+    reverse: boolean = false,
+    limit: number | undefined,
+    merge: (committed: T, toAdd: Uint8Array[], toRemove: Uint8Array[]) => T,
+    map: (vals: Uint8Array[]) => T,
+  ): AsyncIterable<[Uint8Array, T]> {
+    this.assertIsOpen();
+
+    // make a copy of this in case we're running in reverse
+    const uncommittedEntries = [...batch.addEntries];
+    // used to check we're in the right order when comparing between a key and uncommittedEntries
+    let cmpDirection = -1;
+    if (reverse) {
+      cmpDirection = 1;
+      uncommittedEntries.reverse();
+    }
+
+    let uncommittedEntriesIdx = 0;
+    while (uncommittedEntriesIdx < uncommittedEntries.length) {
+      const entry = uncommittedEntries[uncommittedEntriesIdx];
+      // go to the first key in our cache that would be captured by the iterator
+      if (Buffer.compare(entry[0], startKey) !== cmpDirection) {
+        break;
+      }
+      uncommittedEntriesIdx++;
+    }
+
+    let count = 0;
+    // helper to early return if we've reached our limit
+    const checkLimit = typeof limit === 'number' ? () => count < limit : () => true;
+    for await (const [key, values] of iterator) {
+      // yield every key that we have cached that's captured by the iterator
+      while (uncommittedEntriesIdx < uncommittedEntries.length && checkLimit()) {
+        const entry = uncommittedEntries[uncommittedEntriesIdx];
+        if (endKey && Buffer.compare(entry[0], endKey) !== cmpDirection) {
+          break;
+        }
+
+        if (Buffer.compare(entry[0], key) === cmpDirection) {
+          count++;
+          yield [entry[0], map(entry[1])];
+        } else {
+          break;
+        }
+        uncommittedEntriesIdx++;
+      }
+
+      if (!checkLimit()) {
+        // we reached the imposed `limit`
+        break;
+      }
+
+      const toRemove = findInSortedArray(batch.removeEntries, key, singleKeyCmp);
+
+      // at this point we've either exhausted all uncommitted entries,
+      // we reached a key strictly greater/smaller than `key`
+      // or we found the key itself
+      // check if it's the key and use the uncommitted value
+      let toAdd: Uint8Array[] = [];
+      if (
+        uncommittedEntriesIdx < uncommittedEntries.length &&
+        Buffer.compare(uncommittedEntries[uncommittedEntriesIdx][0], key) === 0
+      ) {
+        toAdd = uncommittedEntries[uncommittedEntriesIdx][1];
+        uncommittedEntriesIdx++;
+      }
+
+      if (toRemove && !toRemove[1]) {
+        // we were told to delete this key entirely
+        continue;
+      } else {
+        const mergedValues = merge(values, toAdd, toRemove?.[1] ?? []);
+        if (mergedValues) {
+          count++;
+          yield [key, mergedValues];
+        }
+      }
+    }
+
+    // emit all the uncommitted data that would be captured by this iterator
+    while (uncommittedEntriesIdx < uncommittedEntries.length && checkLimit()) {
+      const entry = uncommittedEntries[uncommittedEntriesIdx];
+      if (endKey && Buffer.compare(entry[0], endKey) !== cmpDirection) {
+        break;
+      }
+      count++;
+      yield [entry[0], map(entry[1])];
+      uncommittedEntriesIdx++;
+    }
+  }
+
+  public async commit() {
+    this.assertIsOpen();
+    this.close();
+    await this.channel.sendMessage(LMDBMessageType.BATCH, {
+      batches: new Map([
+        [Database.DATA, this.dataBatch],
+        [Database.INDEX, this.indexBatch],
+      ]),
+    });
+  }
+}

package/src/stores/index.ts

@@ -0,0 +1 @@
+export * from './l2_tips_store.js';