diff --git a/execution/commitment/commitment.go b/execution/commitment/commitment.go index 6efb209d8ab..fefea9095ba 100644 --- a/execution/commitment/commitment.go +++ b/execution/commitment/commitment.go @@ -121,6 +121,15 @@ type PatriciaContext interface { Storage(plainKey []byte) (*Update, error) } +// BranchNoCopyReader is the optional half of PatriciaContext. A context that can +// hand out branch bytes without copying implements it; one that cannot -- a +// recorder, a mock -- leaves it off and the caller falls back to Branch. The +// result aliases the reader's memory: valid until the next call on the same +// context, and not to be mutated. +type BranchNoCopyReader interface { + BranchNoCopy(prefix []byte) ([]byte, kv.Step, error) +} + type TrieVariant string const ( diff --git a/execution/commitment/commitmentdb/commitment_context.go b/execution/commitment/commitmentdb/commitment_context.go index 431768e17ef..e665368ca3a 100644 --- a/execution/commitment/commitmentdb/commitment_context.go +++ b/execution/commitment/commitmentdb/commitment_context.go @@ -1022,13 +1022,10 @@ func NewTrieContextRo(reader StateReader, stepSize uint64) *TrieContext { } func (sdc *TrieContext) Branch(pref []byte) ([]byte, kv.Step, error) { - enc, step, err := sdc.readDomain(kv.CommitmentDomain, pref) + enc, step, err := sdc.branch(pref) if err != nil { return nil, 0, err } - if sdc.traceW != nil { - fmt.Fprintf(sdc.traceW, "[SDC] Branch read %x => %x\n", pref, enc) - } // The slice from the underlying state cache / getter aliases storage another // commitment worker can recycle, so the bytes have to be owned here. They are // copied into a per-context buffer rather than a fresh allocation: a trie @@ -1046,6 +1043,25 @@ func (sdc *TrieContext) Branch(pref []byte) ([]byte, kv.Step, error) { return sdc.branchBuf, step, nil } +// BranchNoCopy returns the branch bytes uncopied, for a caller that consumes them +// before its next read -- the trie warmup. The result is valid only until the +// next call on this context and must not be mutated; anything that keeps branch +// data uses Branch. +func (sdc *TrieContext) BranchNoCopy(pref []byte) ([]byte, kv.Step, error) { + return sdc.branch(pref) +} + +func (sdc *TrieContext) branch(pref []byte) ([]byte, kv.Step, error) { + enc, step, err := sdc.readDomain(kv.CommitmentDomain, pref) + if err != nil { + return nil, 0, err + } + if sdc.traceW != nil { + fmt.Fprintf(sdc.traceW, "[SDC] Branch read %x => %x\n", pref, enc) + } + return enc, step, nil +} + func (sdc *TrieContext) PutBranch(prefix []byte, data []byte, prevData []byte) error { if sdc.stateReader.WithHistory() { // do not store branches if explicitly operate on history return nil diff --git a/execution/commitment/warmuper.go b/execution/commitment/warmuper.go index 4af7bd3eb82..af1bd6431ce 100644 --- a/execution/commitment/warmuper.go +++ b/execution/commitment/warmuper.go @@ -29,6 +29,7 @@ import ( "github.com/erigontech/erigon/common" "github.com/erigontech/erigon/common/log/v3" + "github.com/erigontech/erigon/db/kv" "github.com/erigontech/erigon/execution/commitment/nibbles" ) @@ -115,6 +116,11 @@ func (w *Warmuper) Start() { return errors.New("warmup trie context factory returned nil PatriciaContext") } + readBranch := trieCtx.Branch + if w, ok := trieCtx.(BranchNoCopyReader); ok { + readBranch = w.BranchNoCopy + } + for { select { case <-w.ctx.Done(): @@ -123,7 +129,7 @@ func (w *Warmuper) Start() { if !ok { return nil } - w.warmupKey(trieCtx, item.hashedKey, item.startDepth) + w.warmupKey(readBranch, item.hashedKey, item.startDepth) w.keysProcessed.Add(1) w.releaseGen(item.gen) } @@ -141,13 +147,15 @@ func (w *Warmuper) Start() { }) } -func (w *Warmuper) warmupKey(trieCtx PatriciaContext, hashedKey []byte, startDepth int) { +// warmupKey descends one key, parsing each branch only to pick the next nibble. +// It keeps nothing, so the caller can pass the non-copying read. +func (w *Warmuper) warmupKey(readBranch func([]byte) ([]byte, kv.Step, error), hashedKey []byte, startDepth int) { depth := startDepth var compactBuf [maxCompactKeyLen]byte for depth <= len(hashedKey) && depth <= w.maxDepth { prefix := nibbles.HexToCompactInto(compactBuf[:], hashedKey[:depth]) - branchData, _, err := trieCtx.Branch(prefix) + branchData, _, err := readBranch(prefix) if err != nil { log.Debug(fmt.Sprintf("[%s][warmup] failed to get branch", w.logPrefix), "prefix", common.Bytes2Hex(prefix), "error", err) diff --git a/execution/commitment/warmuper_test.go b/execution/commitment/warmuper_test.go index 4b1ab31b301..2cf3b9001ff 100644 --- a/execution/commitment/warmuper_test.go +++ b/execution/commitment/warmuper_test.go @@ -17,10 +17,16 @@ package commitment import ( + "bytes" "context" "sync" + "sync/atomic" "testing" "time" + + "github.com/stretchr/testify/require" + + "github.com/erigontech/erigon/db/kv" ) func TestWarmuperFactoryMustNotOutliveCloseAndWait(t *testing.T) { @@ -255,3 +261,55 @@ func TestCloseLeavesWorkChannelOpen(t *testing.T) { default: } } + +// countingBranchCtx records which read path the warmuper took, and holds +// BranchNoCopy to its contract: each call poisons what the previous one returned, +// so a caller that kept the bytes reads garbage and descends differently. +type countingBranchCtx struct { + src []byte + last []byte + owned atomic.Int64 + borrowed atomic.Int64 +} + +func (c *countingBranchCtx) Branch(prefix []byte) ([]byte, kv.Step, error) { + c.owned.Add(1) + return bytes.Clone(c.src), 0, nil +} + +func (c *countingBranchCtx) BranchNoCopy(prefix []byte) ([]byte, kv.Step, error) { + c.borrowed.Add(1) + for i := range c.last { + c.last[i] = 0xff + } + c.last = bytes.Clone(c.src) + return c.last, 0, nil +} + +func (c *countingBranchCtx) PutBranch(prefix, data, prevData []byte) error { return nil } +func (c *countingBranchCtx) Account(plainKey []byte) (*Update, error) { return nil, nil } +func (c *countingBranchCtx) Storage(plainKey []byte) (*Update, error) { return nil, nil } + +// TestWarmuperUsesBranchNoCopy pins the warmuper on the non-copying read. It keeps +// nothing it reads, and it issues several times more branch reads than the fold +// does, so an owning read there is the bulk of the copying. +func TestWarmuperUsesBranchNoCopy(t *testing.T) { + t.Parallel() + // touchMap, afterMap with nibble 0 set, then one cell with no fields. + ctx := &countingBranchCtx{src: []byte{0x00, 0x01, 0x00, 0x01, 0x00}} + w := NewWarmuper(context.Background(), WarmupConfig{ + Enabled: true, + CtxFactory: func(context.Context) (PatriciaContext, func()) { return ctx, nil }, + NumWorkers: 1, + MaxDepth: WarmupMaxDepth, + }) + w.Start() + w.WarmKey([]byte{0, 0, 0, 0}, 0, 0) + require.NoError(t, w.WaitBufferFree(0)) + w.CloseAndWait() + + require.Zero(t, ctx.owned.Load(), "warmup still copies branch bytes it drops immediately") + // One read per nibble of the key, then one more that stops at its end. A + // caller holding a poisoned branch would read 0xff and stop short. + require.Equal(t, int64(5), ctx.borrowed.Load(), "warmup descent read the wrong number of branches") +}