[Cadence 1.0] Optionally check atree storage health before migration

cmd/util/cmd/execution-state-extract/cmd.go

@@ -22,27 +22,28 @@ import (
 )
 
 var (
-	flagExecutionStateDir             string
-	flagOutputDir                     string
-	flagBlockHash                     string
-	flagStateCommitment               string
-	flagDatadir                       string
-	flagChain                         string
-	flagNWorker                       int
-	flagNoMigration                   bool
-	flagNoReport                      bool
-	flagValidateMigration             bool
-	flagAllowPartialStateFromPayloads bool
-	flagSortPayloads                  bool
-	flagPrune                         bool
-	flagLogVerboseValidationError     bool
-	flagDiffMigration                 bool
-	flagLogVerboseDiff                bool
-	flagStagedContractsFile           string
-	flagInputPayloadFileName          string
-	flagOutputPayloadFileName         string
-	flagOutputPayloadByAddresses      string
-	flagMaxAccountSize                uint64
+	flagExecutionStateDir                 string
+	flagOutputDir                         string
+	flagBlockHash                         string
+	flagStateCommitment                   string
+	flagDatadir                           string
+	flagChain                             string
+	flagNWorker                           int
+	flagNoMigration                       bool
+	flagNoReport                          bool
+	flagValidateMigration                 bool
+	flagAllowPartialStateFromPayloads     bool
+	flagSortPayloads                      bool
+	flagPrune                             bool
+	flagLogVerboseValidationError         bool
+	flagDiffMigration                     bool
+	flagLogVerboseDiff                    bool
+	flagCheckStorageHealthBeforeMigration bool
+	flagStagedContractsFile               string
+	flagInputPayloadFileName              string
+	flagOutputPayloadFileName             string
+	flagOutputPayloadByAddresses          string
+	flagMaxAccountSize                    uint64
 )
 
 var Cmd = &cobra.Command{
@@ -92,6 +93,9 @@ func init() {
 	Cmd.Flags().BoolVar(&flagLogVerboseDiff, "log-verbose-diff", false,
 		"log entire Cadence values on diff (requires --diff flag)")
 
+	Cmd.Flags().BoolVar(&flagCheckStorageHealthBeforeMigration, "check-storage-health-before", false,
+		"check (atree) storage health before migration")
+
 	Cmd.Flags().StringVar(&flagStagedContractsFile, "staged-contracts", "",
 		"Staged contracts CSV file")
 
@@ -280,6 +284,10 @@ func run(*cobra.Command, []string) {
 		log.Warn().Msgf("--log-verbose-diff flag is enabled which may increase size of log")
 	}
 
+	if flagCheckStorageHealthBeforeMigration {
+		log.Warn().Msgf("--check-storage-health-before flag is enabled and will increase duration of migration")
+	}
+
 	var inputMsg string
 	if len(flagInputPayloadFileName) > 0 {
 		// Input is payloads
@@ -341,6 +349,7 @@ func run(*cobra.Command, []string) {
 			!flagNoMigration,
 			flagDiffMigration,
 			flagLogVerboseDiff,
+			flagCheckStorageHealthBeforeMigration,
 			chainID,
 			evmContractChange,
 			burnerContractChange,
@@ -362,6 +371,7 @@ func run(*cobra.Command, []string) {
 			!flagNoMigration,
 			flagDiffMigration,
 			flagLogVerboseDiff,
+			flagCheckStorageHealthBeforeMigration,
 			chainID,
 			evmContractChange,
 			burnerContractChange,

cmd/util/cmd/execution-state-extract/execution_state_extract.go

@@ -39,6 +39,7 @@ func extractExecutionState(
 	runMigrations bool,
 	diffMigrations bool,
 	logVerboseDiff bool,
+	checkStorageHealthBeforeMigration bool,
 	chainID flow.ChainID,
 	evmContractChange migrators.EVMContractChange,
 	burnerContractChange migrators.BurnerContractChange,
@@ -114,6 +115,7 @@ func extractExecutionState(
 		runMigrations,
 		diffMigrations,
 		logVerboseDiff,
+		checkStorageHealthBeforeMigration,
 		chainID,
 		evmContractChange,
 		burnerContractChange,
@@ -229,6 +231,7 @@ func extractExecutionStateFromPayloads(
 	runMigrations bool,
 	diffMigrations bool,
 	logVerboseDiff bool,
+	checkStorageHealthBeforeMigration bool,
 	chainID flow.ChainID,
 	evmContractChange migrators.EVMContractChange,
 	burnerContractChange migrators.BurnerContractChange,
@@ -255,6 +258,7 @@ func extractExecutionStateFromPayloads(
 		runMigrations,
 		diffMigrations,
 		logVerboseDiff,
+		checkStorageHealthBeforeMigration,
 		chainID,
 		evmContractChange,
 		burnerContractChange,
@@ -412,6 +416,7 @@ func newMigrations(
 	runMigrations bool,
 	diffMigrations bool,
 	logVerboseDiff bool,
+	checkStorageHealthBeforeMigration bool,
 	chainID flow.ChainID,
 	evmContractChange migrators.EVMContractChange,
 	burnerContractChange migrators.BurnerContractChange,
@@ -434,6 +439,7 @@ func newMigrations(
 		chainID,
 		diffMigrations,
 		logVerboseDiff,
+		checkStorageHealthBeforeMigration,
 		evmContractChange,
 		burnerContractChange,
 		stagedContracts,

cmd/util/cmd/execution-state-extract/execution_state_extract_test.go

@@ -75,6 +75,7 @@ func TestExtractExecutionState(t *testing.T) {
 				false,
 				false,
 				false,
+				false,
 				flow.Emulator,
 				migrations.EVMContractChangeNone,
 				migrations.BurnerContractChangeDeploy,

cmd/util/ledger/migrations/cadence.go

@@ -222,6 +222,7 @@ func NewCadence1ValueMigrations(
 	chainID flow.ChainID,
 	diffMigrations bool,
 	logVerboseDiff bool,
+	checkStorageHealthBeforeMigration bool,
 ) (migrations []NamedMigration) {
 
 	// Populated by CadenceLinkValueMigrator,
@@ -246,33 +247,47 @@ func NewCadence1ValueMigrations(
 		},
 	}
 
-	for _, accountBasedMigration := range []*CadenceBaseMigrator{
-		NewCadence1ValueMigrator(
-			rwf,
-			diffMigrations,
-			logVerboseDiff,
-			errorMessageHandler,
-			contracts,
-			NewCadence1CompositeStaticTypeConverter(chainID),
-			NewCadence1InterfaceStaticTypeConverter(chainID),
-		),
-		NewCadence1LinkValueMigrator(
-			rwf,
-			diffMigrations,
-			logVerboseDiff,
-			errorMessageHandler,
-			contracts,
-			capabilityMapping,
-		),
-		NewCadence1CapabilityValueMigrator(
-			rwf,
-			diffMigrations,
-			logVerboseDiff,
-			errorMessageHandler,
-			contracts,
-			capabilityMapping,
-		),
+	for index, migrationConstructor := range []func(checkStorageHealthBeforeMigration bool) *CadenceBaseMigrator{
+		func(checkStorageHealthBeforeMigration bool) *CadenceBaseMigrator {
+			return NewCadence1ValueMigrator(
+				rwf,
+				diffMigrations,
+				logVerboseDiff,
+				checkStorageHealthBeforeMigration,
+				errorMessageHandler,
+				contracts,
+				NewCadence1CompositeStaticTypeConverter(chainID),
+				NewCadence1InterfaceStaticTypeConverter(chainID),
+			)
+		},
+		func(checkStorageHealthBeforeMigration bool) *CadenceBaseMigrator {
+			return NewCadence1LinkValueMigrator(
+				rwf,
+				diffMigrations,
+				logVerboseDiff,
+				checkStorageHealthBeforeMigration,
+				errorMessageHandler,
+				contracts,
+				capabilityMapping,
+			)
+		},
+		func(checkStorageHealthBeforeMigration bool) *CadenceBaseMigrator {
+			return NewCadence1CapabilityValueMigrator(
+				rwf,
+				diffMigrations,
+				logVerboseDiff,
+				checkStorageHealthBeforeMigration,
+				errorMessageHandler,
+				contracts,
+				capabilityMapping,
+			)
+		},
 	} {
+		accountBasedMigration := migrationConstructor(
+			// Only check storage health before the first migration
+			checkStorageHealthBeforeMigration && index == 0,
+		)
+
 		migrations = append(
 			migrations,
 			NamedMigration{
@@ -357,6 +372,7 @@ func NewCadence1Migrations(
 	chainID flow.ChainID,
 	diffMigrations bool,
 	logVerboseDiff bool,
+	checkStorageHealthBeforeMigration bool,
 	evmContractChange EVMContractChange,
 	burnerContractChange BurnerContractChange,
 	stagedContracts []StagedContract,
@@ -417,6 +433,7 @@ func NewCadence1Migrations(
 			nWorker,
 			chainID,
 			diffMigrations,
+			checkStorageHealthBeforeMigration,
 			logVerboseDiff,
 		)...,
 	)