Skip to content

Commit 60ee2a8

Browse files
brbrrbrbrr
committed
test: consolidate PRs #3657 #3658 #3659 for testing
Merges maksym/trie-migration (which stacks on headstate + statehistory) into a single branch off main for integration testing. - #3657 feat(migration): headstate migration - #3658 feat(migration): statehistory migration - #3659 feat(migration): trie migration
2 parents 345d918 + df52a01 commit 60ee2a8

25 files changed

Lines changed: 3777 additions & 3 deletions

‎core/state/accessors.go‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,21 @@ func HasContract(r db.KeyValueReader, addr *felt.Felt) (bool, error) {
4141
return r.Has(key)
4242
}
4343

44-
func WriteContract(w db.KeyValueWriter, addr *felt.Felt, contract *stateContract) error {
44+
func WriteContract(
45+
w db.KeyValueWriter,
46+
addr *felt.Felt,
47+
nonce, classHash felt.Felt,
48+
deployHeight uint64,
49+
) error {
50+
contract := stateContract{
51+
Nonce: nonce,
52+
ClassHash: classHash,
53+
DeployedHeight: deployHeight,
54+
}
55+
return writeContract(w, addr, &contract)
56+
}
57+
58+
func writeContract(w db.KeyValueWriter, addr *felt.Felt, contract *stateContract) error {
4559
key := db.ContractKey(addr)
4660
data, err := contract.MarshalBinary()
4761
if err != nil {

‎core/state/state.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -376,7 +376,7 @@ func (s *State) flush(
376376
return err
377377
}
378378
} else { // updated
379-
if err := WriteContract(s.batch, &addr, obj.contract); err != nil {
379+
if err := writeContract(s.batch, &addr, obj.contract); err != nil {
380380
return err
381381
}
382382
}

‎migration/headstate/committer.go‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
package headstate
2+
3+
import (
4+
"github.com/NethermindEth/juno/db"
5+
"github.com/NethermindEth/juno/migration/pipeline"
6+
"github.com/NethermindEth/juno/migration/semaphore"
7+
"github.com/NethermindEth/juno/utils/log"
8+
"go.uber.org/zap"
9+
)
10+
11+
type committer struct {
12+
counter counter
13+
logger log.StructuredLogger
14+
batchSemaphore semaphore.ResourceSemaphore[db.Batch]
15+
}
16+
17+
var _ pipeline.State[task, struct{}] = (*committer)(nil)
18+
19+
func newCommitter(
20+
logger log.StructuredLogger,
21+
batchSemaphore semaphore.ResourceSemaphore[db.Batch],
22+
) *committer {
23+
return &committer{
24+
logger: logger,
25+
counter: newCounter(logger, timeLogRate),
26+
batchSemaphore: batchSemaphore,
27+
}
28+
}
29+
30+
func (c *committer) Run(_ int, t task, _ chan<- struct{}) error {
31+
c.logger.Debug(
32+
"writing batch",
33+
zap.Int("completedAddrs", t.completedAddrs),
34+
zap.Int("batchSize", t.batch.Size()),
35+
)
36+
37+
byteSize := uint64(t.batch.Size())
38+
if err := t.batch.Write(); err != nil {
39+
return err
40+
}
41+
42+
c.counter.log(byteSize, t.completedAddrs)
43+
c.batchSemaphore.Put()
44+
return nil
45+
}
46+
47+
func (c *committer) Done(int, chan<- struct{}) error {
48+
return nil
49+
}

‎migration/headstate/counter.go‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
package headstate
2+
3+
import (
4+
"time"
5+
6+
"github.com/NethermindEth/juno/db"
7+
"github.com/NethermindEth/juno/utils/log"
8+
"go.uber.org/zap"
9+
)
10+
11+
type counter struct {
12+
logger log.StructuredLogger
13+
timeLogRate time.Duration
14+
start time.Time
15+
size uint64
16+
completedAddrs uint64
17+
}
18+
19+
func newCounter(logger log.StructuredLogger, timeLogRate time.Duration) counter {
20+
return counter{
21+
logger: logger,
22+
timeLogRate: timeLogRate,
23+
start: time.Now(),
24+
}
25+
}
26+
27+
func (c *counter) log(byteSize uint64, completedAddrs int) {
28+
c.size += byteSize
29+
c.completedAddrs += uint64(completedAddrs)
30+
31+
now := time.Now()
32+
elapsed := now.Sub(c.start).Seconds()
33+
if elapsed > c.timeLogRate.Seconds() {
34+
mbs := float64(c.size) / float64(db.Megabyte)
35+
c.logger.Info(
36+
"write speed",
37+
zap.Float64("MB", mbs),
38+
zap.Float64("MB/s", mbs/elapsed),
39+
zap.Uint64("completedContracts", c.completedAddrs),
40+
zap.Float64("completedContracts/s", float64(c.completedAddrs)/elapsed),
41+
zap.Float64("time", elapsed),
42+
)
43+
c.start = now
44+
c.size = 0
45+
c.completedAddrs = 0
46+
}
47+
}

‎migration/headstate/ingestor.go‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
package headstate
2+
3+
import (
4+
"errors"
5+
"fmt"
6+
7+
"github.com/NethermindEth/juno/core"
8+
"github.com/NethermindEth/juno/core/felt"
9+
"github.com/NethermindEth/juno/core/state"
10+
"github.com/NethermindEth/juno/db"
11+
"github.com/NethermindEth/juno/migration/pipeline"
12+
"github.com/NethermindEth/juno/migration/semaphore"
13+
)
14+
15+
type ingestor struct {
16+
database db.KeyValueReader
17+
batchSemaphore semaphore.ResourceSemaphore[db.Batch]
18+
tasks []task
19+
}
20+
21+
func newIngestor(
22+
database db.KeyValueReader,
23+
batchSemaphore semaphore.ResourceSemaphore[db.Batch],
24+
) *ingestor {
25+
tasks := make([]task, ingestorCount)
26+
for i := range tasks {
27+
tasks[i] = task{batch: batchSemaphore.GetBlocking()}
28+
}
29+
return &ingestor{
30+
database: database,
31+
batchSemaphore: batchSemaphore,
32+
tasks: tasks,
33+
}
34+
}
35+
36+
var _ pipeline.State[felt.Address, task] = (*ingestor)(nil)
37+
38+
func (c *ingestor) Run(index int, addr felt.Address, outputs chan<- task) error {
39+
t := &c.tasks[index]
40+
41+
sizeBefore := t.batch.Size()
42+
if err := c.ingestAddress(t.batch, addr); err != nil {
43+
return err
44+
}
45+
if t.batch.Size() > sizeBefore {
46+
t.completedAddrs++
47+
}
48+
49+
if t.batch.Size() >= targetBatchByteSize {
50+
outputs <- task{batch: t.batch, completedAddrs: t.completedAddrs}
51+
t.completedAddrs = 0
52+
t.batch = c.batchSemaphore.GetBlocking()
53+
}
54+
return nil
55+
}
56+
57+
func (c *ingestor) Done(index int, outputs chan<- task) error {
58+
outputs <- c.tasks[index]
59+
return nil
60+
}
61+
62+
func (c *ingestor) ingestAddress(batch db.Batch, addr felt.Address) error {
63+
addrFelt := felt.Felt(addr)
64+
65+
already, err := state.HasContract(c.database, &addrFelt)
66+
if err != nil {
67+
return fmt.Errorf("HasContract(%s): %w", &addrFelt, err)
68+
}
69+
if already {
70+
return nil
71+
}
72+
73+
classHash, err := core.GetContractClassHash(c.database, &addrFelt)
74+
if err != nil {
75+
return fmt.Errorf("GetContractClassHash(%s): %w", &addrFelt, err)
76+
}
77+
78+
nonce, err := core.GetContractNonce(c.database, &addrFelt)
79+
if err != nil {
80+
if !errors.Is(err, db.ErrKeyNotFound) {
81+
return fmt.Errorf("GetContractNonce(%s): %w", &addrFelt, err)
82+
}
83+
nonce = felt.Zero
84+
}
85+
86+
height, err := core.GetContractDeploymentHeight(c.database, &addrFelt)
87+
if err != nil {
88+
return fmt.Errorf("GetContractDeploymentHeight(%s): %w", &addrFelt, err)
89+
}
90+
91+
if err := state.WriteContract(batch, &addrFelt, nonce, classHash, height); err != nil {
92+
return fmt.Errorf("WriteContract(%s): %w", &addrFelt, err)
93+
}
94+
return nil
95+
}

‎migration/headstate/migrator.go‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
package headstate
2+
3+
import (
4+
"context"
5+
"errors"
6+
"fmt"
7+
"iter"
8+
"time"
9+
10+
"github.com/NethermindEth/juno/blockchain/networks"
11+
"github.com/NethermindEth/juno/core/felt"
12+
"github.com/NethermindEth/juno/db"
13+
"github.com/NethermindEth/juno/db/dbutils"
14+
"github.com/NethermindEth/juno/migration"
15+
"github.com/NethermindEth/juno/migration/pipeline"
16+
"github.com/NethermindEth/juno/migration/semaphore"
17+
"github.com/NethermindEth/juno/utils/log"
18+
)
19+
20+
const (
21+
batchByteSize = 128 * db.Megabyte
22+
targetBatchByteSize = 96 * db.Megabyte
23+
ingestorCount = 4
24+
timeLogRate = 5 * time.Second
25+
)
26+
27+
type task struct {
28+
batch db.Batch
29+
completedAddrs int
30+
}
31+
32+
var (
33+
shouldRerun = []byte{}
34+
shouldNotRerun = []byte(nil)
35+
)
36+
37+
var _ migration.Migration = (*Migrator)(nil)
38+
39+
// Migrator consolidates the deprecated per-field contract layout into a
40+
// single Contract record per address, written via state.WriteContract:
41+
//
42+
// ContractClassHash[addr]
43+
// ContractNonce[addr]
44+
// ContractDeploymentHeight[addr]
45+
// │
46+
// ▼
47+
// Contract[addr] = { ClassHash, Nonce, DeployedHeight }
48+
//
49+
// StorageRoot is left zero — the running node lazily backfills it on the
50+
// contract's first storage write.
51+
//
52+
// Each address discovered in the ContractClassHash bucket is processed by one
53+
// of ingestorCount worker goroutines that read the three old fields into a
54+
// shared db.Batch; a single committer drains batches to disk. Once every
55+
// address has been migrated, the three deprecated buckets are wiped via
56+
// DeleteRange.
57+
//
58+
// Re-run safe: an address whose Contract record already exists is skipped
59+
// (via state.HasContract), and the trailing wipe re-issues DeleteRange over
60+
// the (possibly already empty) ranges.
61+
type Migrator struct{}
62+
63+
func (Migrator) Before([]byte) error {
64+
return nil
65+
}
66+
67+
func (Migrator) Migrate(
68+
ctx context.Context,
69+
database db.KeyValueStore,
70+
_ *networks.Network,
71+
logger log.StructuredLogger,
72+
) ([]byte, error) {
73+
addressesIter, sourceErr := pendingAddresses(database)
74+
res := migrateAddresses(ctx, database, logger, addressesIter)
75+
76+
if err := errors.Join(sourceErr(), res.Err); err != nil {
77+
return shouldRerun, err
78+
}
79+
if !res.IsDone {
80+
if ctxErr := ctx.Err(); ctxErr != nil {
81+
return shouldRerun, ctxErr
82+
}
83+
return shouldRerun, errors.New("headstate migration did not complete")
84+
}
85+
86+
return shouldNotRerun, wipeDeprecatedBuckets(database)
87+
}
88+
89+
func migrateAddresses(
90+
ctx context.Context,
91+
database db.KeyValueStore,
92+
logger log.StructuredLogger,
93+
addresses iter.Seq[felt.Address],
94+
) pipeline.Result {
95+
batchSemaphore := semaphore.New(
96+
ingestorCount+1,
97+
func() db.Batch {
98+
return database.NewBatchWithSize(batchByteSize)
99+
},
100+
)
101+
102+
source := pipeline.Source(addresses)
103+
104+
ingestorPipeline := pipeline.New(
105+
source,
106+
ingestorCount,
107+
newIngestor(database, batchSemaphore),
108+
)
109+
110+
committerPipeline := pipeline.New(
111+
ingestorPipeline,
112+
1,
113+
newCommitter(logger, batchSemaphore),
114+
)
115+
116+
_, wait := committerPipeline.Run(ctx)
117+
return wait()
118+
}
119+
120+
func pendingAddresses(r db.KeyValueReader) (iter.Seq[felt.Address], func() error) {
121+
var iterErr error
122+
seq := func(yield func(felt.Address) bool) {
123+
prefix := db.ContractClassHash.Key()
124+
it, err := r.NewIterator(prefix, true)
125+
if err != nil {
126+
iterErr = err
127+
return
128+
}
129+
defer it.Close()
130+
131+
for valid := it.First(); valid; valid = it.Next() {
132+
key := it.Key()
133+
if len(key) != len(prefix)+felt.Bytes {
134+
iterErr = fmt.Errorf(
135+
"malformed ContractClassHash key: len %d, want %d",
136+
len(key),
137+
len(prefix)+felt.Bytes,
138+
)
139+
return
140+
}
141+
f := felt.FromBytes[felt.Felt](key[len(prefix):])
142+
if !yield(felt.Address(f)) {
143+
return
144+
}
145+
}
146+
}
147+
return seq, func() error { return iterErr }
148+
}
149+
150+
func wipeDeprecatedBuckets(database db.KeyValueStore) error {
151+
for _, bucket := range []db.Bucket{
152+
db.ContractClassHash,
153+
db.ContractNonce,
154+
db.ContractDeploymentHeight,
155+
} {
156+
start := bucket.Key()
157+
end := dbutils.UpperBound(start)
158+
if err := database.DeleteRange(start, end); err != nil {
159+
return err
160+
}
161+
}
162+
return nil
163+
}

0 commit comments

Comments
 (0)