Merge pull request #5438 from onflow/fxamacker/improve-payload-file

Add version, partial state indicator, and checksum to payload file (intermediate migration file)

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

@@ -13,27 +13,29 @@ import (
 	runtimeCommon "github.com/onflow/cadence/runtime/common"
 
 	"github.com/onflow/flow-go/cmd/util/cmd/common"
+	"github.com/onflow/flow-go/cmd/util/ledger/util"
 	"github.com/onflow/flow-go/model/bootstrap"
 	"github.com/onflow/flow-go/model/flow"
 	"github.com/onflow/flow-go/module/metrics"
 	"github.com/onflow/flow-go/storage/badger"
 )
 
 var (
-	flagExecutionStateDir         string
-	flagOutputDir                 string
-	flagBlockHash                 string
-	flagStateCommitment           string
-	flagDatadir                   string
-	flagChain                     string
-	flagNWorker                   int
-	flagNoMigration               bool
-	flagNoReport                  bool
-	flagValidateMigration         bool
-	flagLogVerboseValidationError bool
-	flagInputPayloadFileName      string
-	flagOutputPayloadFileName     string
-	flagOutputPayloadByAddresses  string
+	flagExecutionStateDir             string
+	flagOutputDir                     string
+	flagBlockHash                     string
+	flagStateCommitment               string
+	flagDatadir                       string
+	flagChain                         string
+	flagNWorker                       int
+	flagNoMigration                   bool
+	flagNoReport                      bool
+	flagValidateMigration             bool
+	flagLogVerboseValidationError     bool
+	flagAllowPartialStateFromPayloads bool
+	flagInputPayloadFileName          string
+	flagOutputPayloadFileName         string
+	flagOutputPayloadByAddresses      string
 )
 
 var Cmd = &cobra.Command{
@@ -76,6 +78,9 @@ func init() {
 	Cmd.Flags().BoolVar(&flagLogVerboseValidationError, "log-verbose-validation-error", false,
 		"log entire Cadence values on validation error (atree migration)")
 
+	Cmd.Flags().BoolVar(&flagAllowPartialStateFromPayloads, "allow-partial-state-from-payload-file", false,
+		"allow input payload file containing partial state (e.g. not all accounts)")
+
 	// If specified, the state will consist of payloads from the given input payload file.
 	// If not specified, then the state will be extracted from the latest checkpoint file.
 	// This flag can be used to reduce total duration of migrations when state extraction involves
@@ -166,6 +171,29 @@ func run(*cobra.Command, []string) {
 		if _, err := os.Stat(flagInputPayloadFileName); os.IsNotExist(err) {
 			log.Fatal().Msgf("payload input file %s doesn't exist", flagInputPayloadFileName)
 		}
+
+		partialState, err := util.IsPayloadFilePartialState(flagInputPayloadFileName)
+		if err != nil {
+			log.Fatal().Err(err).Msgf("cannot get flag from payload input file %s", flagInputPayloadFileName)
+		}
+
+		// Check if payload file contains partial state and is allowed by --allow-partial-state-from-payload-file.
+		if !flagAllowPartialStateFromPayloads && partialState {
+			log.Fatal().Msgf("payload input file %s contains partial state, please specify --allow-partial-state-from-payload-file", flagInputPayloadFileName)
+		}
+
+		msg := "input payloads represent "
+		if partialState {
+			msg += "partial state"
+		} else {
+			msg += "complete state"
+		}
+		if flagAllowPartialStateFromPayloads {
+			msg += ", and --allow-partial-state-from-payload-file is specified"
+		} else {
+			msg += ", and --allow-partial-state-from-payload-file is NOT specified"
+		}
+		log.Info().Msg(msg)
 	}
 
 	if len(flagOutputPayloadFileName) > 0 {

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

@@ -120,11 +120,22 @@ func extractExecutionState(
 	if exportPayloads {
 		payloads := newTrie.AllPayloads()
 
+		log.Info().Msgf("sorting %d payloads", len(payloads))
+
+		// Sort payloads to produce deterministic payload file with
+		// same sequence of payloads inside.
+		payloads = util.SortPayloadsByAddress(payloads, nWorker)
+
+		log.Info().Msgf("sorted %d payloads", len(payloads))
+
+		log.Info().Msgf("creating payloads file %s", outputPayloadFile)
+
 		exportedPayloadCount, err := util.CreatePayloadFile(
 			log,
 			outputPayloadFile,
 			payloads,
 			exportPayloadsByAddresses,
+			false, // payloads represents entire state.
 		)
 		if err != nil {
 			return fmt.Errorf("cannot generate payloads file: %w", err)
@@ -204,7 +215,7 @@ func extractExecutionStateFromPayloads(
 	exportPayloadsByAddresses []common.Address,
 ) error {
 
-	payloads, err := util.ReadPayloadFile(log, inputPayloadFile)
+	inputPayloadsFromPartialState, payloads, err := util.ReadPayloadFile(log, inputPayloadFile)
 	if err != nil {
 		return err
 	}
@@ -220,11 +231,23 @@ func extractExecutionStateFromPayloads(
 
 	exportPayloads := len(outputPayloadFile) > 0
 	if exportPayloads {
+
+		log.Info().Msgf("sorting %d payloads", len(payloads))
+
+		// Sort payloads to produce deterministic payload file with
+		// same sequence of payloads inside.
+		payloads = util.SortPayloadsByAddress(payloads, nWorker)
+
+		log.Info().Msgf("sorted %d payloads", len(payloads))
+
+		log.Info().Msgf("creating payloads file %s", outputPayloadFile)
+
 		exportedPayloadCount, err := util.CreatePayloadFile(
 			log,
 			outputPayloadFile,
 			payloads,
 			exportPayloadsByAddresses,
+			inputPayloadsFromPartialState,
 		)
 		if err != nil {
 			return fmt.Errorf("cannot generate payloads file: %w", err)

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

@@ -286,9 +286,10 @@ func TestExtractPayloadsFromExecutionState(t *testing.T) {
 			require.NoError(t, err)
 
 			// Verify exported payloads.
-			payloadsFromFile, err := util.ReadPayloadFile(zerolog.Nop(), outputPayloadFileName)
+			partialState, payloadsFromFile, err := util.ReadPayloadFile(zerolog.Nop(), outputPayloadFileName)
 			require.NoError(t, err)
 			require.Equal(t, len(keysValues), len(payloadsFromFile))
+			require.False(t, partialState)
 
 			for _, payloadFromFile := range payloadsFromFile {
 				k, err := payloadFromFile.Key()
@@ -392,9 +393,10 @@ func TestExtractPayloadsFromExecutionState(t *testing.T) {
 			require.NoError(t, err)
 
 			// Verify exported payloads.
-			payloadsFromFile, err := util.ReadPayloadFile(zerolog.Nop(), outputPayloadFileName)
+			partialState, payloadsFromFile, err := util.ReadPayloadFile(zerolog.Nop(), outputPayloadFileName)
 			require.NoError(t, err)
 			require.Equal(t, len(selectedKeysValues), len(payloadsFromFile))
+			require.True(t, partialState)
 
 			for _, payloadFromFile := range payloadsFromFile {
 				k, err := payloadFromFile.Key()
@@ -441,6 +443,7 @@ func TestExtractStateFromPayloads(t *testing.T) {
 				inputPayloadFileName,
 				payloads,
 				nil,
+				false,
 			)
 			require.NoError(t, err)
 			require.Equal(t, len(payloads), numOfPayloadWritten)
@@ -511,6 +514,7 @@ func TestExtractStateFromPayloads(t *testing.T) {
 				inputPayloadFileName,
 				payloads,
 				nil,
+				false,
 			)
 			require.NoError(t, err)
 			require.Equal(t, len(payloads), numOfPayloadWritten)
@@ -531,9 +535,79 @@ func TestExtractStateFromPayloads(t *testing.T) {
 			require.NoError(t, err)
 
 			// Verify exported payloads.
-			payloadsFromFile, err := util.ReadPayloadFile(zerolog.Nop(), outputPayloadFileName)
+			partialState, payloadsFromFile, err := util.ReadPayloadFile(zerolog.Nop(), outputPayloadFileName)
+			require.NoError(t, err)
+			require.Equal(t, len(keysValues), len(payloadsFromFile))
+			require.False(t, partialState)
+
+			for _, payloadFromFile := range payloadsFromFile {
+				k, err := payloadFromFile.Key()
+				require.NoError(t, err)
+
+				kv, exist := keysValues[k.String()]
+				require.True(t, exist)
+
+				require.Equal(t, kv.value, payloadFromFile.Value())
+			}
+		})
+	})
+
+	t.Run("input is partial state", func(t *testing.T) {
+		withDirs(t, func(_, execdir, outdir string) {
+			size := 10
+
+			inputPayloadFileName := filepath.Join(execdir, payloadFileName)
+			outputPayloadFileName := filepath.Join(outdir, "selected.payload")
+
+			// Generate some data
+			keysValues := make(map[string]keyPair)
+			var payloads []*ledger.Payload
+
+			for i := 0; i < size; i++ {
+				keys, values := getSampleKeyValues(i)
+
+				for j, key := range keys {
+					keysValues[key.String()] = keyPair{
+						key:   key,
+						value: values[j],
+					}
+
+					payloads = append(payloads, ledger.NewPayload(key, values[j]))
+				}
+			}
+
+			// Create input payload file that represents partial state
+			numOfPayloadWritten, err := util.CreatePayloadFile(
+				zerolog.Nop(),
+				inputPayloadFileName,
+				payloads,
+				nil,
+				true,
+			)
+			require.NoError(t, err)
+			require.Equal(t, len(payloads), numOfPayloadWritten)
+
+			// Since input payload file is partial state, --allow-partial-state-from-payload-file must be specified.
+			Cmd.SetArgs([]string{
+				"--execution-state-dir", execdir,
+				"--output-dir", outdir,
+				"--no-migration",
+				"--no-report",
+				"--state-commitment", "",
+				"--input-payload-filename", inputPayloadFileName,
+				"--output-payload-filename", outputPayloadFileName,
+				"--extract-payloads-by-address", "",
+				"--allow-partial-state-from-payload-file",
+				"--chain", flow.Emulator.Chain().String()})
+
+			err = Cmd.Execute()
+			require.NoError(t, err)
+
+			// Verify exported payloads.
+			partialState, payloadsFromFile, err := util.ReadPayloadFile(zerolog.Nop(), outputPayloadFileName)
 			require.NoError(t, err)
 			require.Equal(t, len(keysValues), len(payloadsFromFile))
+			require.True(t, partialState)
 
 			for _, payloadFromFile := range payloadsFromFile {
 				k, err := payloadFromFile.Key()