Merge pull request #6122 from onflow/bastian/optimize-contracts-check…

…ing-migration

Optimize contracts checking migration

cmd/util/ledger/migrations/cadence_values_migration_test.go

@@ -1274,6 +1274,16 @@ func TestProgramParsingError(t *testing.T) {
 	)
 	require.NoError(t, err)
 
+	encodedContractNames, err := environment.EncodeContractNames([]string{contractName})
+	require.NoError(t, err)
+
+	err = registersByAccount.Set(
+		string(testAddress[:]),
+		flow.ContractNamesKey,
+		encodedContractNames,
+	)
+	require.NoError(t, err)
+
 	// Migrate
 
 	// TODO: EVM contract is not deployed in snapshot yet, so can't update it

cmd/util/ledger/migrations/contract_checking_migration.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/onflow/flow-go/cmd/util/ledger/reporters"
 	"github.com/onflow/flow-go/cmd/util/ledger/util/registers"
+	"github.com/onflow/flow-go/fvm/environment"
 	"github.com/onflow/flow-go/model/flow"
 )
 
@@ -30,6 +31,7 @@ func NewContractCheckingMigration(
 	return func(registersByAccount *registers.ByAccount) error {
 
 		reporter := rwf.ReportWriter(contractCheckingReporterName)
+		defer reporter.Close()
 
 		mr, err := NewInterpreterMigrationRuntime(
 			registersByAccount,
@@ -42,6 +44,8 @@ func NewContractCheckingMigration(
 
 		// Gather all contracts
 
+		log.Info().Msg("Gathering contracts ...")
+
 		contractsForPrettyPrinting := make(map[common.Location][]byte, contractCountEstimate)
 
 		type contract struct {
@@ -50,35 +54,49 @@ func NewContractCheckingMigration(
 		}
 		contracts := make([]contract, 0, contractCountEstimate)
 
-		err = registersByAccount.ForEach(func(owner string, key string, value []byte) error {
+		err = registersByAccount.ForEachAccount(func(accountRegisters *registers.AccountRegisters) error {
+			owner := accountRegisters.Owner()
 
-			// Skip payloads that are not contract code
-			contractName := flow.KeyContractName(key)
-			if contractName == "" {
-				return nil
+			encodedContractNames, err := accountRegisters.Get(owner, flow.ContractNamesKey)
+			if err != nil {
+				return err
 			}
 
-			address := common.Address([]byte(owner))
-			code := value
-			location := common.AddressLocation{
-				Address: address,
-				Name:    contractName,
+			contractNames, err := environment.DecodeContractNames(encodedContractNames)
+			if err != nil {
+				return err
 			}
 
-			contracts = append(
-				contracts,
-				contract{
-					location: location,
-					code:     code,
-				},
-			)
+			for _, contractName := range contractNames {
+
+				contractKey := flow.ContractKey(contractName)
+
+				code, err := accountRegisters.Get(owner, contractKey)
+				if err != nil {
+					return err
+				}
+
+				address := common.Address([]byte(owner))
+				location := common.AddressLocation{
+					Address: address,
+					Name:    contractName,
+				}
+
+				contracts = append(
+					contracts,
+					contract{
+						location: location,
+						code:     code,
+					},
+				)
 
-			contractsForPrettyPrinting[location] = code
+				contractsForPrettyPrinting[location] = code
+			}
 
 			return nil
 		})
 		if err != nil {
-			return fmt.Errorf("failed to iterate over registers: %w", err)
+			return fmt.Errorf("failed to get contracts of accounts: %w", err)
 		}
 
 		sort.Slice(contracts, func(i, j int) bool {
@@ -87,6 +105,8 @@ func NewContractCheckingMigration(
 			return a.location.ID() < b.location.ID()
 		})
 
+		log.Info().Msgf("Gathered all contracts (%d)", len(contracts))
+
 		// Check all contracts
 
 		for _, contract := range contracts {
@@ -134,8 +154,6 @@ func NewContractCheckingMigration(
 			}
 		}
 
-		reporter.Close()
-
 		return nil
 	}
 }

fvm/environment/accounts.go

@@ -431,6 +431,20 @@ func (a *StatefulAccounts) setContract(
 	return nil
 }
 
+func EncodeContractNames(contractNames contractNames) ([]byte, error) {
+	var buf bytes.Buffer
+	cborEncoder := cbor.NewEncoder(&buf)
+	err := cborEncoder.Encode(contractNames)
+	if err != nil {
+		return nil, errors.NewEncodingFailuref(
+			err,
+			"cannot encode contract names: %s",
+			contractNames,
+		)
+	}
+	return buf.Bytes(), nil
+}
+
 func (a *StatefulAccounts) setContractNames(
 	contractNames contractNames,
 	address flow.Address,
@@ -443,16 +457,11 @@ func (a *StatefulAccounts) setContractNames(
 	if !ok {
 		return errors.NewAccountNotFoundError(address)
 	}
-	var buf bytes.Buffer
-	cborEncoder := cbor.NewEncoder(&buf)
-	err = cborEncoder.Encode(contractNames)
+
+	newContractNames, err := EncodeContractNames(contractNames)
 	if err != nil {
-		return errors.NewEncodingFailuref(
-			err,
-			"cannot encode contract names: %s",
-			contractNames)
+		return err
 	}
-	newContractNames := buf.Bytes()
 
 	id := flow.ContractNamesRegisterID(address)
 	prevContractNames, err := a.GetValue(id)
@@ -607,20 +616,26 @@ func (a *StatefulAccounts) getContractNames(
 	error,
 ) {
 	// TODO return fatal error if can't fetch
-	encContractNames, err := a.GetValue(flow.ContractNamesRegisterID(address))
+	encodedContractNames, err := a.GetValue(flow.ContractNamesRegisterID(address))
 	if err != nil {
 		return nil, fmt.Errorf("cannot get deployed contract names: %w", err)
 	}
+
+	return DecodeContractNames(encodedContractNames)
+}
+
+func DecodeContractNames(encodedContractNames []byte) ([]string, error) {
 	identifiers := make([]string, 0)
-	if len(encContractNames) > 0 {
-		buf := bytes.NewReader(encContractNames)
+	if len(encodedContractNames) > 0 {
+		buf := bytes.NewReader(encodedContractNames)
 		cborDecoder := cbor.NewDecoder(buf)
-		err = cborDecoder.Decode(&identifiers)
+		err := cborDecoder.Decode(&identifiers)
 		if err != nil {
 			return nil, fmt.Errorf(
 				"cannot decode deployed contract names %x: %w",
-				encContractNames,
-				err)
+				encodedContractNames,
+				err,
+			)
 		}
 	}
 	return identifiers, nil