Merge pull request #5213 from ethereum-optimism/feat/concurrent-iterator-lib

op-chain-ops: Extract concurrent state iterator into util

op-chain-ops/ether/migrate_test.go

 package ether
 
 import (
-	"fmt"
 	"math/big"
 	"math/rand"
 	"testing"
 
+	"github.com/ethereum-optimism/optimism/op-chain-ops/util"
+
 	"github.com/ethereum-optimism/optimism/op-bindings/predeploys"
 
 	"github.com/ethereum-optimism/optimism/op-chain-ops/crossdomain"
@@ -190,7 +191,7 @@ func TestMigrateBalances(t *testing.T) {
 	}
 }
 
-func makeLegacyETH(t *testing.T, totalSupply *big.Int, balances map[common.Address]*big.Int, allowances map[common.Address]common.Address) (*state.StateDB, DBFactory) {
+func makeLegacyETH(t *testing.T, totalSupply *big.Int, balances map[common.Address]*big.Int, allowances map[common.Address]common.Address) (*state.StateDB, util.DBFactory) {
 	memDB := rawdb.NewMemoryDatabase()
 	db, err := state.New(common.Hash{}, state.NewDatabaseWithConfig(memDB, &trie.Config{
 		Preimages: true,
@@ -283,85 +284,6 @@ func TestMigrateBalancesRandomMissing(t *testing.T) {
 	}
 }
 
-func TestPartitionKeyspace(t *testing.T) {
-	tests := []struct {
-		i        int
-		count    int
-		expected [2]common.Hash
-	}{
-		{
-			i:     0,
-			count: 1,
-			expected: [2]common.Hash{
-				common.HexToHash("0x00"),
-				common.HexToHash("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
-			},
-		},
-		{
-			i:     0,
-			count: 2,
-			expected: [2]common.Hash{
-				common.HexToHash("0x00"),
-				common.HexToHash("0x7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
-			},
-		},
-		{
-			i:     1,
-			count: 2,
-			expected: [2]common.Hash{
-				common.HexToHash("0x7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
-				common.HexToHash("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
-			},
-		},
-		{
-			i:     0,
-			count: 3,
-			expected: [2]common.Hash{
-				common.HexToHash("0x00"),
-				common.HexToHash("0x5555555555555555555555555555555555555555555555555555555555555555"),
-			},
-		},
-		{
-			i:     1,
-			count: 3,
-			expected: [2]common.Hash{
-				common.HexToHash("0x5555555555555555555555555555555555555555555555555555555555555555"),
-				common.HexToHash("0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"),
-			},
-		},
-		{
-			i:     2,
-			count: 3,
-			expected: [2]common.Hash{
-				common.HexToHash("0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"),
-				common.HexToHash("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
-			},
-		},
-	}
-	for _, tt := range tests {
-		t.Run(fmt.Sprintf("i %d, count %d", tt.i, tt.count), func(t *testing.T) {
-			start, end := PartitionKeyspace(tt.i, tt.count)
-			require.Equal(t, tt.expected[0], start)
-			require.Equal(t, tt.expected[1], end)
-		})
-	}
-
-	t.Run("panics on invalid i or count", func(t *testing.T) {
-		require.Panics(t, func() {
-			PartitionKeyspace(1, 1)
-		})
-		require.Panics(t, func() {
-			PartitionKeyspace(-1, 1)
-		})
-		require.Panics(t, func() {
-			PartitionKeyspace(0, -1)
-		})
-		require.Panics(t, func() {
-			PartitionKeyspace(-1, -1)
-		})
-	})
-}
-
 func randAddr(t *testing.T) common.Address {
 	var addr common.Address
 	_, err := rand.Read(addr[:])

op-chain-ops/util/state_iterator.go

+package util
+
+import (
+	"fmt"
+	"math/big"
+	"sync"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/log"
+	"github.com/ethereum/go-ethereum/rlp"
+	"github.com/ethereum/go-ethereum/trie"
+)
+
+var (
+	// maxSlot is the maximum possible storage slot.
+	maxSlot = common.HexToHash("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff")
+)
+
+type DBFactory func() (*state.StateDB, error)
+
+type StateCallback func(db *state.StateDB, key, value common.Hash) error
+
+func IterateState(dbFactory DBFactory, address common.Address, cb StateCallback, workers int) error {
+	if workers <= 0 {
+		panic("workers must be greater than 0")
+	}
+
+	// WaitGroup to wait for all workers to finish.
+	var wg sync.WaitGroup
+
+	// Channel to receive errors from each iteration job.
+	errCh := make(chan error, workers)
+	// Channel to cancel all iteration jobs.
+	cancelCh := make(chan struct{})
+
+	worker := func(start, end common.Hash) {
+		// Decrement the WaitGroup when the function returns.
+		defer wg.Done()
+
+		db, err := dbFactory()
+		if err != nil {
+			// Should never happen, so explode if it does.
+			log.Crit("cannot create state db", "err", err)
+		}
+		st, err := db.StorageTrie(address)
+		if err != nil {
+			// Should never happen, so explode if it does.
+			log.Crit("cannot get storage trie", "address", address, "err", err)
+		}
+		// st can be nil if the account doesn't exist.
+		if st == nil {
+			errCh <- fmt.Errorf("account does not exist: %s", address.Hex())
+			return
+		}
+
+		it := trie.NewIterator(st.NodeIterator(start.Bytes()))
+
+		// Below code is largely based on db.ForEachStorage. We can't use that
+		// because it doesn't allow us to specify a start and end key.
+		for it.Next() {
+			select {
+			case <-cancelCh:
+				// If one of the workers encounters an error, cancel all of them.
+				return
+			default:
+				break
+			}
+
+			// Use the raw (i.e., secure hashed) key to check if we've reached
+			// the end of the partition. Use > rather than >= here to account for
+			// the fact that the values returned by PartitionKeys are inclusive.
+			// Duplicate addresses that may be returned by this iteration are
+			// filtered out in the collector.
+			if new(big.Int).SetBytes(it.Key).Cmp(end.Big()) > 0 {
+				return
+			}
+
+			// Skip if the value is empty.
+			rawValue := it.Value
+			if len(rawValue) == 0 {
+				continue
+			}
+
+			// Get the preimage.
+			rawKey := st.GetKey(it.Key)
+			if rawKey == nil {
+				// Should never happen, so explode if it does.
+				log.Crit("cannot get preimage for storage key", "key", it.Key)
+			}
+			key := common.BytesToHash(rawKey)
+
+			// Parse the raw value.
+			_, content, _, err := rlp.Split(rawValue)
+			if err != nil {
+				// Should never happen, so explode if it does.
+				log.Crit("mal-formed data in state: %v", err)
+			}
+
+			value := common.BytesToHash(content)
+
+			// Call the callback with the DB, key, and value. Errors get
+			// bubbled up to the errCh.
+			if err := cb(db, key, value); err != nil {
+				errCh <- err
+				return
+			}
+		}
+	}
+
+	for i := 0; i < workers; i++ {
+		wg.Add(1)
+
+		// Partition the keyspace per worker.
+		start, end := PartitionKeyspace(i, workers)
+
+		// Kick off our worker.
+		go worker(start, end)
+	}
+
+	wg.Wait()
+
+	for len(errCh) > 0 {
+		err := <-errCh
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// PartitionKeyspace divides the key space into partitions by dividing the maximum keyspace
+// by count then multiplying by i. This will leave some slots left over, which we handle below. It
+// returns the start and end keys for the partition as a common.Hash. Note that the returned range
+// of keys is inclusive, i.e., [start, end] NOT [start, end).
+func PartitionKeyspace(i int, count int) (common.Hash, common.Hash) {
+	if i < 0 || count < 0 {
+		panic("i and count must be greater than 0")
+	}
+
+	if i > count-1 {
+		panic("i must be less than count - 1")
+	}
+
+	// Divide the key space into partitions by dividing the key space by the number
+	// of jobs. This will leave some slots left over, which we handle below.
+	partSize := new(big.Int).Div(maxSlot.Big(), big.NewInt(int64(count)))
+
+	start := common.BigToHash(new(big.Int).Mul(big.NewInt(int64(i)), partSize))
+	var end common.Hash
+	if i < count-1 {
+		// If this is not the last partition, use the next partition's start key as the end.
+		end = common.BigToHash(new(big.Int).Mul(big.NewInt(int64(i+1)), partSize))
+	} else {
+		// If this is the last partition, use the max slot as the end.
+		end = maxSlot
+	}
+
+	return start, end
+}

op-chain-ops/util/state_iterator_test.go

+package util
+
+import (
+	crand "crypto/rand"
+	"fmt"
+	"math/rand"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/rawdb"
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/trie"
+	"github.com/stretchr/testify/require"
+)
+
+var testAddr = common.Address{0: 0xff}
+
+func TestStateIteratorWorkers(t *testing.T) {
+	_, factory, _ := setupRandTest(t)
+
+	for i := -1; i <= 0; i++ {
+		require.Panics(t, func() {
+			_ = IterateState(factory, testAddr, func(db *state.StateDB, key, value common.Hash) error {
+				return nil
+			}, i)
+		})
+	}
+}
+
+func TestStateIteratorNonexistentAccount(t *testing.T) {
+	_, factory, _ := setupRandTest(t)
+
+	require.ErrorContains(t, IterateState(factory, common.Address{}, func(db *state.StateDB, key, value common.Hash) error {
+		return nil
+	}, 1), "account does not exist")
+}
+
+func TestStateIteratorRandomOK(t *testing.T) {
+	for i := 0; i < 100; i++ {
+		hashes, factory, workerCount := setupRandTest(t)
+
+		seenHashes := make(map[common.Hash]bool)
+		hashCh := make(chan common.Hash)
+		doneCh := make(chan struct{})
+		go func() {
+			defer close(doneCh)
+			for hash := range hashCh {
+				seenHashes[hash] = true
+			}
+		}()
+
+		require.NoError(t, IterateState(factory, testAddr, func(db *state.StateDB, key, value common.Hash) error {
+			hashCh <- key
+			return nil
+		}, workerCount))
+
+		close(hashCh)
+		<-doneCh
+
+		// Perform a less or equal check here in case of duplicates. The map check below will assert
+		// that all of the hashes are accounted for.
+		require.LessOrEqual(t, len(seenHashes), len(hashes))
+
+		// Every hash we put into state should have been iterated over.
+		for _, hash := range hashes {
+			require.Contains(t, seenHashes, hash)
+		}
+	}
+}
+
+func TestStateIteratorRandomError(t *testing.T) {
+	for i := 0; i < 100; i++ {
+		hashes, factory, workerCount := setupRandTest(t)
+
+		failHash := hashes[rand.Intn(len(hashes))]
+		require.ErrorContains(t, IterateState(factory, testAddr, func(db *state.StateDB, key, value common.Hash) error {
+			if key == failHash {
+				return fmt.Errorf("test error")
+			}
+			return nil
+		}, workerCount), "test error")
+	}
+}
+
+func TestPartitionKeyspace(t *testing.T) {
+	tests := []struct {
+		i        int
+		count    int
+		expected [2]common.Hash
+	}{
+		{
+			i:     0,
+			count: 1,
+			expected: [2]common.Hash{
+				common.HexToHash("0x00"),
+				common.HexToHash("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
+			},
+		},
+		{
+			i:     0,
+			count: 2,
+			expected: [2]common.Hash{
+				common.HexToHash("0x00"),
+				common.HexToHash("0x7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
+			},
+		},
+		{
+			i:     1,
+			count: 2,
+			expected: [2]common.Hash{
+				common.HexToHash("0x7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
+				common.HexToHash("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
+			},
+		},
+		{
+			i:     0,
+			count: 3,
+			expected: [2]common.Hash{
+				common.HexToHash("0x00"),
+				common.HexToHash("0x5555555555555555555555555555555555555555555555555555555555555555"),
+			},
+		},
+		{
+			i:     1,
+			count: 3,
+			expected: [2]common.Hash{
+				common.HexToHash("0x5555555555555555555555555555555555555555555555555555555555555555"),
+				common.HexToHash("0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"),
+			},
+		},
+		{
+			i:     2,
+			count: 3,
+			expected: [2]common.Hash{
+				common.HexToHash("0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"),
+				common.HexToHash("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("i %d, count %d", tt.i, tt.count), func(t *testing.T) {
+			start, end := PartitionKeyspace(tt.i, tt.count)
+			require.Equal(t, tt.expected[0], start)
+			require.Equal(t, tt.expected[1], end)
+		})
+	}
+
+	t.Run("panics on invalid i or count", func(t *testing.T) {
+		require.Panics(t, func() {
+			PartitionKeyspace(1, 1)
+		})
+		require.Panics(t, func() {
+			PartitionKeyspace(-1, 1)
+		})
+		require.Panics(t, func() {
+			PartitionKeyspace(0, -1)
+		})
+		require.Panics(t, func() {
+			PartitionKeyspace(-1, -1)
+		})
+	})
+}
+
+func setupRandTest(t *testing.T) ([]common.Hash, DBFactory, int) {
+	memDB := rawdb.NewMemoryDatabase()
+	db, err := state.New(common.Hash{}, state.NewDatabaseWithConfig(memDB, &trie.Config{
+		Preimages: true,
+		Cache:     1024,
+	}), nil)
+	require.NoError(t, err)
+
+	hashCount := rand.Intn(100)
+	if hashCount == 0 {
+		hashCount = 1
+	}
+
+	hashes := make([]common.Hash, hashCount)
+
+	db.CreateAccount(testAddr)
+
+	for j := 0; j < hashCount; j++ {
+		hashes[j] = randHash(t)
+		db.SetState(testAddr, hashes[j], hashes[j])
+	}
+
+	root, err := db.Commit(false)
+	require.NoError(t, err)
+
+	err = db.Database().TrieDB().Commit(root, true)
+	require.NoError(t, err)
+
+	factory := func() (*state.StateDB, error) {
+		return state.New(root, state.NewDatabaseWithConfig(memDB, &trie.Config{
+			Preimages: true,
+			Cache:     1024,
+		}), nil)
+	}
+
+	workerCount := rand.Intn(64)
+	if workerCount == 0 {
+		workerCount = 1
+	}
+	return hashes, factory, workerCount
+}
+
+func randHash(t *testing.T) common.Hash {
+	var h common.Hash
+	_, err := crand.Read(h[:])
+	require.NoError(t, err)
+	return h
+}