@aztec/kv-store 0.0.1-commit.b655e406 → 0.0.1-commit.b9865e97

package/src/lmdb-v2/factory.ts

@@ -1,23 +1,31 @@
 import { EthAddress } from '@aztec/foundation/eth-address';
-import { type Logger, createLogger } from '@aztec/foundation/log';
-import { DatabaseVersionManager } from '@aztec/stdlib/database-version';
+import { type LoggerBindings, createLogger } from '@aztec/foundation/log';
+import { DatabaseVersionManager, type SchemaVersionMismatchPolicy } from '@aztec/stdlib/database-version/manager';
+import type { DataStoreConfig } from '@aztec/stdlib/kv-store';
 
-import { mkdir, mkdtemp, rm } from 'fs/promises';
+import { copyFile, mkdir, mkdtemp, rm } from 'fs/promises';
 import { tmpdir } from 'os';
 import { join } from 'path';
 
-import type { DataStoreConfig } from '../config.js';
 import { AztecLMDBStoreV2 } from './store.js';
 
 const MAX_READERS = 16;
 
+/** Optional versioning hooks for persistent LMDB stores. */
+export type CreateStoreOptions = {
+  onUpgrade?: (dataDir: string, currentVersion: number, latestVersion: number) => Promise<void>;
+  schemaVersionMismatchPolicy?: SchemaVersionMismatchPolicy;
+};
+
 export async function createStore(
   name: string,
   schemaVersion: number,
   config: DataStoreConfig,
-  log: Logger = createLogger('kv-store:lmdb-v2:' + name),
+  bindings?: LoggerBindings,
+  options: CreateStoreOptions = {},
 ): Promise<AztecLMDBStoreV2> {
-  const { dataDirectory, l1Contracts } = config;
+  const log = createLogger('kv-store:lmdb-v2:' + name, bindings);
+  const { dataDirectory, rollupAddress: rollupFromConfig } = config;
 
   let store: AztecLMDBStoreV2;
   if (typeof dataDirectory !== 'undefined') {
@@ -25,7 +33,7 @@ export async function createStore(
     const subDir = join(dataDirectory, name);
     await mkdir(subDir, { recursive: true });
 
-    const rollupAddress = l1Contracts ? l1Contracts.rollupAddress : EthAddress.ZERO;
+    const rollupAddress = rollupFromConfig ?? EthAddress.ZERO;
 
     // Create a version manager
     const versionManager = new DatabaseVersionManager({
@@ -33,7 +41,9 @@ export async function createStore(
       rollupAddress,
       dataDirectory: subDir,
       onOpen: dbDirectory =>
-        AztecLMDBStoreV2.new(dbDirectory, config.dataStoreMapSizeKb, MAX_READERS, () => Promise.resolve(), log),
+        AztecLMDBStoreV2.new(dbDirectory, config.dataStoreMapSizeKb, MAX_READERS, () => Promise.resolve(), bindings),
+      onUpgrade: options.onUpgrade,
+      schemaVersionMismatchPolicy: options.schemaVersionMismatchPolicy,
     });
 
     log.info(
@@ -41,25 +51,31 @@ export async function createStore(
     );
     [store] = await versionManager.open();
   } else {
-    store = await openTmpStore(name, true, config.dataStoreMapSizeKb, MAX_READERS, log);
+    store = await openTmpStore(name, true, config.dataStoreMapSizeKb, MAX_READERS, bindings);
   }
 
   return store;
 }
 
+/**
+ * Open a persistent on-disk store rooted in OS tmpdir. Caller chooses whether to
+ * auto-remove the directory on close. Uses standard durable LMDB flags (fsync on
+ * every commit) — for full in-memory / NOSYNC semantics use `openEphemeralStore`.
+ */
 export async function openTmpStore(
   name: string,
-  ephemeral: boolean = true,
+  cleanupTmpDir: boolean = true,
   dbMapSizeKb = 10 * 1_024 * 1_024, // 10GB
   maxReaders = MAX_READERS,
-  log: Logger = createLogger('kv-store:lmdb-v2:' + name),
+  bindings?: LoggerBindings,
 ): Promise<AztecLMDBStoreV2> {
+  const log = createLogger('kv-store:lmdb-v2:' + name, bindings);
   const dataDir = await mkdtemp(join(tmpdir(), name + '-'));
   log.debug(`Created temporary data store at: ${dataDir} with size: ${dbMapSizeKb} KB (LMDB v2)`);
 
   // pass a cleanup callback because process.on('beforeExit', cleanup) does not work under Jest
   const cleanup = async () => {
-    if (ephemeral) {
+    if (cleanupTmpDir) {
       try {
         await rm(dataDir, { recursive: true, force: true, maxRetries: 3 });
         log.debug(`Deleted temporary data store: ${dataDir}`);
@@ -71,19 +87,76 @@ export async function openTmpStore(
     }
   };
 
-  // For temporary stores, we don't need to worry about versioning
-  // as they are ephemeral and get cleaned up after use
-  return AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, cleanup, log);
+  return AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, cleanup, bindings, false);
+}
+
+/**
+ * Open a fully-ephemeral store: a fresh tmpdir backing file, LMDB opened with
+ * MDB_NOSYNC | MDB_NOMETASYNC (no fsync, kernel-only writeback), and the directory
+ * unconditionally removed on close. Intended for tests and worker-local scratch
+ * stores where durability is explicitly not desired.
+ */
+export async function openEphemeralStore(
+  name: string,
+  dbMapSizeKb = 10 * 1_024 * 1_024, // 10GB
+  maxReaders = MAX_READERS,
+  bindings?: LoggerBindings,
+): Promise<AztecLMDBStoreV2> {
+  const log = createLogger('kv-store:lmdb-v2:' + name, bindings);
+  const dataDir = await mkdtemp(join(tmpdir(), name + '-'));
+  log.debug(`Created ephemeral data store at: ${dataDir} with size: ${dbMapSizeKb} KB (LMDB v2)`);
+
+  const cleanup = async () => {
+    try {
+      await rm(dataDir, { recursive: true, force: true, maxRetries: 3 });
+      log.debug(`Deleted ephemeral data store: ${dataDir}`);
+    } catch (err) {
+      log.warn(`Failed to delete ephemeral data directory (LMDB v2) ${dataDir}: ${err}`);
+    }
+  };
+
+  return AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, cleanup, bindings, true);
 }
 
 export async function openStoreAt(
   dataDir: string,
   dbMapSizeKb = 10 * 1_024 * 1_024, // 10GB
   maxReaders = MAX_READERS,
-  log: Logger = createLogger('kv-store:lmdb-v2'),
+  bindings?: LoggerBindings,
 ): Promise<AztecLMDBStoreV2> {
+  const log = createLogger('kv-store:lmdb-v2', bindings);
   log.debug(`Opening data store at: ${dataDir} with size: ${dbMapSizeKb} KB (LMDB v2)`);
-  return await AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, undefined, log);
+  return await AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, undefined, bindings);
+}
+
+/**
+ * Open a fully-ephemeral store seeded from an existing `data.mdb` file. Creates a fresh tmpdir,
+ * copies `srcDataMdbPath` into it as `data.mdb`, opens with MDB_NOSYNC | MDB_NOMETASYNC, and
+ * removes the tmpdir on close. Intended for worker clones of a main-thread-built store — the
+ * source file is never written to.
+ */
+export async function cloneEphemeralStoreFrom(
+  srcDataMdbPath: string,
+  name: string,
+  dbMapSizeKb = 10 * 1_024 * 1_024, // 10GB
+  maxReaders = MAX_READERS,
+  bindings?: LoggerBindings,
+): Promise<AztecLMDBStoreV2> {
+  const log = createLogger('kv-store:lmdb-v2:' + name, bindings);
+  const dataDir = await mkdtemp(join(tmpdir(), name + '-'));
+  await copyFile(srcDataMdbPath, join(dataDir, 'data.mdb'));
+  log.debug(`Cloned ephemeral data store at: ${dataDir} from ${srcDataMdbPath} (LMDB v2)`);
+
+  const cleanup = async () => {
+    try {
+      await rm(dataDir, { recursive: true, force: true, maxRetries: 3 });
+      log.debug(`Deleted ephemeral data store: ${dataDir}`);
+    } catch (err) {
+      log.warn(`Failed to delete ephemeral data directory (LMDB v2) ${dataDir}: ${err}`);
+    }
+  };
+
+  return AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, cleanup, bindings, true);
 }
 
 export async function openVersionedStoreAt(
@@ -92,14 +165,15 @@ export async function openVersionedStoreAt(
   rollupAddress: EthAddress,
   dbMapSizeKb = 10 * 1_024 * 1_024, // 10GB
   maxReaders = MAX_READERS,
-  log: Logger = createLogger('kv-store:lmdb-v2'),
+  bindings?: LoggerBindings,
 ): Promise<AztecLMDBStoreV2> {
+  const log = createLogger('kv-store:lmdb-v2', bindings);
   log.debug(`Opening data store at: ${dataDirectory} with size: ${dbMapSizeKb} KB (LMDB v2)`);
   const [store] = await new DatabaseVersionManager({
     schemaVersion,
     rollupAddress,
     dataDirectory,
-    onOpen: dataDir => AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, undefined, log),
+    onOpen: dataDir => AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, undefined, bindings),
   }).open();
   return store;
 }

