Atomic trie height map repair

go.mod

@@ -38,6 +38,7 @@ require (
 	github.com/tyler-smith/go-bip39 v1.1.0
 	github.com/urfave/cli/v2 v2.17.2-0.20221006022127-8f469abc00aa
 	go.uber.org/goleak v1.2.1
+	go.uber.org/mock v0.2.0
 	golang.org/x/crypto v0.14.0
 	golang.org/x/exp v0.0.0-20230626212559-97b1e661b5df
 	golang.org/x/sync v0.3.0
@@ -127,7 +128,6 @@ require (
 	go.opentelemetry.io/otel/sdk v1.11.0 // indirect
 	go.opentelemetry.io/otel/trace v1.11.0 // indirect
 	go.opentelemetry.io/proto/otlp v0.19.0 // indirect
-	go.uber.org/mock v0.2.0 // indirect
 	go.uber.org/multierr v1.10.0 // indirect
 	go.uber.org/zap v1.26.0 // indirect
 	golang.org/x/net v0.17.0 // indirect

go.sum

@@ -55,8 +55,6 @@ github.com/allegro/bigcache v1.2.1-0.20190218064605-e24eb225f156 h1:eMwmnE/GDgah
 github.com/allegro/bigcache v1.2.1-0.20190218064605-e24eb225f156/go.mod h1:Cb/ax3seSYIx7SuZdm2G2xzfwmv3TPSk2ucNfQESPXM=
 github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=
 github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6/go.mod h1:grANhF5doyWs3UAsr3K4I6qtAmlQcZDesFNEHPZAzj8=
-github.com/ava-labs/avalanchego v1.10.17-rc.0.0.20231129204442-3a8cca5f31b9 h1:IDbRSp4QU5gcsC31uxOMc1089nTDhZpr6c+z7pHVNks=
-github.com/ava-labs/avalanchego v1.10.17-rc.0.0.20231129204442-3a8cca5f31b9/go.mod h1:JuMiDdOmbTZfZty/RFszDPpKd1KEy9aivjxbcc1TPZo=
 github.com/ava-labs/avalanchego v1.10.17-rc.9 h1:qMDE56M33lO8nkXMPmycN8uIT1BL42giR46f8j7kbhg=
 github.com/ava-labs/avalanchego v1.10.17-rc.9/go.mod h1:GW0zr0OJfy9Vbki1ChOm47jv0gK+2tnBXP2NpzAiVzE=
 github.com/aymerick/raymond v2.0.3-0.20180322193309-b565731e1464+incompatible/go.mod h1:osfaiScAUVup+UC9Nfq76eWqDhXlp+4UYaA8uhTBO6g=

plugin/evm/atomic_backend.go

@@ -87,6 +87,18 @@ func NewAtomicBackend(
 	db *versiondb.Database, sharedMemory atomic.SharedMemory,
 	bonusBlocks map[uint64]ids.ID, repo AtomicTxRepository,
 	lastAcceptedHeight uint64, lastAcceptedHash common.Hash, commitInterval uint64,
+) (AtomicBackend, error) {
+	return NewAtomicBackendWithBonusBlockRepair(
+		db, sharedMemory, bonusBlocks, nil, repo,
+		lastAcceptedHeight, lastAcceptedHash, commitInterval,
+	)
+}
+
+func NewAtomicBackendWithBonusBlockRepair(
+	db *versiondb.Database, sharedMemory atomic.SharedMemory,
+	bonusBlocks map[uint64]ids.ID, bonusBlocksRlp map[uint64]string,
+	repo AtomicTxRepository,
+	lastAcceptedHeight uint64, lastAcceptedHash common.Hash, commitInterval uint64,
 ) (AtomicBackend, error) {
 	atomicTrieDB := prefixdb.New(atomicTrieDBPrefix, db)
 	metadataDB := prefixdb.New(atomicTrieMetaDBPrefix, db)
@@ -96,6 +108,14 @@ func NewAtomicBackend(
 	if err != nil {
 		return nil, err
 	}
+	if len(bonusBlocksRlp) > 0 {
+		if _, err := atomicTrie.repairAtomicTrie(bonusBlocks, bonusBlocksRlp); err != nil {
+			return nil, err
+		}
+		if err := db.Commit(); err != nil {
+			return nil, err
+		}
+	}
 	atomicBackend := &atomicBackend{
 		codec:            codec,
 		db:               db,

plugin/evm/atomic_trie.go

@@ -36,6 +36,7 @@ var (
 	_                            AtomicTrie = &atomicTrie{}
 	lastCommittedKey                        = []byte("atomicTrieLastCommittedBlock")
 	appliedSharedMemoryCursorKey            = []byte("atomicTrieLastAppliedToSharedMemory")
+	heightMapRepairKey                      = []byte("atomicTrieHeightMapRepair")
 )
 
 // AtomicTrie maintains an index of atomic operations by blockchainIDs for every block
@@ -82,6 +83,10 @@ type AtomicTrie interface {
 
 	// RejectTrie dereferences root from the trieDB, freeing memory.
 	RejectTrie(root common.Hash) error
+
+	// RepairHeightMap repairs the height map of the atomic trie by iterating
+	// over all leaves in the trie and committing the trie at every commit interval.
+	RepairHeightMap(to uint64, iterationDelay time.Duration) (bool, error)
 }
 
 // AtomicTrieIterator is a stateful iterator that iterates the leafs of an AtomicTrie
@@ -94,6 +99,9 @@ type AtomicTrieIterator interface {
 	// returned []byte can be freely modified
 	Key() []byte
 
+	// Value returns the current database value that the iterator is iterating
+	Value() []byte
+
 	// BlockNumber returns the current block number
 	BlockNumber() uint64
 

plugin/evm/atomic_trie_height_map_repair.go

@@ -0,0 +1,121 @@
+// (c) 2020-2021, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package evm
+
+import (
+	"errors"
+	"fmt"
+	"math"
+	"time"
+
+	"github.com/ava-labs/avalanchego/database"
+	"github.com/ava-labs/coreth/core/types"
+	"github.com/ava-labs/coreth/trie/trienode"
+	"github.com/ethereum/go-ethereum/log"
+)
+
+const (
+	repairDone     = math.MaxUint64         // used as a marker for when the height map is repaired
+	iterationDelay = 100 * time.Microsecond // delay between iterations of the repair loop
+)
+
+func (a *atomicTrie) RepairHeightMap(to uint64, iterationDelay time.Duration) (bool, error) {
+	repairFrom, err := database.GetUInt64(a.metadataDB, heightMapRepairKey)
+	switch {
+	case errors.Is(err, database.ErrNotFound):
+		repairFrom = 0 // height map not repaired yet, start at 0
+	case err != nil:
+		return false, err
+	case repairFrom == repairDone:
+		// height map already repaired, nothing to do
+		return false, nil
+	}
+	return true, a.repairHeightMap(repairFrom, to, iterationDelay)
+}
+
+func (a *atomicTrie) repairHeightMap(from, to uint64, iterationDelay time.Duration) error {
+	// open the atomic trie at the last known root with correct height map
+	// correspondance
+	fromRoot, err := getRoot(a.metadataDB, from)
+	if err != nil {
+		return fmt.Errorf("could not get root at height %d: %w", from, err)
+	}
+	hasher, err := a.OpenTrie(fromRoot)
+	if err != nil {
+		return fmt.Errorf("could not open atomic trie at root %s: %w", fromRoot, err)
+	}
+
+	// hashes values inserted in [hasher], and stores the result in the height
+	// map at [commitHeight]. Additionally, it updates the resume marker and
+	// re-opens [hasher] to respect the trie's no use after commit invariant.
+	lastLog := from
+	logEach := uint64(100_000)
+	commitRepairedHeight := func(commitHeight uint64) error {
+		root, nodes := hasher.Commit(false)
+		if nodes != nil {
+			err := a.trieDB.Update(root, types.EmptyRootHash, trienode.NewWithNodeSet(nodes))
+			if err != nil {
+				return err
+			}
+		}
+		err = a.trieDB.Commit(root, false)
+		if err != nil {
+			return err
+		}
+		err = a.metadataDB.Put(database.PackUInt64(commitHeight), root[:])
+		if err != nil {
+			return err
+		}
+		err = database.PutUInt64(a.metadataDB, heightMapRepairKey, commitHeight)
+		if err != nil {
+			return err
+		}
+		if commitHeight >= lastLog+logEach {
+			log.Info("repaired atomic trie height map", "height", commitHeight, "root", root)
+			lastLog = commitHeight
+		}
+		hasher, err = a.OpenTrie(root)
+		return err
+	}
+
+	// iterate over all leaves in the current atomic trie
+	root, _ := a.LastCommitted()
+	it, err := a.Iterator(root, database.PackUInt64(from+1))
+	if err != nil {
+		return fmt.Errorf("could not create iterator for atomic trie at root %s: %w", root, err)
+	}
+
+	var height uint64
+	lastCommit := from
+	for it.Next() {
+		height = it.BlockNumber()
+		if height > to {
+			break
+		}
+
+		for next := lastCommit + a.commitInterval; next < height; next += a.commitInterval {
+			if err := commitRepairedHeight(next); err != nil {
+				return err
+			}
+			lastCommit = next
+		}
+
+		if err := hasher.Update(it.Key(), it.Value()); err != nil {
+			return fmt.Errorf("could not update atomic trie at root %s: %w", root, err)
+		}
+
+		time.Sleep(iterationDelay) // pause to avoid putting a spike of load on the disk
+	}
+	if err := it.Error(); err != nil {
+		return fmt.Errorf("error iterating atomic trie: %w", err)
+	}
+	for next := lastCommit + a.commitInterval; next <= to; next += a.commitInterval {
+		if err := commitRepairedHeight(next); err != nil {
+			return err
+		}
+	}
+
+	// mark height map as repaired
+	return database.PutUInt64(a.metadataDB, heightMapRepairKey, repairDone)
+}

plugin/evm/atomic_trie_height_map_repair_test.go

@@ -0,0 +1,119 @@
+// (c) 2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package evm
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/ava-labs/avalanchego/database"
+	"github.com/ava-labs/avalanchego/database/memdb"
+	"github.com/ava-labs/avalanchego/database/versiondb"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/stretchr/testify/require"
+)
+
+func TestAtomicTrieRepairHeightMap(t *testing.T) {
+	for i, test := range []testAtomicTrieRepairHeightMap{
+		{
+			lastAccepted:  3*testCommitInterval + 5,
+			skipAtomicTxs: func(height uint64) bool { return false },
+		},
+		{
+			lastAccepted:  3 * testCommitInterval,
+			skipAtomicTxs: func(height uint64) bool { return false },
+		},
+		{
+			lastAccepted:  3 * testCommitInterval,
+			skipAtomicTxs: func(height uint64) bool { return height > testCommitInterval && height <= 2*testCommitInterval },
+		},
+		{
+			lastAccepted:  3 * testCommitInterval,
+			skipAtomicTxs: func(height uint64) bool { return height > testCommitInterval+1 },
+		},
+	} {
+		t.Run(fmt.Sprintf("case %d", i), func(t *testing.T) { test.run(t) })
+	}
+}
+
+type testAtomicTrieRepairHeightMap struct {
+	lastAccepted  uint64
+	skipAtomicTxs func(height uint64) bool
+}
+
+func (test testAtomicTrieRepairHeightMap) run(t *testing.T) {
+	require := require.New(t)
+
+	db := versiondb.New(memdb.New())
+	repo, err := NewAtomicTxRepository(db, testTxCodec(), 0, nil, nil, nil)
+	require.NoError(err)
+	atomicBackend, err := NewAtomicBackend(db, testSharedMemory(), nil, repo, 0, common.Hash{}, testCommitInterval)
+	require.NoError(err)
+	atomicTrie := atomicBackend.AtomicTrie().(*atomicTrie)
+
+	heightMap := make(map[uint64]common.Hash)
+	for height := uint64(1); height <= test.lastAccepted; height++ {
+		atomicRequests := testDataImportTx().mustAtomicOps()
+		if test.skipAtomicTxs(height) {
+			atomicRequests = nil
+		}
+		err := indexAtomicTxs(atomicTrie, height, atomicRequests)
+		require.NoError(err)
+		if height%testCommitInterval == 0 {
+			root, _ := atomicTrie.LastCommitted()
+			heightMap[height] = root
+		}
+	}
+
+	// Verify that [atomicTrie] can access each of the expected roots
+	verifyRoots := func(expectZero bool) {
+		for height, hash := range heightMap {
+			root, err := atomicTrie.Root(height)
+			require.NoError(err)
+			if expectZero {
+				require.Zero(root)
+			} else {
+				require.Equal(hash, root)
+			}
+		}
+	}
+	verifyRoots(false)
+
+	// destroy the height map
+	for height := range heightMap {
+		err := atomicTrie.metadataDB.Delete(database.PackUInt64(height))
+		require.NoError(err)
+	}
+	require.NoError(db.Commit())
+	verifyRoots(true)
+
+	// repair the height map
+	testIterationDelay := time.Duration(0) // no need to wait between iterations in tests
+	repaired, err := atomicTrie.RepairHeightMap(test.lastAccepted, testIterationDelay)
+	require.NoError(err)
+	verifyRoots(false)
+	require.True(repaired)
+
+	// partially destroy the height map
+	_, lastHeight := atomicTrie.LastCommitted()
+	err = atomicTrie.metadataDB.Delete(database.PackUInt64(lastHeight))
+	require.NoError(err)
+	err = atomicTrie.metadataDB.Put(
+		heightMapRepairKey,
+		database.PackUInt64(lastHeight-testCommitInterval),
+	)
+	require.NoError(err)
+
+	// repair the height map
+	repaired, err = atomicTrie.RepairHeightMap(test.lastAccepted, testIterationDelay)
+	require.NoError(err)
+	verifyRoots(false)
+	require.True(repaired)
+
+	// try to repair the height map again
+	repaired, err = atomicTrie.RepairHeightMap(test.lastAccepted, testIterationDelay)
+	require.NoError(err)
+	require.False(repaired)
+}

plugin/evm/atomic_trie_iterator.go

@@ -112,3 +112,11 @@ func (a *atomicTrieIterator) AtomicOps() *atomic.Requests {
 func (a *atomicTrieIterator) Key() []byte {
 	return a.key
 }
+
+// Value returns the current database value that the iterator is iterating
+func (a *atomicTrieIterator) Value() []byte {
+	if a.trieIterator == nil {
+		return nil
+	}
+	return a.trieIterator.Value
+}