fix(x/swingset): migration upgrade handler to fix state-sync

golang/cosmos/app/app.go

@@ -810,14 +810,107 @@ func NewAgoricApp(
 	return app
 }
 
+type swingStoreMigrationEventHandler struct {
+	swingStore sdk.KVStore
+}
+
+func (eventHandler swingStoreMigrationEventHandler) OnExportStarted(height uint64, retrieveSwingStoreExport func() error) error {
+	return retrieveSwingStoreExport()
+}
+
+func (eventHandler swingStoreMigrationEventHandler) OnExportRetrieved(provider swingsetkeeper.SwingStoreExportProvider) (err error) {
+	exportDataReader, err := provider.GetExportDataReader()
+	if err != nil {
+		return err
+	}
+	defer exportDataReader.Close()
+
+	var hasExportData bool
+
+	for {
+		entry, err := exportDataReader.Read()
+		if err == io.EOF {
+			break
+		} else if err != nil {
+			return err
+		}
+		hasExportData = true
+		if !entry.HasValue() {
+			return fmt.Errorf("no value for export data key %s", entry.Key())
+		}
+		eventHandler.swingStore.Set([]byte(entry.Key()), []byte(entry.StringValue()))
+	}
+	if !hasExportData {
+		return fmt.Errorf("export data had no entries")
+	}
+	return nil
+}
+
 // upgrade11Handler performs standard upgrade actions plus custom actions for upgrade-11.
 func upgrade11Handler(app *GaiaApp, targetUpgrade string) func(sdk.Context, upgradetypes.Plan, module.VersionMap) (module.VersionMap, error) {
 	return func(ctx sdk.Context, plan upgradetypes.Plan, fromVm module.VersionMap) (module.VersionMap, error) {
 		app.CheckControllerInited(false)
 		// Record the plan to send to SwingSet
 		app.upgradePlan = &plan
 
-		// TODO: Migrate x/vstorage swingStore to x/swingset SwingStore
+		// Perform swing-store migrations. We do this in the app upgrade handler
+		// since it involves multiple modules (x/vstorage and x/swingset) which
+		// don't strictly have a version change on their own.
+
+		// We are at the begining of the upgrade block, so all stores are commited
+		// as of the end of the previous block
+		savedBlockHeight := uint64(ctx.BlockHeight() - 1)
+
+		// First, repair swing-store metadata in case this node was previously
+		// initialized from a state-sync snapshot. This is done with a check on the
+		// block height to catch early any hangover related mismatch.
+		// Only entries related to missing historical metadata are imported, but we
+		// don't know what these look like here, so we provide it all.
+		getSwingStoreExportDataFromVstorage := func() (reader agorictypes.KVEntryReader, err error) {
+			return agorictypes.NewVstorageDataEntriesReader(
+				app.VstorageKeeper.ExportStorageFromPrefix(ctx, swingsetkeeper.StoragePathSwingStore),
+			), nil
+		}
+
+		// We're not restoring any artifact to swing-store, nor have any to provide
+		readNoArtifact := func() (artifact swingsettypes.SwingStoreArtifact, err error) {
+			return artifact, io.EOF
+		}
+
+		err := app.SwingStoreExportsHandler.RestoreExport(
+			swingsetkeeper.SwingStoreExportProvider{
+				BlockHeight:         savedBlockHeight,
+				GetExportDataReader: getSwingStoreExportDataFromVstorage,
+				ReadNextArtifact:    readNoArtifact,
+			},
+			swingsetkeeper.SwingStoreRestoreOptions{
+				ArtifactMode:   swingsetkeeper.SwingStoreArtifactModeNone,
+				ExportDataMode: swingsetkeeper.SwingStoreExportDataModeRepairMetadata,
+			},
+		)
+		if err != nil {
+			return nil, err
+		}
+
+		// Then migrate the swing-store shadow copy:
+		// 1. Remove the swing-store "export data" shadow-copy entries from vstorage.
+		// 2. Export swing-store "export-data" (as of the previous block) through a
+		//    handler that writes every entry into the swingset module's new Store.
+		app.VstorageKeeper.RemoveEntriesWithPrefix(ctx, swingsetkeeper.StoragePathSwingStore)
+		err = app.SwingStoreExportsHandler.InitiateExport(
+			savedBlockHeight,
+			swingStoreMigrationEventHandler{swingStore: app.SwingSetKeeper.GetSwingStore(ctx)},
+			swingsetkeeper.SwingStoreExportOptions{
+				ArtifactMode:   swingsetkeeper.SwingStoreArtifactModeNone,
+				ExportDataMode: swingsetkeeper.SwingStoreExportDataModeAll,
+			},
+		)
+		if err == nil {
+			err = swingsetkeeper.WaitUntilSwingStoreExportDone()
+		}
+		if err != nil {
+			return nil, err
+		}
 
 		// Always run module migrations
 		mvm, err := app.mm.RunMigrations(ctx, app.configurator, fromVm)

golang/cosmos/x/swingset/keeper/extension_snapshotter.go

@@ -125,8 +125,8 @@ func (snapshotter *ExtensionSnapshotter) InitiateSnapshot(height int64) error {
 	blockHeight := uint64(height)
 
 	return snapshotter.swingStoreExportsHandler.InitiateExport(blockHeight, snapshotter, SwingStoreExportOptions{
-		ExportMode:        SwingStoreExportModeCurrent,
-		IncludeExportData: false,
+		ArtifactMode:   SwingStoreArtifactModeReplay,
+		ExportDataMode: SwingStoreExportDataModeSkip,
 	})
 }
 
@@ -304,6 +304,6 @@ func (snapshotter *ExtensionSnapshotter) RestoreExtension(blockHeight uint64, fo
 
 	return snapshotter.swingStoreExportsHandler.RestoreExport(
 		SwingStoreExportProvider{BlockHeight: blockHeight, GetExportDataReader: getExportDataReader, ReadNextArtifact: readNextArtifact},
-		SwingStoreRestoreOptions{IncludeHistorical: false},
+		SwingStoreRestoreOptions{ArtifactMode: SwingStoreArtifactModeReplay, ExportDataMode: SwingStoreExportDataModeAll},
 	)
 }

golang/cosmos/x/swingset/keeper/swing_store_exports_handler.go

@@ -72,7 +72,7 @@ import (
 // - OnExportRetrieved reads the export using the provider.
 //
 // Restoring a swing-store export does not have similar non-blocking requirements.
-// The component simply invokes swingStoreExportHandler.RestoreExport with a
+// The component simply invokes swingStoreExportsHandler.RestoreExport with a
 // SwingStoreExportProvider representing the swing-store export to
 // be restored, and RestoreExport will consume it and block until the JS side
 // has completed the restore before returning.
@@ -157,44 +157,81 @@ type swingStoreRestoreExportAction struct {
 	Args        [1]swingStoreImportOptions `json:"args"`
 }
 
-// SwingStoreExportModeCurrent represents the minimal set of artifacts needed
-// to operate a node.
-const SwingStoreExportModeCurrent = "current"
+const (
+	// SwingStoreArtifactModeNone means that no artifacts are part of the
+	// export / import.
+	SwingStoreArtifactModeNone = "none"
 
-// SwingStoreExportModeArchival represents the set of all artifacts needed to
-// not lose any historical state.
-const SwingStoreExportModeArchival = "archival"
+	// SwingStoreArtifactModeOperational represents the minimal set of artifacts
+	// needed to operate a node.
+	SwingStoreArtifactModeOperational = "operational"
 
-// SwingStoreExportModeDebug represents the maximal set of artifacts available
-// in the JS swing-store, including any kept around for debugging purposed only
-// (like previous XS heap snapshots)
-const SwingStoreExportModeDebug = "debug"
+	// SwingStoreArtifactModeReplay represents the set of artifacts needed to
+	// replay the current incarnation of every vat.
+	SwingStoreArtifactModeReplay = "replay"
+
+	// SwingStoreArtifactModeArchival represents the set of all artifacts
+	// providing all available historical state.
+	SwingStoreArtifactModeArchival = "archival"
+
+	// SwingStoreArtifactModeDebug represents the maximal set of artifacts
+	// available in the JS swing-store, including any kept around for debugging
+	// purposes only (like previous XS heap snapshots)
+	SwingStoreArtifactModeDebug = "debug"
+)
+
+const (
+	// SwingStoreExportDataModeSkip indicates "export data" should be excluded from
+	// an export. ArtifactMode cannot be "none" in this case.
+	SwingStoreExportDataModeSkip = "skip"
+
+	// SwingStoreExportDataModeRepairMetadata indicates the "export data" should be
+	// used to repair the metadata of an existing swing-store for an import
+	// operation. ArtifactMode must be "none" in this case.
+	SwingStoreExportDataModeRepairMetadata = "repair-metadata"
+
+	// SwingStoreExportDataModeAll indicates "export data" should be part of the
+	// export or import. For import, ArtifactMode cannot be "none".
+	SwingStoreExportDataModeAll = "all"
+)
 
 // SwingStoreExportOptions are configurable options provided to the JS swing-store export
 type SwingStoreExportOptions struct {
-	// The export mode can be "current", "archival" or "debug" (SwingStoreExportMode* const)
-	// See packages/cosmic-swingset/src/export-kernel-db.js initiateSwingStoreExport and
-	// packages/swing-store/src/swingStore.js makeSwingStoreExporter
-	ExportMode string `json:"exportMode,omitempty"`
-	// A flag indicating whether "export data" should be part of the swing-store export
-	// If false, the resulting SwingStoreExportProvider's GetExportDataReader
-	// will return nil
-	IncludeExportData bool `json:"includeExportData,omitempty"`
+	// ArtifactMode controls the set of artifacts that should be included in the
+	// swing-store export. Any SwingStoreArtifactMode* const value can be used
+	// (None, Operational, Replay, Archival, Debug).
+	// See packages/cosmic-swingset/src/export-kernel-db.js initiateSwingStoreExport
+	ArtifactMode string `json:"artifactMode,omitempty"`
+	// ExportDataMode selects whether to include "export data" in the swing-store
+	// export or not. Use the value SwingStoreExportDataModeSkip or
+	// SwingStoreExportDataModeAll. If "skip", the reader returned by
+	// SwingStoreExportProvider's GetExportDataReader will be nil.
+	ExportDataMode string `json:"exportDataMode,omitempty"`
 }
 
 // SwingStoreRestoreOptions are configurable options provided to the JS swing-store import
 type SwingStoreRestoreOptions struct {
-	// A flag indicating whether the swing-store import should attempt to load
-	// all historical artifacts available from the export provider
-	IncludeHistorical bool `json:"includeHistorical,omitempty"`
+	// ArtifactMode controls the set of artifacts that should be restored in
+	// swing-store. Any SwingStoreArtifactMode* const value can be used
+	// (None, Operational, Replay, Archival, Debug).
+	// See packages/cosmic-swingset/src/import-kernel-db.js performStateSyncImport
+	ArtifactMode string `json:"artifactMode,omitempty"`
+	// ExportDataMode selects the purpose of the restore, to recreate a
+	// swing-store (SwingStoreExportDataModeAll), or just to import missing
+	// metadata (SwingStoreExportDataModeRepairMetadata).
+	// If RepairMetadata, ArtifactMode should be SwingStoreArtifactModeNone.
+	// If All, ArtifactMode must be at least SwingStoreArtifactModeOperational.
+	ExportDataMode string `json:"exportDataMode,omitempty"`
 }
 
 type swingStoreImportOptions struct {
 	// ExportDir is the directory created by RestoreExport that JS swing-store
 	// should import from.
 	ExportDir string `json:"exportDir"`
-	// IncludeHistorical is a copy of SwingStoreRestoreOptions.IncludeHistorical
-	IncludeHistorical bool `json:"includeHistorical,omitempty"`
+	// ArtifactMode is a copy of SwingStoreRestoreOptions.ArtifactMode
+	ArtifactMode string `json:"artifactMode,omitempty"`
+	// ExportDataMode is a copy of SwingStoreRestoreOptions.ExportDataMode
+	ExportDataMode string `json:"exportDataMode,omitempty"`
 }
 
 var disallowedArtifactNameChar = regexp.MustCompile(`[^-_.a-zA-Z0-9]`)
@@ -781,8 +818,9 @@ func (exportsHandler SwingStoreExportsHandler) RestoreExport(provider SwingStore
 		BlockHeight: blockHeight,
 		Request:     restoreRequest,
 		Args: [1]swingStoreImportOptions{{
-			ExportDir:         exportDir,
-			IncludeHistorical: restoreOptions.IncludeHistorical,
+			ExportDir:      exportDir,
+			ArtifactMode:   restoreOptions.ArtifactMode,
+			ExportDataMode: restoreOptions.ExportDataMode,
 		}},
 	}
 

golang/cosmos/x/vstorage/keeper/keeper.go

@@ -181,6 +181,54 @@ func (k Keeper) ImportStorage(ctx sdk.Context, entries []*types.DataEntry) {
 	}
 }
 
+func getEncodedKeysWithPrefixFromIterator(iterator sdk.Iterator, prefix string) [][]byte {
+	keys := make([][]byte, 0)
+	defer iterator.Close()
+	for ; iterator.Valid(); iterator.Next() {
+		key := iterator.Key()
+		path := types.EncodedKeyToPath(key)
+		if strings.HasPrefix(path, prefix) {
+			keys = append(keys, key)
+		}
+	}
+	return keys
+}
+
+// RemoveEntriesWithPrefix removes all storage entries starting with the
+// supplied pathPrefix, which may not be empty.
+// It has the same effect as listing children of the prefix and removing each
+// descendant recursively.
+func (k Keeper) RemoveEntriesWithPrefix(ctx sdk.Context, pathPrefix string) {
+	store := ctx.KVStore(k.storeKey)
+
+	if len(pathPrefix) == 0 {
+		panic("cannot remove all content")
+	}
+	if err := types.ValidatePath(pathPrefix); err != nil {
+		panic(err)
+	}
+	descendantPrefix := pathPrefix + types.PathSeparator
+
+	// since vstorage encodes keys with a prefix indicating the number of path
+	// elements, we cannot use a simple prefix iterator.
+	// Instead we iterate over the whole vstorage content and check
+	// whether each entry matches the descendantPrefix. This choice assumes most
+	// entries will be deleted. An alternative implementation would be to
+	// recursively list all children under the descendantPrefix, and delete them.
+
+	iterator := sdk.KVStorePrefixIterator(store, nil)
+
+	keys := getEncodedKeysWithPrefixFromIterator(iterator, descendantPrefix)
+
+	for _, key := range keys {
+		store.Delete(key)
+	}
+
+	// Update the prefix entry itself with SetStorage, which will effectively
+	// delete it and all necessary ancestors.
+	k.SetStorage(ctx, agoric.NewKVEntryWithNoValue(pathPrefix))
+}
+
 func (k Keeper) EmitChange(ctx sdk.Context, change *ProposedChange) {
 	if change.NewValue == change.ValueFromLastBlock {
 		// No change.

golang/cosmos/x/vstorage/keeper/keeper_test.go

@@ -185,7 +185,25 @@ func TestStorage(t *testing.T) {
 		t.Errorf("got export %q, want %q", got, expectedKey2Export)
 	}
 
+	keeper.RemoveEntriesWithPrefix(ctx, "key2.child2")
+	if keeper.HasEntry(ctx, "key2") {
+		t.Errorf("got leftover entries for key2 after removal")
+	}
+	expectedRemainingExport := []*types.DataEntry{
+		{Path: "alpha2", Value: "value2"},
+		{Path: "beta3", Value: "value3"},
+		{Path: "inited", Value: ""},
+	}
+	gotRemainingExport := keeper.ExportStorage(ctx)
+	if !reflect.DeepEqual(gotRemainingExport, expectedRemainingExport) {
+		t.Errorf("got remaining export %q, want %q", expectedRemainingExport, expectedRemainingExport)
+	}
+
 	keeper.ImportStorage(ctx, gotExport)
+	gotExport = keeper.ExportStorage(ctx)
+	if !reflect.DeepEqual(gotExport, expectedExport) {
+		t.Errorf("got export %q after import, want %q", gotExport, expectedExport)
+	}
 }
 
 func TestStorageNotify(t *testing.T) {