fix: Do not call Remove during Walk (cosmos#19833)

Co-authored-by: Aleksandr Bezobchuk <[email protected]>
crypto-org-chain · Mar 24, 2024 · ea7bdd1 · ea7bdd1
1 parent 4599439
commit ea7bdd1
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 15 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -98,6 +98,7 @@ Every module contains its own CHANGELOG.md. Please refer to the module you are i
 * (baseapp) [#18551](https://github.com/cosmos/cosmos-sdk/pull/18551) Fix SelectTxForProposal the calculation method of tx bytes size is inconsistent with CometBFT
 * (server) [#18994](https://github.com/cosmos/cosmos-sdk/pull/18994) Update server context directly rather than a reference to a sub-object
 * (crypto) [#19691](https://github.com/cosmos/cosmos-sdk/pull/19691) Fix tx sign doesn't throw an error when incorrect Ledger is used.
+* [#19833](https://github.com/cosmos/cosmos-sdk/pull/19833) Fix some places in which we call Remove inside a Walk.
 
 ### API Breaking Changes
 

diff --git a/x/feegrant/keeper/keeper.go b/x/feegrant/keeper/keeper.go
@@ -302,16 +302,15 @@ func (k Keeper) RemoveExpiredAllowances(ctx context.Context, limit int) error {
 	rng := collections.NewPrefixUntilTripleRange[time.Time, sdk.AccAddress, sdk.AccAddress](exp)
 	count := 0
 
+	keysToRemove := []collections.Triple[time.Time, sdk.AccAddress, sdk.AccAddress]{}
 	err := k.FeeAllowanceQueue.Walk(ctx, rng, func(key collections.Triple[time.Time, sdk.AccAddress, sdk.AccAddress], value bool) (stop bool, err error) {
 		grantee, granter := key.K2(), key.K3()
 
 		if err := k.FeeAllowance.Remove(ctx, collections.Join(grantee, granter)); err != nil {
 			return true, err
 		}
 
-		if err := k.FeeAllowanceQueue.Remove(ctx, key); err != nil {
-			return true, err
-		}
+		keysToRemove = append(keysToRemove, key)
 
 		// limit the amount of iterations to avoid taking too much time
 		count++
@@ -325,5 +324,11 @@ func (k Keeper) RemoveExpiredAllowances(ctx context.Context, limit int) error {
 		return err
 	}
 
+	for _, key := range keysToRemove {
+		if err := k.FeeAllowanceQueue.Remove(ctx, key); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
diff --git a/x/gov/keeper/tally.go b/x/gov/keeper/tally.go
@@ -247,6 +247,7 @@ func defaultCalculateVoteResultsAndVotingPower(
 
 	// iterate over all votes, tally up the voting power of each validator
 	rng := collections.NewPrefixedPairRange[uint64, sdk.AccAddress](proposalID)
+	votesToRemove := []collections.Pair[uint64, sdk.AccAddress]{}
 	if err := k.Votes.Walk(ctx, rng, func(key collections.Pair[uint64, sdk.AccAddress], vote v1.Vote) (bool, error) {
 		// if validator, just record it in the map
 		voter, err := k.authKeeper.AddressCodec().StringToBytes(vote.Voter)
@@ -292,11 +293,19 @@ func defaultCalculateVoteResultsAndVotingPower(
 			return false, err
 		}
 
-		return false, k.Votes.Remove(ctx, collections.Join(vote.ProposalId, sdk.AccAddress(voter)))
+		votesToRemove = append(votesToRemove, key)
+		return false, nil
 	}); err != nil {
 		return math.LegacyDec{}, nil, err
 	}
 
+	// remove all votes from store
+	for _, key := range votesToRemove {
+		if err := k.Votes.Remove(ctx, key); err != nil {
+			return math.LegacyDec{}, nil, err
+		}
+	}
+
 	// iterate over the validators again to tally their voting power
 	for _, val := range validators {
 		if len(val.Vote) == 0 {

diff --git a/x/slashing/keeper/signing_info.go b/x/slashing/keeper/signing_info.go
@@ -182,13 +182,7 @@ func (k Keeper) DeleteMissedBlockBitmap(ctx context.Context, addr sdk.ConsAddres
 	}
 
 	rng := collections.NewPrefixedPairRange[[]byte, uint64](addr.Bytes())
-	return k.ValidatorMissedBlockBitmap.Walk(ctx, rng, func(key collections.Pair[[]byte, uint64], value []byte) (bool, error) {
-		err := k.ValidatorMissedBlockBitmap.Remove(ctx, key)
-		if err != nil {
-			return true, err
-		}
-		return false, nil
-	})
+	return k.ValidatorMissedBlockBitmap.Clear(ctx, rng)
 }
 
 // IterateMissedBlockBitmap iterates over a validator's signed blocks window

diff --git a/x/staking/keeper/cons_pubkey.go b/x/staking/keeper/cons_pubkey.go
@@ -202,9 +202,7 @@ func (k Keeper) deleteConsKeyIndexKey(ctx context.Context, valAddr sdk.ValAddres
 		StartInclusive(collections.Join(valAddr.Bytes(), time.Time{})).
 		EndInclusive(collections.Join(valAddr.Bytes(), ts))
 
-	return k.ValidatorConsensusKeyRotationRecordIndexKey.Walk(ctx, rng, func(key collections.Pair[[]byte, time.Time]) (stop bool, err error) {
-		return false, k.ValidatorConsensusKeyRotationRecordIndexKey.Remove(ctx, key)
-	})
+	return k.ValidatorConsensusKeyRotationRecordIndexKey.Clear(ctx, rng)
 }
 
 // getAndRemoveAllMaturedRotatedKeys returns all matured valaddresses.
@@ -213,14 +211,23 @@ func (k Keeper) getAndRemoveAllMaturedRotatedKeys(ctx context.Context, matureTim
 
 	// get an iterator for all timeslices from time 0 until the current HeaderInfo time
 	rng := new(collections.Range[time.Time]).EndInclusive(matureTime)
+	keysToRemove := []time.Time{}
 	err := k.ValidatorConsensusKeyRotationRecordQueue.Walk(ctx, rng, func(key time.Time, value types.ValAddrsOfRotatedConsKeys) (stop bool, err error) {
 		valAddrs = append(valAddrs, value.Addresses...)
-		return false, k.ValidatorConsensusKeyRotationRecordQueue.Remove(ctx, key)
+		keysToRemove = append(keysToRemove, key)
+		return false, nil
 	})
 	if err != nil {
 		return nil, err
 	}
 
+	// remove all the keys from the list
+	for _, key := range keysToRemove {
+		if err := k.ValidatorConsensusKeyRotationRecordQueue.Remove(ctx, key); err != nil {
+			return nil, err
+		}
+	}
+
 	return valAddrs, nil
 }