Batch Merkle Update (#432)

- Added re-entrant lock to SMT
- Fixed issue in SMT.getLeaf(...)
- Added batchUpdate(...) to MerkleTree and SMT implementation

packages/core/package.json

@@ -29,6 +29,7 @@
   },
   "dependencies": {
     "abstract-leveldown": "^6.0.3",
+    "async-lock": "^1.2.2",
     "async-mutex": "^0.1.3",
     "axios": "^0.19.0",
     "bn.js": "^4.11.8",

packages/core/src/app/block-production/merkle-tree.ts

@@ -1,20 +1,24 @@
 /* External Imports */
-import { Mutex } from 'async-mutex'
+import * as AsyncLock from 'async-lock'
+import * as domain from 'domain'
 import * as assert from 'assert'
 
+/* Internal Imports */
 import {
   BIG_ENDIAN,
   BigNumber,
   DB,
   HashFunction,
   MerkleTreeInclusionProof,
   MerkleTreeNode,
+  MerkleUpdate,
   ONE,
   SparseMerkleTree,
   TWO,
   ZERO,
 } from '../../types'
 import { keccak256 } from '../utils/crypto'
+import { runInDomain } from '../utils'
 
 /**
  * SparseMerkleTree implementation assuming a 256-bit hash algorithm is used.
@@ -25,7 +29,10 @@ export class SparseMerkleTreeImpl implements SparseMerkleTree {
 
   private root: MerkleTreeNode
   private zeroHashes: Buffer[]
-  private readonly treeMutex: Mutex = new Mutex()
+  private static readonly lockKey: string = 'lock'
+  private readonly treeLock: AsyncLock = new AsyncLock({
+    domainReentrant: true,
+  })
   private readonly hashBuffer: Buffer = Buffer.alloc(64)
 
   constructor(
@@ -46,7 +53,7 @@ export class SparseMerkleTreeImpl implements SparseMerkleTree {
   }
 
   public async getLeaf(leafKey: BigNumber, rootHash?: Buffer): Promise<Buffer> {
-    return this.treeMutex.runExclusive(async () => {
+    return this.treeLock.acquire(SparseMerkleTreeImpl.lockKey, async () => {
       if (!!rootHash && !rootHash.equals(this.root.hash)) {
         return undefined
       }
@@ -55,7 +62,7 @@ export class SparseMerkleTreeImpl implements SparseMerkleTree {
       if (!nodesInPath || !nodesInPath.length) {
         return undefined
       }
-      const leaf: MerkleTreeNode = nodesInPath[nodesInPath.length]
+      const leaf: MerkleTreeNode = nodesInPath[nodesInPath.length - 1]
       // Will only match if we were able to traverse all the way to the leaf
       return leaf.key.equals(leafKey) ? leaf.value : undefined
     })
@@ -69,7 +76,7 @@ export class SparseMerkleTreeImpl implements SparseMerkleTree {
       return false
     }
 
-    return this.treeMutex.runExclusive(async () => {
+    return this.treeLock.acquire(SparseMerkleTreeImpl.lockKey, async () => {
       const leafHash: Buffer = this.hashFunction(inclusionProof.value)
       if (!!(await this.getNode(leafHash, inclusionProof.key))) {
         return true
@@ -126,72 +133,95 @@ export class SparseMerkleTreeImpl implements SparseMerkleTree {
     })
   }
 
-  public async update(leafKey: BigNumber, leafValue: Buffer): Promise<boolean> {
-    let nodesToUpdate: MerkleTreeNode[] = await this.treeMutex.runExclusive(
-      () => {
-        return this.getNodesInPath(leafKey)
-      }
-    )
+  public async update(
+    leafKey: BigNumber,
+    leafValue: Buffer,
+    d?: domain.Domain
+  ): Promise<boolean> {
+    return runInDomain(d, async () => {
+      return this.treeLock.acquire(SparseMerkleTreeImpl.lockKey, async () => {
+        let nodesToUpdate: MerkleTreeNode[] = await this.getNodesInPath(leafKey)
+
+        if (!nodesToUpdate) {
+          return false
+        } else if (nodesToUpdate.length !== this.height) {
+          if (
+            !(await this.verifyAndStorePartiallyEmptyPath(
+              leafKey,
+              nodesToUpdate.length
+            ))
+          ) {
+            return false
+          }
+          nodesToUpdate = await this.getNodesInPath(leafKey)
+        }
 
-    if (!nodesToUpdate) {
-      return false
-    } else if (nodesToUpdate.length !== this.height) {
-      if (
-        !(await this.verifyAndStorePartiallyEmptyPath(
-          leafKey,
-          nodesToUpdate.length
-        ))
-      ) {
-        return false
-      }
-      nodesToUpdate = await this.getNodesInPath(leafKey)
-    }
+        const leaf: MerkleTreeNode = nodesToUpdate[nodesToUpdate.length - 1]
+        const idsToDelete: Buffer[] = [this.getNodeID(leaf)]
+        leaf.hash = this.hashFunction(leafValue)
+        leaf.value = leafValue
+
+        let updatedChild: MerkleTreeNode = leaf
+        let depth: number = nodesToUpdate.length - 2 // -2 because this array also contains the leaf
+
+        // Iteratively update all nodes from the leaf-pointer node up to the root
+        for (; depth >= 0; depth--) {
+          idsToDelete.push(this.getNodeID(nodesToUpdate[depth]))
+          updatedChild = this.updateNode(
+            nodesToUpdate[depth],
+            updatedChild,
+            leafKey,
+            depth
+          )
+        }
 
-    return this.treeMutex.runExclusive(async () => {
-      // Have to check to make sure nodesToUpdate didn't change while we didn't have the lock.
-      const prevLeaf = nodesToUpdate[nodesToUpdate.length - 1]
-      const leafStillExists: MerkleTreeNode = await this.getNode(
-        prevLeaf.hash,
-        prevLeaf.key
-      )
-      if (!leafStillExists) {
-        nodesToUpdate = await this.getNodesInPath(leafKey)
-      }
+        await Promise.all([
+          ...nodesToUpdate.map((n) => this.db.put(this.getNodeID(n), n.value)),
+          ...idsToDelete.map((id) => this.db.del(id)),
+        ])
 
-      const leaf: MerkleTreeNode = nodesToUpdate[nodesToUpdate.length - 1]
-      const idsToDelete: Buffer[] = [this.getNodeID(leaf)]
-      leaf.hash = this.hashFunction(leafValue)
-      leaf.value = leafValue
-
-      let updatedChild: MerkleTreeNode = leaf
-      let depth: number = nodesToUpdate.length - 2 // -2 because this array also contains the leaf
-
-      // Iteratively update all nodes from the leaf-pointer node up to the root
-      for (; depth >= 0; depth--) {
-        idsToDelete.push(this.getNodeID(nodesToUpdate[depth]))
-        updatedChild = this.updateNode(
-          nodesToUpdate[depth],
-          updatedChild,
-          leafKey,
-          depth
-        )
-      }
+        this.root = nodesToUpdate[0]
+        return true
+      })
+    })
+  }
 
-      await Promise.all([
-        ...nodesToUpdate.map((n) => this.db.put(this.getNodeID(n), n.value)),
-        ...idsToDelete.map((id) => this.db.del(id)),
-      ])
+  public async batchUpdate(updates: MerkleUpdate[]): Promise<boolean> {
+    const d: domain.Domain = domain.create()
+
+    return runInDomain(d, () => {
+      return this.treeLock.acquire(SparseMerkleTreeImpl.lockKey, async () => {
+        for (const update of updates) {
+          if (
+            !(await this.verifyAndStore({
+              rootHash: this.root.hash,
+              key: update.key,
+              value: update.oldValue,
+              siblings: update.oldValueProofSiblings,
+            }))
+          ) {
+            return false
+          }
+        }
 
-      this.root = nodesToUpdate[0]
-      return true
+        for (const update of updates) {
+          if (!(await this.update(update.key, update.newValue, d))) {
+            throw Error(
+              "Verify and Store worked but update didn't! This should never happen!"
+            )
+          }
+        }
+
+        return true
+      })
     })
   }
 
   public async getMerkleProof(
     leafKey: BigNumber,
     leafValue: Buffer
   ): Promise<MerkleTreeInclusionProof> {
-    return this.treeMutex.runExclusive(async () => {
+    return this.treeLock.acquire(SparseMerkleTreeImpl.lockKey, async () => {
       if (!this.root || !this.root.hash || !this.root.value) {
         return undefined
       }

packages/core/src/app/utils/misc.ts

@@ -1,4 +1,5 @@
 /* External Imports */
+import * as domain from 'domain'
 
 /* Internal Imports */
 import { Transaction } from '../../types'
@@ -190,3 +191,18 @@ export const except = <T>(list: T[], element: T): T[] => {
     return item !== element
   })
 }
