From 4315480e09824f3c8b50d30948e7327612869bbf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?P=C3=A9ter=20Szil=C3=A1gyi?= <peterke@gmail.com>
Date: Wed, 24 Jan 2018 16:31:26 +0200
Subject: [PATCH] core, trie: fix duplicate references, sanity check on
 shutdown

---
 core/blockchain.go |  3 +++
 trie/database.go   | 29 ++++++++++++++++-------------
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/core/blockchain.go b/core/blockchain.go
index 40b4a9934b..3a5264e022 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -669,6 +669,9 @@ func (bc *BlockChain) Stop() {
 	for !bc.triegc.Empty() {
 		bc.triedb.Dereference(bc.triegc.PopItem().(common.Hash), common.Hash{})
 	}
+	if size := bc.triedb.Size(); size != 0 {
+		log.Error("Dangling trie nodes after full cleanup")
+	}
 	log.Info("Blockchain manager stopped")
 }
 
diff --git a/trie/database.go b/trie/database.go
index a0878f2850..2e36377b5c 100644
--- a/trie/database.go
+++ b/trie/database.go
@@ -61,9 +61,9 @@ type Database struct {
 // cachedNode is all the information we know about a single cached node in the
 // memory database write layer.
 type cachedNode struct {
-	blob     []byte              // Cached data block of the trie node
-	parents  int                 // Number of live nodes referencing this one
-	children map[common.Hash]int // Children referenced by this nodes
+	blob     []byte                   // Cached data block of the trie node
+	parents  int                      // Number of live nodes referencing this one
+	children map[common.Hash]struct{} // Children referenced by this nodes
 }
 
 // NewDatabase creates a new trie database to store ephemeral trie content before
@@ -72,7 +72,7 @@ func NewDatabase(diskdb ethdb.Database) *Database {
 	return &Database{
 		diskdb: diskdb,
 		nodes: map[common.Hash]*cachedNode{
-			common.Hash{}: {children: make(map[common.Hash]int)},
+			common.Hash{}: {children: make(map[common.Hash]struct{})},
 		},
 		preimages: make(map[common.Hash][]byte),
 	}
@@ -99,7 +99,7 @@ func (db *Database) insert(hash common.Hash, blob []byte) {
 	}
 	db.nodes[hash] = &cachedNode{
 		blob:     common.CopyBytes(blob),
-		children: make(map[common.Hash]int),
+		children: make(map[common.Hash]struct{}),
 	}
 	db.size += common.StorageSize(common.HashLength + len(blob))
 }
@@ -185,8 +185,12 @@ func (db *Database) reference(child common.Hash, parent common.Hash) {
 	if !ok {
 		return
 	}
+	// If the reference already exists, don't duplicate it
+	if _, ok = db.nodes[parent].children[child]; ok {
+		return
+	}
 	node.parents++
-	db.nodes[parent].children[child]++
+	db.nodes[parent].children[child] = struct{}{}
 }
 
 // Dereference removes an existing reference from a parent node to a child node.
@@ -200,17 +204,16 @@ func (db *Database) Dereference(child common.Hash, parent common.Hash) {
 	db.gcnodes += uint64(nodes - len(db.nodes))
 	db.gcsize += storage - db.size
 	db.gctime += time.Since(start)
+
+	log.Debug("Dereferenced trie from memory database", "nodes", nodes-len(db.nodes), "size", storage-db.size, "time", time.Since(start),
+		"gcnodes", db.gcnodes, "gcsize", db.gcsize, "gctime", db.gctime, "livenodes", len(db.nodes), "livesize", db.size)
 }
 
 // dereference is the private locked version of Dereference.
 func (db *Database) dereference(child common.Hash, parent common.Hash) {
-	// Dereference the parent-child
-	node := db.nodes[parent]
+	// Dereference the parent-child relationship
+	delete(db.nodes[parent].children, child)
 
-	node.children[child]--
-	if node.children[child] == 0 {
-		delete(node.children, child)
-	}
 	// If the node does not exist, it's a previously comitted node.
 	node, ok := db.nodes[child]
 	if !ok {
@@ -271,7 +274,7 @@ func (db *Database) Commit(node common.Hash) error {
 	db.preimages = make(map[common.Hash][]byte)
 	db.uncache(node)
 
-	log.Debug("Persisted trie from memory database", "nodes", nodes-len(db.nodes), "size", storage-db.size, "time", time.Since(start),
+	log.Info("Persisted trie from memory database", "nodes", nodes-len(db.nodes), "size", storage-db.size, "time", time.Since(start),
 		"gcnodes", db.gcnodes, "gcsize", db.gcsize, "gctime", db.gctime, "livenodes", len(db.nodes), "livesize", db.size)
 
 	// Reset the garbage collection statistics