integrate new migration

onflow · Apr 23, 2024 · f9798b7 · f9798b7
1 parent a792258
commit f9798b7
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 29 deletions.
diff --git a/cmd/util/cmd/execution-state-extract/cmd.go b/cmd/util/cmd/execution-state-extract/cmd.go
@@ -39,7 +39,7 @@ var (
 	flagInputPayloadFileName               string
 	flagOutputPayloadFileName              string
 	flagOutputPayloadByAddresses           string
-	flagFixSlabWithBrokenReferences        bool
+	flagFixSlabsWithBrokenReferences       bool
 )
 
 var Cmd = &cobra.Command{
@@ -124,7 +124,7 @@ func init() {
 		"extract payloads of addresses (comma separated hex-encoded addresses) to file specified by output-payload-filename",
 	)
 
-	Cmd.Flags().BoolVar(&flagFixSlabWithBrokenReferences, "fix-testnet-slabs-with-broken-references", false,
+	Cmd.Flags().BoolVar(&flagFixSlabsWithBrokenReferences, "fix-testnet-slabs-with-broken-references", false,
 		"fix slabs with broken references in testnet")
 }
 
@@ -303,9 +303,8 @@ func run(*cobra.Command, []string) {
 
 	log.Info().Msgf("state extraction plan: %s, %s", inputMsg, outputMsg)
 
-
 	chainID := flow.ChainID(flagChain)
-	fixSlabWithBrokenReferences := chainID == flow.Testnet && flagFixSlabWithBrokenReferences
+	fixSlabsWithBrokenReferences := chainID == flow.Testnet && flagFixSlabsWithBrokenReferences
 
 	var err error
 	if len(flagInputPayloadFileName) > 0 {
@@ -318,7 +317,7 @@ func run(*cobra.Command, []string) {
 			flagInputPayloadFileName,
 			flagOutputPayloadFileName,
 			exportedAddresses,
-			fixSlabWithBrokenReferences,
+			fixSlabsWithBrokenReferences,
 		)
 	} else {
 		err = extractExecutionState(
@@ -330,7 +329,7 @@ func run(*cobra.Command, []string) {
 			!flagNoMigration,
 			flagOutputPayloadFileName,
 			exportedAddresses,
-			fixSlabWithBrokenReferences,
+			fixSlabsWithBrokenReferences,
 		)
 	}
 

diff --git a/cmd/util/cmd/execution-state-extract/execution_state_extract.go b/cmd/util/cmd/execution-state-extract/execution_state_extract.go
@@ -39,7 +39,7 @@ func extractExecutionState(
 	runMigrations bool,
 	outputPayloadFile string,
 	exportPayloadsByAddresses []common.Address,
-	fixSlabWithBrokenReferences bool,
+	fixSlabsWithBrokenReferences bool,
 ) error {
 
 	log.Info().Msg("init WAL")
@@ -90,7 +90,13 @@ func extractExecutionState(
 		<-compactor.Done()
 	}()
 
-	migrations := newMigrations(log, dir, nWorker, runMigrations)
+	migrations := newMigrations(
+		log,
+		dir,
+		nWorker,
+		runMigrations,
+		fixSlabsWithBrokenReferences,
+	)
 
 	newState := ledger.State(targetHash)
 
@@ -214,7 +220,7 @@ func extractExecutionStateFromPayloads(
 	inputPayloadFile string,
 	outputPayloadFile string,
 	exportPayloadsByAddresses []common.Address,
-	fixSlabWithBrokenReferences bool,
+	fixSlabsWithBrokenReferences bool,
 ) error {
 
 	inputPayloadsFromPartialState, payloads, err := util.ReadPayloadFile(log, inputPayloadFile)
@@ -224,7 +230,13 @@ func extractExecutionStateFromPayloads(
 
 	log.Info().Msgf("read %d payloads", len(payloads))
 
-	migrations := newMigrations(log, dir, nWorker, runMigrations)
+	migrations := newMigrations(
+		log,
+		dir,
+		nWorker,
+		runMigrations,
+		fixSlabsWithBrokenReferences,
+	)
 
 	payloads, err = migratePayloads(log, payloads, migrations)
 	if err != nil {
@@ -353,32 +365,52 @@ func newMigrations(
 	dir string,
 	nWorker int, // number of concurrent worker to migation payloads
 	runMigrations bool,
+	fixSlabsWithBrokenReferences bool,
 ) []ledger.Migration {
 	if runMigrations {
+
 		rwf := reporters.NewReportFileWriterFactory(dir, log)
 
-		migrations := []ledger.Migration{
+		var accountBasedMigrations []migrators.AccountBasedMigration
+
+		if fixSlabsWithBrokenReferences {
+			accountBasedMigrations = append(
+				accountBasedMigrations,
+				migrators.NewFixBrokenReferencesInSlabsMigration(
+					rwf,
+					migrators.TestnetAccountsWithBrokenSlabReferences,
+				),
+			)
+		}
+
+		accountBasedMigrations = append(
+			accountBasedMigrations,
+			migrators.NewFixBrokenReferencesInSlabsMigration(
+				rwf,
+				migrators.TestnetAccountsWithBrokenSlabReferences,
+			),
+			migrators.NewAtreeRegisterMigrator(
+				rwf,
+				flagValidateMigration,
+				flagLogVerboseValidationError,
+				flagContinueMigrationOnValidationError,
+				flagCheckStorageHealthBeforeMigration,
+				flagCheckStorageHealthAfterMigration,
+			),
+
+			&migrators.DeduplicateContractNamesMigration{},
+
+			// This will fix storage used discrepancies caused by the previous migrations
+			&migrators.AccountUsageMigrator{},
+		)
+
+		return []ledger.Migration{
 			migrators.CreateAccountBasedMigration(
 				log,
 				nWorker,
-				[]migrators.AccountBasedMigration{
-					migrators.NewAtreeRegisterMigrator(
-						rwf,
-						flagValidateMigration,
-						flagLogVerboseValidationError,
-						flagContinueMigrationOnValidationError,
-						flagCheckStorageHealthBeforeMigration,
-						flagCheckStorageHealthAfterMigration,
-					),
-
-					&migrators.DeduplicateContractNamesMigration{},
-
-					// This will fix storage used discrepancies caused by the previous migrations
-					&migrators.AccountUsageMigrator{},
-				}),
+				accountBasedMigrations,
+			),
 		}
-
-		return migrations
 	}
 
 	return nil

diff --git a/cmd/util/ledger/migrations/account_based_migration.go b/cmd/util/ledger/migrations/account_based_migration.go
@@ -300,7 +300,7 @@ func MigrateGroupConcurrently(
 	return migrated, nil
 }
 
-var testnetAccountsWithBrokenSlabReferences = func() map[common.Address]struct{} {
+var TestnetAccountsWithBrokenSlabReferences = func() map[common.Address]struct{} {
 	testnetAddresses := map[common.Address]struct{}{
 		mustHexToAddress("434a1f199a7ae3ba"): {},
 		mustHexToAddress("454c9991c2b8d947"): {},