fix(x/gov): Do not call Remove during Walk (#24460)

Co-authored-by: Tyler <48813565+technicallyty@users.noreply.github.com>

CHANGELOG.md

@@ -83,6 +83,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### Bug Fixes
 
+* (x/gov)[#24460](https://github.com/cosmos/cosmos-sdk/pull/24460) Do not call Remove during Walk in defaultCalculateVoteResultsAndVotingPower.
 * (baseapp) [24261](https://github.com/cosmos/cosmos-sdk/pull/24261) Fix post handler error always results in code 1
 * (server) [#24068](https://github.com/cosmos/cosmos-sdk/pull/24068) Allow align block header with skip check header in grpc server.
 * (x/gov) [#24044](https://github.com/cosmos/cosmos-sdk/pull/24044) Fix some places in which we call Remove inside a Walk (x/gov).

x/gov/keeper/tally.go

@@ -38,6 +38,7 @@ func defaultCalculateVoteResultsAndVotingPower(
 	results[v1.OptionNoWithVeto] = math.LegacyZeroDec()
 
 	rng := collections.NewPrefixedPairRange[uint64, sdk.AccAddress](proposal.Id)
+	votesToRemove := []collections.Pair[uint64, sdk.AccAddress]{}
 	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)
@@ -81,12 +82,20 @@ 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
 	})
 	if err != nil {
 		return math.LegacyZeroDec(), nil, fmt.Errorf("error while iterating delegations: %w", err)
 	}
 
+	// remove all votes from store
+	for _, key := range votesToRemove {
+		if err := k.Votes.Remove(ctx, key); err != nil {
+			return math.LegacyDec{}, nil, fmt.Errorf("error while removing vote (%d/%s): %w", key.K1(), key.K2(), err)
+		}
+	}
+
 	// iterate over the validators again to tally their voting power
 	for _, val := range validators {
 		if len(val.Vote) == 0 {

x/gov/keeper/tally_test.go

@@ -0,0 +1,115 @@
+package keeper_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"cosmossdk.io/collections"
+	"cosmossdk.io/math"
+
+	"github.com/cosmos/cosmos-sdk/codec/address"
+	simtestutil "github.com/cosmos/cosmos-sdk/testutil/sims"
+	v1 "github.com/cosmos/cosmos-sdk/x/gov/types/v1"
+)
+
+func TestVoteRemovalAfterTally(t *testing.T) {
+	govKeeper, authKeeper, bankKeeper, stakingKeeper, _, _, ctx := setupGovKeeper(t)
+	authKeeper.EXPECT().AddressCodec().Return(address.NewBech32Codec("cosmos")).AnyTimes()
+	stakingKeeper.EXPECT().ValidatorAddressCodec().Return(address.NewBech32Codec("cosmos")).AnyTimes()
+
+	addrs := simtestutil.AddTestAddrsIncremental(bankKeeper, stakingKeeper, ctx, 3, math.NewInt(30000000))
+
+	// Create a test proposal
+	tp := TestProposal
+	proposal, err := govKeeper.SubmitProposal(ctx, tp, "", "test", "summary", addrs[0], false)
+	require.NoError(t, err)
+	proposalID := proposal.Id
+
+	// Activate voting period
+	proposal.Status = v1.StatusVotingPeriod
+	require.NoError(t, govKeeper.SetProposal(ctx, proposal))
+
+	// Add votes from different addresses
+	require.NoError(t, govKeeper.AddVote(ctx, proposalID, addrs[0], v1.NewNonSplitVoteOption(v1.OptionYes), ""))
+	require.NoError(t, govKeeper.AddVote(ctx, proposalID, addrs[1], v1.NewNonSplitVoteOption(v1.OptionNo), ""))
+	require.NoError(t, govKeeper.AddVote(ctx, proposalID, addrs[2], v1.NewNonSplitVoteOption(v1.OptionAbstain), ""))
+
+	// verify votes were added to state
+	for i, addr := range addrs {
+		vote, err := govKeeper.Votes.Get(ctx, collections.Join(proposalID, addr))
+		require.NoError(t, err, "Vote for address %d should exist before tally", i)
+		require.NotNil(t, vote, "Vote for address %d should not be nil before tally", i)
+	}
+
+	// tally the proposal
+	proposal, err = govKeeper.Proposals.Get(ctx, proposalID)
+	require.NoError(t, err)
+	_, _, _, err = govKeeper.Tally(ctx, proposal)
+	require.NoError(t, err)
+
+	// votes should be deleted.
+	for i, addr := range addrs {
+		_, err := govKeeper.Votes.Get(ctx, collections.Join(proposalID, addr))
+		require.Error(t, err, "Vote for address %d should be removed after tally", i)
+		require.ErrorIs(t, err, collections.ErrNotFound, "Error should be ErrNotFound for address %d after tally", i)
+	}
+}
+
+// TestMultipleProposalsVoteRemoval verifies that votes for one proposal are removed
+// while votes for another proposal are preserved during tallying
+func TestMultipleProposalsVoteRemoval(t *testing.T) {
+	govKeeper, authKeeper, bankKeeper, stakingKeeper, _, _, ctx := setupGovKeeper(t)
+	authKeeper.EXPECT().AddressCodec().Return(address.NewBech32Codec("cosmos")).AnyTimes()
+	stakingKeeper.EXPECT().ValidatorAddressCodec().Return(address.NewBech32Codec("cosmos")).AnyTimes()
+
+	addrs := simtestutil.AddTestAddrsIncremental(bankKeeper, stakingKeeper, ctx, 2, math.NewInt(30000000))
+
+	tp := TestProposal
+	proposal1, err := govKeeper.SubmitProposal(ctx, tp, "", "test1", "summary", addrs[0], false)
+	require.NoError(t, err)
+	proposal1ID := proposal1.Id
+
+	proposal2, err := govKeeper.SubmitProposal(ctx, tp, "", "test2", "summary", addrs[0], false)
+	require.NoError(t, err)
+	proposal2ID := proposal2.Id
+
+	// activate both proposals
+	proposal1.Status = v1.StatusVotingPeriod
+	require.NoError(t, govKeeper.SetProposal(ctx, proposal1))
+	proposal2.Status = v1.StatusVotingPeriod
+	require.NoError(t, govKeeper.SetProposal(ctx, proposal2))
+
+	// add some votes for both proposals
+	require.NoError(t, govKeeper.AddVote(ctx, proposal1ID, addrs[0], v1.NewNonSplitVoteOption(v1.OptionYes), ""))
+	require.NoError(t, govKeeper.AddVote(ctx, proposal1ID, addrs[1], v1.NewNonSplitVoteOption(v1.OptionNo), ""))
+
+	require.NoError(t, govKeeper.AddVote(ctx, proposal2ID, addrs[0], v1.NewNonSplitVoteOption(v1.OptionNo), ""))
+	require.NoError(t, govKeeper.AddVote(ctx, proposal2ID, addrs[1], v1.NewNonSplitVoteOption(v1.OptionYes), ""))
+
+	// votes should eixst
+	vote1Addr0, err := govKeeper.Votes.Get(ctx, collections.Join(proposal1ID, addrs[0]))
+	require.NoError(t, err)
+	require.Equal(t, v1.OptionYes, vote1Addr0.Options[0].Option)
+	vote2Addr0, err := govKeeper.Votes.Get(ctx, collections.Join(proposal2ID, addrs[0]))
+	require.NoError(t, err)
+	require.Equal(t, v1.OptionNo, vote2Addr0.Options[0].Option)
+
+	// only tally proposal1
+	proposal1, err = govKeeper.Proposals.Get(ctx, proposal1ID)
+	require.NoError(t, err)
+	_, _, _, err = govKeeper.Tally(ctx, proposal1)
+	require.NoError(t, err)
+
+	// check votes
+	for _, addr := range addrs {
+		// proposal1 votes should be deleted
+		_, err := govKeeper.Votes.Get(ctx, collections.Join(proposal1ID, addr))
+		require.Error(t, err)
+		require.ErrorIs(t, err, collections.ErrNotFound)
+
+		// proposal2 votes should still exist.
+		_, err = govKeeper.Votes.Get(ctx, collections.Join(proposal2ID, addr))
+		require.NoError(t, err)
+	}
+}