From c03e94cc0bef3f48d7f525ef22d88ed0b55657f1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?P=C3=A9ter=20Szil=C3=A1gyi?= <peterke@gmail.com>
Date: Mon, 22 Jan 2018 15:45:07 +0200
Subject: [PATCH] core: garbage collect side fork tries too

---
 core/blockchain.go      | 20 +++++++++++++++---
 core/blockchain_test.go | 46 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/core/blockchain.go b/core/blockchain.go
index 1f01de38c4..8b02b5c576 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -42,6 +42,7 @@ import (
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
 	"github.com/hashicorp/golang-lru"
+	"gopkg.in/karalabe/cookiejar.v2/collections/prque"
 )
 
 var (
@@ -89,6 +90,7 @@ type BlockChain struct {
 
 	diskdb ethdb.Database // Low level persistent database to store final content in
 	triedb *trie.Database // High level ephemeral database to store non-final tries in
+	triegc *prque.Prque   // Priority queue mapping block numbers to tries to gc
 
 	hc            *HeaderChain
 	rmLogsFeed    event.Feed
@@ -152,6 +154,7 @@ func NewBlockChain(diskdb ethdb.Database, cacheConfig *CacheConfig, chainConfig
 		cacheConfig:  cacheConfig,
 		diskdb:       diskdb,
 		triedb:       triedb,
+		triegc:       prque.New(),
 		stateCache:   state.NewDatabase(triedb),
 		quit:         make(chan struct{}),
 		bodyCache:    bodyCache,
@@ -659,6 +662,9 @@ func (bc *BlockChain) Stop() {
 			log.Error("Failed to commit recent state trie", "err", err)
 		}
 	}
+	for !bc.triegc.Empty() {
+		bc.triedb.Dereference(bc.triegc.PopItem().(common.Hash), common.Hash{})
+	}
 	log.Info("Blockchain manager stopped")
 }
 
@@ -869,8 +875,10 @@ func (bc *BlockChain) WriteBlockAndState(block *types.Block, receipts []*types.R
 	}
 
 	bc.triedb.Reference(root, common.Hash{}) // metadata reference to keep trie alive
+	bc.triegc.Push(root, -float32(block.NumberU64()))
+
 	if current := block.NumberU64(); current > triesInMemory {
-		// Find the next state trie we need to get rid of or commit
+		// Find the next state trie we need to commit
 		header := bc.GetHeaderByNumber(current - triesInMemory)
 		chosen := header.Number.Uint64()
 
@@ -882,8 +890,14 @@ func (bc *BlockChain) WriteBlockAndState(block *types.Block, receipts []*types.R
 			bc.procTime = 0
 		}
 		// Garbage collect anything below our required write retention
-		bc.triedb.Dereference(header.Root, common.Hash{})
-
+		for !bc.triegc.Empty() {
+			root, number := bc.triegc.Pop()
+			if uint64(-number) > chosen {
+				bc.triegc.Push(root, number)
+				break
+			}
+			bc.triedb.Dereference(root.(common.Hash), common.Hash{})
+		}
 		if current%10000 == 0 {
 			log.Warn("Current trie pruning state", "size", bc.triedb.Size(), "elapsed", bc.procTime)
 		}
diff --git a/core/blockchain_test.go b/core/blockchain_test.go
index cab4ecadf3..c9287bcb7b 100644
--- a/core/blockchain_test.go
+++ b/core/blockchain_test.go
@@ -1245,3 +1245,49 @@ func TestBlockchainHeaderchainReorgConsistency(t *testing.T) {
 		}
 	}
 }
+
+// Tests that importing small side forks doesn't leave junk in the trie database
+// cache (which would eventually cause memory issues).
+func TestTrieForkGC(t *testing.T) {
+	// Generate a canonical chain to act as the main dataset
+	engine := ethash.NewFaker()
+
+	db, _ := ethdb.NewMemDatabase()
+	genesis := new(Genesis).MustCommit(db)
+	blocks, _ := GenerateChain(params.TestChainConfig, genesis, engine, db, 2*triesInMemory, func(i int, b *BlockGen) { b.SetCoinbase(common.Address{1}) })
+
+	// Generate a bunch of fork blocks, each side forking from the canonical chain
+	forks := make([]*types.Block, len(blocks))
+	for i := 0; i < len(forks); i++ {
+		parent := genesis
+		if i > 0 {
+			parent = blocks[i-1]
+		}
+		fork, _ := GenerateChain(params.TestChainConfig, parent, engine, db, 1, func(i int, b *BlockGen) { b.SetCoinbase(common.Address{2}) })
+		forks[i] = fork[0]
+	}
+	// Import the canonical and fork chain side by side, forcing the trie cache to cache both
+	diskdb, _ := ethdb.NewMemDatabase()
+	new(Genesis).MustCommit(diskdb)
+
+	chain, err := NewBlockChain(diskdb, nil, params.TestChainConfig, engine, vm.Config{})
+	if err != nil {
+		t.Fatalf("failed to create tester chain: %v", err)
+	}
+	for i := 0; i < len(blocks); i++ {
+		if _, err := chain.InsertChain(blocks[i : i+1]); err != nil {
+			t.Fatalf("block %d: failed to insert into chain: %v", i, err)
+		}
+		if _, err := chain.InsertChain(forks[i : i+1]); err != nil {
+			t.Fatalf("fork %d: failed to insert into chain: %v", i, err)
+		}
+	}
+	// Dereference all the recent tries and ensure no past trie is left in
+	for i := 0; i < triesInMemory; i++ {
+		chain.triedb.Dereference(blocks[len(blocks)-1-i].Root(), common.Hash{})
+		chain.triedb.Dereference(forks[len(blocks)-1-i].Root(), common.Hash{})
+	}
+	if len(chain.triedb.Nodes()) > 0 {
+		t.Fatalf("stale tries still alive after garbase collection")
+	}
+}