package/src/lmdb-v2/map.ts

@@ -3,8 +3,8 @@ import { Encoder } from 'msgpackr';
 import type { Key, Range, Value } from '../interfaces/common.js';
 import type { AztecAsyncMap } from '../interfaces/map.js';
 import type { ReadTransaction } from './read_transaction.js';
-// eslint-disable-next-line import/no-cycle
-import { type AztecLMDBStoreV2, execInReadTx, execInWriteTx } from './store.js';
+import type { AztecLMDBStoreV2 } from './store.js';
+import { execInReadTx, execInWriteTx } from './tx-helpers.js';
 import { deserializeKey, maxKey, minKey, serializeKey } from './utils.js';
 
 export class LMDBMap<K extends Key, V extends Value> implements AztecAsyncMap<K, V> {

package/src/lmdb-v2/multi_map.ts

@@ -1,9 +1,11 @@
 import { Encoder } from 'msgpackr/pack';
+import { MAXIMUM_KEY, toBufferKey } from 'ordered-binary';
 
 import type { Key, Range, Value } from '../interfaces/common.js';
 import type { AztecAsyncMultiMap } from '../interfaces/multi_map.js';
 import type { ReadTransaction } from './read_transaction.js';
-import { type AztecLMDBStoreV2, execInReadTx, execInWriteTx } from './store.js';
+import type { AztecLMDBStoreV2 } from './store.js';
+import { execInReadTx, execInWriteTx } from './tx-helpers.js';
 import { deserializeKey, maxKey, minKey, serializeKey } from './utils.js';
 
 export class LMDBMultiMap<K extends Key, V extends Value> implements AztecAsyncMultiMap<K, V> {
@@ -138,4 +140,11 @@ export class LMDBMultiMap<K extends Key, V extends Value> implements AztecAsyncM
       yield this.encoder.unpack(value);
     }
   }
