diff --git a/types/mempool/priority_nonce.go b/types/mempool/priority_nonce.go index 216cd8148cae..f081e2b413db 100644 --- a/types/mempool/priority_nonce.go +++ b/types/mempool/priority_nonce.go @@ -351,9 +351,13 @@ func (i *PriorityNonceIterator[C]) Tx() sdk.Tx { // // NOTE: It is not safe to use this iterator while removing transactions from // the underlying mempool. -func (mp *PriorityNonceMempool[C]) Select(_ context.Context, _ [][]byte) Iterator { +func (mp *PriorityNonceMempool[C]) Select(ctx context.Context, txs [][]byte) Iterator { mp.mtx.Lock() defer mp.mtx.Unlock() + return mp.doSelect(ctx, txs) +} + +func (mp *PriorityNonceMempool[C]) doSelect(_ context.Context, _ [][]byte) Iterator { if mp.priorityIndex.Len() == 0 { return nil } @@ -369,21 +373,11 @@ func (mp *PriorityNonceMempool[C]) Select(_ context.Context, _ [][]byte) Iterato } // SelectBy will hold the mutex during the iteration, callback returns if continue. -func (mp *PriorityNonceMempool[C]) SelectBy(_ context.Context, _ [][]byte, callback func(sdk.Tx) bool) { +func (mp *PriorityNonceMempool[C]) SelectBy(ctx context.Context, txs [][]byte, callback func(sdk.Tx) bool) { mp.mtx.Lock() defer mp.mtx.Unlock() - if mp.priorityIndex.Len() == 0 { - return - } - - mp.reorderPriorityTies() - - iterator := &PriorityNonceIterator[C]{ - mempool: mp, - senderCursors: make(map[string]*skiplist.Element), - } - iter := iterator.iteratePriority() + iter := mp.doSelect(ctx, txs) for iter != nil && callback(iter.Tx()) { iter = iter.Next() } diff --git a/types/mempool/sender_nonce.go b/types/mempool/sender_nonce.go index 371e701d280a..00f554f26216 100644 --- a/types/mempool/sender_nonce.go +++ b/types/mempool/sender_nonce.go @@ -159,9 +159,13 @@ func (snm *SenderNonceMempool) Insert(_ context.Context, tx sdk.Tx) error { // // NOTE: It is not safe to use this iterator while removing transactions from // the underlying mempool. -func (snm *SenderNonceMempool) Select(_ context.Context, _ [][]byte) Iterator { +func (snm *SenderNonceMempool) Select(ctx context.Context, txs [][]byte) Iterator { snm.mtx.Lock() defer snm.mtx.Unlock() + return snm.doSelect(ctx, txs) +} + +func (snm *SenderNonceMempool) doSelect(_ context.Context, _ [][]byte) Iterator { var senders []string senderCursors := make(map[string]*skiplist.Element) @@ -190,34 +194,11 @@ func (snm *SenderNonceMempool) Select(_ context.Context, _ [][]byte) Iterator { } // SelectBy will hold the mutex during the iteration, callback returns if continue. -func (snm *SenderNonceMempool) SelectBy(_ context.Context, _ [][]byte, callback func(sdk.Tx) bool) { +func (snm *SenderNonceMempool) SelectBy(ctx context.Context, txs [][]byte, callback func(sdk.Tx) bool) { snm.mtx.Lock() defer snm.mtx.Unlock() - var senders []string - - senderCursors := make(map[string]*skiplist.Element) - orderedSenders := skiplist.New(skiplist.String) - - // #nosec - for s := range snm.senders { - orderedSenders.Set(s, s) - } - - s := orderedSenders.Front() - for s != nil { - sender := s.Value.(string) - senders = append(senders, sender) - senderCursors[sender] = snm.senders[sender].Front() - s = s.Next() - } - - iterator := &senderNonceMempoolIterator{ - senders: senders, - rnd: snm.rnd, - senderCursors: senderCursors, - } - iter := iterator.Next() + iter := snm.doSelect(ctx, txs) for iter != nil && callback(iter.Tx()) { iter = iter.Next() }