chore: self review

Author: MaksymMalicki

migration/statehistory/class_hash_ingestor.go

@@ -0,0 +1,120 @@
+package statehistory
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/NethermindEth/juno/core/felt"
+	"github.com/NethermindEth/juno/core/state"
+	"github.com/NethermindEth/juno/db"
+	"github.com/NethermindEth/juno/db/dbutils"
+	"github.com/NethermindEth/juno/migration/pipeline"
+	"github.com/NethermindEth/juno/migration/semaphore"
+)
+
+type classHashIngestor struct {
+	baseIngestor
+}
+
+var _ pipeline.State[*felt.Felt, task] = (*classHashIngestor)(nil)
+
+func newClassHashIngestor(
+	ctx context.Context,
+	sem semaphore.ResourceSemaphore[db.Batch],
+	database db.KeyValueReader,
+) *classHashIngestor {
+	return &classHashIngestor{baseIngestor: newBaseIngestor(ctx, sem, database)}
+}
+
+func (i *classHashIngestor) Run(index int, addr *felt.Felt, outputs chan<- task) error {
+	t := &i.tasks[index]
+
+	deprecatedPrefix := db.DeprecatedContractClassHashHistoryKey(addr)
+	contract, err := state.GetContract(i.database, addr)
+	if err != nil {
+		return fmt.Errorf("class-hash: GetContract(%s): %w", addr, err)
+	}
+
+	deployKey := db.ContractClassHashHistoryAtBlockKey(addr, contract.DeployedHeight)
+	deployEntryExists, err := i.database.Has(deployKey)
+	if err != nil {
+		return fmt.Errorf("class-hash: Has(deploy entry): %w", err)
+	}
+
+	depIt, err := i.database.NewIterator(deprecatedPrefix, true)
+	if err != nil {
+		return fmt.Errorf("class-hash: open deprecated iter(%s): %w", addr, err)
+	}
+	defer depIt.Close()
+
+	if !depIt.First() {
+		if deployEntryExists {
+			return nil
+		}
+		err = state.WriteClassHashHistory(
+			t.batch,
+			addr,
+			contract.DeployedHeight,
+			&contract.ClassHash,
+		)
+		if err != nil {
+			return err
+		}
+		t.completedAddrs++
+		t.entryCount++
+		return i.flush(t, outputs)
+	}
+
+	rawValue, err := depIt.Value()
+	if err != nil {
+		return fmt.Errorf("class-hash: read first value(%s): %w", addr, err)
+	}
+	deployClassHash := felt.FromBytes[felt.Felt](rawValue)
+	if err := state.WriteClassHashHistory(
+		t.batch,
+		addr,
+		contract.DeployedHeight,
+		&deployClassHash,
+	); err != nil {
+		return err
+	}
+	t.entryCount++
+	if err := i.flush(t, outputs); err != nil {
+		return err
+	}
+
+	// Shift-up loop: each block in the deprecated history gets the *next*
+	// entry's value (since in the old layout the value at block B was the
+	// value before B's write). The final block gets the head class hash.
+	for {
+		block, err := parseBlockKey(depIt.Key(), deprecatedPrefix)
+		if err != nil {
+			return fmt.Errorf("class-hash(%s): %w", addr, err)
+		}
+		hasNext := depIt.Next()
+		historyValue := contract.ClassHash
+		if hasNext {
+			rawValue, err := depIt.Value()
+			if err != nil {
+				return fmt.Errorf("class-hash(%s): %w", addr, err)
+			}
+			historyValue = felt.FromBytes[felt.Felt](rawValue)
+		}
+		if err := state.WriteClassHashHistory(t.batch, addr, block, &historyValue); err != nil {
+			return err
+		}
+		t.entryCount++
+		if err := i.flush(t, outputs); err != nil {
+			return err
+		}
+		if !hasNext {
+			break
+		}
+	}
+
+	if err := t.batch.DeleteRange(deprecatedPrefix, dbutils.UpperBound(deprecatedPrefix)); err != nil {
+		return fmt.Errorf("class-hash: DeleteRange deprecated(%s): %w", addr, err)
+	}
+	t.completedAddrs++
+	return nil
+}

migration/statehistory/committer.go

@@ -29,9 +29,11 @@ func newCommitter(
 }
 
 func (c *committer) Run(_ int, t task, _ chan<- struct{}) error {
+	defer c.batchSemaphore.Put()
+
 	c.logger.Debug(
 		"writing batch",
-		zap.Int("addrCount", t.addrCount),
+		zap.Int("completedAddrs", t.completedAddrs),
 		zap.Int("entryCount", t.entryCount),
 		zap.Int("batchSize", t.batch.Size()),
 	)
@@ -41,8 +43,7 @@ func (c *committer) Run(_ int, t task, _ chan<- struct{}) error {
 		return err
 	}
 
-	c.counter.log(byteSize, t.addrCount, t.entryCount)
-	c.batchSemaphore.Put()
+	c.counter.log(byteSize, t.completedAddrs, t.entryCount)
 	return nil
 }
 

migration/statehistory/counter.go

