From 8d3f8bb3d5df784cfbb699878f8903adfed8549d Mon Sep 17 00:00:00 2001
From: vsmk98 <vsaimk@gmail.com>
Date: Fri, 14 Sep 2018 03:59:10 +0000
Subject: [PATCH] Raft: private state tridb commit added

---
 core/blockchain.go | 18 ++++++++++++++++--
 miner/worker.go    |  2 +-
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/core/blockchain.go b/core/blockchain.go
index e18ef7e29..f9e74c5c8 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -903,7 +903,7 @@ func (bc *BlockChain) WriteBlockWithoutState(block *types.Block, td *big.Int) (e
 }
 
 // WriteBlockWithState writes the block and all associated state to the database.
-func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.Receipt, state *state.StateDB) (status WriteStatus, err error) {
+func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.Receipt, state , privateState *state.StateDB) (status WriteStatus, err error) {
 	bc.wg.Add(1)
 	defer bc.wg.Done()
 
@@ -929,16 +929,30 @@ func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.
 	rawdb.WriteBlock(batch, block)
 
 	root, err := state.Commit(bc.chainConfig.IsEIP158(block.Number()))
+
 	if err != nil {
 		return NonStatTy, err
 	}
 	triedb := bc.stateCache.TrieDB()
 
+	// Explicit commit for privateStateTriedb to handle Raft
+	if privateState != nil {
+		privateRoot, err := privateState.Commit(bc.chainConfig.IsEIP158(block.Number()))
+		if err != nil {
+			return NonStatTy, err
+		}
+		privateTriedb := bc.privateStateCache.TrieDB()
+		if err := privateTriedb.Commit(privateRoot, false); err != nil {
+			return NonStatTy, err
+		}
+	}
+
 	// If we're running an archive node, always flush
 	if bc.cacheConfig.Disabled {
 		if err := triedb.Commit(root, false); err != nil {
 			return NonStatTy, err
 		}
+
 	} else {
 		// Full but not archive node, do proper garbage collection
 		triedb.Reference(root, common.Hash{}) // metadata reference to keep trie alive
@@ -1234,7 +1248,7 @@ func (bc *BlockChain) insertChain(chain types.Blocks) (int, []interface{}, []*ty
 		proctime := time.Since(bstart)
 
 		// Write the block to the chain and get the status.
-		status, err := bc.WriteBlockWithState(block, allReceipts, state)
+		status, err := bc.WriteBlockWithState(block, allReceipts, state, privateState)
 		if err != nil {
 			return i, events, coalescedLogs, err
 		}
diff --git a/miner/worker.go b/miner/worker.go
index d4603c411..8e169f5d3 100644
--- a/miner/worker.go
+++ b/miner/worker.go
@@ -339,7 +339,7 @@ func (self *worker) wait() {
 			core.WritePrivateStateRoot(self.chainDb, block.Root(), privateStateRoot)
 			allReceipts := mergeReceipts(work.receipts, work.privateReceipts)
 
-			stat, err := self.chain.WriteBlockWithState(block, allReceipts, work.state)
+			stat, err := self.chain.WriteBlockWithState(block, allReceipts, work.state, nil)
 			if err != nil {
 				log.Error("Failed writWriteBlockAndStating block to chain", "err", err)
 				continue