From 316fbf3e85e86e8e3d5ea7a1ae3dbaa3da1aae90 Mon Sep 17 00:00:00 2001
From: Mathieu Hofman <mathieu@agoric.com>
Date: Mon, 24 Jul 2023 23:53:02 +0000
Subject: [PATCH] feat(cosmos): fix and migrate swing-store

---
 golang/cosmos/app/app.go | 95 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 94 insertions(+), 1 deletion(-)

diff --git a/golang/cosmos/app/app.go b/golang/cosmos/app/app.go
index c6be511c249a..52c6cd043a2c 100644
--- a/golang/cosmos/app/app.go
+++ b/golang/cosmos/app/app.go
@@ -107,6 +107,7 @@ import (
 	gaiaappparams "github.com/Agoric/agoric-sdk/golang/cosmos/app/params"
 
 	appante "github.com/Agoric/agoric-sdk/golang/cosmos/ante"
+	agorictypes "github.com/Agoric/agoric-sdk/golang/cosmos/types"
 	"github.com/Agoric/agoric-sdk/golang/cosmos/vm"
 	"github.com/Agoric/agoric-sdk/golang/cosmos/x/lien"
 	"github.com/Agoric/agoric-sdk/golang/cosmos/x/swingset"
@@ -803,6 +804,42 @@ 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) {
@@ -810,7 +847,63 @@ func upgrade11Handler(app *GaiaApp, targetUpgrade string) func(sdk.Context, upgr
 		// 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 involved multile 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 state-synced
+		// This is done with a check on the block height to catch any hangover early
+		// 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,
+				ReadArtifact:        readNoArtifact,
+			},
+			swingsetkeeper.SwingStoreRestoreOptions{
+				ArtifactMode:   swingsetkeeper.SwingStoreArtifactModeNone,
+				ExportDataMode: swingsetkeeper.SwingStoreRepairMetadataExportDataMode,
+			},
+		)
+		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)