chore: unwinding work

Author: LHerskind

yarn-project/archiver/src/archiver/archiver.ts

@@ -561,7 +561,17 @@ enum Operation {
  * store would need to include otherwise while exposing fewer functions and logic directly to the archiver.
  */
 class ArchiverStoreHelper
-  implements Omit<ArchiverDataStore, 'addLogs' | 'addContractClasses' | 'addContractInstances' | 'addFunctions'>
+  implements
+    Omit<
+      ArchiverDataStore,
+      | 'addLogs'
+      | 'deleteLogs'
+      | 'addContractClasses'
+      | 'deleteContractClasses'
+      | 'addContractInstances'
+      | 'deleteContractInstances'
+      | 'addFunctions'
+    >
 {
   #log = createDebugLogger('aztec:archiver:block-helper');
 
@@ -699,7 +709,8 @@ class ArchiverStoreHelper
           await this.#updateDeployedContractInstances(blockLogs, block.data.number, Operation.Delete);
         }),
       )),
-      await this.store.unwindBlocks(from, blocksToUnwind),
+      this.store.deleteLogs(blocks.map(b => b.data)),
+      this.store.unwindBlocks(from, blocksToUnwind),
     ].every(Boolean);
   }
 

yarn-project/archiver/src/archiver/archiver_store.ts

