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

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`

package/src/stores/l2_tips_store.ts

@@ -1,79 +1,88 @@
-import type {
-  L2BlockId,
-  L2BlockStreamEvent,
-  L2BlockStreamEventHandler,
-  L2BlockStreamLocalDataProvider,
-  L2BlockTag,
-  L2Tips,
-} from '@aztec/stdlib/block';
+import { BlockNumber, CheckpointNumber } from '@aztec/foundation/branded-types';
+import { type L2BlockTag, L2TipsStoreBase } from '@aztec/stdlib/block';
+import { PublishedCheckpoint } from '@aztec/stdlib/checkpoint';
 
 import type { AztecAsyncMap } from '../interfaces/map.js';
 import type { AztecAsyncKVStore } from '../interfaces/store.js';
 
-/** Stores currently synced L2 tips and unfinalized block hashes. */
-export class L2TipsKVStore implements L2BlockStreamEventHandler, L2BlockStreamLocalDataProvider {
-  private readonly l2TipsStore: AztecAsyncMap<L2BlockTag, number>;
-  private readonly l2BlockHashesStore: AztecAsyncMap<number, string>;
+/**
+ * Persistent implementation of L2 tips store backed by a KV store.
+ * Used by nodes that need to persist chain state across restarts.
+ */
+export class L2TipsKVStore extends L2TipsStoreBase {
+  private readonly l2TipsStore: AztecAsyncMap<L2BlockTag, BlockNumber>;
+  private readonly l2BlockHashesStore: AztecAsyncMap<BlockNumber, string>;
+  private readonly l2BlockNumberToCheckpointNumberStore: AztecAsyncMap<BlockNumber, CheckpointNumber>;
+  private readonly l2CheckpointStore: AztecAsyncMap<CheckpointNumber, Buffer>;
 
-  constructor(store: AztecAsyncKVStore, namespace: string) {
+  constructor(
+    private store: AztecAsyncKVStore,
+    namespace: string,
+  ) {
+    super();
     this.l2TipsStore = store.openMap([namespace, 'l2_tips'].join('_'));
     this.l2BlockHashesStore = store.openMap([namespace, 'l2_block_hashes'].join('_'));
+    this.l2BlockNumberToCheckpointNumberStore = store.openMap(
+      [namespace, 'l2_block_number_to_checkpoint_number'].join('_'),
+    );
+    this.l2CheckpointStore = store.openMap([namespace, 'l2_checkpoint_store'].join('_'));
   }
 
-  public getL2BlockHash(number: number): Promise<string | undefined> {
-    return this.l2BlockHashesStore.getAsync(number);
+  protected getTip(tag: L2BlockTag): Promise<BlockNumber | undefined> {
+    return this.l2TipsStore.getAsync(tag);
   }
 
-  public async getL2Tips(): Promise<L2Tips> {
-    return {
-      latest: await this.getL2Tip('latest'),
-      finalized: await this.getL2Tip('finalized'),
-      proven: await this.getL2Tip('proven'),
-    };
+  protected setTip(tag: L2BlockTag, blockNumber: BlockNumber): Promise<void> {
+    return this.l2TipsStore.set(tag, blockNumber);
   }
 
-  private async getL2Tip(tag: L2BlockTag): Promise<L2BlockId> {
-    const blockNumber = await this.l2TipsStore.getAsync(tag);
-    if (blockNumber === undefined || blockNumber === 0) {
-      return { number: 0, hash: undefined };
-    }
-    const blockHash = await this.l2BlockHashesStore.getAsync(blockNumber);
-    if (!blockHash) {
-      throw new Error(`Block hash not found for block number ${blockNumber}`);
+  protected getStoredBlockHash(blockNumber: BlockNumber): Promise<string | undefined> {
+    return this.l2BlockHashesStore.getAsync(blockNumber);
+  }
+
+  protected setBlockHash(blockNumber: BlockNumber, hash: string): Promise<void> {
+    return this.l2BlockHashesStore.set(blockNumber, hash);
+  }
+
+  protected async deleteBlockHashesBefore(blockNumber: BlockNumber): Promise<void> {
+    for await (const key of this.l2BlockHashesStore.keysAsync({ end: blockNumber })) {
+      await this.l2BlockHashesStore.delete(key);
     }
+  }
+
+  protected getCheckpointNumberForBlock(blockNumber: BlockNumber): Promise<CheckpointNumber | undefined> {
+    return this.l2BlockNumberToCheckpointNumberStore.getAsync(blockNumber);
+  }
 
-    return { number: blockNumber, hash: blockHash };
+  protected setCheckpointNumberForBlock(blockNumber: BlockNumber, checkpointNumber: CheckpointNumber): Promise<void> {
+    return this.l2BlockNumberToCheckpointNumberStore.set(blockNumber, checkpointNumber);
   }
 
-  public async handleBlockStreamEvent(event: L2BlockStreamEvent): Promise<void> {
-    switch (event.type) {
-      case 'blocks-added': {
-        const blocks = event.blocks.map(b => b.block);
-        for (const block of blocks) {
-          await this.l2BlockHashesStore.set(block.number, (await block.hash()).toString());
-        }
-        await this.l2TipsStore.set('latest', blocks.at(-1)!.number);
-        break;
-      }
-      case 'chain-pruned':
-        await this.saveTag('latest', event.block);
-        break;
-      case 'chain-proven':
-        await this.saveTag('proven', event.block);
-        break;
-      case 'chain-finalized':
-        await this.saveTag('finalized', event.block);
-        for await (const key of this.l2BlockHashesStore.keysAsync({ end: event.block.number })) {
-          await this.l2BlockHashesStore.delete(key);
-        }
-        break;
+  protected async deleteBlockToCheckpointBefore(blockNumber: BlockNumber): Promise<void> {
+    for await (const key of this.l2BlockNumberToCheckpointNumberStore.keysAsync({ end: blockNumber })) {
+      await this.l2BlockNumberToCheckpointNumberStore.delete(key);
     }
   }
 
-  private async saveTag(name: L2BlockTag, block: L2BlockId) {
-    await this.l2TipsStore.set(name, block.number);
-    if (block.hash) {
-      await this.l2BlockHashesStore.set(block.number, block.hash);
+  protected async getCheckpoint(checkpointNumber: CheckpointNumber): Promise<PublishedCheckpoint | undefined> {
+    const buffer = await this.l2CheckpointStore.getAsync(checkpointNumber);
+    if (!buffer) {
+      return undefined;
     }
+    return PublishedCheckpoint.fromBuffer(buffer);
+  }
+
+  protected saveCheckpointData(checkpoint: PublishedCheckpoint): Promise<void> {
+    return this.l2CheckpointStore.set(checkpoint.checkpoint.number, checkpoint.toBuffer());
+  }
+
+  protected async deleteCheckpointsBefore(checkpointNumber: CheckpointNumber): Promise<void> {
+    for await (const key of this.l2CheckpointStore.keysAsync({ end: checkpointNumber })) {
+      await this.l2CheckpointStore.delete(key);
+    }
+  }
+
+  protected runInTransaction<T>(fn: () => Promise<T>): Promise<T> {
+    return this.store.transactionAsync(fn);
   }
 }

package/src/utils.ts

@@ -1,40 +1,98 @@
-import type { EthAddress } from '@aztec/foundation/eth-address';
+import { EthAddress } from '@aztec/foundation/eth-address';
 import type { Logger } from '@aztec/foundation/log';
+import { DatabaseVersion } from '@aztec/stdlib/database-version/version';
 
 import type { AztecAsyncSingleton, AztecSingleton } from './interfaces/singleton.js';
 import type { AztecAsyncKVStore, AztecKVStore } from './interfaces/store.js';
 import { isSyncStore } from './interfaces/utils.js';
 
 /**
- * Clears the store if the rollup address does not match the one stored in the database.
- * This is to prevent data from being accidentally shared between different rollup instances.
+ * Clears the store if the schema version or rollup address does not match the one stored in the database.
+ * Also clears if migrating from an older store format that didn't track schema version.
+ * This is to prevent data from being accidentally mixed up between different rollup instances or schema versions.
  * @param store - The store to check
+ * @param targetSchemaVersion - The current schema version
  * @param rollupAddress - The ETH address of the rollup contract
- * @returns A promise that resolves when the store is cleared, or rejects if the rollup address does not match
+ * @param log - Optional logger
+ * @returns The store (cleared if necessary)
  */
-export async function initStoreForRollup<T extends AztecKVStore | AztecAsyncKVStore>(
+export async function initStoreForRollupAndSchemaVersion<T extends AztecKVStore | AztecAsyncKVStore>(
   store: T,
-  rollupAddress: EthAddress,
+  schemaVersion: number | undefined,
+  rollupAddress: EthAddress | undefined,
   log?: Logger,
 ): Promise<T> {
-  if (!rollupAddress) {
-    throw new Error('Rollup address is required');
-  }
-  const rollupAddressValue = store.openSingleton<ReturnType<EthAddress['toString']>>('rollupAddress');
-  const rollupAddressString = rollupAddress.toString();
-  const storedRollupAddressString = isSyncStore(store)
-    ? (rollupAddressValue as AztecSingleton<ReturnType<EthAddress['toString']>>).get()
-    : await (rollupAddressValue as AztecAsyncSingleton<ReturnType<EthAddress['toString']>>).getAsync();
-
-  if (typeof storedRollupAddressString !== 'undefined' && storedRollupAddressString !== rollupAddressString) {
-    log?.warn(`Rollup address mismatch. Clearing entire database...`, {
-      expected: rollupAddressString,
-      found: storedRollupAddressString,
-    });
+  const targetSchemaVersion = schemaVersion ?? 0;
+  const targetRollupAddress = rollupAddress ?? EthAddress.ZERO;
+  const targetDatabaseVersion = new DatabaseVersion(targetSchemaVersion, targetRollupAddress);
+
+  // DB version: database schema version + rollup address combined)
+  const dbVersion = store.openSingleton<string>('dbVersion');
+
+  const storedDatabaseVersion = isSyncStore(store)
+    ? (dbVersion as AztecSingleton<string>).get()
+    : await (dbVersion as AztecAsyncSingleton<string>).getAsync();
+
+  // Check if this is an old format store (has rollupAddress singleton but no dbVersion)
+  const oldRollupSingleton = store.openSingleton<string>('rollupAddress');
+  const hasOldFormat = isSyncStore(store)
+    ? !storedDatabaseVersion && !!(oldRollupSingleton as AztecSingleton<string>).get()
+    : !storedDatabaseVersion && !!(await (oldRollupSingleton as AztecAsyncSingleton<string>).getAsync());
 
+  if (
+    hasOldFormat ||
+    doesStoreNeedToBeCleared(
+      targetDatabaseVersion,
+      storedDatabaseVersion,
+      targetSchemaVersion,
+      targetRollupAddress,
+      log,
+    )
+  ) {
+    if (hasOldFormat) {
+      log?.warn('Detected old store format without dbVersion, clearing database');
+    }
     await store.clear();
   }
 
-  await rollupAddressValue.set(rollupAddressString);
+  await dbVersion.set(targetDatabaseVersion.toBuffer().toString('utf-8'));
+
   return store;
 }
+
+function doesStoreNeedToBeCleared(
+  targetDatabaseVersion: DatabaseVersion,
+  storedDatabaseVersion: string | undefined,
+  targetSchemaVersion: number,
+  targetRollupAddress: EthAddress,
+  log?: Logger,
+) {
+  if (storedDatabaseVersion) {
+    try {
+      const storedVersion = DatabaseVersion.fromBuffer(Buffer.from(storedDatabaseVersion, 'utf-8'));
+      const cmp = storedVersion.cmp(targetDatabaseVersion);
+
+      if (cmp === undefined) {
+        log?.warn('Rollup address changed, clearing database', {
+          stored: storedVersion.rollupAddress.toString(),
+          current: targetRollupAddress.toString(),
+        });
+        return true;
+      }
+
+      if (cmp !== 0) {
+        log?.warn('Schema version changed, clearing database', {
+          stored: storedVersion.schemaVersion,
+          current: targetSchemaVersion,
+        });
+
+        return true;
+      }
+    } catch (err) {
+      log?.warn('Failed to parse stored version, clearing database', { err });
+      return true;
+    }
+  }
+
+  return false;
+}