+
+/**
+ * Runs the provided function in the provided domain if one is provided.
+ * If a domain is falsy, this will create a domain for this function to run in.
+ * @param d The domain in which this function should run.
+ * @param func The function to run.
+ * @returns The result of the function to be run
+ */
+export const runInDomain = async (
+  d: domain.Domain,
+  func: () => any
+): Promise<any> => {
+  const domainToUse: domain.Domain = !!d ? d : domain.create()
+  return domainToUse.run(func)
+}

packages/core/src/types/utils/merkle-tree.interface.ts → packages/core/src/types/block-production/merkle-tree.interface.ts

@@ -13,6 +13,13 @@ export interface MerkleTreeInclusionProof {
   siblings: Buffer[]
 }
 
+export interface MerkleUpdate {
+  key: BigNumber
+  oldValue: Buffer
+  oldValueProofSiblings: Buffer[]
+  newValue: Buffer
+}
+
 export interface MerkleTree {
   /**
    * Gets the root hash for this tree.
@@ -31,6 +38,18 @@ export interface MerkleTree {
    */
   update(leafKey: BigNumber, leafValue: Buffer): Promise<boolean>
 
+  /**
+   * Updates the provided keys in the Merkle Tree in an atomic fashion
+   * including all ancestor hashes that result from these modifications.
+   *
+   * Note: It is known that applying one update invalidates the proof for the next
+   * update, which should be accounted for within this method.
+   *
+   * @param updates The updates to execute
+   * @return true if the update succeeded, false if we're missing the intermediate nodes / siblings required for this
+   */
+  batchUpdate(updates: MerkleUpdate[]): Promise<boolean>
+
   /**
    * Gets a Merkle proof for the provided leaf value at the provided key in the tree.
    *

packages/core/src/types/utils/index.ts

@@ -1,2 +1,2 @@
 export * from './hash-algorithms'
-export * from './merkle-tree.interface'
+export * from '../block-production/merkle-tree.interface'