@aztec/kv-store 0.0.0-test.0

package/src/lmdb-v2/map.ts

@@ -0,0 +1,233 @@
+import { Encoder } from 'msgpackr';
+
+import type { Key, Range } from '../interfaces/common.js';
+import type { AztecAsyncMap, AztecAsyncMultiMap } from '../interfaces/map.js';
+import type { ReadTransaction } from './read_transaction.js';
+import { type AztecLMDBStoreV2, execInReadTx, execInWriteTx } from './store.js';
+import { deserializeKey, maxKey, minKey, serializeKey } from './utils.js';
+
+export class LMDBMap<K extends Key, V> implements AztecAsyncMap<K, V> {
+  private prefix: string;
+  private encoder = new Encoder();
+
+  constructor(private store: AztecLMDBStoreV2, name: string) {
+    this.prefix = `map:${name}`;
+  }
+  /**
+   * Sets the value at the given key.
+   * @param key - The key to set the value at
+   * @param val - The value to set
+   */
+  set(key: K, val: V): Promise<void> {
+    return execInWriteTx(this.store, tx => tx.set(serializeKey(this.prefix, key), this.encoder.pack(val)));
+  }
+
+  /**
+   * Sets the value at the given key if it does not already exist.
+   * @param key - The key to set the value at
+   * @param val - The value to set
+   */
+  setIfNotExists(key: K, val: V): Promise<boolean> {
+    return execInWriteTx(this.store, async tx => {
+      const strKey = serializeKey(this.prefix, key);
+      const exists = !!(await tx.get(strKey));
+      if (!exists) {
+        await tx.set(strKey, this.encoder.pack(val));
+        return true;
+      }
+      return false;
+    });
+  }
+
+  /**
+   * Deletes the value at the given key.
+   * @param key - The key to delete the value at
+   */
+  delete(key: K): Promise<void> {
+    return execInWriteTx(this.store, tx => tx.remove(serializeKey(this.prefix, key)));
+  }
+
+  getAsync(key: K): Promise<V | undefined> {
+    return execInReadTx(this.store, async tx => {
+      const val = await tx.get(serializeKey(this.prefix, key));
+      return val ? this.encoder.unpack(val) : undefined;
+    });
+  }
+
+  hasAsync(key: K): Promise<boolean> {
+    return execInReadTx(this.store, async tx => !!(await tx.get(serializeKey(this.prefix, key))));
+  }
+
+  /**
+   * Iterates over the map's key-value entries in the key's natural order
+   * @param range - The range of keys to iterate over
+   */
+  async *entriesAsync(range?: Range<K>): AsyncIterableIterator<[K, V]> {
+    const reverse = range?.reverse ?? false;
+    const startKey = range?.start ? serializeKey(this.prefix, range.start) : minKey(this.prefix);
+
+    const endKey = range?.end ? serializeKey(this.prefix, range.end) : reverse ? maxKey(this.prefix) : undefined;
+
+    let tx: ReadTransaction | undefined = this.store.getCurrentWriteTx();
+    const shouldClose = !tx;
+    tx ??= this.store.getReadTx();
+
+    try {
+      for await (const [key, val] of tx.iterate(
+        reverse ? endKey! : startKey,
+        reverse ? startKey : endKey,
+        reverse,
+        range?.limit,
+      )) {
+        const deserializedKey = deserializeKey<K>(this.prefix, key);
+        if (!deserializedKey) {
+          break;
+        }
+        yield [deserializedKey, this.encoder.unpack(val)];
+      }
+    } finally {
+      if (shouldClose) {
+        tx.close();
+      }
+    }
+  }
+
+  /**
+   * Iterates over the map's values in the key's natural order
+   * @param range - The range of keys to iterate over
+   */
+  async *valuesAsync(range?: Range<K>): AsyncIterableIterator<V> {
+    for await (const [_, value] of this.entriesAsync(range)) {
+      yield value;
+    }
+  }
+
+  /**
+   * Iterates over the map's keys in the key's natural order
+   * @param range - The range of keys to iterate over
+   */
+  async *keysAsync(range?: Range<K>): AsyncIterableIterator<K> {
+    for await (const [key, _] of this.entriesAsync(range)) {
+      yield key;
+    }
+  }
+}
+
+export class LMDBMultiMap<K extends Key, V> implements AztecAsyncMultiMap<K, V> {
+  private prefix: string;
+  private encoder = new Encoder();
+  constructor(private store: AztecLMDBStoreV2, name: string) {
+    this.prefix = `multimap:${name}`;
+  }
+
+  /**
+   * Sets the value at the given key.
+   * @param key - The key to set the value at
+   * @param val - The value to set
+   */
+  set(key: K, val: V): Promise<void> {
+    return execInWriteTx(this.store, tx => tx.setIndex(serializeKey(this.prefix, key), this.encoder.pack(val)));
+  }
+
+  /**
+   * Sets the value at the given key if it does not already exist.
+   * @param key - The key to set the value at
+   * @param val - The value to set
+   */
+  setIfNotExists(key: K, val: V): Promise<boolean> {
+    return execInWriteTx(this.store, async tx => {
+      const exists = !!(await this.getAsync(key));
+      if (!exists) {
+        await tx.setIndex(serializeKey(this.prefix, key), this.encoder.pack(val));
+        return true;
+      }
+      return false;
+    });
+  }
+
+  /**
+   * Deletes the value at the given key.
+   * @param key - The key to delete the value at
+   */
+  delete(key: K): Promise<void> {
+    return execInWriteTx(this.store, tx => tx.removeIndex(serializeKey(this.prefix, key)));
+  }
+
+  getAsync(key: K): Promise<V | undefined> {
+    return execInReadTx(this.store, async tx => {
+      const val = await tx.getIndex(serializeKey(this.prefix, key));
+      return val.length > 0 ? this.encoder.unpack(val[0]) : undefined;
+    });
+  }
+
+  hasAsync(key: K): Promise<boolean> {
+    return execInReadTx(this.store, async tx => (await tx.getIndex(serializeKey(this.prefix, key))).length > 0);
+  }
+
+  /**
+   * Iterates over the map's key-value entries in the key's natural order
+   * @param range - The range of keys to iterate over
+   */
+  async *entriesAsync(range?: Range<K>): AsyncIterableIterator<[K, V]> {
+    const reverse = range?.reverse ?? false;
+    const startKey = range?.start ? serializeKey(this.prefix, range.start) : minKey(this.prefix);
+    const endKey = range?.end ? serializeKey(this.prefix, range.end) : reverse ? maxKey(this.prefix) : undefined;
+
+    let tx: ReadTransaction | undefined = this.store.getCurrentWriteTx();
+    const shouldClose = !tx;
+    tx ??= this.store.getReadTx();
+
+    try {
+      for await (const [key, vals] of tx.iterateIndex(
+        reverse ? endKey! : startKey,
+        reverse ? startKey : endKey,
+        reverse,
+        range?.limit,
+      )) {
+        const deserializedKey = deserializeKey<K>(this.prefix, key);
+        if (!deserializedKey) {
+          break;
+        }
+
+        for (const val of vals) {
+          yield [deserializedKey, this.encoder.unpack(val)];
+        }
+      }
+    } finally {
+      if (shouldClose) {
+        tx.close();
+      }
+    }
+  }
+
+  /**
+   * Iterates over the map's values in the key's natural order
+   * @param range - The range of keys to iterate over
+   */
+  async *valuesAsync(range?: Range<K>): AsyncIterableIterator<V> {
+    for await (const [_, value] of this.entriesAsync(range)) {
+      yield value;
+    }
+  }
+
+  /**
+   * Iterates over the map's keys in the key's natural order
+   * @param range - The range of keys to iterate over
+   */
+  async *keysAsync(range?: Range<K>): AsyncIterableIterator<K> {
+    for await (const [key, _] of this.entriesAsync(range)) {
+      yield key;
+    }
+  }
+
+  deleteValue(key: K, val: V | undefined): Promise<void> {
+    return execInWriteTx(this.store, tx => tx.removeIndex(serializeKey(this.prefix, key), this.encoder.pack(val)));
+  }
+
+  async *getValuesAsync(key: K): AsyncIterableIterator<V> {
+    const values = await execInReadTx(this.store, tx => tx.getIndex(serializeKey(this.prefix, key)));
+    for (const value of values) {
+      yield this.encoder.unpack(value);
+    }
+  }
+}

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;
+    });
+  }
+}