From 38a55abb0b6d5e2dc0f3dfdf716ad217018a3cf5 Mon Sep 17 00:00:00 2001
From: Elizabeth Engelman <emengelman@gmail.com>
Date: Fri, 8 Feb 2019 15:15:36 -0600
Subject: [PATCH] Make sure that statediff has been processed before pruning

---
 cmd/geth/config.go      |  4 ++++
 core/blockchain.go      | 34 +++++++++++++++++++++++++++-------
 core/blockchain_test.go | 12 ++++++------
 eth/backend.go          |  2 +-
 4 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/cmd/geth/config.go b/cmd/geth/config.go
index ec97070c7..fe916d37a 100644
--- a/cmd/geth/config.go
+++ b/cmd/geth/config.go
@@ -158,6 +158,10 @@ func makeFullNode(ctx *cli.Context) *node.Node {
 		cfg.Eth.ConstantinopleOverride = new(big.Int).SetUint64(ctx.GlobalUint64(utils.ConstantinopleOverrideFlag.Name))
 	}
 
+	if ctx.GlobalBool(utils.StateDiffFlag.Name) {
+		cfg.Eth.StateDiff = true
+	}
+
 	utils.RegisterEthService(stack, &cfg.Eth)
 
 	if ctx.GlobalBool(utils.StateDiffFlag.Name) {
diff --git a/core/blockchain.go b/core/blockchain.go
index f0dba102c..912074ef8 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -95,7 +95,7 @@ type CacheConfig struct {
 	TrieDirtyLimit      int           // Memory limit (MB) at which to start flushing dirty trie nodes to disk
 	TrieDirtyDisabled   bool          // Whether to disable trie write caching and GC altogether (archive node)
 	TrieTimeLimit       time.Duration // Time limit after which to flush the current in-memory trie to disk
-	ProcessingStateDiffs bool          // Whether statediffs processing should be taken into a account before a trie is pruned
+	ProcessStateDiffs bool            // Whether statediffs processing should be taken into a account before a trie is pruned
 }
 
 // BlockChain represents the canonical chain given a database with a genesis
@@ -937,8 +937,17 @@ func (bc *BlockChain) WriteBlockWithoutState(block *types.Block, td *big.Int) (e
 }
 
 func (bc *BlockChain) AddToStateDiffProcessedCollection(hash common.Hash) {
-	count := bc.stateDiffsProcessed[hash]
-	bc.stateDiffsProcessed[hash] = count + 1
+	count, ok := bc.stateDiffsProcessed[hash]
+	if count > 1 {
+		log.Error("count is too high", "count", count, "hash", hash.Hex())
+	}
+
+	if ok {
+		count++
+		bc.stateDiffsProcessed[hash] = count
+	} else {
+		bc.stateDiffsProcessed[hash] = 1
+	}
 }
 
 // WriteBlockWithState writes the block and all associated state to the database.
@@ -1025,16 +1034,27 @@ func (bc *BlockChain) writeBlockWithState(block *types.Block, receipts []*types.
 					bc.triegc.Push(root, number)
 					break
 				}
-
-				if bc.cacheConfig.ProcessingStateDiffs {
-					if !bc.allowedRootToBeDereferenced(root.(common.Hash)) {
+				if bc.cacheConfig.ProcessStateDiffs {
+					count, ok := bc.stateDiffsProcessed[root.(common.Hash)]
+					//if we haven't processed the statediff for a given state root and it's child, don't dereference it yet
+					if !ok {
+						log.Info("Current root NOT found root in stateDiffsProcessed", "root", root.(common.Hash).Hex())
+						bc.triegc.Push(root, number)
+						break
+					}
+					if count < 2 {
+						log.Info("Current root has not yet been processed for it's child", "root", root.(common.Hash).Hex())
 						bc.triegc.Push(root, number)
 						break
 					} else {
+						log.Warn("Current root found in stateDiffsProcessed collection with a count of 2, okay to dereference",
+							"root", root.(common.Hash).Hex(),
+							"blockNumber", uint64(-number),
+							"size of stateDiffsProcessed", len(bc.stateDiffsProcessed))
 						delete(bc.stateDiffsProcessed, root.(common.Hash))
 					}
 				}
-
+				log.Info("DEREFERENCING", "root", root.(common.Hash).Hex())
 				triedb.Dereference(root.(common.Hash))
 			}
 		}
diff --git a/core/blockchain_test.go b/core/blockchain_test.go
index 682bab8d9..22b6df734 100644
--- a/core/blockchain_test.go
+++ b/core/blockchain_test.go
@@ -1816,11 +1816,11 @@ func TestProcessingStateDiffs(t *testing.T) {
 	defaultTrieDirtyCache := 256
 	defaultTrieTimeout := 60 * time.Minute
 	cacheConfig := &CacheConfig{
-		Disabled:             false,
-		TrieCleanLimit:       defaultTrieCleanCache,
-		TrieDirtyLimit:       defaultTrieDirtyCache,
-		TrieTimeLimit:        defaultTrieTimeout,
-		ProcessingStateDiffs: true,
+		Disabled:          false,
+		TrieCleanLimit:    defaultTrieCleanCache,
+		TrieDirtyLimit:    defaultTrieDirtyCache,
+		TrieTimeLimit:     defaultTrieTimeout,
+		ProcessStateDiffs: true,
 	}
 	db := ethdb.NewMemDatabase()
 	genesis := new(Genesis).MustCommit(db)
@@ -1883,7 +1883,7 @@ func TestProcessingStateDiffs(t *testing.T) {
 	}
 }
 
-func containsRootHash(collection []common.Hash, hash common.Hash) bool {
+func containsRootHash(collection []common.Hash, hash common.Hash) bool{
 	for _, n := range collection {
 		if n == hash {
 			return true
diff --git a/eth/backend.go b/eth/backend.go
index 130cb44ef..a64d96347 100644
--- a/eth/backend.go
+++ b/eth/backend.go
@@ -171,7 +171,7 @@ func New(ctx *node.ServiceContext, config *Config) (*Ethereum, error) {
 			TrieDirtyLimit:      config.TrieDirtyCache,
 			TrieDirtyDisabled:   config.NoPruning,
 			TrieTimeLimit:       config.TrieTimeout,
-			ProcessingStateDiffs: config.StateDiff,
+			ProcessStateDiffs: config.StateDiff,
 		}
 	)
 	eth.blockchain, err = core.NewBlockChain(chainDb, cacheConfig, chainConfig, eth.engine, vmConfig, eth.shouldPreserve)