@@ -84,6 +84,7 @@ export interface ArchiverDataStore {
    * @returns True if the operation is successful.
    */
   addLogs(blocks: L2Block[]): Promise<boolean>;
+  deleteLogs(blocks: L2Block[]): Promise<boolean>;
 
   /**
    * Append L1 to L2 messages to the store.
@@ -170,6 +171,8 @@ export interface ArchiverDataStore {
    */
   addContractClasses(data: ContractClassPublic[], blockNumber: number): Promise<boolean>;
 
+  deleteContractClasses(data: ContractClassPublic[], blockNumber: number): Promise<boolean>;
+
   /**
    * Returns a contract class given its id, or undefined if not exists.
    * @param id - Id of the contract class.
@@ -183,6 +186,7 @@ export interface ArchiverDataStore {
    * @returns True if the operation is successful.
    */
   addContractInstances(data: ContractInstanceWithAddress[], blockNumber: number): Promise<boolean>;
+  deleteContractInstances(data: ContractInstanceWithAddress[], blockNumber: number): Promise<boolean>;
 
   /**
    * Adds private functions to a contract class.

yarn-project/archiver/src/archiver/archiver_store_test_suite.ts

@@ -52,6 +52,20 @@ export function describeArchiverDataStore(testName: string, getStore: () => Arch
       });
     });
 
+    describe('unwindBlocks', () => {
+      it('unwinding blocks will remove blocks from the chain', async () => {
+        await store.addBlocks(blocks);
+        const blockNumber = await store.getSynchedL2BlockNumber();
+
+        expect(await store.getBlocks(blockNumber, 1)).toEqual([blocks[blocks.length - 1]]);
+
+        await store.unwindBlocks(blockNumber, 1);
+
+        expect(await store.getSynchedL2BlockNumber()).toBe(blockNumber - 1);
+        expect(await store.getBlocks(blockNumber, 1)).toEqual([]);
+      });
+    });
+
     describe('getBlocks', () => {
       beforeEach(async () => {
         await store.addBlocks(blocks);
@@ -120,12 +134,32 @@ export function describeArchiverDataStore(testName: string, getStore: () => Arch
       });
     });
 
+    describe('deleteLogs', () => {
+      it('deletes encrypted & unencrypted logs', async () => {
+        const block = blocks[0].data;
+        await store.addBlocks([blocks[0]]);
+        await expect(store.addLogs([block])).resolves.toEqual(true);
+
+        expect((await store.getLogs(1, 1, LogType.NOTEENCRYPTED))[0]).toEqual(block.body.noteEncryptedLogs);
+        expect((await store.getLogs(1, 1, LogType.ENCRYPTED))[0]).toEqual(block.body.encryptedLogs);
+        expect((await store.getLogs(1, 1, LogType.UNENCRYPTED))[0]).toEqual(block.body.unencryptedLogs);
+
+        // This one is a pain for memory as we would never want to just delete memory in the middle.
+        await store.deleteLogs([block]);
+
+        expect((await store.getLogs(1, 1, LogType.NOTEENCRYPTED))[0]).toEqual(undefined);
+        expect((await store.getLogs(1, 1, LogType.ENCRYPTED))[0]).toEqual(undefined);
+        expect((await store.getLogs(1, 1, LogType.UNENCRYPTED))[0]).toEqual(undefined);
+      });
+    });
+
     describe.each([
       ['note_encrypted', LogType.NOTEENCRYPTED],
       ['encrypted', LogType.ENCRYPTED],
       ['unencrypted', LogType.UNENCRYPTED],
     ])('getLogs (%s)', (_, logType) => {
       beforeEach(async () => {
+        await store.addBlocks(blocks);
         await store.addLogs(blocks.map(b => b.data));
       });
 
@@ -167,6 +201,24 @@ export function describeArchiverDataStore(testName: string, getStore: () => Arch
       it('returns undefined if tx is not found', async () => {
         await expect(store.getTxEffect(new TxHash(Fr.random().toBuffer()))).resolves.toBeUndefined();
       });
+
+      it.each([
+        () => blocks[0].data.body.txEffects[0],
+        () => blocks[9].data.body.txEffects[3],
+        () => blocks[3].data.body.txEffects[1],
+        () => blocks[5].data.body.txEffects[2],
+        () => blocks[1].data.body.txEffects[0],
+      ])('tries to retrieves a previously stored transaction after deleted', async getExpectedTx => {
+        await store.unwindBlocks(blocks.length, blocks.length);
+
+        const expectedTx = getExpectedTx();
+        const actualTx = await store.getTxEffect(expectedTx.txHash);
+        expect(actualTx).toEqual(undefined);
+      });
+
+      it('returns undefined if tx is not found', async () => {
+        await expect(store.getTxEffect(new TxHash(Fr.random().toBuffer()))).resolves.toBeUndefined();
+      });
     });
 
     describe('L1 to L2 Messages', () => {
@@ -237,6 +289,11 @@ export function describeArchiverDataStore(testName: string, getStore: () => Arch
       it('returns undefined if contract instance is not found', async () => {
         await expect(store.getContractInstance(AztecAddress.random())).resolves.toBeUndefined();
       });
+
+      it('returns undefined if previously stored contract instances was deleted', async () => {
+        await store.deleteContractInstances([contractInstance], blockNum);
+        await expect(store.getContractInstance(contractInstance.address)).resolves.toBeUndefined();
+      });
     });
 
     describe('contractClasses', () => {
@@ -252,6 +309,17 @@ export function describeArchiverDataStore(testName: string, getStore: () => Arch
         await expect(store.getContractClass(contractClass.id)).resolves.toMatchObject(contractClass);
       });
 
+      it('returns undefined if the initial deployed contract class was deleted', async () => {
+        await store.deleteContractClasses([contractClass], blockNum);
+        await expect(store.getContractClass(contractClass.id)).resolves.toBeUndefined();
+      });
+
+      it('returns contract class if later "deployment" class was deleted', async () => {
+        await store.addContractClasses([contractClass], blockNum + 1);
+        await store.deleteContractClasses([contractClass], blockNum + 1);
+        await expect(store.getContractClass(contractClass.id)).resolves.toMatchObject(contractClass);
+      });
+
       it('returns undefined if contract class is not found', async () => {
         await expect(store.getContractClass(Fr.random())).resolves.toBeUndefined();
       });
@@ -304,6 +372,24 @@ export function describeArchiverDataStore(testName: string, getStore: () => Arch
         await store.addLogs(blocks.map(b => b.data));
       });
 
+      it('no logs returned if deleted ("txHash" filter param is respected variant)', async () => {
+        // get random tx
+        const targetBlockIndex = randomInt(numBlocks);
+        const targetTxIndex = randomInt(txsPerBlock);
+        const targetTxHash = blocks[targetBlockIndex].data.body.txEffects[targetTxIndex].txHash;
+
+        await Promise.all([
+          store.unwindBlocks(blocks.length, blocks.length),
+          store.deleteLogs(blocks.map(b => b.data)),
+        ]);
+
+        const response = await store.getUnencryptedLogs({ txHash: targetTxHash });
+        const logs = response.logs;
+
+        expect(response.maxLogsHit).toBeFalsy();
+        expect(logs.length).toEqual(0);
+      });
+
       it('"txHash" filter param is respected', async () => {
         // get random tx
         const targetBlockIndex = randomInt(numBlocks);

yarn-project/archiver/src/archiver/kv_archiver_store/block_store.ts

@@ -89,7 +89,7 @@ export class BlockStore {
       // We should only allow deleting the very last ye, otherwise we can really get some messy shit.
       const last = this.getSynchedL2BlockNumber();
       if (from != last) {
-        throw new Error(`Can only remove the tip`);
+        throw new Error(`Can only remove from the tip`);
       }
 
       for (let i = 0; i < blocksToUnwind; i++) {

yarn-project/archiver/src/archiver/kv_archiver_store/contract_class_store.ts

@@ -3,6 +3,7 @@ import { BufferReader, numToUInt8, serializeToBuffer } from '@aztec/foundation/s
 import { type AztecKVStore, type AztecMap } from '@aztec/kv-store';
 import {
   type ContractClassPublic,
+  type ContractClassPublicWithBlockNumber,
   type ExecutablePrivateFunctionWithMembershipProof,
   type UnconstrainedFunctionWithMembershipProof,
 } from '@aztec/types/contracts';
@@ -17,8 +18,18 @@ export class ContractClassStore {
     this.#contractClasses = db.openMap('archiver_contract_classes');
   }
 
-  addContractClass(contractClass: ContractClassPublic): Promise<void> {
-    return this.#contractClasses.set(contractClass.id.toString(), serializeContractClassPublic(contractClass));
+  async addContractClass(contractClass: ContractClassPublic, blockNumber: number): Promise<void> {
+    await this.#contractClasses.setIfNotExists(
+      contractClass.id.toString(),
+      serializeContractClassPublic({ ...contractClass, l2BlockNumber: blockNumber }),
+    );
+  }
+
+  async deleteContractClasses(contractClass: ContractClassPublic, blockNumber: number): Promise<void> {
+    const restoredContractClass = this.#contractClasses.get(contractClass.id.toString());
+    if (restoredContractClass && deserializeContractClassPublic(restoredContractClass).l2BlockNumber >= blockNumber) {
+      await this.#contractClasses.delete(contractClass.id.toString());
+    }
   }
 
   getContractClass(id: Fr): ContractClassPublic | undefined {
@@ -44,7 +55,7 @@ export class ContractClassStore {
       const existingClass = deserializeContractClassPublic(existingClassBuffer);
       const { privateFunctions: existingPrivateFns, unconstrainedFunctions: existingUnconstrainedFns } = existingClass;
 
-      const updatedClass: Omit<ContractClassPublic, 'id'> = {
+      const updatedClass: Omit<ContractClassPublicWithBlockNumber, 'id'> = {
         ...existingClass,
         privateFunctions: [
           ...existingPrivateFns,
@@ -63,8 +74,9 @@ export class ContractClassStore {
   }
 }
 
-function serializeContractClassPublic(contractClass: Omit<ContractClassPublic, 'id'>): Buffer {
+function serializeContractClassPublic(contractClass: Omit<ContractClassPublicWithBlockNumber, 'id'>): Buffer {
   return serializeToBuffer(
+    contractClass.l2BlockNumber,
     numToUInt8(contractClass.version),
     contractClass.artifactHash,
     contractClass.publicFunctions.length,
@@ -108,9 +120,10 @@ function serializeUnconstrainedFunction(fn: UnconstrainedFunctionWithMembershipP
   );
 }
 
-function deserializeContractClassPublic(buffer: Buffer): Omit<ContractClassPublic, 'id'> {
+function deserializeContractClassPublic(buffer: Buffer): Omit<ContractClassPublicWithBlockNumber, 'id'> {
   const reader = BufferReader.asReader(buffer);
   return {
+    l2BlockNumber: reader.readNumber(),
     version: reader.readUInt8() as 1,
     artifactHash: reader.readObject(Fr),
     publicFunctions: reader.readVector({

yarn-project/archiver/src/archiver/kv_archiver_store/contract_instance_store.ts

@@ -19,6 +19,10 @@ export class ContractInstanceStore {
     );
   }
 
+  deleteContractInstance(contractInstance: ContractInstanceWithAddress): Promise<void> {
+    return this.#contractInstances.delete(contractInstance.address.toString());
+  }
+
   getContractInstance(address: AztecAddress): ContractInstanceWithAddress | undefined {
     const contractInstance = this.#contractInstances.get(address.toString());
     return contractInstance && SerializableContractInstance.fromBuffer(contractInstance).withAddress(address);

yarn-project/archiver/src/archiver/kv_archiver_store/kv_archiver_store.ts

@@ -74,8 +74,14 @@ export class KVArchiverDataStore implements ArchiverDataStore {
     return Promise.resolve(this.#contractInstanceStore.getContractInstance(address));
   }
 
-  async addContractClasses(data: ContractClassPublic[], _blockNumber: number): Promise<boolean> {
-    return (await Promise.all(data.map(c => this.#contractClassStore.addContractClass(c)))).every(Boolean);
+  async addContractClasses(data: ContractClassPublic[], blockNumber: number): Promise<boolean> {
+    return (await Promise.all(data.map(c => this.#contractClassStore.addContractClass(c, blockNumber)))).every(Boolean);
+  }
+
+  async deleteContractClasses(data: ContractClassPublic[], blockNumber: number): Promise<boolean> {
+    return (await Promise.all(data.map(c => this.#contractClassStore.deleteContractClasses(c, blockNumber)))).every(
+      Boolean,
+    );
   }
 
   addFunctions(
@@ -90,6 +96,10 @@ export class KVArchiverDataStore implements ArchiverDataStore {
     return (await Promise.all(data.map(c => this.#contractInstanceStore.addContractInstance(c)))).every(Boolean);
   }
 
+  async deleteContractInstances(data: ContractInstanceWithAddress[], _blockNumber: number): Promise<boolean> {
+    return (await Promise.all(data.map(c => this.#contractInstanceStore.deleteContractInstance(c)))).every(Boolean);
+  }
+
   /**
    * Append new blocks to the store's list.
    * @param blocks - The L2 blocks to be added to the store and the last processed L1 block.
@@ -153,6 +163,10 @@ export class KVArchiverDataStore implements ArchiverDataStore {
     return this.#logStore.addLogs(blocks);
   }
 
+  deleteLogs(blocks: L2Block[]): Promise<boolean> {
+    return this.#logStore.deleteLogs(blocks);
+  }
+
   /**
    * Append L1 to L2 messages to the store.
    * @param messages - The L1 to L2 messages to be added to the store and the last processed L1 block.

yarn-project/archiver/src/archiver/kv_archiver_store/log_store.ts

@@ -44,17 +44,21 @@ export class LogStore {
   addLogs(blocks: L2Block[]): Promise<boolean> {
     return this.db.transaction(() => {
       blocks.forEach(block => {
-        if (block.body.noteEncryptedLogs) {
-          void this.#noteEncryptedLogs.set(block.number, block.body.noteEncryptedLogs.toBuffer());
-        }
+        void this.#noteEncryptedLogs.set(block.number, block.body.noteEncryptedLogs.toBuffer());
+        void this.#encryptedLogs.set(block.number, block.body.encryptedLogs.toBuffer());
+        void this.#unencryptedLogs.set(block.number, block.body.unencryptedLogs.toBuffer());
+      });
 
-        if (block.body.encryptedLogs) {
-          void this.#encryptedLogs.set(block.number, block.body.encryptedLogs.toBuffer());
-        }
+      return true;
+    });
+  }
 
-        if (block.body.unencryptedLogs) {
-          void this.#unencryptedLogs.set(block.number, block.body.unencryptedLogs.toBuffer());
-        }
+  deleteLogs(blocks: L2Block[]): Promise<boolean> {
+    return this.db.transaction(() => {
+      blocks.forEach(block => {
+        void this.#noteEncryptedLogs.delete(block.number);
+        void this.#encryptedLogs.delete(block.number);
+        void this.#unencryptedLogs.delete(block.number);
       });
 
       return true;