From ce5f95b6a1bb00d6c2aec79384cf96e149ebef14 Mon Sep 17 00:00:00 2001 From: ucwong Date: Fri, 28 Mar 2025 06:31:02 +0800 Subject: [PATCH] trie memmory optimized --- trie/committer.go | 19 +++----- trie/errors.go | 6 +++ trie/hasher.go | 57 ++++++++++++---------- trie/iterator_test.go | 4 +- trie/node.go | 46 ++++++++++++++++-- trie/secure_trie.go | 6 +++ trie/trie.go | 107 +++++++++++++++++++++++++++++++++--------- 7 files changed, 178 insertions(+), 67 deletions(-) diff --git a/trie/committer.go b/trie/committer.go index e460499df8..b5c0181444 100644 --- a/trie/committer.go +++ b/trie/committer.go @@ -86,9 +86,6 @@ func (c *committer) commit(n node, db *Database) (node, error) { // Commit children, then parent, and remove the dirty flag. switch cn := n.(type) { case *shortNode: - // Commit child - collapsed := cn.copy() - // If the child is fullnode, recursively commit. // Otherwise it can only be hashNode or valueNode. if _, ok := cn.Val.(*fullNode); ok { @@ -96,28 +93,26 @@ func (c *committer) commit(n node, db *Database) (node, error) { if err != nil { return nil, err } - collapsed.Val = childV + cn.Val = childV } // The key needs to be copied, since we're delivering it to database - collapsed.Key = hexToCompact(cn.Key) - hashedNode := c.store(collapsed, db) + cn.Key = hexToCompact(cn.Key) + hashedNode := c.store(cn, db) if hn, ok := hashedNode.(hashNode); ok { return hn, nil } - return collapsed, nil + return cn, nil case *fullNode: hashedKids, err := c.commitChildren(cn, db) if err != nil { return nil, err } - collapsed := cn.copy() - collapsed.Children = hashedKids - - hashedNode := c.store(collapsed, db) + cn.Children = hashedKids + hashedNode := c.store(cn, db) if hn, ok := hashedNode.(hashNode); ok { return hn, nil } - return collapsed, nil + return cn, nil case hashNode: return cn, nil default: diff --git a/trie/errors.go b/trie/errors.go index 19eec38bde..fdd8987e1a 100644 --- a/trie/errors.go +++ b/trie/errors.go @@ -17,11 +17,17 @@ package trie import ( + "errors" "fmt" "github.com/CortexFoundation/CortexTheseus/common" ) +// ErrCommitted is returned when an already committed trie is requested for usage. +// The potential usages can be `Get`, `Update`, `Delete`, `NodeIterator`, `Prove` +// and so on. +var ErrCommitted = errors.New("trie is already committed") + // MissingNodeError is returned by the trie functions (TryGet, Update, TryDelete) // in the case where a trie node is not present in the local database. It contains // information necessary for retrieving the missing node. diff --git a/trie/hasher.go b/trie/hasher.go index 8676c05d03..1ca03a4318 100644 --- a/trie/hasher.go +++ b/trie/hasher.go @@ -55,44 +55,45 @@ func returnHasherToPool(h *hasher) { // hash collapses a node down into a hash node, also returning a copy of the // original node initialized with the computed hash to replace the original one. -func (h *hasher) hash(n node, force bool) (hashed node, cached node) { +func (h *hasher) hash(n node, force bool) node { // Return the cached hash if it's available if hash, _ := n.cache(); hash != nil { - return hash, n + return hash } // Trie not processed yet, walk the children switch n := n.(type) { case *shortNode: - collapsed, cached := h.hashShortNodeChildren(n) + collapsed := h.hashShortNodeChildren(n) hashed := h.shortnodeToHash(collapsed, force) // We need to retain the possibly _not_ hashed node, in case it was too // small to be hashed if hn, ok := hashed.(hashNode); ok { - cached.flags.hash = hn + n.flags.hash = hn } else { - cached.flags.hash = nil + n.flags.hash = nil } - return hashed, cached + return hashed case *fullNode: - collapsed, cached := h.hashFullNodeChildren(n) - hashed = h.fullnodeToHash(collapsed, force) + collapsed := h.hashFullNodeChildren(n) + hashed := h.fullnodeToHash(collapsed, force) if hn, ok := hashed.(hashNode); ok { - cached.flags.hash = hn + n.flags.hash = hn } else { - cached.flags.hash = nil + n.flags.hash = nil } - return hashed, cached + return hashed default: // Value and hash nodes don't have children, so they're left as were - return n, n + return n } } // hashShortNodeChildren collapses the short node. The returned collapsed node // holds a live reference to the Key, and must not be modified. -func (h *hasher) hashShortNodeChildren(n *shortNode) (collapsed, cached *shortNode) { +func (h *hasher) hashShortNodeChildren(n *shortNode) *shortNode { // Hash the short node's child, caching the newly hashed subtree - collapsed, cached = n.copy(), n.copy() + //collapsed, cached = n.copy(), n.copy() + var collapsed shortNode // Previously, we did copy this one. We don't seem to need to actually // do that, since we don't overwrite/reuse keys // cached.Key = common.CopyBytes(n.Key) @@ -100,15 +101,16 @@ func (h *hasher) hashShortNodeChildren(n *shortNode) (collapsed, cached *shortNo // Unless the child is a valuenode or hashnode, hash it switch n.Val.(type) { case *fullNode, *shortNode: - collapsed.Val, cached.Val = h.hash(n.Val, false) + collapsed.Val = h.hash(n.Val, false) + default: + collapsed.Val = n.Val } - return collapsed, cached + return &collapsed } -func (h *hasher) hashFullNodeChildren(n *fullNode) (collapsed *fullNode, cached *fullNode) { +func (h *hasher) hashFullNodeChildren(n *fullNode) *fullNode { // Hash the full node's children, caching the newly hashed subtrees - cached = n.copy() - collapsed = n.copy() + var children [17]node if h.parallel { var wg sync.WaitGroup wg.Add(16) @@ -116,9 +118,9 @@ func (h *hasher) hashFullNodeChildren(n *fullNode) (collapsed *fullNode, cached go func(i int) { hasher := newHasher(false) if child := n.Children[i]; child != nil { - collapsed.Children[i], cached.Children[i] = hasher.hash(child, false) + children[i] = hasher.hash(child, false) } else { - collapsed.Children[i] = nilValueNode + children[i] = nilValueNode } returnHasherToPool(hasher) wg.Done() @@ -128,13 +130,16 @@ func (h *hasher) hashFullNodeChildren(n *fullNode) (collapsed *fullNode, cached } else { for i := 0; i < 16; i++ { if child := n.Children[i]; child != nil { - collapsed.Children[i], cached.Children[i] = h.hash(child, false) + children[i] = h.hash(child, false) } else { - collapsed.Children[i] = nilValueNode + children[i] = nilValueNode } } } - return collapsed, cached + if n.Children[16] != nil { + children[16] = n.Children[16] + } + return &fullNode{flags: nodeFlag{}, Children: children} } // shortnodeToHash creates a hashNode from a shortNode. The supplied shortnode @@ -203,10 +208,10 @@ func (h *hasher) hashDataTo(dst, data []byte) { func (h *hasher) proofHash(original node) (collapsed, hashed node) { switch n := original.(type) { case *shortNode: - sn, _ := h.hashShortNodeChildren(n) + sn := h.hashShortNodeChildren(n) return sn, h.shortnodeToHash(sn, false) case *fullNode: - fn, _ := h.hashFullNodeChildren(n) + fn := h.hashFullNodeChildren(n) return fn, h.fullnodeToHash(fn, false) default: // Value and hash nodes don't have children, so they're left as were diff --git a/trie/iterator_test.go b/trie/iterator_test.go index d7b77e14d1..87bd2661e9 100644 --- a/trie/iterator_test.go +++ b/trie/iterator_test.go @@ -59,7 +59,7 @@ func TestIterator(t *testing.T) { all[val.k] = val.v trie.Update([]byte(val.k), []byte(val.v)) } - trie.Commit(nil) + //trie.Commit(nil) found := make(map[string]string) it := NewIterator(trie.NodeIterator(nil)) @@ -123,7 +123,7 @@ func TestNodeIteratorCoverage(t *testing.T) { // Gather all the node hashes found by the iterator hashes := make(map[common.Hash]struct{}) - for it := trie.NodeIterator(nil); it.Next(true); { + for it := trie.MustNodeIterator(nil); it.Next(true); { if it.Hash() != (common.Hash{}) { hashes[it.Hash()] = struct{}{} } diff --git a/trie/node.go b/trie/node.go index 4d34c93d53..1f655e0367 100644 --- a/trie/node.go +++ b/trie/node.go @@ -79,15 +79,19 @@ func (n *fullNode) EncodeRLP(w io.Writer) error { return eb.Flush() } -func (n *fullNode) copy() *fullNode { copy := *n; return © } -func (n *shortNode) copy() *shortNode { copy := *n; return © } - // nodeFlag contains caching-related metadata about a node. type nodeFlag struct { hash hashNode // cached hash of the node (may be nil) dirty bool // whether the node has changes that must be written to the database } +func (n nodeFlag) copy() nodeFlag { + return nodeFlag{ + hash: common.CopyBytes(n.hash), + dirty: n.dirty, + } +} + func (n *fullNode) cache() (hashNode, bool) { return n.flags.hash, n.flags.dirty } func (n *shortNode) cache() (hashNode, bool) { return n.flags.hash, n.flags.dirty } func (n hashNode) cache() (hashNode, bool) { return nil, true } @@ -130,6 +134,16 @@ func mustDecodeNode(hash, buf []byte) node { return n } +// mustDecodeNodeUnsafe is a wrapper of decodeNodeUnsafe and panic if any error is +// encountered. +func mustDecodeNodeUnsafe(hash, buf []byte) node { + n, err := decodeNodeUnsafe(hash, buf) + if err != nil { + panic(fmt.Sprintf("node %x: %v", hash, err)) + } + return n +} + // decodeNode parses the RLP encoding of a trie node. func decodeNode(hash, buf []byte) (node, error) { if len(buf) == 0 { @@ -151,6 +165,29 @@ func decodeNode(hash, buf []byte) (node, error) { } } +// decodeNodeUnsafe parses the RLP encoding of a trie node. The passed byte slice +// will be directly referenced by node without bytes deep copy, so the input MUST +// not be changed after. +func decodeNodeUnsafe(hash, buf []byte) (node, error) { + if len(buf) == 0 { + return nil, io.ErrUnexpectedEOF + } + elems, _, err := rlp.SplitList(buf) + if err != nil { + return nil, fmt.Errorf("decode error: %v", err) + } + switch c, _ := rlp.CountValues(elems); c { + case 2: + n, err := decodeShort(hash, elems) + return n, wrapError(err, "short") + case 17: + n, err := decodeFull(hash, elems) + return n, wrapError(err, "full") + default: + return nil, fmt.Errorf("invalid number of list elements: %v", c) + } +} + func decodeShort(hash, elems []byte) (node, error) { kbuf, rest, err := rlp.SplitString(elems) if err != nil { @@ -207,7 +244,8 @@ func decodeRef(buf []byte) (node, []byte, error) { err := fmt.Errorf("oversized embedded node (size is %d bytes, want size < %d)", size, hashLen) return nil, buf, err } - n, err := decodeNode(nil, buf) + //n, err := decodeNode(nil, buf) + n, err := decodeNodeUnsafe(nil, buf) return n, rest, err case kind == rlp.String && len(val) == 0: // empty node diff --git a/trie/secure_trie.go b/trie/secure_trie.go index 2f40f9bc42..7616ec20fe 100644 --- a/trie/secure_trie.go +++ b/trie/secure_trie.go @@ -273,6 +273,12 @@ func (t *StateTrie) NodeIterator(start []byte) NodeIterator { return t.trie.NodeIterator(start) } +// MustNodeIterator is a wrapper of NodeIterator and will omit any encountered +// error but just print out an error message. +func (t *StateTrie) MustNodeIterator(start []byte) NodeIterator { + return t.trie.MustNodeIterator(start) +} + // hashKey returns the hash of key as an ephemeral buffer. // The caller must not hold onto the return value because it will become // invalid on the next call to hashKey or secKey. diff --git a/trie/trie.go b/trie/trie.go index d7e0f267c8..472900637c 100644 --- a/trie/trie.go +++ b/trie/trie.go @@ -45,11 +45,17 @@ type Trie struct { owner common.Hash root node + // Flag whether the commit operation is already performed. If so the + // trie is not usable(latest states is invisible). + committed bool + // Keep track of the number leaves which have been inserted since the last // hashing operation. This number will not directly map to the number of - // actually unhashed nodes + // actually unhashed nodes. unhashed int + // uncommitted is the number of updates since last commit. + uncommitted int // tracer is the state diff tracer can be used to track newly added/deleted // trie node. It will be reset after each commit operation. tracer *tracer @@ -63,10 +69,12 @@ func (t *Trie) newFlag() nodeFlag { // Copy returns a copy of Trie. func (t *Trie) Copy() *Trie { return &Trie{ - db: t.db, - root: t.root, - unhashed: t.unhashed, - tracer: t.tracer.copy(), + db: t.db, + root: copyNode(t.root), + committed: t.committed, + unhashed: t.unhashed, + uncommitted: t.uncommitted, + tracer: t.tracer.copy(), } } @@ -98,9 +106,9 @@ func New(id *ID, db *Database) (*Trie, error) { // It's only used by range prover. func newWithRootNode(root node) *Trie { return &Trie{ - root: root, - //tracer: newTracer(), - db: NewDatabase(rawdb.NewMemoryDatabase()), + root: root, + tracer: newTracer(), + db: NewDatabase(rawdb.NewMemoryDatabase()), } } @@ -120,12 +128,20 @@ func (t *Trie) MustNodeIterator(start []byte) NodeIterator { // NodeIterator returns an iterator that returns nodes of the trie. Iteration starts at // the key after the given start key. func (t *Trie) NodeIterator(start []byte) NodeIterator { + // Short circuit if the trie is already committed and not usable. + if t.committed { + //return nil + } return newNodeIterator(t, start) } // Get returns the value for key stored in the trie. // The value bytes must not be modified by the caller. func (t *Trie) Get(key []byte) []byte { + // Short circuit if the trie is already committed and not usable. + if t.committed { + //return nil + } res, err := t.TryGet(key) if err != nil { log.Error(fmt.Sprintf("Unhandled trie error: %v", err)) @@ -157,14 +173,12 @@ func (t *Trie) tryGet(origNode node, key []byte, pos int) (value []byte, newnode } value, newnode, didResolve, err = t.tryGet(n.Val, key, pos+len(n.Key)) if err == nil && didResolve { - n = n.copy() n.Val = newnode } return value, n, didResolve, err case *fullNode: value, newnode, didResolve, err = t.tryGet(n.Children[key[pos]], key, pos+1) if err == nil && didResolve { - n = n.copy() n.Children[key[pos]] = newnode } return value, n, didResolve, err @@ -183,6 +197,10 @@ func (t *Trie) tryGet(origNode node, key []byte, pos int) (value []byte, newnode // TryGetNode attempts to retrieve a trie node by compact-encoded path. It is not // possible to use keybyte-encoding as the path might contain odd nibbles. func (t *Trie) TryGetNode(path []byte) ([]byte, int, error) { + // Short circuit if the trie is already committed and not usable. + if t.committed { + //return nil, 0, ErrCommitted + } item, newroot, resolved, err := t.tryGetNode(t.root, compactToHex(path), 0) if err != nil { return nil, resolved, err @@ -228,7 +246,6 @@ func (t *Trie) tryGetNode(origNode node, path []byte, pos int) (item []byte, new } item, newnode, resolved, err = t.tryGetNode(n.Val, path, pos+len(n.Key)) if err == nil && resolved > 0 { - n = n.copy() n.Val = newnode } return item, n, resolved, err @@ -236,7 +253,6 @@ func (t *Trie) tryGetNode(origNode node, path []byte, pos int) (item []byte, new case *fullNode: item, newnode, resolved, err = t.tryGetNode(n.Children[path[pos]], path, pos+1) if err == nil && resolved > 0 { - n = n.copy() n.Children[path[pos]] = newnode } return item, n, resolved, err @@ -261,6 +277,10 @@ func (t *Trie) tryGetNode(origNode node, path []byte, pos int) (item []byte, new // The value bytes must not be modified by the caller while they are // stored in the trie. func (t *Trie) Update(key, value []byte) { + // Short circuit if the trie is already committed and not usable. + if t.committed { + //return + } if err := t.TryUpdate(key, value); err != nil { log.Error(fmt.Sprintf("Unhandled trie error: %v", err)) } @@ -276,6 +296,7 @@ func (t *Trie) Update(key, value []byte) { // If a node was not found in the database, a MissingNodeError is returned. func (t *Trie) TryUpdate(key, value []byte) error { t.unhashed++ + t.uncommitted++ k := keybytesToHex(key) if len(value) != 0 { _, n, err := t.insert(t.root, nil, k, valueNode(value)) @@ -340,7 +361,6 @@ func (t *Trie) insert(n node, prefix, key []byte, value node) (bool, node, error if !dirty || err != nil { return false, n, err } - n = n.copy() n.flags = t.newFlag() n.Children[key[0]] = nn return true, n, nil @@ -374,6 +394,12 @@ func (t *Trie) insert(n node, prefix, key []byte, value node) (bool, node, error // Delete removes any existing value for key from the trie. func (t *Trie) Delete(key []byte) { + // Short circuit if the trie is already committed and not usable. + if t.committed { + //return + } + t.uncommitted++ + t.unhashed++ if err := t.TryDelete(key); err != nil { log.Error(fmt.Sprintf("Unhandled trie error: %v", err)) } @@ -440,7 +466,6 @@ func (t *Trie) delete(n node, prefix, key []byte) (bool, node, error) { if !dirty || err != nil { return false, n, err } - n = n.copy() n.flags = t.newFlag() n.Children[key[0]] = nn @@ -533,6 +558,36 @@ func concat(s1 []byte, s2 ...byte) []byte { return r } +// copyNode deep-copies the supplied node along with its children recursively. +func copyNode(n node) node { + switch n := (n).(type) { + case nil: + return nil + case valueNode: + return valueNode(common.CopyBytes(n)) + + case *shortNode: + return &shortNode{ + flags: n.flags.copy(), + Key: common.CopyBytes(n.Key), + Val: copyNode(n.Val), + } + case *fullNode: + var children [17]node + for i, cn := range n.Children { + children[i] = copyNode(cn) + } + return &fullNode{ + flags: n.flags.copy(), + Children: children, + } + case hashNode: + return n + default: + panic(fmt.Sprintf("%T: unknown node type", n)) + } +} + func (t *Trie) resolve(n node, prefix []byte) (node, error) { if n, ok := n.(hashNode); ok { return t.resolveHash(n, prefix) @@ -560,17 +615,19 @@ func (t *Trie) resolveBlob(n hashNode, prefix []byte) ([]byte, error) { // Hash returns the root hash of the trie. It does not write to the // database and can be used even if the trie doesn't have one. func (t *Trie) Hash() common.Hash { - hash, cached, _ := t.hashRoot() - t.root = cached - return common.BytesToHash(hash.(hashNode)) + return common.BytesToHash(t.hashRoot().(hashNode)) } // Commit writes all nodes to the trie's memory database, tracking the internal // and external (for account tries) references. func (t *Trie) Commit(onleaf LeafCallback) (root common.Hash, err error) { + defer func() { + t.committed = true + }() if t.db == nil { panic("commit called on trie with nil database") } + defer t.tracer.reset() if t.root == nil { @@ -612,25 +669,29 @@ func (t *Trie) Commit(onleaf LeafCallback) (root common.Hash, err error) { return common.Hash{}, err } t.root = newRoot + t.uncommitted = 0 return rootHash, nil } // hashRoot calculates the root hash of the given trie -func (t *Trie) hashRoot() (node, node, error) { +func (t *Trie) hashRoot() node { if t.root == nil { - return hashNode(types.EmptyRootHash.Bytes()), nil, nil + return hashNode(types.EmptyRootHash.Bytes()) } // If the number of changes is below 100, we let one thread handle it h := newHasher(t.unhashed >= 100) - defer returnHasherToPool(h) - hashed, cached := h.hash(t.root, true) - t.unhashed = 0 - return hashed, cached, nil + defer func() { + returnHasherToPool(h) + t.unhashed = 0 + }() + return h.hash(t.root, true) } // Reset drops the referenced root node and cleans all internal state. func (t *Trie) Reset() { t.root = nil t.unhashed = 0 + t.uncommitted = 0 t.tracer.reset() + t.committed = false }