From e8d54831f9a9e57e1850a1e10629e6fd40a3dfc5 Mon Sep 17 00:00:00 2001
From: mmsqe <mavis@crypto.com>
Date: Fri, 15 Mar 2024 09:13:48 +0800
Subject: [PATCH] align signer extraction adapter for mempool remove

this change should have been included in 545258668753e4defc2aa3912362a42df6a7c652

add test

update doc

Apply suggestions from code review
---
 CHANGELOG.md                         |  2 +
 types/mempool/priority_nonce.go      |  5 +-
 types/mempool/priority_nonce_test.go | 81 ++++++++++++++++++++++++++++
 3 files changed, 85 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7bace8e0a9bd..26a3d1975725 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -38,6 +38,8 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ## [Unreleased]
 
+* (types) [#19759](https://github.com/cosmos/cosmos-sdk/pull/19759) Align SignerExtractionAdapter in PriorityNonceMempool Remove.
+
 ## [v0.50.5](https://github.com/cosmos/cosmos-sdk/releases/tag/v0.50.5) - 2024-03-12
 
 ### Features
diff --git a/types/mempool/priority_nonce.go b/types/mempool/priority_nonce.go
index bf996b9a6878..f0df79e70882 100644
--- a/types/mempool/priority_nonce.go
+++ b/types/mempool/priority_nonce.go
@@ -9,7 +9,6 @@ import (
 	"github.com/huandu/skiplist"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	"github.com/cosmos/cosmos-sdk/x/auth/signing"
 )
 
 var (
@@ -432,7 +431,7 @@ func (mp *PriorityNonceMempool[C]) CountTx() int {
 func (mp *PriorityNonceMempool[C]) Remove(tx sdk.Tx) error {
 	mp.mtx.Lock()
 	defer mp.mtx.Unlock()
-	sigs, err := tx.(signing.SigVerifiableTx).GetSignaturesV2()
+	sigs, err := mp.cfg.SignerExtractor.GetSigners(tx)
 	if err != nil {
 		return err
 	}
@@ -441,7 +440,7 @@ func (mp *PriorityNonceMempool[C]) Remove(tx sdk.Tx) error {
 	}
 
 	sig := sigs[0]
-	sender := sdk.AccAddress(sig.PubKey.Address()).String()
+	sender := sig.Signer.String()
 	nonce := sig.Sequence
 
 	scoreKey := txMeta[C]{nonce: nonce, sender: sender}
diff --git a/types/mempool/priority_nonce_test.go b/types/mempool/priority_nonce_test.go
index 462b213201cc..99127f80b889 100644
--- a/types/mempool/priority_nonce_test.go
+++ b/types/mempool/priority_nonce_test.go
@@ -60,6 +60,87 @@ func TestOutOfOrder(t *testing.T) {
 	require.Error(t, validateOrder(rmtxs))
 }
 
+type signerExtractionAdapter struct {
+	UseOld bool
+}
+
+func (a signerExtractionAdapter) GetSigners(tx sdk.Tx) ([]mempool.SignerData, error) {
+	if !a.UseOld {
+		return mempool.NewDefaultSignerExtractionAdapter().GetSigners(tx)
+	}
+	sigs, err := tx.(signing.SigVerifiableTx).GetSignaturesV2()
+	if err != nil {
+		return nil, err
+	}
+	signerData := make([]mempool.SignerData, len(sigs))
+	for _, sig := range sigs {
+		signerData = append(signerData, mempool.SignerData{
+			Signer:   sig.PubKey.Address().Bytes(),
+			Sequence: sig.Sequence,
+		})
+	}
+	return signerData, nil
+}
+
+func (s *MempoolTestSuite) TestPriorityNonceTxOrderWithAdapter() {
+	t := s.T()
+	ctx := sdk.NewContext(nil, false, log.NewNopLogger())
+	accounts := simtypes.RandomAccounts(rand.New(rand.NewSource(0)), 5)
+	sa := accounts[0].Address
+	sb := accounts[1].Address
+
+	tests := []struct {
+		txs   []txSpec
+		order []int
+		fail  bool
+	}{
+		{
+			txs: []txSpec{
+				{p: 21, n: 4, a: sa},
+				{p: 8, n: 3, a: sa},
+				{p: 6, n: 2, a: sa},
+				{p: 15, n: 1, a: sb},
+				{p: 20, n: 1, a: sa},
+			},
+			order: []int{4, 3, 2, 1, 0},
+		},
+	}
+	for i, tt := range tests {
+		t.Run(fmt.Sprintf("case %d", i), func(t *testing.T) {
+			adapter := signerExtractionAdapter{}
+			pool := mempool.NewPriorityMempool(mempool.PriorityNonceMempoolConfig[int64]{
+				TxPriority:      mempool.NewDefaultTxPriority(),
+				SignerExtractor: adapter,
+			})
+
+			// create test txs and insert into mempool
+			for i, ts := range tt.txs {
+				tx := testTx{id: i, priority: int64(ts.p), nonce: uint64(ts.n), address: ts.a}
+				c := ctx.WithPriority(tx.priority)
+				err := pool.Insert(c, tx)
+				require.NoError(t, err)
+			}
+
+			orderedTxs := fetchTxs(pool.Select(ctx, nil), 1000)
+
+			var txOrder []int
+			for _, tx := range orderedTxs {
+				txOrder = append(txOrder, tx.(testTx).id)
+			}
+
+			require.Equal(t, tt.order, txOrder)
+			require.NoError(t, validateOrder(orderedTxs))
+
+			adapter.UseOld = true
+			for _, tx := range orderedTxs {
+				require.NoError(t, pool.Remove(tx))
+			}
+
+			require.NoError(t, mempool.IsEmpty[int64](pool))
+		})
+	}
+}
+
 func (s *MempoolTestSuite) TestPriorityNonceTxOrder() {
 	t := s.T()
 	ctx := sdk.NewContext(nil, cmtproto.Header{}, false, log.NewNopLogger())