+
+  getValueCountAsync(key: K): Promise<number> {
+    const startKey = serializeKey(this.prefix, key);
+    const endKey = toBufferKey([this.prefix, key, MAXIMUM_KEY]);
+
+    return execInReadTx(this.store, tx => tx.countEntriesIndex(startKey, endKey, false));
+  }
 }

package/src/lmdb-v2/read_transaction.ts

@@ -66,20 +66,20 @@ export class ReadTransaction {
   ): 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;
-
+    let cursor: number | undefined;
     try {
+      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,
+      });
+
+      cursor = response.cursor ?? undefined;
+      let entries = response.entries;
+      let done = typeof cursor !== 'number';
+      let count = 0;
       // 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) {
@@ -125,17 +125,17 @@ export class ReadTransaction {
   async #countEntries(db: string, startKey: Uint8Array, endKey: Uint8Array, reverse: boolean): Promise<number> {
     this.assertIsOpen();
 
-    const response = await this.channel.sendMessage(LMDBMessageType.START_CURSOR, {
-      key: startKey,
-      reverse,
-      count: 0,
-      onePage: false,
-      db,
-    });
-
-    const cursor = response.cursor;
-
+    let cursor: number | undefined;
     try {
+      const response = await this.channel.sendMessage(LMDBMessageType.START_CURSOR, {
+        key: startKey,
+        reverse,
+        count: 0,
+        onePage: false,
+        db,
+      });
+
+      cursor = response.cursor ?? undefined;
       if (!cursor) {
         return 0;
       }

package/src/lmdb-v2/singleton.ts

@@ -1,8 +1,8 @@
 import { Encoder } from 'msgpackr';
 
 import type { AztecAsyncSingleton } from '../interfaces/singleton.js';
-// eslint-disable-next-line import/no-cycle
-import { type AztecLMDBStoreV2, execInReadTx, execInWriteTx } from './store.js';
+import type { AztecLMDBStoreV2 } from './store.js';
+import { execInReadTx, execInWriteTx } from './tx-helpers.js';
 import { serializeKey } from './utils.js';
 
 export class LMDBSingleValue<T> implements AztecAsyncSingleton<T> {

package/src/lmdb-v2/store.ts

@@ -1,4 +1,4 @@
-import { type Logger, createLogger } from '@aztec/foundation/log';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import { Semaphore, SerialQueue } from '@aztec/foundation/queue';
 import { MsgpackChannel, NativeLMDBStore } from '@aztec/native';
 
@@ -13,9 +13,7 @@ import type { AztecAsyncMultiMap } from '../interfaces/multi_map.js';
 import type { AztecAsyncSet } from '../interfaces/set.js';
 import type { AztecAsyncSingleton } from '../interfaces/singleton.js';
 import type { AztecAsyncKVStore } from '../interfaces/store.js';
-// eslint-disable-next-line import/no-cycle
 import { LMDBArray } from './array.js';
-// eslint-disable-next-line import/no-cycle
 import { LMDBMap } from './map.js';
 import {
   Database,
@@ -27,10 +25,11 @@ import {
 import { LMDBMultiMap } from './multi_map.js';
 import { ReadTransaction } from './read_transaction.js';
 import { LMDBSet } from './set.js';
-// eslint-disable-next-line import/no-cycle
 import { LMDBSingleValue } from './singleton.js';
 import { WriteTransaction } from './write_transaction.js';
 
+export { execInReadTx, execInWriteTx } from './tx-helpers.js';
+
 export class AztecLMDBStoreV2 implements AztecAsyncKVStore, LMDBMessageChannel {
   private open = false;
   private channel: MsgpackChannel<LMDBMessageType, LMDBRequestBody, LMDBResponseBody>;
@@ -44,9 +43,10 @@ export class AztecLMDBStoreV2 implements AztecAsyncKVStore, LMDBMessageChannel {
     maxReaders: number,
     private log: Logger,
     private cleanup?: () => Promise<void>,
+    ephemeral: boolean = false,
   ) {
     this.log.info(`Starting data store with maxReaders ${maxReaders}`);
-    this.channel = new MsgpackChannel(new NativeLMDBStore(dataDir, mapSize, maxReaders));
+    this.channel = new MsgpackChannel(new NativeLMDBStore(dataDir, mapSize, maxReaders, ephemeral));
     // leave one reader to always be available for regular, atomic, reads
     this.availableCursors = new Semaphore(maxReaders - 1);
   }
@@ -76,9 +76,11 @@ export class AztecLMDBStoreV2 implements AztecAsyncKVStore, LMDBMessageChannel {
     dbMapSizeKb: number = 10 * 1024 * 1024,
     maxReaders: number = 16,
     cleanup?: () => Promise<void>,
-    log = createLogger('kv-store:lmdb-v2'),
+    bindings?: LoggerBindings,
+    ephemeral: boolean = false,
   ) {
-    const db = new AztecLMDBStoreV2(dataDir, dbMapSizeKb, maxReaders, log, cleanup);
+    const log = createLogger('kv-store:lmdb-v2', bindings);
+    const db = new AztecLMDBStoreV2(dataDir, dbMapSizeKb, maxReaders, log, cleanup, ephemeral);
     await db.start();
     return db;
   }
@@ -217,29 +219,3 @@ export class AztecLMDBStoreV2 implements AztecAsyncKVStore, LMDBMessageChannel {
     };
   }
 }
-
-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/tx-helpers.ts

@@ -0,0 +1,29 @@
+import type { ReadTransaction } from './read_transaction.js';
+import type { AztecLMDBStoreV2 } from './store.js';
+import type { WriteTransaction } from './write_transaction.js';
+
+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

@@ -2,124 +2,6 @@ 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]);
 }

package/src/lmdb-v2/write_transaction.ts

@@ -1,16 +1,16 @@
-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';
+} from '@aztec/foundation/array';
+
+import { type Batch, Database, LMDBMessageType } from './message.js';
+import { ReadTransaction } from './read_transaction.js';
+import { keyCmp, singleKeyCmp } from './utils.js';
 
 export class WriteTransaction extends ReadTransaction {
   // exposed for tests
@@ -44,7 +44,7 @@ export class WriteTransaction extends ReadTransaction {
   remove(key: Uint8Array): Promise<void> {
     const removeEntryIndex = findIndexInSortedArray(this.dataBatch.removeEntries, key, singleKeyCmp);
     if (removeEntryIndex === -1) {
-      this.dataBatch.removeEntries.push([key, null]);
+      insertIntoSortedArray(this.dataBatch.removeEntries, [key, null], keyCmp);
     }
 
     const addEntryIndex = findIndexInSortedArray(this.dataBatch.addEntries, key, singleKeyCmp);
@@ -218,8 +218,9 @@ export class WriteTransaction extends ReadTransaction {
   ): AsyncIterable<[Uint8Array, T]> {
     this.assertIsOpen();
 
-    // make a copy of this in case we're running in reverse
+    // Snapshot both add and remove entries at the start of iteration to ensure consistency
     const uncommittedEntries = [...batch.addEntries];
+    const removeEntries = [...batch.removeEntries];
     // used to check we're in the right order when comparing between a key and uncommittedEntries
     let cmpDirection = -1;
     if (reverse) {
@@ -262,7 +263,7 @@ export class WriteTransaction extends ReadTransaction {
         break;
       }
 
-      const toRemove = findInSortedArray(batch.removeEntries, key, singleKeyCmp);
+      const toRemove = findInSortedArray(removeEntries, key, singleKeyCmp);
 
       // at this point we've either exhausted all uncommitted entries,
       // we reached a key strictly greater/smaller than `key`