fix removeTx to push next queued evm tx (#191)

sei-protocol · Mar 26, 2024 · 2eba403 · 2eba403
1 parent 4a085eb
commit 2eba403
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 3 deletions.
diff --git a/internal/mempool/priority_queue.go b/internal/mempool/priority_queue.go
@@ -121,6 +121,8 @@ func (pq *TxPriorityQueue) removeQueuedEvmTxUnsafe(tx *WrappedTx) {
 				pq.evmQueue[tx.evmAddress] = append(queue[:i], queue[i+1:]...)
 				if len(pq.evmQueue[tx.evmAddress]) == 0 {
 					delete(pq.evmQueue, tx.evmAddress)
+				} else {
+					heap.Push(pq, pq.evmQueue[tx.evmAddress][0])
 				}
 				break
 			}
@@ -195,9 +197,6 @@ func (pq *TxPriorityQueue) popTxUnsafe() *WrappedTx {
 	}
 
 	pq.removeQueuedEvmTxUnsafe(tx)
-	if len(pq.evmQueue[tx.evmAddress]) > 0 {
-		heap.Push(pq, pq.evmQueue[tx.evmAddress][0])
-	}
 
 	return tx
 }

diff --git a/internal/mempool/priority_queue_test.go b/internal/mempool/priority_queue_test.go
@@ -90,6 +90,13 @@ func TestTxPriorityQueue_PriorityAndNonceOrdering(t *testing.T) {
 			},
 			expectedOutput: []int64{12, 13, 11},
 		},
+		{
+			name: "OneItem",
+			inputTxs: []*WrappedTx{
+				{sender: "14", isEVM: true, evmAddress: "0xabc", evmNonce: 1, priority: 10},
+			},
+			expectedOutput: []int64{14},
+		},
 	}
 
 	for _, tc := range testCases {
@@ -263,6 +270,31 @@ func TestTxPriorityQueue_GetEvictableTxs(t *testing.T) {
 	}
 }
 
+func TestTxPriorityQueue_RemoveTxEvm(t *testing.T) {
+	pq := NewTxPriorityQueue()
+
+	tx1 := &WrappedTx{
+		priority:   1,
+		isEVM:      true,
+		evmAddress: "0xabc",
+		evmNonce:   1,
+	}
+	tx2 := &WrappedTx{
+		priority:   1,
+		isEVM:      true,
+		evmAddress: "0xabc",
+		evmNonce:   2,
+	}
+
+	pq.PushTx(tx1)
+	pq.PushTx(tx2)
+
+	pq.RemoveTx(tx1)
+
+	result := pq.PopTx()
+	require.Equal(t, tx2, result)
+}
+
 func TestTxPriorityQueue_RemoveTx(t *testing.T) {
 	pq := NewTxPriorityQueue()
 	rng := rand.New(rand.NewSource(time.Now().UnixNano()))