fix(mempool): prevent memory leaks in PriorityNonceMempool Remove

CHANGELOG.md

@@ -126,6 +126,7 @@ This patch release contains only minor dependency bumps.
 * (baseapp) [#25862](https://github.com/cosmos/cosmos-sdk/pull/25862) Skip running validateBasic for rechecking txs. (Backport of https://github.com/cosmos/cosmos-sdk/pull/20208).
 * (blockstm) [25883](https://github.com/cosmos/cosmos-sdk/pull/25883) Re-use decoded tx object in pre-estimates.
 * (blockstm) [#25788](https://github.com/cosmos/cosmos-sdk/pull/25788) Only validate transactions that's executed at lease once.
+* (type/mempool) [#25858] (https://github.com/cosmos/cosmos-sdk/pull/25858) fix(mempool): prevent memory leaks in PriorityNonceMempool Remove.
 * (blockstm) [#25767](https://github.com/cosmos/cosmos-sdk/pull/25767) Optimize block-stm MVMemory with bitmap index.
 
 ### Bug Fixes

types/mempool/priority_nonce.go

@@ -485,8 +485,14 @@ func (mp *PriorityNonceMempool[C]) Remove(tx sdk.Tx) error {
 
 	mp.priorityIndex.Remove(tk)
 	senderTxs.Remove(tk)
+	if senderTxs.Len() == 0 {
+		delete(mp.senderIndices, sender)
+	}
 	delete(mp.scores, scoreKey)
 	mp.priorityCounts[score.priority]--
+	if mp.priorityCounts[score.priority] == 0 {
+		delete(mp.priorityCounts, score.priority)
+	}
 
 	return nil
 }

types/mempool/priority_nonce_test.go

@@ -840,6 +840,44 @@ func TestPriorityNonceMempool_NextSenderTx(t *testing.T) {
 	require.Equal(t, txs[0], tx)
 }
 
+func TestPriorityNonceMempool_CleanupOnRemove(t *testing.T) {
+	accounts := simtypes.RandomAccounts(rand.New(rand.NewSource(0)), 3)
+	ctx := sdk.NewContext(nil, cmtproto.Header{}, false, log.NewNopLogger())
+	sa := accounts[0].Address
+	sb := accounts[1].Address
+	sc := accounts[2].Address
+
+	mp := mempool.DefaultPriorityMempool()
+
+	txs := []testTx{
+		{priority: 10, nonce: 1, address: sa},
+		{priority: 10, nonce: 2, address: sa},
+		{priority: 20, nonce: 1, address: sb},
+		{priority: 30, nonce: 1, address: sc},
+	}
+
+	for _, tx := range txs {
+		c := ctx.WithPriority(tx.priority)
+		require.NoError(t, mp.Insert(c, tx))
+	}
+
+	require.Equal(t, 4, mp.CountTx())
+
+	require.NoError(t, mp.Remove(txs[0]))
+	require.NoError(t, mp.Remove(txs[1]))
+
+	require.NotPanics(t, func() {
+		require.Nil(t, mp.NextSenderTx(sa.String()))
+	})
+
+	require.NoError(t, mp.Remove(txs[2]))
+
+	require.NoError(t, mp.Remove(txs[3]))
+
+	require.Equal(t, 0, mp.CountTx())
+	require.NoError(t, mempool.IsEmpty[int64](mp))
+}
+
 func TestNextSenderTx_ConcurrentAccess(t *testing.T) {
 	ctx := sdk.NewContext(nil, cmtproto.Header{}, false, log.NewNopLogger())
 	account := simtypes.RandomAccounts(rand.New(rand.NewSource(2)), 1)[0].Address