Merge pull request #1035 from ethereumjs/checkpoint-trie-cache

Trie: cache CheckpointTrie in memory

Author: holgerd77

packages/trie/src/checkpointDb.ts

@@ -0,0 +1,152 @@
+import { LevelUp } from 'levelup'
+import { DB, BatchDBOp, ENCODING_OPTS } from './db'
+
+export type Checkpoint = {
+  // We cannot use a Buffer => Buffer map directly. If you create two Buffers with the same internal value,
+  // then when setting a value on the Map, it actually creates two indices.
+  keyValueMap: Map<string, Buffer | null>
+  root: Buffer
+}
+
+/**
+ * DB is a thin wrapper around the underlying levelup db,
+ * which validates inputs and sets encoding type.
+ */
+export class CheckpointDB extends DB {
+  public checkpoints: Checkpoint[]
+
+  /**
+   * Initialize a DB instance. If `leveldb` is not provided, DB
+   * defaults to an [in-memory store](https://github.com/Level/memdown).
+   * @param leveldb - An abstract-leveldown compliant store
+   */
+  constructor(leveldb?: LevelUp) {
+    super(leveldb)
+    // Roots of trie at the moment of checkpoint
+    this.checkpoints = []
+  }
+
+  /**
+   * Is the DB during a checkpoint phase?
+   */
+  get isCheckpoint() {
+    return this.checkpoints.length > 0
+  }
+
+  /**
+   * Adds a new checkpoint to the stack
+   * @param root
+   */
+  checkpoint(root: Buffer) {
+    this.checkpoints.push({ keyValueMap: new Map<string, Buffer>(), root })
+  }
+
+  /**
+   * Commits the latest checkpoint
+   */
+  async commit() {
+    const { keyValueMap } = this.checkpoints.pop()!
+    if (!this.isCheckpoint) {
+      // This was the final checkpoint, we should now commit and flush everything to disk
+      const batchOp: BatchDBOp[] = []
+      keyValueMap.forEach(function (value, key) {
+        if (value === null) {
+          batchOp.push({
+            type: 'del',
+            key: Buffer.from(key, 'binary'),
+          })
+        } else {
+          batchOp.push({
+            type: 'put',
+            key: Buffer.from(key, 'binary'),
+            value,
+          })
+        }
+      })
+      await this.batch(batchOp)
+    } else {
+      // dump everything into the current (higher level) cache
+      const currentKeyValueMap = this.checkpoints[this.checkpoints.length - 1].keyValueMap
+      keyValueMap.forEach((value, key) => currentKeyValueMap.set(key, value))
+    }
+  }
+
+  /**
+   * Reverts the latest checkpoint
+   */
+  async revert() {
+    const { root } = this.checkpoints.pop()!
+    return root
+  }
+
+  /**
+   * Retrieves a raw value from leveldb.
+   * @param key
+   * @returns A Promise that resolves to `Buffer` if a value is found or `null` if no value is found.
+   */
+  async get(key: Buffer): Promise<Buffer | null> {
+    // Lookup the value in our cache. We return the latest checkpointed value (which should be the value on disk)
+    for (let index = this.checkpoints.length - 1; index >= 0; index--) {
+      const value = this.checkpoints[index].keyValueMap.get(key.toString('binary'))
+      if (value !== undefined) {
+        return value
+      }
+    }
+    // Nothing has been found in cache, look up from disk
+
+    const value = await super.get(key)
+    if (this.isCheckpoint) {
+      // Since we are a checkpoint, put this value in cache, so future `get` calls will not look the key up again from disk.
+      this.checkpoints[this.checkpoints.length - 1].keyValueMap.set(key.toString('binary'), value)
+    }
+
+    return value
+  }
+
+  /**
+   * Writes a value directly to leveldb.
+   * @param key The key as a `Buffer`
+   * @param value The value to be stored
+   */
+  async put(key: Buffer, val: Buffer): Promise<void> {
+    if (this.isCheckpoint) {
+      // put value in cache
+      this.checkpoints[this.checkpoints.length - 1].keyValueMap.set(key.toString('binary'), val)
+    } else {
+      await super.put(key, val)
+    }
+  }
+
+  /**
+   * Removes a raw value in the underlying leveldb.
+   * @param keys
+   */
+  async del(key: Buffer): Promise<void> {
+    if (this.isCheckpoint) {
+      // delete the value in the current cache
+      this.checkpoints[this.checkpoints.length - 1].keyValueMap.set(key.toString('binary'), null)
+    } else {
+      // delete the value on disk
+      await this._leveldb.del(key, ENCODING_OPTS)
+    }
+  }
+
+  /**
+   * Performs a batch operation on db.
+   * @param opStack A stack of levelup operations
+   */
+  async batch(opStack: BatchDBOp[]): Promise<void> {
+    if (this.isCheckpoint) {
+      for (const op of opStack) {
+        if (op.type === 'put') {
+          await this.put(op.key, op.value)
+          // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+        } else if (op.type === 'del') {
+          await this.del(op.key)
+        }
+      }
+    } else {
+      await super.batch(opStack)
+    }
+  }
+}

packages/trie/src/checkpointTrie.ts

@@ -1,18 +1,22 @@
 import { Trie as BaseTrie } from './baseTrie'
