Merge pull request #6977 from onflow/fxamacker/update-storage-health-…

…check

Add  support for account format v2 in util `check-storage`

cmd/util/cmd/check-storage/cmd.go

@@ -2,7 +2,9 @@ package check_storage
 
 import (
 	"context"
+	"fmt"
 
+	"github.com/onflow/atree"
 	"github.com/onflow/cadence/interpreter"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
@@ -22,12 +24,14 @@ import (
 )
 
 var (
-	flagPayloads        string
-	flagState           string
-	flagStateCommitment string
-	flagOutputDirectory string
-	flagChain           string
-	flagNWorker         int
+	flagPayloads           string
+	flagState              string
+	flagStateCommitment    string
+	flagOutputDirectory    string
+	flagChain              string
+	flagNWorker            int
+	flagHasAccountFormatV1 bool
+	flagHasAccountFormatV2 bool
 )
 
 var (
@@ -94,6 +98,20 @@ func init() {
 		"Chain name",
 	)
 	_ = Cmd.MarkFlagRequired("chain")
+
+	Cmd.Flags().BoolVar(
+		&flagHasAccountFormatV1,
+		"account-format-v1",
+		false,
+		"State contains accounts in v1 format",
+	)
+
+	Cmd.Flags().BoolVar(
+		&flagHasAccountFormatV2,
+		"account-format-v2",
+		false,
+		"State contains accounts in v2 format",
+	)
 }
 
 func run(*cobra.Command, []string) {
@@ -110,6 +128,9 @@ func run(*cobra.Command, []string) {
 	if flagState != "" && flagStateCommitment == "" {
 		log.Fatal().Msg("--state-commitment must be provided when --state is provided")
 	}
+	if !flagHasAccountFormatV1 && !flagHasAccountFormatV2 {
+		log.Fatal().Msg("both of or one of --account-format-v1 and --account-format-v2 must be true")
+	}
 
 	// Get EVM account by chain
 	evmAccount = systemcontracts.SystemContractsForChain(chainID).EVMStorage.Address
@@ -323,7 +344,30 @@ func checkAccountStorageHealth(accountRegisters *registers.AccountRegisters, nWo
 	// Check atree storage health
 
 	ledger := &registers.ReadOnlyLedger{Registers: accountRegisters}
-	storage := runtime.NewStorage(ledger, nil, runtime.StorageConfig{})
+	var config runtime.StorageConfig
+	if flagHasAccountFormatV2 {
+		config.StorageFormatV2Enabled = true
+	}
+	storage := runtime.NewStorage(ledger, nil, config)
+
+	// Check account format against specified flags.
+	err = checkAccountFormat(
+		ledger,
+		address,
+		flagHasAccountFormatV1,
+		flagHasAccountFormatV2,
+	)
+	if err != nil {
+		issues = append(
+			issues,
+			accountStorageIssue{
+				Address: address.Hex(),
+				Kind:    storageErrorKindString[storageFormatErrorKind],
+				Msg:     err.Error(),
+			},
+		)
+		return issues
+	}
 
 	inter, err := interpreter.NewInterpreter(
 		nil,
@@ -367,6 +411,7 @@ const (
 	otherErrorKind storageErrorKind = iota
 	cadenceAtreeStorageErrorKind
 	evmAtreeStorageErrorKind
+	storageFormatErrorKind
 )
 
 var storageErrorKindString = map[storageErrorKind]string{
@@ -380,3 +425,65 @@ type accountStorageIssue struct {
 	Kind    string
 	Msg     string
 }
+
+func hasDomainRegister(ledger atree.Ledger, address common.Address) (bool, error) {
+	for _, domain := range common.AllStorageDomains {
+		value, err := ledger.GetValue(address[:], []byte(domain.Identifier()))
+		if err != nil {
+			return false, err
+		}
+		if len(value) > 0 {
+			return true, nil
+		}
+	}
+
+	return false, nil
+}
+
+func hasAccountRegister(ledger atree.Ledger, address common.Address) (bool, error) {
+	value, err := ledger.GetValue(address[:], []byte(runtime.AccountStorageKey))
+	if err != nil {
+		return false, err
+	}
+	return len(value) > 0, nil
+}
+
+func checkAccountFormat(
+	ledger atree.Ledger,
+	address common.Address,
+	expectV1 bool,
+	expectV2 bool,
+) error {
+	// Skip empty address because it doesn't have any account or domain registers.
+	if len(address) == 0 || address == common.ZeroAddress {
+		return nil
+	}
+
+	foundDomainRegister, err := hasDomainRegister(ledger, address)
+	if err != nil {
+		return err
+	}
+
+	foundAccountRegister, err := hasAccountRegister(ledger, address)
+	if err != nil {
+		return err
+	}
+
+	if !foundAccountRegister && !foundDomainRegister {
+		return fmt.Errorf("found neither domain nor account registers")
+	}
+
+	if foundAccountRegister && foundDomainRegister {
+		return fmt.Errorf("found both domain and account registers")
+	}
+
+	if foundAccountRegister && !expectV2 {
+		return fmt.Errorf("found account in format v2 while only expect account in format v1")
+	}
+
+	if foundDomainRegister && !expectV1 {
+		return fmt.Errorf("found account in format v1 while only expect account in format v2")
+	}
+
+	return nil
+}

cmd/util/cmd/check-storage/evm_account_storage_health.go

@@ -70,7 +70,7 @@ func checkCadenceAtreeRegistersInEVMAccount(
 ) []accountStorageIssue {
 	var issues []accountStorageIssue
 
-	storage := runtime.NewStorage(ledger, nil, runtime.StorageConfig{})
+	storage := runtime.NewStorage(ledger, nil, runtime.StorageConfig{StorageFormatV2Enabled: true})
 
 	inter, err := interpreter.NewInterpreter(
 		nil,

cmd/util/cmd/checkpoint-collect-stats/account_stats.go

@@ -33,12 +33,11 @@ func (format accountFormat) MarshalJSON() ([]byte, error) {
 
 type AccountStats struct {
 	stats
-	FormatV1Count      int            `json:"account_format_v1_count"`
-	FormatV2Count      int            `json:"account_format_v2_count"`
-	AccountsInFormatV2 []string       `json:"accounts_in_format_v2,omitempty"`
-	ServiceAccount     *AccountInfo   `json:"service_account,omitempty"`
-	EVMAccount         *AccountInfo   `json:"evm_account,omitempty"`
-	TopN               []*AccountInfo `json:"largest_accounts"`
+	FormatV1Count  int            `json:"account_format_v1_count"`
+	FormatV2Count  int            `json:"account_format_v2_count"`
+	ServiceAccount *AccountInfo   `json:"service_account,omitempty"`
+	EVMAccount     *AccountInfo   `json:"evm_account,omitempty"`
+	TopN           []*AccountInfo `json:"largest_accounts"`
 }
 
 type AccountInfo struct {
@@ -55,7 +54,6 @@ func getAccountStatus(
 	accountsSlice := make([]*AccountInfo, 0, len(accounts))
 	accountSizesSlice := make([]float64, 0, len(accounts))
 
-	var accountsInFormatV2 []string
 	var accountFormatV1Count, accountFormatV2Count int
 
 	for _, acct := range accounts {
@@ -68,7 +66,6 @@ func getAccountStatus(
 
 		case accountFormatV2:
 			accountFormatV2Count++
-			accountsInFormatV2 = append(accountsInFormatV2, acct.Address)
 
 		default:
 			if acct.Address != "" {
@@ -82,23 +79,19 @@ func getAccountStatus(
 		return cmp.Compare(b.PayloadSize, a.PayloadSize)
 	})
 
-	// Sort accounts in format v2
-	slices.SortFunc(accountsInFormatV2, cmp.Compare)
-
 	stats := getValueStats(accountSizesSlice, percentiles)
 
 	evmAccountAddress := systemcontracts.SystemContractsForChain(chainID).EVMStorage.Address
 
 	serviceAccountAddress := serviceAccountAddressForChain(chainID)
 
 	return AccountStats{
-		stats:              stats,
-		FormatV1Count:      accountFormatV1Count,
-		FormatV2Count:      accountFormatV2Count,
-		AccountsInFormatV2: accountsInFormatV2,
-		ServiceAccount:     accounts[string(serviceAccountAddress[:])],
-		EVMAccount:         accounts[string(evmAccountAddress[:])],
-		TopN:               accountsSlice[:flagTopN],
+		stats:          stats,
+		FormatV1Count:  accountFormatV1Count,
+		FormatV2Count:  accountFormatV2Count,
+		ServiceAccount: accounts[string(serviceAccountAddress[:])],
+		EVMAccount:     accounts[string(evmAccountAddress[:])],
+		TopN:           accountsSlice[:flagTopN],
 	}
 }
 

cmd/util/cmd/diff-states/cmd.go

@@ -338,14 +338,6 @@ func diffAccount(
 	acctsToSkip []string,
 ) (err error) {
 
-	if accountRegisters1.Count() != accountRegisters2.Count() {
-		rw.Write(countDiff{
-			Owner:  owner,
-			State1: accountRegisters1.Count(),
-			State2: accountRegisters2.Count(),
-		})
-	}
-
 	diffValues := flagAlwaysDiffValues
 
 	err = accountRegisters1.ForEach(func(owner, key string, value1 []byte) error {