trie: unload old cached nodes

trie/hasher.go

@@ -29,6 +29,7 @@ import (
 type hasher struct {
 	tmp                  *bytes.Buffer
 	sha                  hash.Hash
+	cachegen, cachelimit uint16
 }
 
 // hashers live in a global pool.
@@ -38,8 +39,10 @@ var hasherPool = sync.Pool{
 	},
 }
 
-func newHasher() *hasher {
-	return hasherPool.Get().(*hasher)
+func newHasher(cachegen, cachelimit uint16) *hasher {
+	h := hasherPool.Get().(*hasher)
+	h.cachegen, h.cachelimit = cachegen, cachelimit
+	return h
 }
 
 func returnHasherToPool(h *hasher) {
@@ -50,9 +53,17 @@ func returnHasherToPool(h *hasher) {
 // original node initialzied with the computed hash to replace the original one.
 func (h *hasher) hash(n node, db DatabaseWriter, force bool) (node, node, error) {
 	// If we're not storing the node, just hashing, use avaialble cached data
-	if hash, dirty := n.cache(); hash != nil && (db == nil || !dirty) {
+	if hash, dirty := n.cache(); hash != nil {
+		if db == nil {
 			return hash, n, nil
 		}
+		if n.canUnload(h.cachegen, h.cachelimit) {
+			return hash, hash, nil
+		}
+		if !dirty {
+			return hash, n, nil
+		}
+	}
 	// Trie not processed yet or needs storage, walk the children
 	collapsed, cached, err := h.hashChildren(n, db)
 	if err != nil {
@@ -110,8 +121,7 @@ func (h *hasher) hashChildren(original node, db DatabaseWriter) (node, node, err
 
 	case *fullNode:
 		// Hash the full node's children, caching the newly hashed subtrees
-		collapsed := n.copy()
-		cached := &fullNode{dirty: n.dirty}
+		collapsed, cached := n.copy(), n.copy()
 
 		for i := 0; i < 16; i++ {
 			if n.Children[i] != nil {

trie/node.go

@@ -30,19 +30,18 @@ var indices = []string{"0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "a", "b
 type node interface {
 	fstring(string) string
 	cache() (hashNode, bool)
+	canUnload(cachegen, cachelimit uint16) bool
 }
 
 type (
 	fullNode struct {
 		Children [17]node // Actual trie node data to encode/decode (needs custom encoder)
-		hash     hashNode // Cached hash of the node to prevent rehashing (may be nil)
-		dirty    bool     // Cached flag whether the node's new or already stored
+		nodeFlag
 	}
 	shortNode struct {
 		Key []byte
 		Val node
-		hash  hashNode // Cached hash of the node to prevent rehashing (may be nil)
-		dirty bool     // Cached flag whether the node's new or already stored
+		nodeFlag
 	}
 	hashNode  []byte
 	valueNode []byte
@@ -56,9 +55,28 @@ func (n *fullNode) EncodeRLP(w io.Writer) error {
 func (n *fullNode) copy() *fullNode   { copy := *n; return &copy }
 func (n *shortNode) copy() *shortNode { copy := *n; return &copy }
 
-// Cache accessors to retrieve precalculated values (avoid lengthy type switches).
-func (n *fullNode) cache() (hashNode, bool)  { return n.hash, n.dirty }
-func (n *shortNode) cache() (hashNode, bool) { return n.hash, n.dirty }
+// nodeFlag contains caching-related metadata about a node.
+type nodeFlag struct {
+	hash  hashNode // cached hash of the node (may be nil)
+	gen   uint16   // cache generation counter
+	dirty bool     // whether the node has changes that must be written to the database
+}
+
+// canUnload tells whether a node can be unloaded.
+func (n nodeFlag) canUnload(cachegen, cachelimit uint16) bool {
+	var dist uint16
+	if n.gen > cachegen {
+		dist = n.gen - cachegen
+	} else {
+		dist = cachegen - n.gen
+	}
+	return !n.dirty && dist > cachelimit
+}
+
+func (n hashNode) canUnload(uint16, uint16) bool  { return false }
+func (n valueNode) canUnload(uint16, uint16) bool { return false }
+
+func (n nodeFlag) cache() (hashNode, bool)  { return n.hash, n.dirty }
 func (n hashNode) cache() (hashNode, bool)  { return nil, true }
 func (n valueNode) cache() (hashNode, bool) { return nil, true }
 
@@ -123,6 +141,7 @@ func decodeShort(hash, buf, elems []byte) (node, error) {
 	if err != nil {
 		return nil, err
 	}
+	flag := nodeFlag{hash: hash}
 	key := compactDecode(kbuf)
 	if key[len(key)-1] == 16 {
 		// value node
@@ -130,17 +149,17 @@ func decodeShort(hash, buf, elems []byte) (node, error) {
 		if err != nil {
 			return nil, fmt.Errorf("invalid value node: %v", err)
 		}
-		return &shortNode{key, append(valueNode{}, val...), hash, false}, nil
+		return &shortNode{key, append(valueNode{}, val...), flag}, nil
 	}
 	r, _, err := decodeRef(rest)
 	if err != nil {
 		return nil, wrapError(err, "val")
 	}
-	return &shortNode{key, r, hash, false}, nil
+	return &shortNode{key, r, flag}, nil
 }
 
 func decodeFull(hash, buf, elems []byte) (*fullNode, error) {
-	n := &fullNode{hash: hash}
+	n := &fullNode{nodeFlag: nodeFlag{hash: hash}}
 	for i := 0; i < 16; i++ {
 		cld, rest, err := decodeRef(elems)
 		if err != nil {

trie/proof.go

@@ -70,7 +70,7 @@ func (t *Trie) Prove(key []byte) []rlp.RawValue {
 			panic(fmt.Sprintf("%T: invalid node: %v", tn, tn))
 		}
 	}
-	hasher := newHasher()
+	hasher := newHasher(0, 0)
 	proof := make([]rlp.RawValue, 0, len(nodes))
 	for i, n := range nodes {
 		// Don't bother checking for errors here since hasher panics

trie/secure_trie.go

@@ -191,7 +191,7 @@ func (t *SecureTrie) secKey(key []byte) []byte {
 // The caller must not hold onto the return value because it will become
 // invalid on the next call to hashKey or secKey.
 func (t *SecureTrie) hashKey(key []byte) []byte {
-	h := newHasher()
+	h := newHasher(0, 0)
 	h.sha.Reset()
 	h.sha.Write(key)
 	buf := h.sha.Sum(t.hashKeyBuf[:0])

trie/trie.go

@@ -62,6 +62,23 @@ type Trie struct {
 	root         node
 	db           Database
 	originalRoot common.Hash
+
+	// Cache generation values.
+	// cachegen increase by one with each commit operation.
+	// new nodes are tagged with the current generation and unloaded
+	// when their generation is older than than cachegen-cachelimit.
+	cachegen, cachelimit uint16
+}
+
+// SetCacheLimit sets the number of 'cache generations' to keep.
+// A cache generations is created by a call to Commit.
+func (t *Trie) SetCacheLimit(l uint16) {
+	t.cachelimit = l
+}
+
+// newFlag returns the cache flag value for a newly created node.
+func (t *Trie) newFlag() nodeFlag {
+	return nodeFlag{dirty: true, gen: t.cachegen}
 }
 
 // New creates a trie with an existing root node from db.
@@ -206,10 +223,10 @@ func (t *Trie) insert(n node, prefix, key []byte, value node) (bool, node, error
 			if !dirty || err != nil {
 				return false, n, err
 			}
-			return true, &shortNode{n.Key, nn, nil, true}, nil
+			return true, &shortNode{n.Key, nn, t.newFlag()}, nil
 		}
 		// Otherwise branch out at the index where they differ.
-		branch := &fullNode{dirty: true}
+		branch := &fullNode{nodeFlag: t.newFlag()}
 		var err error
 		_, branch.Children[n.Key[matchlen]], err = t.insert(nil, append(prefix, n.Key[:matchlen+1]...), n.Key[matchlen+1:], n.Val)
 		if err != nil {
@@ -224,7 +241,7 @@ func (t *Trie) insert(n node, prefix, key []byte, value node) (bool, node, error
 			return true, branch, nil
 		}
 		// Otherwise, replace it with a short node leading up to the branch.
-		return true, &shortNode{key[:matchlen], branch, nil, true}, nil
+		return true, &shortNode{key[:matchlen], branch, t.newFlag()}, nil
 
 	case *fullNode:
 		dirty, nn, err := t.insert(n.Children[key[0]], append(prefix, key[0]), key[1:], value)
@@ -236,7 +253,7 @@ func (t *Trie) insert(n node, prefix, key []byte, value node) (bool, node, error
 		return true, n, nil
 
 	case nil:
-		return true, &shortNode{key, value, nil, true}, nil
+		return true, &shortNode{key, value, t.newFlag()}, nil
 
 	case hashNode:
 		// We've hit a part of the trie that isn't loaded yet. Load
@@ -305,9 +322,9 @@ func (t *Trie) delete(n node, prefix, key []byte) (bool, node, error) {
 			// always creates a new slice) instead of append to
 			// avoid modifying n.Key since it might be shared with
 			// other nodes.
-			return true, &shortNode{concat(n.Key, child.Key...), child.Val, nil, true}, nil
+			return true, &shortNode{concat(n.Key, child.Key...), child.Val, t.newFlag()}, nil
 		default:
-			return true, &shortNode{n.Key, child, nil, true}, nil
+			return true, &shortNode{n.Key, child, t.newFlag()}, nil
 		}
 
 	case *fullNode:
@@ -352,12 +369,12 @@ func (t *Trie) delete(n node, prefix, key []byte) (bool, node, error) {
 				}
 				if cnode, ok := cnode.(*shortNode); ok {
 					k := append([]byte{byte(pos)}, cnode.Key...)
-					return true, &shortNode{k, cnode.Val, nil, true}, nil
+					return true, &shortNode{k, cnode.Val, t.newFlag()}, nil
 				}
 			}
 			// Otherwise, n is replaced by a one-nibble short node
 			// containing the child.
-			return true, &shortNode{[]byte{byte(pos)}, n.Children[pos], nil, true}, nil
+			return true, &shortNode{[]byte{byte(pos)}, n.Children[pos], t.newFlag()}, nil
 		}
 		// n still contains at least two values and cannot be reduced.
 		return true, n, nil
@@ -453,6 +470,7 @@ func (t *Trie) CommitTo(db DatabaseWriter) (root common.Hash, err error) {
 		return (common.Hash{}), err
 	}
 	t.root = cached
+	t.cachegen++
 	return common.BytesToHash(hash.(hashNode)), nil
 }
 
@@ -460,7 +478,7 @@ func (t *Trie) hashRoot(db DatabaseWriter) (node, node, error) {
 	if t.root == nil {
 		return hashNode(emptyRoot.Bytes()), nil, nil
 	}
-	h := newHasher()
+	h := newHasher(t.cachegen, t.cachelimit)
 	defer returnHasherToPool(h)
 	return h.hash(t.root, db, true)
 }