+import { CheckpointDB } from './checkpointDb'
 
 /**
  * Adds checkpointing to the {@link BaseTrie}
  */
 export class CheckpointTrie extends BaseTrie {
+  db: CheckpointDB
+
   constructor(...args: any) {
     super(...args)
+    this.db = new CheckpointDB(...args)
   }
 
   /**
    * Is the trie during a checkpoint phase?
    */
   get isCheckpoint() {
-    return this.db.checkpoints.length > 0
+    return this.db.isCheckpoint
   }
 
   /**
@@ -34,7 +38,7 @@ export class CheckpointTrie extends BaseTrie {
     }
 
     await this.lock.wait()
-    this.db.commit()
+    await this.db.commit()
     this.lock.signal()
   }
 
@@ -61,7 +65,7 @@ export class CheckpointTrie extends BaseTrie {
     const db = this.db.copy()
     const trie = new CheckpointTrie(db._leveldb, this.root)
     if (includeCheckpoints && this.isCheckpoint) {
-      trie.db.checkpoints = this.db.checkpoints.slice()
+      trie.db.checkpoints = [...this.db.checkpoints]
     }
     return trie
   }

packages/trie/src/db.ts

@@ -3,11 +3,6 @@ const level = require('level-mem')
 
 export const ENCODING_OPTS = { keyEncoding: 'binary', valueEncoding: 'binary' }
 
-export type Checkpoint = {
-  root: Buffer
-  revertOps: BatchDBOp[]
-}
-
 export type BatchDBOp = PutBatch | DelBatch
 export interface PutBatch {
   type: 'put'
@@ -24,8 +19,6 @@ export interface DelBatch {
  * which validates inputs and sets encoding type.
  */
 export class DB {
-  public checkpoints: Checkpoint[]
-
   _leveldb: LevelUp
 
   /**
@@ -34,53 +27,9 @@ export class DB {
    * @param leveldb - An abstract-leveldown compliant store
    */
   constructor(leveldb?: LevelUp) {
-    // Roots of trie at the moment of checkpoint
-    this.checkpoints = []
-
     this._leveldb = leveldb || level()
   }
 
-  /**
-   * Is the DB during a checkpoint phase?
-   */
-  get isCheckpoint() {
-    return this.checkpoints.length > 0
-  }
-
-  /**
-   * Adds a new checkpoint to the stack
-   * @param root
-   */
-  checkpoint(root: Buffer) {
-    this.checkpoints.push({ root, revertOps: [] })
-  }
-
-  /**
-   * Commits the latest checkpoint
-   */
-  commit() {
-    const { root, revertOps } = this.checkpoints.pop()!
-    // On nested checkpoints put the revertOps on the parent
-    // stack in case there is a revert
-    if (this.isCheckpoint) {
-      this.checkpoints[this.checkpoints.length - 1].revertOps.concat(revertOps)
-    }
-    return root
-  }
-
-  /**
-   * Reverts the latest checkpoint
-   */
-  async revert() {
-    const { root, revertOps } = this.checkpoints.pop()!
-    await this.batch(revertOps.reverse())
-    return root
-  }
-
-  private addCPRevertOperation(op: BatchDBOp) {
-    this.checkpoints[this.checkpoints.length - 1].revertOps.push(op)
-  }
-
   /**
    * Retrieves a raw value from leveldb.
    * @param key
@@ -106,75 +55,23 @@ export class DB {
    * @param value The value to be stored
    */
   async put(key: Buffer, val: Buffer): Promise<void> {
-    const revertOps: BatchDBOp[] = []
-    // In CP mode check for an old value to be put
-    // on the revert ops stack
-    if (this.isCheckpoint) {
-      const oldValue = await this.get(key)
-      if (oldValue !== null) {
-        revertOps.push({
-          type: 'put',
-          key: key,
-          value: oldValue,
-        })
-      }
-    }
     await this._leveldb.put(key, val, ENCODING_OPTS)
-    // In CP mode add del to the revert ops stack
-    if (this.isCheckpoint) {
-      revertOps.push({
-        type: 'del',
-        key: key,
-      })
-      for (const revertOp of revertOps) {
-        this.addCPRevertOperation(revertOp)
-      }
-    }
   }
 
   /**
    * Removes a raw value in the underlying leveldb.
    * @param keys
    */
   async del(key: Buffer): Promise<void> {
-    const revertOps: BatchDBOp[] = []
-    // In CP mode check for an old value to be put
-    // on the revert ops stack
-    if (this.isCheckpoint) {
-      const oldValue = await this.get(key)
-      if (oldValue !== null) {
-        revertOps.push({
-          type: 'put',
-          key: key,
-          value: oldValue,
-        })
-      }
-    }
     await this._leveldb.del(key, ENCODING_OPTS)
-    if (this.isCheckpoint) {
-      for (const revertOp of revertOps) {
-        this.addCPRevertOperation(revertOp)
-      }
-    }
   }
 
   /**
    * Performs a batch operation on db.
    * @param opStack A stack of levelup operations
    */
   async batch(opStack: BatchDBOp[]): Promise<void> {
-    if (this.isCheckpoint) {
-      for (const op of opStack) {
-        if (op.type === 'put') {
-          await this.put(op.key, op.value)
-          // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
-        } else if (op.type === 'del') {
-          await this.del(op.key)
-        }
-      }
-    } else {
-      await this._leveldb.batch(opStack, ENCODING_OPTS)
-    }
+    await this._leveldb.batch(opStack, ENCODING_OPTS)
   }
 
   /**

packages/trie/src/secure.ts

@@ -89,8 +89,11 @@ export class SecureTrie extends CheckpointTrie {
    * @param includeCheckpoints - If true and during a checkpoint, the copy will contain the checkpointing metadata and will use the same scratch as underlying db.
    */
   copy(includeCheckpoints = true): SecureTrie {
-    const trie = super.copy(includeCheckpoints)
-    const db = trie.db.copy()
-    return new SecureTrie(db._leveldb, this.root)
+    const db = this.db.copy()
+    const secureTrie = new SecureTrie(db._leveldb, this.root)
+    if (includeCheckpoints && this.isCheckpoint) {
+      secureTrie.db.checkpoints = [...this.db.checkpoints]
+    }
+    return secureTrie
   }
 }