modify store outside of iterator

protocol/app/upgrades/v8.0/migrate_accountplus_test.go

@@ -6,7 +6,6 @@ import (
 	"cosmossdk.io/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	testapp "github.com/dydxprotocol/v4-chain/protocol/testutil/app"
-	"github.com/dydxprotocol/v4-chain/protocol/x/accountplus/types"
 	accountplustypes "github.com/dydxprotocol/v4-chain/protocol/x/accountplus/types"
 	"github.com/stretchr/testify/suite"
 )
@@ -67,7 +66,7 @@ func (s *UpgradeTestSuite) TestUpgrade_MigrateAccountplusAccountState() {
 
 		// Check that the new prefixed key exists
 		bzNew := prefixStore.Get(accAddress)
-		var actualAccountState types.AccountState
+		var actualAccountState accountplustypes.AccountState
 		s.tApp.App.AccountPlusKeeper.GetCdc().MustUnmarshal(bzNew, &actualAccountState)
 		expectedAccountState := accountplustypes.AccountState{
 			Address: addr,

protocol/app/upgrades/v8.0/upgrade.go

@@ -20,6 +20,11 @@ func migrateAccountplusAccountState(ctx sdk.Context, k accountpluskeeper.Keeper)
 	iterator := storetypes.KVStorePrefixIterator(store, nil)
 	defer iterator.Close()
 
+	var keysToDelete [][]byte
+	var accountStatesToSet []struct {
+		address      sdk.AccAddress
+		accountState accountplustypes.AccountState
+	}
 	for ; iterator.Valid(); iterator.Next() {
 		key := iterator.Key()
 
@@ -34,10 +39,21 @@ func migrateAccountplusAccountState(ctx sdk.Context, k accountpluskeeper.Keeper)
 			panic(fmt.Sprintf("failed to unmarshal AccountState for key %X: %s", key, err))
 		}
 
-		// SetAccountState stores with prefix
-		k.SetAccountState(ctx, key, accountState)
+		accountStatesToSet = append(accountStatesToSet, struct {
+			address      sdk.AccAddress
+			accountState accountplustypes.AccountState
+		}{key, accountState})
+
+		keysToDelete = append(keysToDelete, key)
+	}
+
+	// Set prefixed keys
+	for _, item := range accountStatesToSet {
+		k.SetAccountState(ctx, item.address, item.accountState)
+	}
 
-		// Delete unprefixed key
+	// Delete unprefixed keys
+	for _, key := range keysToDelete {
 		store.Delete(key)
 	}