@@ -9,13 +9,13 @@ import (
 )
 
 type counter struct {
-	logger      log.StructuredLogger
-	timeLogRate time.Duration
-	phaseName   string
-	start       time.Time
-	size        uint64
-	addrCount   uint64
-	entryCount  uint64
+	logger         log.StructuredLogger
+	timeLogRate    time.Duration
+	phaseName      string
+	start          time.Time
+	size           uint64
+	completedAddrs uint64
+	entryCount     uint64
 }
 
 func newCounter(logger log.StructuredLogger, timeLogRate time.Duration, phaseName string) counter {
@@ -27,29 +27,29 @@ func newCounter(logger log.StructuredLogger, timeLogRate time.Duration, phaseNam
 	}
 }
 
-func (c *counter) log(byteSize uint64, addrCount, entryCount int) {
+func (c *counter) log(byteSize uint64, completedAddrs, entryCount int) {
 	c.size += byteSize
-	c.addrCount += uint64(addrCount)
+	c.completedAddrs += uint64(completedAddrs)
 	c.entryCount += uint64(entryCount)
 
 	now := time.Now()
 	elapsed := now.Sub(c.start).Seconds()
-	if elapsed > float64(c.timeLogRate.Seconds()) {
+	if elapsed > c.timeLogRate.Seconds() {
 		mbs := float64(c.size) / float64(db.Megabyte)
 		c.logger.Info(
 			"write speed",
 			zap.String("phase", c.phaseName),
 			zap.Float64("MB", mbs),
 			zap.Float64("MB/s", mbs/elapsed),
-			zap.Uint64("contracts", c.addrCount),
-			zap.Float64("contracts/s", float64(c.addrCount)/elapsed),
+			zap.Uint64("completedContracts", c.completedAddrs),
+			zap.Float64("completedContracts/s", float64(c.completedAddrs)/elapsed),
 			zap.Uint64("entries", c.entryCount),
 			zap.Float64("entries/s", float64(c.entryCount)/elapsed),
 			zap.Float64("time", elapsed),
 		)
 		c.start = now
 		c.size = 0
-		c.addrCount = 0
+		c.completedAddrs = 0
 		c.entryCount = 0
 	}
 }

migration/statehistory/ingestor.go

@@ -1,57 +1,63 @@
 package statehistory
 
 import (
-	"github.com/NethermindEth/juno/core/felt"
+	"context"
+
 	"github.com/NethermindEth/juno/db"
-	"github.com/NethermindEth/juno/migration/pipeline"
 	"github.com/NethermindEth/juno/migration/semaphore"
 )
 
-type FlushBatchFn func(t *task)
-
-type ingestor struct {
+type baseIngestor struct {
+	ctx            context.Context
 	batchSemaphore semaphore.ResourceSemaphore[db.Batch]
 	database       db.KeyValueReader
 	tasks          []task
-	transform      func(db.KeyValueReader, *task, felt.Address, FlushBatchFn) error
 }
 
-var _ pipeline.State[felt.Address, task] = (*ingestor)(nil)
-
-func newIngestor(
+// newBaseIngestor pre-allocates one batch per ingestor slot. The semaphore is
+// created with capacity ingestorCount+1 immediately before this call, so the
+// acquires cannot block — using GetBlocking keeps the constructor signature
+// error-free.
+func newBaseIngestor(
+	ctx context.Context,
 	sem semaphore.ResourceSemaphore[db.Batch],
 	database db.KeyValueReader,
-	transform func(db.KeyValueReader, *task, felt.Address, FlushBatchFn) error,
-) *ingestor {
+) baseIngestor {
 	tasks := make([]task, ingestorCount)
 	for i := range tasks {
 		tasks[i] = task{batch: sem.GetBlocking()}
 	}
-	return &ingestor{batchSemaphore: sem, database: database, tasks: tasks, transform: transform}
-}
-
-func (p *ingestor) Run(index int, addr felt.Address, outputs chan<- task) error {
-	t := &p.tasks[index]
-	flush := func(t *task) {
-		if t.batch.Size() < targetBatchByteSize {
-			return
-		}
-		outputs <- *t
-		*t = task{batch: p.batchSemaphore.GetBlocking()}
+	return baseIngestor{
+		ctx:            ctx,
+		batchSemaphore: sem,
+		database:       database,
+		tasks:          tasks,
 	}
+}
 
-	if err := p.transform(p.database, t, addr, flush); err != nil {
-		return err
+// flush emits the current task downstream when its batch hits target size and
+// acquires a fresh batch. The ctx-aware select on the channel send is the
+// snappy cancellation point. The semaphore acquire uses GetBlocking — it is
+// guaranteed to unblock within one committer iteration because the committer's
+// deferred Put always runs.
+func (b *baseIngestor) flush(t *task, outputs chan<- task) error {
+	if t.batch.Size() < targetBatchByteSize {
+		return nil
 	}
-
-	if t.batch.Size() >= targetBatchByteSize {
-		outputs <- *t
-		*t = task{batch: p.batchSemaphore.GetBlocking()}
+	select {
+	case <-b.ctx.Done():
+		return b.ctx.Err()
+	case outputs <- *t:
 	}
+	*t = task{batch: b.batchSemaphore.GetBlocking()}
 	return nil
 }
 
-func (p *ingestor) Done(index int, outputs chan<- task) error {
-	outputs <- p.tasks[index]
+func (b *baseIngestor) Done(index int, outputs chan<- task) error {
+	select {
+	case <-b.ctx.Done():
+		return b.ctx.Err()
+	case outputs <- b.tasks[index]:
+	}
 	return nil
 }