From fba3d29c06e12898e22edbdab4c64900fcefd1f0 Mon Sep 17 00:00:00 2001
From: Maddiaa0 <47148561+Maddiaa0@users.noreply.github.com>
Date: Fri, 13 Mar 2026 18:16:30 +0000
Subject: [PATCH] fix(kv-store): make LMDB clear and drop operations atomic
 across sub-databases

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 yarn-project/kv-store/src/lmdb/store.test.ts | 124 +++++++++++++++++++
 yarn-project/kv-store/src/lmdb/store.ts      |  20 +--
 2 files changed, 136 insertions(+), 8 deletions(-)
 create mode 100644 yarn-project/kv-store/src/lmdb/store.test.ts

diff --git a/yarn-project/kv-store/src/lmdb/store.test.ts b/yarn-project/kv-store/src/lmdb/store.test.ts
new file mode 100644
index 000000000000..8ea74548eef2
--- /dev/null
+++ b/yarn-project/kv-store/src/lmdb/store.test.ts
@@ -0,0 +1,124 @@
+import { toArray } from '@aztec/foundation/iterable';
+
+import { expect } from 'chai';
+
+import { openTmpStore } from './index.js';
+import type { AztecLmdbStore } from './store.js';
+
+describe('AztecLmdbStore', () => {
+  let store: AztecLmdbStore;
+
+  beforeEach(() => {
+    store = openTmpStore(true);
+  });
+
+  afterEach(async () => {
+    try {
+      await store.delete();
+    } catch {
+      // Store may already be deleted by the test
+    }
+  });
+
+  describe('clear', () => {
+    it('removes all data from maps', async () => {
+      const map = store.openMap<string, string>('test-map');
+      await map.set('key1', 'value1');
+      await map.set('key2', 'value2');
+
+      expect(map.get('key1')).to.equal('value1');
+      expect(map.get('key2')).to.equal('value2');
+
+      await store.clear();
+
+      expect(map.get('key1')).to.be.undefined;
+      expect(map.get('key2')).to.be.undefined;
+    });
+
+    it('removes all data from multimaps', async () => {
+      const multiMap = store.openMultiMap<string, string>('test-multimap');
+      await multiMap.set('key1', 'value1');
+      await multiMap.set('key1', 'value2');
+
+      const valuesBefore = await toArray(multiMap.getValues('key1'));
+      expect(valuesBefore.length).to.be.greaterThan(0);
+
+      await store.clear();
+
+      const valuesAfter = await toArray(multiMap.getValues('key1'));
+      expect(valuesAfter).to.deep.equal([]);
+    });
+
+    it('removes all data from singletons', async () => {
+      const singleton = store.openSingleton<string>('test-singleton');
+      await singleton.set('hello');
+
+      expect(singleton.get()).to.equal('hello');
+
+      await store.clear();
+
+      expect(singleton.get()).to.be.undefined;
+    });
+
+    it('removes all data from all sub-databases at once', async () => {
+      const map = store.openMap<string, string>('test-map');
+      const multiMap = store.openMultiMap<string, string>('test-multimap');
+      const singleton = store.openSingleton<string>('test-singleton');
+      const counter = store.openCounter<string>('test-counter');
+      const set = store.openSet<string>('test-set');
+
+      await map.set('mk', 'mv');
+      await multiMap.set('mmk', 'mmv');
+      await singleton.set('sv');
+      await counter.update('ck', 5);
+      await set.add('sk');
+
+      await store.clear();
+
+      expect(map.get('mk')).to.be.undefined;
+      expect(await toArray(multiMap.getValues('mmk'))).to.deep.equal([]);
+      expect(singleton.get()).to.be.undefined;
+      expect(counter.get('ck')).to.equal(0);
+      expect(set.has('sk')).to.equal(false);
+    });
+
+    it('allows writing new data after clear', async () => {
+      const map = store.openMap<string, string>('test-map');
+      await map.set('key1', 'value1');
+
+      await store.clear();
+
+      await map.set('key2', 'value2');
+      expect(map.get('key1')).to.be.undefined;
+      expect(map.get('key2')).to.equal('value2');
+    });
+  });
+
+  describe('drop', () => {
+    it('completes without error after populating all sub-databases', async () => {
+      const map = store.openMap<string, string>('test-map');
+      const multiMap = store.openMultiMap<string, string>('test-multimap');
+      const singleton = store.openSingleton<string>('test-singleton');
+
+      await map.set('mk', 'mv');
+      await multiMap.set('mmk', 'mmv');
+      await singleton.set('sv');
+
+      await store.drop();
+    });
+  });
+
+  describe('delete', () => {
+    it('drops and closes the store without error', async () => {
+      const map = store.openMap<string, string>('test-map');
+      const multiMap = store.openMultiMap<string, string>('test-multimap');
+      const singleton = store.openSingleton<string>('test-singleton');
+
+      await map.set('mk', 'mv');
+      await multiMap.set('mmk', 'mmv');
+      await singleton.set('sv');
+
+      await store.delete();
+    });
+  });
+});
diff --git a/yarn-project/kv-store/src/lmdb/store.ts b/yarn-project/kv-store/src/lmdb/store.ts
index 4ab029ede280..47226942e9cb 100644
--- a/yarn-project/kv-store/src/lmdb/store.ts
+++ b/yarn-project/kv-store/src/lmdb/store.ts
@@ -147,21 +147,25 @@ export class AztecLmdbStore implements AztecKVStore, AztecAsyncKVStore {
   }
 
   /**
-   * Clears all entries in the store & sub DBs.
+   * Clears all entries in the store & sub DBs atomically within a single transaction.
    */
   async clear() {
-    await this.#data.clearAsync();
-    await this.#multiMapData.clearAsync();
-    await this.#rootDb.clearAsync();
+    await this.#rootDb.transaction(async () => {
+      await this.#data.clearAsync();
+      await this.#multiMapData.clearAsync();
+      await this.#rootDb.clearAsync();
+    });
   }
 
   /**
-   * Drops the database & sub DBs.
+   * Drops the database & sub DBs atomically within a single transaction.
    */
   async drop() {
-    await this.#data.drop();
-    await this.#multiMapData.drop();
-    await this.#rootDb.drop();
+    await this.#rootDb.transaction(async () => {
+      await this.#data.drop();
+      await this.#multiMapData.drop();
+      await this.#rootDb.drop();
+    });
   }
 
   /**