Improve how StateStore handles multiple snapshots (#8768)

* Improve how StateStore handles multiple snapshots * Ensure that the snapshot id limit is not exceeded * Refactor the limit check Co-authored-by: shuse2 <[email protected]> * Update unit test --------- Co-authored-by: shuse2 <[email protected]>
LiskArchive · Jul 31, 2023 · 970b641 · 970b641
1 parent 137515d
commit 970b641
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/elements/lisk-chain/src/state_store/state_store.ts b/elements/lisk-chain/src/state_store/state_store.ts
@@ -38,11 +38,13 @@ export class StateStore {
 	private readonly _prefix: Buffer;
 	private _cache: CacheDB;
 	private _snapshot: CacheDB | undefined;
+	private _latestSnapshotId: number;
 
 	public constructor(db: DatabaseReader, prefix?: Buffer, cache?: CacheDB) {
 		this._db = db;
 		this._prefix = prefix ?? DB_KEY_STATE_STORE;
 		this._cache = cache ?? new CacheDB();
+		this._latestSnapshotId = -1;
 	}
 
 	// TODO: Remove accepting number for subStorePrefix
@@ -222,14 +224,25 @@ export class StateStore {
 	// createSnapshot follows the same interface as stateDB. However, it does not support multi snapshot.
 	public createSnapshot(): number {
 		this._snapshot = this._cache.copy();
-		return 0;
+		this._latestSnapshotId += 1;
+
+		if (this._latestSnapshotId === Number.MAX_SAFE_INTEGER) {
+			this._latestSnapshotId = 0;
+		}
+
+		return this._latestSnapshotId;
 	}
 
 	// restoreSnapshot does not support multi-snapshot. Therefore, id is not used.
-	public restoreSnapshot(_id: number): void {
+	public restoreSnapshot(id: number): void {
+		if (id !== this._latestSnapshotId) {
+			throw new Error('Invalid snapshot ID. Cannot revert to an older snapshot.');
+		}
+
 		if (!this._snapshot) {
 			throw new Error('Snapshot must be taken first before reverting');
 		}
+
 		this._cache = this._snapshot;
 		this._snapshot = undefined;
 	}

diff --git a/elements/lisk-chain/test/unit/state_store/state_store.spec.ts b/elements/lisk-chain/test/unit/state_store/state_store.spec.ts
@@ -313,6 +313,30 @@ describe('state store', () => {
 
 			await expect(subStore.get(Buffer.from([0]))).rejects.toThrow(NotFoundError);
 			await expect(subStore.get(existingKey)).resolves.toEqual(existingValue);
+			expect(subStore['_snapshot']).toBeUndefined();
+			expect(subStore['_latestSnapshotId']).toBe(id);
+		});
+
+		it('should throw an error when restoring with an invalid snapshot ID', () => {
+			const subStore = stateStore.getStore(moduleID, storePrefix);
+
+			expect(() => subStore.restoreSnapshot(100)).toThrow(
+				'Invalid snapshot ID. Cannot revert to an older snapshot.',
+			);
+		});
+
+		it('should throw an error when restoring without taking a snapshot first', () => {
+			const subStore = stateStore.getStore(moduleID, storePrefix);
+
+			expect(() => subStore.restoreSnapshot(0)).toThrow(
+				'Invalid snapshot ID. Cannot revert to an older snapshot.',
+			);
+		});
+
+		it('should reset the snapshot id to 0 when the limit is reached', () => {
+			stateStore['_latestSnapshotId'] = Number.MAX_SAFE_INTEGER - 1;
+			const snapshotId = stateStore.createSnapshot();
+			expect(snapshotId).toBe(0);
 		});
 	});