@aztec/kv-store 0.71.0 → 0.73.0

package/src/lmdb-v2/message.ts

@@ -0,0 +1,146 @@
+export enum Database {
+  DATA = 'data',
+  INDEX = 'index',
+}
+
+export const CURSOR_PAGE_SIZE = 10;
+
+export enum LMDBMessageType {
+  OPEN_DATABASE = 100,
+  GET,
+  HAS,
+
+  START_CURSOR,
+  ADVANCE_CURSOR,
+  CLOSE_CURSOR,
+
+  BATCH,
+
+  STATS,
+
+  CLOSE,
+}
+
+type Key = Uint8Array;
+type Value = Uint8Array;
+type OptionalValues = Array<Value[] | null>;
+type KeyOptionalValues = [Key, null | Array<Value>];
+type KeyValues = [Key, Value[]];
+
+interface OpenDatabaseRequest {
+  db: string;
+  uniqueKeys?: boolean;
+}
+
+interface GetRequest {
+  keys: Key[];
+  db: string;
+}
+
+interface GetResponse {
+  values: OptionalValues;
+}
+
+interface HasRequest {
+  entries: KeyOptionalValues[];
+  db: string;
+}
+
+interface StartCursorRequest {
+  key: Key;
+  reverse: boolean;
+  count: number | null;
+  onePage: boolean | null;
+  db: string;
+}
+
+interface AdvanceCursorRequest {
+  cursor: number;
+  count: number | null;
+}
+
+interface CloseCursorRequest {
+  cursor: number;
+}
+
+export interface Batch {
+  addEntries: Array<KeyValues>;
+  removeEntries: Array<KeyOptionalValues>;
+}
+
+interface BatchRequest {
+  batches: Map<string, Batch>;
+}
+
+export type LMDBRequestBody = {
+  [LMDBMessageType.OPEN_DATABASE]: OpenDatabaseRequest;
+
+  [LMDBMessageType.GET]: GetRequest;
+  [LMDBMessageType.HAS]: HasRequest;
+
+  [LMDBMessageType.START_CURSOR]: StartCursorRequest;
+  [LMDBMessageType.ADVANCE_CURSOR]: AdvanceCursorRequest;
+  [LMDBMessageType.CLOSE_CURSOR]: CloseCursorRequest;
+
+  [LMDBMessageType.BATCH]: BatchRequest;
+
+  [LMDBMessageType.STATS]: void;
+
+  [LMDBMessageType.CLOSE]: void;
+};
+
+interface GetResponse {
+  values: OptionalValues;
+}
+
+interface HasResponse {
+  exists: boolean[];
+}
+
+interface StartCursorResponse {
+  cursor: number | null;
+  entries: Array<KeyValues>;
+}
+
+interface AdvanceCursorResponse {
+  entries: Array<KeyValues>;
+  done: boolean;
+}
+
+interface BatchResponse {
+  durationNs: number;
+}
+
+interface BoolResponse {
+  ok: true;
+}
+
+interface StatsResponse {
+  stats: Array<{
+    name: string;
+    numDataItems: bigint | number;
+    totalUsedSize: bigint | number;
+  }>;
+  dbMapSizeBytes: bigint | number;
+}
+
+export type LMDBResponseBody = {
+  [LMDBMessageType.OPEN_DATABASE]: BoolResponse;
+
+  [LMDBMessageType.GET]: GetResponse;
+  [LMDBMessageType.HAS]: HasResponse;
+
+  [LMDBMessageType.START_CURSOR]: StartCursorResponse;
+  [LMDBMessageType.ADVANCE_CURSOR]: AdvanceCursorResponse;
+  [LMDBMessageType.CLOSE_CURSOR]: BoolResponse;
+
+  [LMDBMessageType.BATCH]: BatchResponse;
+
+  [LMDBMessageType.STATS]: StatsResponse;
+
+  [LMDBMessageType.CLOSE]: BoolResponse;
+};
+
+export interface LMDBMessageChannel {
+  sendMessage<T extends LMDBMessageType>(msgType: T, body: LMDBRequestBody[T]): Promise<LMDBResponseBody[T]>;
+}

package/src/lmdb-v2/read_transaction.ts

