From ed36c23480adeed2fc3738c1c6f27a6b0ce924a1 Mon Sep 17 00:00:00 2001 From: Ignacio Hagopian Date: Thu, 31 Aug 2023 12:29:58 -0300 Subject: [PATCH] core/state: rewrite a new optimized keyValueMigrator (#256) * trie/utils: add helper to calculate code tree indices * core/state: rewrite optimized version of keyValueMigrator Signed-off-by: Ignacio Hagopian * trie/verkle: remove uint256 allocs (#257) Signed-off-by: Ignacio Hagopian --------- Signed-off-by: Ignacio Hagopian --- core/state_processor.go | 179 ++++++++++++++++++++++++++++------------ trie/utils/verkle.go | 72 ++++++++-------- 2 files changed, 165 insertions(+), 86 deletions(-) diff --git a/core/state_processor.go b/core/state_processor.go index c66c0049d5..8dd4de5436 100644 --- a/core/state_processor.go +++ b/core/state_processor.go @@ -25,6 +25,8 @@ import ( "io" "math/big" "os" + "runtime" + "sync" "time" "github.com/ethereum/go-ethereum/common" @@ -170,7 +172,7 @@ func (p *StateProcessor) Process(block *types.Block, statedb *state.StateDB, cfg // mkv will be assiting in the collection of up to maxMovedCount key values to be migrated to the VKT. // It has internal caches to do efficient MPT->VKT key calculations, which will be discarded after // this function. - mkv := &keyValueMigrator{vktLeafData: make(map[string]*verkle.BatchNewLeafNodeData)} + mkv := newKeyValueMigrator() // move maxCount accounts into the verkle tree, starting with the // slots from the previous account. count := 0 @@ -297,8 +299,17 @@ func (p *StateProcessor) Process(block *types.Block, statedb *state.StateDB, cfg } migrdb.SetCurrentPreimageOffset(preimageSeek) - log.Info("Collected and prepared key values from base tree", "count", count, "duration", time.Since(now), "last account", statedb.Database().GetCurrentAccountHash()) + log.Info("Collected key values from base tree", "count", count, "duration", time.Since(now), "last account", statedb.Database().GetCurrentAccountHash()) + // Take all the collected key-values and prepare the new leaf values. + // This fires a background routine that will start doing the work that + // migrateCollectedKeyValues() will use to insert into the tree. + // + // TODO: Now both prepare() and migrateCollectedKeyValues() are next to each other, but + // after we fix an existing bug, we can call prepare() before the block execution and + // let it do the work in the background. After the block execution and finalization + // finish, we can call migrateCollectedKeyValues() which should already find everything ready. + mkv.prepare() now = time.Now() if err := mkv.migrateCollectedKeyValues(tt.Overlay()); err != nil { return nil, nil, 0, fmt.Errorf("could not migrate key values: %w", err) @@ -380,30 +391,60 @@ func ApplyTransaction(config *params.ChainConfig, bc ChainContext, author *commo return applyTransaction(msg, config, gp, statedb, header.Number, header.Hash(), tx, usedGas, vmenv) } -// keyValueMigrator is a helper struct that collects key-values from the base tree. -// The walk is done in account order, so **we assume** the APIs hold this invariant. This is -// useful to be smart about caching banderwagon.Points to make VKT key calculations faster. -type keyValueMigrator struct { - currAddr []byte - currAddrPoint *verkle.Point +var zeroTreeIndex uint256.Int - vktLeafData map[string]*verkle.BatchNewLeafNodeData +// keyValueMigrator is a helper module that collects key-values from the overlay-tree migration for Verkle Trees. +// It assumes that the walk of the base tree is done in address-order, so it exploit that fact to +// collect the key-values in a way that is efficient. +type keyValueMigrator struct { + // leafData contains the values for the future leaf for a particular VKT branch. + leafData []migratedKeyValue + + // When prepare() is called, it will start a background routine that will process the leafData + // saving the result in newLeaves to be used by migrateCollectedKeyValues(). The background + // routine signals that it is done by closing processingReady. + processingReady chan struct{} + newLeaves []verkle.LeafNode + prepareErr error +} + +func newKeyValueMigrator() *keyValueMigrator { + // We do initialize the VKT config since prepare() might indirectly make multiple GetConfig() calls + // in different goroutines when we never called GetConfig() before, causing a race considering the way + // that `config` is designed in go-verkle. + // TODO: jsign as a fix for this in the PR where we move to a file-less precomp, since it allows safe + // concurrent calls to GetConfig(). When that gets merged, we can remove this line. + _ = verkle.GetConfig() + return &keyValueMigrator{ + processingReady: make(chan struct{}), + leafData: make([]migratedKeyValue, 0, 10_000), + } +} + +type migratedKeyValue struct { + branchKey branchKey + leafNodeData verkle.BatchNewLeafNodeData +} +type branchKey struct { + addr common.Address + treeIndex uint256.Int +} + +func newBranchKey(addr []byte, treeIndex *uint256.Int) branchKey { + var sk branchKey + copy(sk.addr[:], addr) + sk.treeIndex = *treeIndex + return sk } func (kvm *keyValueMigrator) addStorageSlot(addr []byte, slotNumber []byte, slotValue []byte) { - addrPoint := kvm.getAddrPoint(addr) - - vktKey := tutils.GetTreeKeyStorageSlotWithEvaluatedAddress(addrPoint, slotNumber) - leafNodeData := kvm.getOrInitLeafNodeData(vktKey) - - leafNodeData.Values[vktKey[verkle.StemSize]] = slotValue + treeIndex, subIndex := tutils.GetTreeKeyStorageSlotTreeIndexes(slotNumber) + leafNodeData := kvm.getOrInitLeafNodeData(newBranchKey(addr, treeIndex)) + leafNodeData.Values[subIndex] = slotValue } func (kvm *keyValueMigrator) addAccount(addr []byte, acc *types.StateAccount) { - addrPoint := kvm.getAddrPoint(addr) - - vktKey := tutils.GetTreeKeyVersionWithEvaluatedAddress(addrPoint) - leafNodeData := kvm.getOrInitLeafNodeData(vktKey) + leafNodeData := kvm.getOrInitLeafNodeData(newBranchKey(addr, &zeroTreeIndex)) var version [verkle.LeafValueSize]byte leafNodeData.Values[tutils.VersionLeafKey] = version[:] @@ -419,16 +460,10 @@ func (kvm *keyValueMigrator) addAccount(addr []byte, acc *types.StateAccount) { leafNodeData.Values[tutils.NonceLeafKey] = nonce[:] leafNodeData.Values[tutils.CodeKeccakLeafKey] = acc.CodeHash[:] - - // Code size is ignored here. If this isn't an EOA, the tree-walk will call - // addAccountCode with this information. } func (kvm *keyValueMigrator) addAccountCode(addr []byte, codeSize uint64, chunks []byte) { - addrPoint := kvm.getAddrPoint(addr) - - vktKey := tutils.GetTreeKeyVersionWithEvaluatedAddress(addrPoint) - leafNodeData := kvm.getOrInitLeafNodeData(vktKey) + leafNodeData := kvm.getOrInitLeafNodeData(newBranchKey(addr, &zeroTreeIndex)) // Save the code size. var codeSizeBytes [verkle.LeafValueSize]byte @@ -442,8 +477,8 @@ func (kvm *keyValueMigrator) addAccountCode(addr []byte, codeSize uint64, chunks // Potential further chunks, have their own leaf nodes. for i := 128; i < len(chunks)/32; { - vktKey := tutils.GetTreeKeyCodeChunkWithEvaluatedAddress(addrPoint, uint256.NewInt(uint64(i))) - leafNodeData := kvm.getOrInitLeafNodeData(vktKey) + treeIndex, _ := tutils.GetTreeKeyCodeChunkIndices(uint256.NewInt(uint64(i))) + leafNodeData := kvm.getOrInitLeafNodeData(newBranchKey(addr, treeIndex)) j := i for ; (j-i) < 256 && j < len(chunks)/32; j++ { @@ -453,41 +488,79 @@ func (kvm *keyValueMigrator) addAccountCode(addr []byte, codeSize uint64, chunks } } -func (kvm *keyValueMigrator) getAddrPoint(addr []byte) *verkle.Point { - if bytes.Equal(addr, kvm.currAddr) { - return kvm.currAddrPoint +func (kvm *keyValueMigrator) getOrInitLeafNodeData(bk branchKey) *verkle.BatchNewLeafNodeData { + // Remember that keyValueMigration receives actions ordered by (address, subtreeIndex). + // This means that we can assume that the last element of leafData is the one that we + // are looking for, or that we need to create a new one. + if len(kvm.leafData) == 0 || kvm.leafData[len(kvm.leafData)-1].branchKey != bk { + kvm.leafData = append(kvm.leafData, migratedKeyValue{ + branchKey: bk, + leafNodeData: verkle.BatchNewLeafNodeData{ + Stem: nil, // It will be calculated in the prepare() phase, since it's CPU heavy. + Values: make(map[byte][]byte), + }, + }) } - kvm.currAddr = addr - kvm.currAddrPoint = tutils.EvaluateAddressPoint(addr) - return kvm.currAddrPoint + return &kvm.leafData[len(kvm.leafData)-1].leafNodeData } -func (kvm *keyValueMigrator) getOrInitLeafNodeData(stem []byte) *verkle.BatchNewLeafNodeData { - stemStr := string(stem) - if _, ok := kvm.vktLeafData[stemStr]; !ok { - kvm.vktLeafData[stemStr] = &verkle.BatchNewLeafNodeData{ - Stem: stem[:verkle.StemSize], - Values: make(map[byte][]byte), +func (kvm *keyValueMigrator) prepare() { + // We fire a background routine to process the leafData and save the result in newLeaves. + // The background routine signals that it is done by closing processingReady. + go func() { + // Step 1: We split kvm.leafData in numBatches batches, and we process each batch in a separate goroutine. + // This fills each leafNodeData.Stem with the correct value. + var wg sync.WaitGroup + batchNum := runtime.NumCPU() + batchSize := (len(kvm.leafData) + batchNum - 1) / batchNum + for i := 0; i < len(kvm.leafData); i += batchSize { + start := i + end := i + batchSize + if end > len(kvm.leafData) { + end = len(kvm.leafData) + } + wg.Add(1) + + batch := kvm.leafData[start:end] + go func() { + defer wg.Done() + var currAddr common.Address + var currPoint *verkle.Point + for i := range batch { + if batch[i].branchKey.addr != currAddr { + currAddr = batch[i].branchKey.addr + currPoint = tutils.EvaluateAddressPoint(currAddr[:]) + } + stem := tutils.GetTreeKeyWithEvaluatedAddess(currPoint, &batch[i].branchKey.treeIndex, 0) + stem = stem[:verkle.StemSize] + batch[i].leafNodeData.Stem = stem + } + }() } - } - return kvm.vktLeafData[stemStr] + wg.Wait() + + // Step 2: Now that we have all stems (i.e: tree keys) calcualted, we can create the new leaves. + nodeValues := make([]verkle.BatchNewLeafNodeData, len(kvm.leafData)) + for i := range kvm.leafData { + nodeValues[i] = kvm.leafData[i].leafNodeData + } + + // Create all leaves in batch mode so we can optimize cryptography operations. + kvm.newLeaves, kvm.prepareErr = verkle.BatchNewLeafNode(nodeValues) + close(kvm.processingReady) + }() } func (kvm *keyValueMigrator) migrateCollectedKeyValues(tree *trie.VerkleTrie) error { - // Transform the map into a slice. - nodeValues := make([]verkle.BatchNewLeafNodeData, 0, len(kvm.vktLeafData)) - for _, vld := range kvm.vktLeafData { - nodeValues = append(nodeValues, *vld) - } - - // Create all leaves in batch mode so we can optimize cryptography operations. - newLeaves, err := verkle.BatchNewLeafNode(nodeValues) - if err != nil { - return fmt.Errorf("failed to batch-create new leaf nodes") + now := time.Now() + <-kvm.processingReady + if kvm.prepareErr != nil { + return fmt.Errorf("failed to prepare key values: %w", kvm.prepareErr) } + log.Info("Prepared key values from base tree", "duration", time.Since(now)) // Insert into the tree. - if err := tree.InsertMigratedLeaves(newLeaves); err != nil { + if err := tree.InsertMigratedLeaves(kvm.newLeaves); err != nil { return fmt.Errorf("failed to insert migrated leaves: %w", err) } diff --git a/trie/utils/verkle.go b/trie/utils/verkle.go index 85e479b641..07949ec65e 100644 --- a/trie/utils/verkle.go +++ b/trie/utils/verkle.go @@ -17,7 +17,7 @@ package utils import ( - "math/big" + "encoding/binary" "sync" "github.com/crate-crypto/go-ipa/bandersnatch/fr" @@ -34,18 +34,14 @@ const ( ) var ( - zero = uint256.NewInt(0) - HeaderStorageOffset = uint256.NewInt(64) - CodeOffset = uint256.NewInt(128) - MainStorageOffset = new(uint256.Int).Lsh(uint256.NewInt(256), 31) - VerkleNodeWidth = uint256.NewInt(256) - codeStorageDelta = uint256.NewInt(0).Sub(CodeOffset, HeaderStorageOffset) - - // BigInt versions of the above. - headerStorageOffsetBig = HeaderStorageOffset.ToBig() - mainStorageOffsetBig = MainStorageOffset.ToBig() - verkleNodeWidthBig = VerkleNodeWidth.ToBig() - codeStorageDeltaBig = codeStorageDelta.ToBig() + zero = uint256.NewInt(0) + VerkleNodeWidthLog2 = 8 + HeaderStorageOffset = uint256.NewInt(64) + mainStorageOffsetLshVerkleNodeWidth = new(uint256.Int).Lsh(uint256.NewInt(256), 31-uint(VerkleNodeWidthLog2)) + CodeOffset = uint256.NewInt(128) + MainStorageOffset = new(uint256.Int).Lsh(uint256.NewInt(256), 31) + VerkleNodeWidth = uint256.NewInt(256) + codeStorageDelta = uint256.NewInt(0).Sub(CodeOffset, HeaderStorageOffset) getTreePolyIndex0Point *verkle.Point ) @@ -164,6 +160,11 @@ func GetTreeKeyCodeSize(address []byte) []byte { } func GetTreeKeyCodeChunk(address []byte, chunk *uint256.Int) []byte { + treeIndex, subIndex := GetTreeKeyCodeChunkIndices(chunk) + return GetTreeKey(address, treeIndex, subIndex) +} + +func GetTreeKeyCodeChunkIndices(chunk *uint256.Int) (*uint256.Int, byte) { chunkOffset := new(uint256.Int).Add(CodeOffset, chunk) treeIndex := new(uint256.Int).Div(chunkOffset, VerkleNodeWidth) subIndexMod := new(uint256.Int).Mod(chunkOffset, VerkleNodeWidth) @@ -171,7 +172,7 @@ func GetTreeKeyCodeChunk(address []byte, chunk *uint256.Int) []byte { if len(subIndexMod) != 0 { subIndex = byte(subIndexMod[0]) } - return GetTreeKey(address, treeIndex, subIndex) + return treeIndex, subIndex } func GetTreeKeyCodeChunkWithEvaluatedAddress(addressPoint *verkle.Point, chunk *uint256.Int) []byte { @@ -230,8 +231,8 @@ func GetTreeKeyWithEvaluatedAddess(evaluated *verkle.Point, treeIndex *uint256.I // little-endian, 32-byte aligned treeIndex var index [32]byte - for i, b := range treeIndex.Bytes() { - index[len(treeIndex.Bytes())-1-i] = b + for i := 0; i < len(treeIndex); i++ { + binary.LittleEndian.PutUint64(index[i*8:(i+1)*8], treeIndex[i]) } verkle.FromLEBytes(&poly[3], index[:16]) verkle.FromLEBytes(&poly[4], index[16:]) @@ -274,22 +275,27 @@ func GetTreeKeyStorageSlotWithEvaluatedAddress(evaluated *verkle.Point, storageK } func GetTreeKeyStorageSlotTreeIndexes(storageKey []byte) (*uint256.Int, byte) { - // Note that `pos` must be a big.Int and not a uint256.Int, because the subsequent - // arithmetics operations could overflow. (e.g: imagine if storageKey is 2^256-1) - pos := new(big.Int).SetBytes(storageKey) - if pos.Cmp(codeStorageDeltaBig) < 0 { - pos.Add(headerStorageOffsetBig, pos) - } else { - pos.Add(mainStorageOffsetBig, pos) - } - treeIndex, overflow := uint256.FromBig(big.NewInt(0).Div(pos, verkleNodeWidthBig)) - if overflow { // Must never happen considering the EIP definition. - panic("tree index overflow") - } - // calculate the sub_index, i.e. the index in the stem tree. - // Because the modulus is 256, it's the last byte of treeIndex - posBytes := pos.Bytes() - subIndex := posBytes[len(posBytes)-1] + var pos uint256.Int + pos.SetBytes(storageKey) - return treeIndex, subIndex + // If the storage slot is in the header, we need to add the header offset. + if pos.Cmp(codeStorageDelta) < 0 { + // This addition is always safe; it can't ever overflow since pos