optimize migrations

Author: keithsue

x/dlc/keeper/migrations.go

@@ -1,11 +1,9 @@
 package keeper
 
 import (
-	storetypes "cosmossdk.io/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
 	v2 "github.com/sideprotocol/side/x/dlc/migrations/v2"
-	tsstypes "github.com/sideprotocol/side/x/tss/types"
 )
 
 // Migrator is a struct for handling in-place store migrations
@@ -20,5 +18,5 @@ func NewMigrator(keeper Keeper) Migrator {
 
 // Migrate1to2 migrates from version 1 to 2
 func (m Migrator) Migrate1to2(ctx sdk.Context) error {
-	return v2.MigrateStore(ctx, m.keeper.storeKey, storetypes.NewKVStoreKey(tsstypes.StoreKey), m.keeper.cdc)
+	return v2.MigrateStore(ctx, m.keeper.storeKey, m.keeper.tssKeeper, m.keeper.cdc)
 }

x/dlc/migrations/v2/store.go

@@ -13,9 +13,9 @@ import (
 
 // MigrateStore migrates the x/dlc module state from the consensus version 1 to
 // version 2
-func MigrateStore(ctx sdk.Context, storeKey storetypes.StoreKey, tssStoreKey storetypes.StoreKey, cdc codec.BinaryCodec) error {
+func MigrateStore(ctx sdk.Context, storeKey storetypes.StoreKey, tssKeeper types.TSSKeeper, cdc codec.BinaryCodec) error {
 	migrateDLCEvents(ctx, storeKey, cdc)
-	migrateDCMsAndOracles(ctx, storeKey, tssStoreKey, cdc)
+	migrateDCMsAndOracles(ctx, storeKey, tssKeeper, cdc)
 
 	return nil
 }
@@ -37,34 +37,50 @@ func migrateDLCEvents(ctx sdk.Context, storeKey storetypes.StoreKey, cdc codec.B
 }
 
 // migrateDCMsAndOracles performs the DCMs and oracles migration
-func migrateDCMsAndOracles(ctx sdk.Context, storeKey storetypes.StoreKey, tssStoreKey storetypes.StoreKey, cdc codec.BinaryCodec) {
-	tssStore := ctx.KVStore(tssStoreKey)
-
-	iterator := storetypes.KVStorePrefixIterator(tssStore, tsstypes.DKGRequestKeyPrefix)
-	defer iterator.Close()
-
-	for ; iterator.Valid(); iterator.Next() {
-		var dkgRequest tsstypes.DKGRequest
-		cdc.MustUnmarshal(iterator.Value(), &dkgRequest)
-
-		if dkgRequest.Status != tsstypes.DKGStatus_DKG_STATUS_COMPLETED {
-			continue
+func migrateDCMsAndOracles(ctx sdk.Context, storeKey storetypes.StoreKey, tssKeeper types.TSSKeeper, cdc codec.BinaryCodec) {
+	// tssStore := ctx.KVStore(tssStoreKey)
+
+	// iterator := storetypes.KVStorePrefixIterator(tssStore, tsstypes.DKGRequestKeyPrefix)
+	// defer iterator.Close()
+
+	// for ; iterator.Valid(); iterator.Next() {
+	// 	var dkgRequest tsstypes.DKGRequest
+	// 	cdc.MustUnmarshal(iterator.Value(), &dkgRequest)
+
+	// 	if dkgRequest.Status != tsstypes.DKGStatus_DKG_STATUS_COMPLETED {
+	// 		continue
+	// 	}
+
+	// 	bz := tssStore.Get(tsstypes.DKGCompletionKey(dkgRequest.Id, dkgRequest.Participants[0]))
+	// 	var dkgCompletion tsstypes.DKGCompletion
+	// 	cdc.MustUnmarshal(bz, &dkgCompletion)
+
+	// 	// DCM or oracle pub key
+	// 	pubKey := dkgCompletion.PubKeys[0]
+	// 	pubKeyBz, _ := hex.DecodeString(pubKey)
+
+	// 	if dkgRequest.Type == types.DKG_TYPE_DCM {
+	// 		updateDCM(ctx, storeKey, dkgRequest.Id, pubKeyBz, cdc)
+	// 	} else if dkgRequest.Type == types.DKG_TYPE_NONCE {
+	// 		updateOracle(ctx, storeKey, dkgRequest.Id, pubKeyBz, cdc)
+	// 	}
+	// }
+
+	tssKeeper.IterateDKGRequests(ctx, func(req *tsstypes.DKGRequest) (stop bool) {
+		if req.Status == tsstypes.DKGStatus_DKG_STATUS_COMPLETED {
+			// dcm or oracle pub key
+			pubKey := tssKeeper.GetDKGPubKeys(ctx, req.Id)[0]
+			pubKeyBz, _ := hex.DecodeString(pubKey)
+
+			if req.Type == types.DKG_TYPE_DCM {
+				updateDCM(ctx, storeKey, req.Id, pubKeyBz, cdc)
+			} else if req.Type == types.DKG_TYPE_NONCE {
+				updateOracle(ctx, storeKey, req.Id, pubKeyBz, cdc)
+			}
 		}
 
-		bz := tssStore.Get(tsstypes.DKGCompletionKey(dkgRequest.Id, dkgRequest.Participants[0]))
-		var dkgCompletion tsstypes.DKGCompletion
-		cdc.MustUnmarshal(bz, &dkgCompletion)
-
-		// DCM or oracle pub key
-		pubKey := dkgCompletion.PubKeys[0]
-		pubKeyBz, _ := hex.DecodeString(pubKey)
-
-		if dkgRequest.Type == types.DKG_TYPE_DCM {
-			updateDCM(ctx, storeKey, dkgRequest.Id, pubKeyBz, cdc)
-		} else if dkgRequest.Type == types.DKG_TYPE_NONCE {
-			updateOracle(ctx, storeKey, dkgRequest.Id, pubKeyBz, cdc)
-		}
-	}
+		return false
+	})
 }
 
 // updateDCM updates the given dcm

x/dlc/types/expected_keepers.go

@@ -29,6 +29,8 @@ type TSSKeeper interface {
 
 	HasDKGRequest(ctx sdk.Context, id uint64) bool
 	GetDKGRequest(ctx sdk.Context, id uint64) *tsstypes.DKGRequest
+	GetDKGPubKeys(ctx sdk.Context, id uint64) []string
+	IterateDKGRequests(ctx sdk.Context, cb func(req *tsstypes.DKGRequest) (stop bool))
 
 	InitiateDKG(ctx sdk.Context, module string, ty string, intent int32, participants []string, threshold uint32, batchSize uint32) *tsstypes.DKGRequest
 	InitiateSigningRequest(ctx sdk.Context, module string, scopedId string, ty tsstypes.SigningType, intent int32, pubKey string, sigHashes []string, options *tsstypes.SigningOptions) *tsstypes.SigningRequest