@@ -0,0 +1,116 @@
+import { CURSOR_PAGE_SIZE, Database, type LMDBMessageChannel, LMDBMessageType } from './message.js';
+
+export class ReadTransaction {
+  protected open = true;
+
+  constructor(protected channel: LMDBMessageChannel) {}
+
+  public close(): void {
+    if (!this.open) {
+      return;
+    }
+    this.open = false;
+  }
+
+  protected assertIsOpen() {
+    if (!this.open) {
+      throw new Error('Transaction is closed');
+    }
+  }
+
+  public async get(key: Uint8Array): Promise<Uint8Array | undefined> {
+    this.assertIsOpen();
+    const response = await this.channel.sendMessage(LMDBMessageType.GET, { keys: [key], db: Database.DATA });
+    return response.values[0]?.[0] ?? undefined;
+  }
+
+  public async getIndex(key: Uint8Array): Promise<Uint8Array[]> {
+    this.assertIsOpen();
+    const response = await this.channel.sendMessage(LMDBMessageType.GET, { keys: [key], db: Database.INDEX });
+    return response.values[0] ?? [];
+  }
+
+  public async *iterate(
+    startKey: Uint8Array,
+    endKey?: Uint8Array,
+    reverse = false,
+    limit?: number,
+  ): AsyncIterable<[Uint8Array, Uint8Array]> {
+    yield* this.#iterate(Database.DATA, startKey, endKey, reverse, limit, vals => vals[0]);
+  }
+
+  public async *iterateIndex(
+    startKey: Uint8Array,
+    endKey?: Uint8Array,
+    reverse = false,
+    limit?: number,
+  ): AsyncIterable<[Uint8Array, Uint8Array[]]> {
+    yield* this.#iterate(Database.INDEX, startKey, endKey, reverse, limit, vals => vals);
+  }
+
+  async *#iterate<T>(
+    db: string,
+    startKey: Uint8Array,
+    endKey: Uint8Array | undefined,
+    reverse: boolean,
+    limit: number | undefined,
+    map: (val: Uint8Array[]) => T,
+  ): AsyncIterable<[Uint8Array, T]> {
+    this.assertIsOpen();
+
+    const response = await this.channel.sendMessage(LMDBMessageType.START_CURSOR, {
+      key: startKey,
+      reverse,
+      count: typeof limit === 'number' ? Math.min(limit, CURSOR_PAGE_SIZE) : CURSOR_PAGE_SIZE,
+      onePage: typeof limit === 'number' && limit < CURSOR_PAGE_SIZE,
+      db,
+    });
+
+    const cursor = response.cursor;
+    let entries = response.entries;
+    let done = typeof cursor !== 'number';
+    let count = 0;
+
+    try {
+      // emit the first page and any subsequent pages in a while loop
+      // NB: end contition is in the middle of the while loop
+      while (entries.length > 0) {
+        for (const [key, values] of entries) {
+          if (typeof limit === 'number' && count >= limit) {
+            done = true;
+            break;
+          }
+
+          if (endKey) {
+            const cmp = Buffer.compare(key, endKey);
+            if ((!reverse && cmp >= 0) || (reverse && cmp <= 0)) {
+              done = true;
+              break;
+            }
+          }
+
+          count++;
+          yield [key, map(values)];
+        }
+
+        // cursor is null if DB returned everything in the first page
+        if (typeof cursor !== 'number' || done) {
+          break;
+        }
+
+        const response = await this.channel.sendMessage(LMDBMessageType.ADVANCE_CURSOR, {
+          cursor,
+          count: CURSOR_PAGE_SIZE,
+        });
+
+        done = response.done;
+        entries = response.entries;
+      }
+    } finally {
+      // we might not have anything to close
+      if (typeof cursor === 'number') {
+        await this.channel.sendMessage(LMDBMessageType.CLOSE_CURSOR, { cursor });
+      }
+    }
+  }
+}

package/src/lmdb-v2/singleton.ts

@@ -0,0 +1,34 @@
+import { Encoder } from 'msgpackr';
+
+import { type AztecAsyncSingleton } from '../interfaces/singleton.js';
+import { type AztecLMDBStoreV2, execInReadTx, execInWriteTx } from './store.js';
+import { serializeKey } from './utils.js';
+
+export class LMDBSingleValue<T> implements AztecAsyncSingleton<T> {
+  private key: Uint8Array;
+  private encoder = new Encoder();
+  constructor(private store: AztecLMDBStoreV2, name: string) {
+    this.key = serializeKey(`singleton:${name}`, 'value');
+  }
+
+  getAsync(): Promise<T | undefined> {
+    return execInReadTx(this.store, async tx => {
+      const val = await tx.get(this.key);
+      return val ? this.encoder.unpack(val) : undefined;
+    });
+  }
+
+  set(val: T): Promise<boolean> {
+    return execInWriteTx(this.store, async tx => {
+      await tx.set(this.key, this.encoder.pack(val));
+      return true;
+    });
+  }
+
+  delete(): Promise<boolean> {
+    return execInWriteTx(this.store, async tx => {
+      await tx.remove(this.key);
+      return true;
+    });
+  }
+}

package/src/lmdb-v2/store.ts

@@ -0,0 +1,210 @@
+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 { 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 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);
+  }
+
+  private async start() {
+    this.writerQueue.start();
+
+    await this.sendMessage(LMDBMessageType.OPEN_DATABASE, {
+      db: Database.DATA,
+      uniqueKeys: true,
+    });
+
+    await this.sendMessage(LMDBMessageType.OPEN_DATABASE, {
+      db: Database.INDEX,
+      uniqueKeys: false,
+    });
+  }
+
+  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 {
+    return new ReadTransaction(this);
+  }
+
+  public getCurrentWriteTx(): WriteTransaction | undefined {
+    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> {
+    throw new Error('Not implemented');
+  }
+
+  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> {
+    // 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 });
+    this.log.verbose(`Deleted database files at ${this.dataDir}`);
+    await this.cleanup?.();
+  }
+
+  async close() {
+    await this.writerQueue.cancel();
+    await this.sendMessage(LMDBMessageType.CLOSE, undefined);
+  }
+
+  public async sendMessage<T extends LMDBMessageType>(
+    msgType: T,
+    body: LMDBRequestBody[T],
+  ): Promise<LMDBResponseBody[T]> {
+    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;
+}