Merge pull request #5916 from onflow/fxamacker/optimize-commit-and-pr…

…eload-for-cadence-migration

Use faster atree funcs to speedup Cadence 1.0 migration

cmd/util/ledger/migrations/account_storage_migration.go

@@ -41,7 +41,7 @@ func NewAccountStorageMigration(
 		}
 
 		// Commit the changes
-		err = storage.Commit(inter, false)
+		err = storage.NondeterministicCommit(inter, false)
 		if err != nil {
 			return fmt.Errorf("failed to commit changes: %w", err)
 		}

cmd/util/ledger/migrations/cadence_value_diff.go

@@ -77,44 +77,53 @@ type CadenceValueDiffReporter struct {
 	chainID        flow.ChainID
 	reportWriter   reporters.ReportWriter
 	verboseLogging bool
+	nWorkers       int
 }
 
 func NewCadenceValueDiffReporter(
 	address common.Address,
 	chainID flow.ChainID,
 	rw reporters.ReportWriter,
 	verboseLogging bool,
+	nWorkers int,
 ) *CadenceValueDiffReporter {
 	return &CadenceValueDiffReporter{
 		address:        address,
 		chainID:        chainID,
 		reportWriter:   rw,
 		verboseLogging: verboseLogging,
+		nWorkers:       nWorkers,
 	}
 }
 
 func (dr *CadenceValueDiffReporter) newStorageRuntime(
 	payloads []*ledger.Payload,
 ) (
 	*InterpreterMigrationRuntime,
+	registers.Registers,
 	error,
 ) {
 	registersByAccount, err := registers.NewByAccountFromPayloads(payloads)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	// TODO: maybe make read-only again
-	return NewInterpreterMigrationRuntime(
+	runtimeInterface, err := NewInterpreterMigrationRuntime(
 		registersByAccount,
 		dr.chainID,
 		InterpreterMigrationRuntimeConfig{},
 	)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return runtimeInterface, registersByAccount, nil
 }
 
 func (dr *CadenceValueDiffReporter) DiffStates(oldPayloads, newPayloads []*ledger.Payload, domains []string) {
 	// Create all the runtime components we need for comparing Cadence values.
-	oldRuntime, err := dr.newStorageRuntime(oldPayloads)
+	oldRuntime, oldRegs, err := dr.newStorageRuntime(oldPayloads)
 	if err != nil {
 		dr.reportWriter.Write(
 			diffError{
@@ -125,7 +134,18 @@ func (dr *CadenceValueDiffReporter) DiffStates(oldPayloads, newPayloads []*ledge
 		return
 	}
 
-	newRuntime, err := dr.newStorageRuntime(newPayloads)
+	err = loadAtreeSlabsInStorage(oldRuntime.Storage, oldRegs, dr.nWorkers)
+	if err != nil {
+		dr.reportWriter.Write(
+			diffError{
+				Address: dr.address.Hex(),
+				Kind:    diffErrorKindString[abortErrorKind],
+				Msg:     fmt.Sprintf("failed to preload old state atree payloads: %s", err),
+			})
+		return
+	}
+
+	newRuntime, newRegs, err := dr.newStorageRuntime(newPayloads)
 	if err != nil {
 		dr.reportWriter.Write(
 			diffError{
@@ -136,6 +156,17 @@ func (dr *CadenceValueDiffReporter) DiffStates(oldPayloads, newPayloads []*ledge
 		return
 	}
 
+	err = loadAtreeSlabsInStorage(newRuntime.Storage, newRegs, dr.nWorkers)
+	if err != nil {
+		dr.reportWriter.Write(
+			diffError{
+				Address: dr.address.Hex(),
+				Kind:    diffErrorKindString[abortErrorKind],
+				Msg:     fmt.Sprintf("failed to preload new state atree payloads: %s", err),
+			})
+		return
+	}
+
 	// Iterate through all domains and compare cadence values.
 	for _, domain := range domains {
 		dr.diffStorageDomain(oldRuntime, newRuntime, domain)

cmd/util/ledger/migrations/cadence_value_diff_test.go

@@ -2,6 +2,7 @@ package migrations
 
 import (
 	"fmt"
+	"runtime"
 	"strconv"
 	"testing"
 
@@ -30,7 +31,7 @@ func TestDiffCadenceValues(t *testing.T) {
 
 		writer := &testReportWriter{}
 
-		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, true)
+		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, true, runtime.NumCPU())
 
 		diffReporter.DiffStates(
 			createTestPayloads(t, address, domain),
@@ -46,7 +47,7 @@ func TestDiffCadenceValues(t *testing.T) {
 
 		writer := &testReportWriter{}
 
-		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, true)
+		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, true, runtime.NumCPU())
 
 		diffReporter.DiffStates(
 			createTestPayloads(t, address, domain),
@@ -67,7 +68,7 @@ func TestDiffCadenceValues(t *testing.T) {
 
 		writer := &testReportWriter{}
 
-		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, true)
+		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, true, runtime.NumCPU())
 
 		diffReporter.DiffStates(
 			createTestPayloads(t, address, domain),
@@ -94,7 +95,7 @@ func TestDiffCadenceValues(t *testing.T) {
 
 		writer := &testReportWriter{}
 
-		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, true)
+		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, true, runtime.NumCPU())
 
 		diffReporter.DiffStates(
 			createStorageMapPayloads(t, address, domain, []string{"0", "1"}, []interpreter.Value{interpreter.UInt64Value(0), interpreter.UInt64Value(0)}),
@@ -121,7 +122,7 @@ func TestDiffCadenceValues(t *testing.T) {
 
 		writer := &testReportWriter{}
 
-		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, false)
+		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, false, runtime.NumCPU())
 
 		diffReporter.DiffStates(
 			createStorageMapPayloads(t, address, domain, []string{"0", "1"}, []interpreter.Value{interpreter.UInt64Value(100), interpreter.UInt64Value(101)}),
@@ -148,7 +149,7 @@ func TestDiffCadenceValues(t *testing.T) {
 
 		writer := &testReportWriter{}
 
-		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, false)
+		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, false, runtime.NumCPU())
 
 		diffReporter.DiffStates(
 			createStorageMapPayloads(t, address, domain, []string{"0", "1"}, []interpreter.Value{interpreter.UInt64Value(100), interpreter.UInt64Value(101)}),
@@ -174,7 +175,7 @@ func TestDiffCadenceValues(t *testing.T) {
 
 		writer := &testReportWriter{}
 
-		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, false)
+		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, false, runtime.NumCPU())
 
 		createPayloads := func(arrayValues []interpreter.Value) []*ledger.Payload {
 
@@ -229,7 +230,7 @@ func TestDiffCadenceValues(t *testing.T) {
 				),
 			)
 
-			err = mr.Storage.Commit(mr.Interpreter, false)
+			err = mr.Storage.NondeterministicCommit(mr.Interpreter, false)
 			require.NoError(t, err)
 
 			// finalize the transaction
@@ -284,7 +285,7 @@ func TestDiffCadenceValues(t *testing.T) {
 
 		writer := &testReportWriter{}
 
-		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, false)
+		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, false, runtime.NumCPU())
 
 		createPayloads := func(dictValues []interpreter.Value) []*ledger.Payload {
 
@@ -340,7 +341,7 @@ func TestDiffCadenceValues(t *testing.T) {
 				),
 			)
 
-			err = mr.Storage.Commit(mr.Interpreter, false)
+			err = mr.Storage.NondeterministicCommit(mr.Interpreter, false)
 			require.NoError(t, err)
 
 			// finalize the transaction
@@ -401,7 +402,7 @@ func TestDiffCadenceValues(t *testing.T) {
 
 		writer := &testReportWriter{}
 
-		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, false)
+		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, false, runtime.NumCPU())
 
 		createPayloads := func(compositeFields []string, compositeValues []interpreter.Value) []*ledger.Payload {
 
@@ -462,7 +463,7 @@ func TestDiffCadenceValues(t *testing.T) {
 				),
 			)
 
-			err = mr.Storage.Commit(mr.Interpreter, false)
+			err = mr.Storage.NondeterministicCommit(mr.Interpreter, false)
 			require.NoError(t, err)
 
 			// finalize the transaction
@@ -529,7 +530,7 @@ func TestDiffCadenceValues(t *testing.T) {
 
 		writer := &testReportWriter{}
 
-		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, true)
+		diffReporter := NewCadenceValueDiffReporter(address, flow.Emulator, writer, true, runtime.NumCPU())
 
 		createPayloads := func(compositeFields []string, compositeValues []interpreter.Value) []*ledger.Payload {
 
@@ -590,7 +591,7 @@ func TestDiffCadenceValues(t *testing.T) {
 				),
 			)
 
-			err = mr.Storage.Commit(mr.Interpreter, false)
+			err = mr.Storage.NondeterministicCommit(mr.Interpreter, false)
 			require.NoError(t, err)
 
 			// finalize the transaction
@@ -701,7 +702,7 @@ func createStorageMapPayloads(t *testing.T, address common.Address, domain strin
 		)
 	}
 
-	err = mr.Storage.Commit(mr.Interpreter, false)
+	err = mr.Storage.NondeterministicCommit(mr.Interpreter, false)
 	require.NoError(t, err)
 
 	// finalize the transaction
@@ -850,7 +851,7 @@ func createTestPayloads(t *testing.T, address common.Address, domain string) []*
 		),
 	)
 
-	err = mr.Storage.Commit(mr.Interpreter, false)
+	err = mr.Storage.NondeterministicCommit(mr.Interpreter, false)
 	require.NoError(t, err)
 
 	// finalize the transaction

cmd/util/ledger/migrations/cadence_values_migration.go

@@ -46,6 +46,7 @@ type CadenceBaseMigration struct {
 	errorMessageHandler               *errorMessageHandler
 	programs                          map[runtime.Location]*interpreter.Program
 	chainID                           flow.ChainID
+	nWorkers                          int
 }
 
 var _ AccountBasedMigration = (*CadenceBaseMigration)(nil)
@@ -65,9 +66,10 @@ func (m *CadenceBaseMigration) Close() error {
 func (m *CadenceBaseMigration) InitMigration(
 	log zerolog.Logger,
 	_ *registers.ByAccount,
-	_ int,
+	nWorkers int,
 ) error {
 	m.log = log.With().Str("migration", m.name).Logger()
+	m.nWorkers = nWorkers
 
 	// During the migration, we only provide already checked programs,
 	// no parsing/checking of contracts is expected.
@@ -120,7 +122,7 @@ func (m *CadenceBaseMigration) MigrateAccount(
 	var storageHealthErrorBefore error
 	if m.checkStorageHealthBeforeMigration {
 
-		storageHealthErrorBefore = checkStorageHealth(address, storage, accountRegisters)
+		storageHealthErrorBefore = checkStorageHealth(address, storage, accountRegisters, m.nWorkers)
 		if storageHealthErrorBefore != nil {
 			m.log.Warn().
 				Err(storageHealthErrorBefore).
@@ -204,6 +206,7 @@ func (m *CadenceBaseMigration) MigrateAccount(
 			m.chainID,
 			m.diffReporter,
 			m.logVerboseDiff,
+			m.nWorkers,
 		)
 
 		accountDiffReporter.DiffStates(

cmd/util/ledger/migrations/cadence_values_migration_test.go

@@ -791,7 +791,7 @@ func TestProgramParsingError(t *testing.T) {
 		capabilityValue,
 	)
 
-	err = storage.Commit(runtime.Interpreter, false)
+	err = storage.NondeterministicCommit(runtime.Interpreter, false)
 	require.NoError(t, err)
 
 	// finalize the transaction
@@ -930,7 +930,7 @@ func TestCoreContractUsage(t *testing.T) {
 			capabilityValue,
 		)
 
-		err = storage.Commit(runtime.Interpreter, false)
+		err = storage.NondeterministicCommit(runtime.Interpreter, false)
 		require.NoError(t, err)
 
 		// finalize the transaction

cmd/util/ledger/migrations/filter_unreferenced_slabs_migration.go

@@ -86,7 +86,7 @@ func (m *FilterUnreferencedSlabsMigration) MigrateAccount(
 		nil,
 	)
 
-	err := checkStorageHealth(address, storage, accountRegisters)
+	err := checkStorageHealth(address, storage, accountRegisters, m.nWorkers)
 	if err == nil {
 		return nil
 	}

cmd/util/ledger/migrations/filter_unreferenced_slabs_migration_test.go

@@ -123,7 +123,7 @@ func TestFilterUnreferencedSlabs(t *testing.T) {
 		dict1,
 	)
 
-	err = storage.Commit(inter, false)
+	err = storage.NondeterministicCommit(inter, false)
 	require.NoError(t, err)
 
 	oldPayloads := make([]*ledger.Payload, 0, len(payloads))

cmd/util/ledger/migrations/fix_broken_data_migration.go

@@ -78,7 +78,7 @@ func (m *FixSlabsWithBrokenReferencesMigration) MigrateAccount(
 	storage := migrationRuntime.Storage
 
 	// Load all atree registers in storage
-	err := loadAtreeSlabsInStorage(storage, accountRegisters)
+	err := loadAtreeSlabsInStorage(storage, accountRegisters, m.nWorkers)
 	if err != nil {
 		return err
 	}
@@ -116,7 +116,7 @@ func (m *FixSlabsWithBrokenReferencesMigration) MigrateAccount(
 
 	m.mergeBrokenPayloads(brokenPayloads)
 
-	err = storage.FastCommit(m.nWorkers)
+	err = storage.NondeterministicFastCommit(m.nWorkers)
 	if err != nil {
 		return fmt.Errorf("failed to commit storage: %w", err)
 	}

cmd/util/ledger/migrations/staged_contracts_migration_test.go

@@ -707,7 +707,7 @@ func TestStagedContractsMigration(t *testing.T) {
 			interpreter.NewUnmeteredStringValue("Also in the same storage path prefix"),
 		)
 
-		err = mr.Storage.Commit(mr.Interpreter, false)
+		err = mr.Storage.NondeterministicCommit(mr.Interpreter, false)
 		require.NoError(t, err)
 
 		result, err := mr.TransactionState.FinalizeMainTransaction()