From 3916f60deb2ce05cbd64641675bfeb5e419c835f Mon Sep 17 00:00:00 2001 From: Jerry Fan Date: Fri, 18 Oct 2024 17:29:06 -0400 Subject: [PATCH 1/4] add wallet when transfer to subaccount --- .../__tests__/handlers/transfer-handler.test.ts | 15 +++++++++++---- .../scripts/handlers/dydx_transfer_handler.sql | 5 +++++ 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/indexer/services/ender/__tests__/handlers/transfer-handler.test.ts b/indexer/services/ender/__tests__/handlers/transfer-handler.test.ts index 4f7f959070..cd8c7e8005 100644 --- a/indexer/services/ender/__tests__/handlers/transfer-handler.test.ts +++ b/indexer/services/ender/__tests__/handlers/transfer-handler.test.ts @@ -311,7 +311,7 @@ describe('transferHandler', () => { expect(wallet).toEqual(defaultWallet); }); - it('creates new deposit for previously non-existent subaccount', async () => { + it('creates new deposit for previously non-existent subaccount (also non-existent recipient wallet)', async () => { const transactionIndex: number = 0; const depositEvent: TransferEventV1 = defaultDepositEvent; @@ -348,16 +348,23 @@ describe('transferHandler', () => { newTransfer, asset, ); - // Confirm the wallet was created - const wallet: WalletFromDatabase | undefined = await WalletTable.findById( + // Confirm the wallet was created for the sender and recipient + const walletSender: WalletFromDatabase | undefined = await WalletTable.findById( defaultWalletAddress, ); + const walletRecipient: WalletFromDatabase | undefined = await WalletTable.findById( + defaultDepositEvent.recipient!.subaccountId!.owner, + ); const newRecipientSubaccount: SubaccountFromDatabase | undefined = await SubaccountTable.findById( defaultRecipientSubaccountId, ); expect(newRecipientSubaccount).toBeDefined(); - expect(wallet).toEqual(defaultWallet); + expect(walletSender).toEqual(defaultWallet); + expect(walletRecipient).toEqual({ + ...defaultWallet, + address: defaultDepositEvent.recipient!.subaccountId!.owner, + }); }); it('creates new withdrawal for existing subaccount', async () => { diff --git a/indexer/services/ender/src/scripts/handlers/dydx_transfer_handler.sql b/indexer/services/ender/src/scripts/handlers/dydx_transfer_handler.sql index 396a0075f8..8e9b251548 100644 --- a/indexer/services/ender/src/scripts/handlers/dydx_transfer_handler.sql +++ b/indexer/services/ender/src/scripts/handlers/dydx_transfer_handler.sql @@ -46,6 +46,11 @@ BEGIN SET "updatedAtHeight" = recipient_subaccount_record."updatedAtHeight", "updatedAt" = recipient_subaccount_record."updatedAt"; + + recipient_wallet_record."address" = event_data->'recipient'->'subaccountId'->>'owner'; + recipient_wallet_record."totalTradingRewards" = '0'; + recipient_wallet_record."totalVolume" = '0'; + INSERT INTO wallets VALUES (recipient_wallet_record.*) ON CONFLICT DO NOTHING; END IF; IF event_data->'sender'->'subaccountId' IS NOT NULL THEN From d80f69c3baffbf5ee06b50aabeaefae2a5acfb03 Mon Sep 17 00:00:00 2001 From: Jerry Fan Date: Mon, 21 Oct 2024 17:02:08 -0400 Subject: [PATCH 2/4] add prefix to AccountState keys in accountplus keeper --- protocol/app/upgrades/v8.0/upgrade.go | 55 ++++++++++++++ protocol/app/upgrades/v8.0/upgrade_test.go | 73 +++++++++++++++++++ protocol/x/accountplus/keeper/keeper.go | 39 +++++----- protocol/x/accountplus/keeper/keeper_test.go | 46 ++++++++++++ .../x/accountplus/keeper/timestampnonce.go | 11 +-- protocol/x/accountplus/types/keys.go | 5 ++ 6 files changed, 204 insertions(+), 25 deletions(-) create mode 100644 protocol/app/upgrades/v8.0/upgrade.go create mode 100644 protocol/app/upgrades/v8.0/upgrade_test.go diff --git a/protocol/app/upgrades/v8.0/upgrade.go b/protocol/app/upgrades/v8.0/upgrade.go new file mode 100644 index 0000000000..568b033fdb --- /dev/null +++ b/protocol/app/upgrades/v8.0/upgrade.go @@ -0,0 +1,55 @@ +package v_8_0 + +import ( + "bytes" + "fmt" + + storetypes "cosmossdk.io/store/types" + sdk "github.com/cosmos/cosmos-sdk/types" + accountpluskeeper "github.com/dydxprotocol/v4-chain/protocol/x/accountplus/keeper" + accountplustypes "github.com/dydxprotocol/v4-chain/protocol/x/accountplus/types" +) + +// Migrate accountplus AccountState states in kvstore from non-prefixed keys to prefixed keys +func migrateAccountplusAccountState(ctx sdk.Context, k accountpluskeeper.Keeper) { + ctx.Logger().Info("Migrating accountplus module AccountState in kvstore from non-prefixed keys to prefixed keys") + + store := ctx.KVStore(k.GetStoreKey()) + + var keysToDelete [][]byte + + // Create an iterator over all keys without prefixes. At the time of migration, the only + // key/value pairs are un-prefixed address/AccountState pairs. + iterator := storetypes.KVStorePrefixIterator(store, nil) + defer iterator.Close() + + for ; iterator.Valid(); iterator.Next() { + key := iterator.Key() + + // Double check that no keys have prefix + if bytes.HasPrefix(key, []byte(accountplustypes.AccountStateKeyPrefix)) { + panic(fmt.Sprintf("unexpected key with prefix %X found during migration", accountplustypes.AccountStateKeyPrefix)) + } + + value := iterator.Value() + var accountState accountplustypes.AccountState + if err := k.GetCdc().Unmarshal(value, &accountState); err != nil { + panic(fmt.Sprintf("failed to unmarshal AccountState for key %X: %s", key, err)) + } + + // keeper method SetAccountState stores with prefix + k.SetAccountState(ctx, key, accountState) + + keysToDelete = append(keysToDelete, key) + } + + // Delete old keys + ctx.Logger().Info("Deleteding old accountplus AccountState keys") + for _, key := range keysToDelete { + store.Delete(key) + } + + ctx.Logger().Info("Successfully migrated accountplus AccountState keys") +} + +// TODO: Scaffolding for upgrade: https://linear.app/dydx/issue/OTE-886/v8-upgrade-handler-scaffold diff --git a/protocol/app/upgrades/v8.0/upgrade_test.go b/protocol/app/upgrades/v8.0/upgrade_test.go new file mode 100644 index 0000000000..c4fe1ec095 --- /dev/null +++ b/protocol/app/upgrades/v8.0/upgrade_test.go @@ -0,0 +1,73 @@ +package v_8_0 + +import ( + "testing" + + "cosmossdk.io/store/prefix" + sdk "github.com/cosmos/cosmos-sdk/types" + + // upgrade "github.com/dydxprotocol/v4-chain/protocol/app/upgrades/v8_0" + testapp "github.com/dydxprotocol/v4-chain/protocol/testutil/app" + accountplustypes "github.com/dydxprotocol/v4-chain/protocol/x/accountplus/types" + "github.com/stretchr/testify/suite" +) + +type UpgradeTestSuite struct { + suite.Suite + + tApp *testapp.TestApp + Ctx sdk.Context +} + +func TestUpgradeTestSuite(t *testing.T) { + suite.Run(t, new(UpgradeTestSuite)) +} + +func (s *UpgradeTestSuite) SetupTest() { + s.tApp = testapp.NewTestAppBuilder(s.T()).Build() + s.Ctx = s.tApp.InitChain() +} + +func (s *UpgradeTestSuite) TestUpgrade_MigrateAccountplusAccountState() { + ctx := s.Ctx + store := ctx.KVStore(s.tApp.App.AccountPlusKeeper.GetStoreKey()) + prefixStore := prefix.NewStore(store, []byte(accountplustypes.AccountStateKeyPrefix)) + + // Create some AccountState with no prefixes + addresses := []string{"address1", "address2", "address3"} + for _, addr := range addresses { + accAddress := sdk.AccAddress([]byte(addr)) + accountState := accountplustypes.AccountState{ + Address: addr, + TimestampNonceDetails: accountplustypes.TimestampNonceDetails{ + TimestampNonces: []uint64{1, 2, 3}, + MaxEjectedNonce: 0, + }, + } + bz := s.tApp.App.AccountPlusKeeper.GetCdc().MustMarshal(&accountState) + store.Set(accAddress, bz) + } + + // Verify unprefixed keys were successfully created + for _, addr := range addresses { + accAddress := sdk.AccAddress([]byte(addr)) + bz := store.Get(accAddress) + s.Require().NotNil(bz, "Unprefixed key not created for %s", addr) + } + + // Migrate + migrateAccountplusAccountState(ctx, s.tApp.App.AccountPlusKeeper) + + // Verify that unprefixed keys are deleted and prefixed keys exist + for _, addr := range addresses { + accAddress := sdk.AccAddress([]byte(addr)) + + // Check that the old key no longer exists + bzOld := store.Get(accAddress) + s.Require().Nil(bzOld, "Unprefixed AccountState should be deleted for %s", addr) + + // Check that the new prefixed key exists + bzNew := prefixStore.Get(accAddress) + s.Require().NotNil(bzNew, "Prefixed AccountState should exist for %s", addr) + } +} diff --git a/protocol/x/accountplus/keeper/keeper.go b/protocol/x/accountplus/keeper/keeper.go index d63297e19f..bd0bf5cd56 100644 --- a/protocol/x/accountplus/keeper/keeper.go +++ b/protocol/x/accountplus/keeper/keeper.go @@ -1,12 +1,11 @@ package keeper import ( - "bytes" - "errors" "fmt" "strings" "cosmossdk.io/log" + "cosmossdk.io/store/prefix" storetypes "cosmossdk.io/store/types" "github.com/cosmos/cosmos-sdk/codec" sdk "github.com/cosmos/cosmos-sdk/types" @@ -38,6 +37,14 @@ func NewKeeper( } } +func (k Keeper) GetStoreKey() storetypes.StoreKey { + return k.storeKey +} + +func (k Keeper) GetCdc() codec.BinaryCodec { + return k.cdc +} + func (k Keeper) Logger(ctx sdk.Context) log.Logger { return ctx.Logger().With(log.ModuleKey, fmt.Sprintf("x/%s", types.ModuleName)) } @@ -46,26 +53,22 @@ func (k Keeper) Logger(ctx sdk.Context) log.Logger { func (k Keeper) InitializeForGenesis(ctx sdk.Context) { } -// Get all account details pairs in store +// Get all AccountStates from kvstore func (k Keeper) GetAllAccountStates(ctx sdk.Context) ([]types.AccountState, error) { store := ctx.KVStore(k.storeKey) + prefixStore := prefix.NewStore(store, []byte(types.AccountStateKeyPrefix)) - iterator := storetypes.KVStorePrefixIterator(store, nil) + iterator := prefixStore.Iterator(nil, nil) defer iterator.Close() accounts := []types.AccountState{} for ; iterator.Valid(); iterator.Next() { - key := iterator.Key() - - // Temporary workaround to exclude smart account kv pairs. - if bytes.HasPrefix(key, []byte(types.SmartAccountKeyPrefix)) { - continue + var accountState types.AccountState + err := k.cdc.Unmarshal(iterator.Value(), &accountState) + if err != nil { + return nil, fmt.Errorf("failed to unmarshal account state: %w", err) } - accountState, found := k.GetAccountState(ctx, key) - if !found { - return accounts, errors.New("Could not get account state for address: " + sdk.AccAddress(iterator.Key()).String()) - } accounts = append(accounts, accountState) } @@ -144,7 +147,7 @@ func (k Keeper) GetAllAuthenticatorData(ctx sdk.Context) ([]types.AuthenticatorD return accountAuthenticators, nil } -func GetAccountPlusStateWithTimestampNonceDetails( +func AccountStateFromTimestampNonceDetails( address sdk.AccAddress, tsNonce uint64, ) types.AccountState { @@ -162,8 +165,8 @@ func (k Keeper) GetAccountState( ctx sdk.Context, address sdk.AccAddress, ) (types.AccountState, bool) { - store := ctx.KVStore(k.storeKey) - bz := store.Get(address.Bytes()) + prefixStore := prefix.NewStore(ctx.KVStore(k.storeKey), []byte(types.AccountStateKeyPrefix)) + bz := prefixStore.Get(address.Bytes()) if bz == nil { return types.AccountState{}, false } @@ -185,9 +188,9 @@ func (k Keeper) SetAccountState( address sdk.AccAddress, accountState types.AccountState, ) { - store := ctx.KVStore(k.storeKey) + prefixStore := prefix.NewStore(ctx.KVStore(k.storeKey), []byte(types.AccountStateKeyPrefix)) bz := k.cdc.MustMarshal(&accountState) - store.Set(address.Bytes(), bz) + prefixStore.Set(address.Bytes(), bz) } func (k Keeper) HasAuthority(authority string) bool { diff --git a/protocol/x/accountplus/keeper/keeper_test.go b/protocol/x/accountplus/keeper/keeper_test.go index bb0d2d0bbb..060d87befe 100644 --- a/protocol/x/accountplus/keeper/keeper_test.go +++ b/protocol/x/accountplus/keeper/keeper_test.go @@ -53,6 +53,52 @@ func (s *KeeperTestSuite) SetupTest() { ) } +func (s *KeeperTestSuite) TestKeeper_Set_Get_GetAllAccountState() { + ctx := s.Ctx + address := "address1" + + // SetAccountState + s.tApp.App.AccountPlusKeeper.SetAccountState( + ctx, + sdk.AccAddress([]byte(address)), + types.AccountState{ + Address: address, + TimestampNonceDetails: types.TimestampNonceDetails{ + TimestampNonces: []uint64{1, 2, 3}, + MaxEjectedNonce: 0, + }, + }, + ) + + // GetAccountState + _, found := s.tApp.App.AccountPlusKeeper.GetAccountState(ctx, sdk.AccAddress([]byte(address))) + s.Require().True(found, "Account state not found") + + // GetAllAccountStates + accountStates, err := s.tApp.App.AccountPlusKeeper.GetAllAccountStates(ctx) + s.Require().NoError(err, "Should not error when getting all account states") + s.Require().Equal(len(accountStates), 1, "Incorrect number of AccountStates retrieved") + + // Add one more AccountState and check GetAllAccountStates + // SetAccountState + address2 := "address2" + s.tApp.App.AccountPlusKeeper.SetAccountState( + ctx, + sdk.AccAddress([]byte(address2)), + types.AccountState{ + Address: address2, + TimestampNonceDetails: types.TimestampNonceDetails{ + TimestampNonces: []uint64{1, 2, 3}, + MaxEjectedNonce: 0, + }, + }, + ) + + accountStates2, err := s.tApp.App.AccountPlusKeeper.GetAllAccountStates(ctx) + s.Require().NoError(err, "Should not error when getting all account states") + s.Require().Equal(len(accountStates2), 2, "Incorrect number of AccountStates retrieved") +} + func (s *KeeperTestSuite) TestKeeper_AddAuthenticator() { ctx := s.Ctx diff --git a/protocol/x/accountplus/keeper/timestampnonce.go b/protocol/x/accountplus/keeper/timestampnonce.go index e23a0196a6..dad8e25a16 100644 --- a/protocol/x/accountplus/keeper/timestampnonce.go +++ b/protocol/x/accountplus/keeper/timestampnonce.go @@ -41,7 +41,7 @@ func (k Keeper) ProcessTimestampNonce(ctx sdk.Context, acc sdk.AccountI, tsNonce accountState, found := k.GetAccountState(ctx, address) if !found { // initialize accountplus state with ts nonce details - k.SetAccountState(ctx, address, GetAccountPlusStateWithTimestampNonceDetails(address, tsNonce)) + k.SetAccountState(ctx, address, AccountStateFromTimestampNonceDetails(address, tsNonce)) } else { EjectStaleTimestampNonces(&accountState, blockTs) tsNonceAccepted := AttemptTimestampNonceUpdate(tsNonce, &accountState) @@ -56,9 +56,10 @@ func (k Keeper) ProcessTimestampNonce(ctx sdk.Context, acc sdk.AccountI, tsNonce // Inplace eject all stale timestamps. func EjectStaleTimestampNonces(accountState *types.AccountState, referenceTs uint64) { tsNonceDetails := &accountState.TimestampNonceDetails + oldestAllowedTs := referenceTs - MaxTimeInPastMs var newTsNonces []uint64 for _, tsNonce := range tsNonceDetails.TimestampNonces { - if tsNonce >= referenceTs-MaxTimeInPastMs { + if tsNonce >= oldestAllowedTs { newTsNonces = append(newTsNonces, tsNonce) } else { if tsNonce > tsNonceDetails.MaxEjectedNonce { @@ -115,9 +116,5 @@ func isLargerThanSmallestValue(value uint64, values []uint64) (bool, int) { } } - if value > values[minIndex] { - return true, minIndex - } else { - return false, minIndex - } + return value > values[minIndex], minIndex } diff --git a/protocol/x/accountplus/types/keys.go b/protocol/x/accountplus/types/keys.go index 3844749bf3..2e0d0edd91 100644 --- a/protocol/x/accountplus/types/keys.go +++ b/protocol/x/accountplus/types/keys.go @@ -17,6 +17,11 @@ const ( StoreKey = ModuleName ) +// Prefix for account state. +const ( + AccountStateKeyPrefix = "AS/" +) + // Below key prefixes are for smart account implementation. const ( // SmartAccountKeyPrefix is the prefix key for all smart account store state. From b2370d263d2763fd5fc7287ef606a2584ca64c57 Mon Sep 17 00:00:00 2001 From: Jerry Fan Date: Wed, 23 Oct 2024 13:42:27 -0400 Subject: [PATCH 3/4] add additional assertions in unit tests --- ...de_test.go => migrate_accountplus_test.go} | 13 ++++- protocol/x/accountplus/keeper/keeper.go | 8 ++-- protocol/x/accountplus/keeper/keeper_test.go | 48 ++++++++++--------- 3 files changed, 42 insertions(+), 27 deletions(-) rename protocol/app/upgrades/v8.0/{upgrade_test.go => migrate_accountplus_test.go} (78%) diff --git a/protocol/app/upgrades/v8.0/upgrade_test.go b/protocol/app/upgrades/v8.0/migrate_accountplus_test.go similarity index 78% rename from protocol/app/upgrades/v8.0/upgrade_test.go rename to protocol/app/upgrades/v8.0/migrate_accountplus_test.go index 92e28aa8c8..c107c63c6e 100644 --- a/protocol/app/upgrades/v8.0/upgrade_test.go +++ b/protocol/app/upgrades/v8.0/migrate_accountplus_test.go @@ -6,6 +6,7 @@ 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" ) @@ -17,7 +18,7 @@ type UpgradeTestSuite struct { Ctx sdk.Context } -func TestUpgradeTestSuite(t *testing.T) { +func TestMigrateAccountplusAccountState(t *testing.T) { suite.Run(t, new(UpgradeTestSuite)) } @@ -66,6 +67,16 @@ func (s *UpgradeTestSuite) TestUpgrade_MigrateAccountplusAccountState() { // Check that the new prefixed key exists bzNew := prefixStore.Get(accAddress) + var actualAccountState types.AccountState + s.tApp.App.AccountPlusKeeper.GetCdc().MustUnmarshal(bzNew, &actualAccountState) + expectedAccountState := accountplustypes.AccountState{ + Address: addr, + TimestampNonceDetails: accountplustypes.TimestampNonceDetails{ + TimestampNonces: []uint64{1, 2, 3}, + MaxEjectedNonce: 0, + }, + } s.Require().NotNil(bzNew, "Prefixed AccountState should exist for %s", addr) + s.Require().Equal(expectedAccountState, actualAccountState, "Incorrect AccountState after migration for %s", addr) } } diff --git a/protocol/x/accountplus/keeper/keeper.go b/protocol/x/accountplus/keeper/keeper.go index bd0bf5cd56..b09d435d82 100644 --- a/protocol/x/accountplus/keeper/keeper.go +++ b/protocol/x/accountplus/keeper/keeper.go @@ -55,10 +55,10 @@ func (k Keeper) InitializeForGenesis(ctx sdk.Context) { // Get all AccountStates from kvstore func (k Keeper) GetAllAccountStates(ctx sdk.Context) ([]types.AccountState, error) { - store := ctx.KVStore(k.storeKey) - prefixStore := prefix.NewStore(store, []byte(types.AccountStateKeyPrefix)) - - iterator := prefixStore.Iterator(nil, nil) + iterator := storetypes.KVStorePrefixIterator( + ctx.KVStore(k.storeKey), + []byte(types.AccountStateKeyPrefix), + ) defer iterator.Close() accounts := []types.AccountState{} diff --git a/protocol/x/accountplus/keeper/keeper_test.go b/protocol/x/accountplus/keeper/keeper_test.go index 060d87befe..58e0101c66 100644 --- a/protocol/x/accountplus/keeper/keeper_test.go +++ b/protocol/x/accountplus/keeper/keeper_test.go @@ -55,48 +55,52 @@ func (s *KeeperTestSuite) SetupTest() { func (s *KeeperTestSuite) TestKeeper_Set_Get_GetAllAccountState() { ctx := s.Ctx - address := "address1" + + accountState1 := types.AccountState{ + Address: "address1", + TimestampNonceDetails: types.TimestampNonceDetails{ + TimestampNonces: []uint64{1, 2, 3}, + MaxEjectedNonce: 0, + }, + } + + accountState2 := types.AccountState{ + Address: "address2", + TimestampNonceDetails: types.TimestampNonceDetails{ + TimestampNonces: []uint64{1, 2, 3}, + MaxEjectedNonce: 0, + }, + } // SetAccountState s.tApp.App.AccountPlusKeeper.SetAccountState( ctx, - sdk.AccAddress([]byte(address)), - types.AccountState{ - Address: address, - TimestampNonceDetails: types.TimestampNonceDetails{ - TimestampNonces: []uint64{1, 2, 3}, - MaxEjectedNonce: 0, - }, - }, + sdk.AccAddress([]byte(accountState1.Address)), + accountState1, ) // GetAccountState - _, found := s.tApp.App.AccountPlusKeeper.GetAccountState(ctx, sdk.AccAddress([]byte(address))) + _, found := s.tApp.App.AccountPlusKeeper.GetAccountState(ctx, sdk.AccAddress([]byte(accountState1.Address))) s.Require().True(found, "Account state not found") // GetAllAccountStates accountStates, err := s.tApp.App.AccountPlusKeeper.GetAllAccountStates(ctx) s.Require().NoError(err, "Should not error when getting all account states") s.Require().Equal(len(accountStates), 1, "Incorrect number of AccountStates retrieved") + s.Require().Equal(accountStates[0], accountState1, "Incorrect AccountState retrieved") // Add one more AccountState and check GetAllAccountStates - // SetAccountState - address2 := "address2" s.tApp.App.AccountPlusKeeper.SetAccountState( ctx, - sdk.AccAddress([]byte(address2)), - types.AccountState{ - Address: address2, - TimestampNonceDetails: types.TimestampNonceDetails{ - TimestampNonces: []uint64{1, 2, 3}, - MaxEjectedNonce: 0, - }, - }, + sdk.AccAddress([]byte(accountState2.Address)), + accountState2, ) - accountStates2, err := s.tApp.App.AccountPlusKeeper.GetAllAccountStates(ctx) + accountStates, err = s.tApp.App.AccountPlusKeeper.GetAllAccountStates(ctx) s.Require().NoError(err, "Should not error when getting all account states") - s.Require().Equal(len(accountStates2), 2, "Incorrect number of AccountStates retrieved") + s.Require().Equal(len(accountStates), 2, "Incorrect number of AccountStates retrieved") + s.Require().Contains(accountStates, accountState1, "Retrieved AccountStates does not contain accountState1") + s.Require().Contains(accountStates, accountState2, "Retrieved AccountStates does not contain accountState2") } func (s *KeeperTestSuite) TestKeeper_AddAuthenticator() { From 3d9b80a8fc3bea83296f7263a53fba06c969ffc8 Mon Sep 17 00:00:00 2001 From: Jerry Fan Date: Wed, 23 Oct 2024 15:05:26 -0400 Subject: [PATCH 4/4] modify store outside of iterator --- .../upgrades/v8.0/migrate_accountplus_test.go | 3 +-- protocol/app/upgrades/v8.0/upgrade.go | 22 ++++++++++++++++--- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/protocol/app/upgrades/v8.0/migrate_accountplus_test.go b/protocol/app/upgrades/v8.0/migrate_accountplus_test.go index c107c63c6e..88be8a4994 100644 --- a/protocol/app/upgrades/v8.0/migrate_accountplus_test.go +++ b/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, diff --git a/protocol/app/upgrades/v8.0/upgrade.go b/protocol/app/upgrades/v8.0/upgrade.go index 9010906bfb..12707b727b 100644 --- a/protocol/app/upgrades/v8.0/upgrade.go +++ b/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) }