Add heapIndex with safety check (#213)

* add heapIndex with safety check

* cleanup

* comment out for perf test

* add back perf improvement

* fix nil test

* Use write-lock in (*TxPriorityQueue).ReapMax funcs (#209)

ReapMaxBytesMaxGas and ReapMaxTxs funcs in TxPriorityQueue claim
> Transactions returned are not removed from the mempool transaction
> store or indexes.

However, they use a priority queue to accomplish the claim
> Transaction are retrieved in priority order.

This is accomplished by popping all items out of the whole heap, and
then pushing then back in sequentially. A copy of the heap cannot be
obtained otherwise. Both of the mentioned functions use a read-lock
(RLock) when doing this. This results in a potential scenario where
multiple executions of the ReapMax can be started in parallel, and
both would be popping items out of the priority queue.

In practice, this can be abused by executing the `unconfirmed_txs` RPC
call repeatedly. Based on our observations, running it multiple times
per millisecond results in multiple threads picking it up at the same
time. Such a scenario can be obtained via the WebSocket interface, and
spamming `unconfirmed_txs` calls there. The behavior that happens is a
`Panic in WSJSONRPC handler` when a queue item unexpectedly disappears
for `mempool.(*TxPriorityQueue).Swap`.
(`runtime error: index out of range [0] with length 0`)

This can additionally lead to a `CONSENSUS FAILURE!!!` if the race
condition occurs for `internal/consensus.(*State).finalizeCommit`
when it tries to do `mempool.(*TxPriorityQueue).RemoveTx`, but
the ReapMax has already removed all elements from the underlying
heap. (`runtime error: index out of range [-1]`)

This commit switches the lock type to a write-lock (Lock) to ensure
no parallel modifications take place. This commit additionally updates
the tests to allow parallel execution of the func calls in testing,
as to prevent regressions (in case someone wants to downgrade the locks
without considering the implications from the underlying heap usage).

---------

Co-authored-by: Valters Jansons <sigv@users.noreply.github.com>

Author: stevenlanders

internal/mempool/mempool.go

@@ -417,8 +417,8 @@ func (txmp *TxMempool) Flush() {
 //   - Transactions returned are not removed from the mempool transaction
 //     store or indexes.
 func (txmp *TxMempool) ReapMaxBytesMaxGas(maxBytes, maxGas int64) types.Txs {
-	txmp.mtx.RLock()
-	defer txmp.mtx.RUnlock()
+	txmp.mtx.Lock()
+	defer txmp.mtx.Unlock()
 
 	var (
 		totalGas  int64
@@ -458,8 +458,8 @@ func (txmp *TxMempool) ReapMaxBytesMaxGas(maxBytes, maxGas int64) types.Txs {
 //   - Transactions returned are not removed from the mempool transaction
 //     store or indexes.
 func (txmp *TxMempool) ReapMaxTxs(max int) types.Txs {
-	txmp.mtx.RLock()
-	defer txmp.mtx.RUnlock()
+	txmp.mtx.Lock()
+	defer txmp.mtx.Unlock()
 
 	wTxs := txmp.priorityIndex.PeekTxs(max)
 	txs := make([]types.Tx, 0, len(wTxs))

internal/mempool/mempool_test.go

@@ -379,27 +379,43 @@ func TestTxMempool_ReapMaxBytesMaxGas(t *testing.T) {
 		require.Equal(t, priorities[:len(reapedPriorities)], reapedPriorities)
 	}
 
+	var wg sync.WaitGroup
+
 	// reap by gas capacity only
-	reapedTxs := txmp.ReapMaxBytesMaxGas(-1, 50)
-	ensurePrioritized(reapedTxs)
-	require.Equal(t, len(tTxs), txmp.Size())
-	require.Equal(t, int64(5690), txmp.SizeBytes())
-	require.Len(t, reapedTxs, 50)
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		reapedTxs := txmp.ReapMaxBytesMaxGas(-1, 50)
+		ensurePrioritized(reapedTxs)
+		require.Equal(t, len(tTxs), txmp.Size())
+		require.Equal(t, int64(5690), txmp.SizeBytes())
+		require.Len(t, reapedTxs, 50)
+	}()
 
 	// reap by transaction bytes only
-	reapedTxs = txmp.ReapMaxBytesMaxGas(1000, -1)
-	ensurePrioritized(reapedTxs)
-	require.Equal(t, len(tTxs), txmp.Size())
-	require.Equal(t, int64(5690), txmp.SizeBytes())
-	require.GreaterOrEqual(t, len(reapedTxs), 16)
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		reapedTxs := txmp.ReapMaxBytesMaxGas(1000, -1)
+		ensurePrioritized(reapedTxs)
+		require.Equal(t, len(tTxs), txmp.Size())
+		require.Equal(t, int64(5690), txmp.SizeBytes())
+		require.GreaterOrEqual(t, len(reapedTxs), 16)
+	}()
 
 	// Reap by both transaction bytes and gas, where the size yields 31 reaped
 	// transactions and the gas limit reaps 25 transactions.
-	reapedTxs = txmp.ReapMaxBytesMaxGas(1500, 30)
-	ensurePrioritized(reapedTxs)
-	require.Equal(t, len(tTxs), txmp.Size())
-	require.Equal(t, int64(5690), txmp.SizeBytes())
-	require.Len(t, reapedTxs, 25)
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		reapedTxs := txmp.ReapMaxBytesMaxGas(1500, 30)
+		ensurePrioritized(reapedTxs)
+		require.Equal(t, len(tTxs), txmp.Size())
+		require.Equal(t, int64(5690), txmp.SizeBytes())
+		require.Len(t, reapedTxs, 25)
+	}()
+
+	wg.Wait()
 }
 
 func TestTxMempool_ReapMaxTxs(t *testing.T) {
@@ -438,26 +454,42 @@ func TestTxMempool_ReapMaxTxs(t *testing.T) {
 		require.Equal(t, priorities[:len(reapedPriorities)], reapedPriorities)
 	}
 
+	var wg sync.WaitGroup
+
 	// reap all transactions
-	reapedTxs := txmp.ReapMaxTxs(-1)
-	ensurePrioritized(reapedTxs)
-	require.Equal(t, len(tTxs), txmp.Size())
-	require.Equal(t, int64(5690), txmp.SizeBytes())
-	require.Len(t, reapedTxs, len(tTxs))
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		reapedTxs := txmp.ReapMaxTxs(-1)
+		ensurePrioritized(reapedTxs)
+		require.Equal(t, len(tTxs), txmp.Size())
+		require.Equal(t, int64(5690), txmp.SizeBytes())
+		require.Len(t, reapedTxs, len(tTxs))
+	}()
 
 	// reap a single transaction
-	reapedTxs = txmp.ReapMaxTxs(1)
-	ensurePrioritized(reapedTxs)
-	require.Equal(t, len(tTxs), txmp.Size())
-	require.Equal(t, int64(5690), txmp.SizeBytes())
-	require.Len(t, reapedTxs, 1)
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		reapedTxs := txmp.ReapMaxTxs(1)
+		ensurePrioritized(reapedTxs)
+		require.Equal(t, len(tTxs), txmp.Size())
+		require.Equal(t, int64(5690), txmp.SizeBytes())
+		require.Len(t, reapedTxs, 1)
+	}()
 
 	// reap half of the transactions
-	reapedTxs = txmp.ReapMaxTxs(len(tTxs) / 2)
-	ensurePrioritized(reapedTxs)
-	require.Equal(t, len(tTxs), txmp.Size())
-	require.Equal(t, int64(5690), txmp.SizeBytes())
-	require.Len(t, reapedTxs, len(tTxs)/2)
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		reapedTxs := txmp.ReapMaxTxs(len(tTxs) / 2)
+		ensurePrioritized(reapedTxs)
+		require.Equal(t, len(tTxs), txmp.Size())
+		require.Equal(t, int64(5690), txmp.SizeBytes())
+		require.Len(t, reapedTxs, len(tTxs)/2)
+	}()
+
+	wg.Wait()
 }
 
 func TestTxMempool_CheckTxExceedsMaxSize(t *testing.T) {

internal/mempool/priority_queue.go

@@ -178,6 +178,12 @@ func (pq *TxPriorityQueue) removeQueuedEvmTxUnsafe(tx *WrappedTx) (removedIdx in
 }
 
 func (pq *TxPriorityQueue) findTxIndexUnsafe(tx *WrappedTx) (int, bool) {
+	// safety check for race situation where heapIndex is out of range of txs
+	if tx.heapIndex >= 0 && tx.heapIndex < len(pq.txs) && pq.txs[tx.heapIndex].tx.Key() == tx.tx.Key() {
+		return tx.heapIndex, true
+	}
+
+	// heap index isn't trustable here, so attempt to find it
 	for i, t := range pq.txs {
 		if t.tx.Key() == tx.tx.Key() {
 			return i, true
@@ -443,7 +449,9 @@ func (pq *TxPriorityQueue) PeekTxs(max int) []*WrappedTx {
 //
 // NOTE: A caller should never call Push. Use PushTx instead.
 func (pq *TxPriorityQueue) Push(x interface{}) {
+	n := len(pq.txs)
 	item := x.(*WrappedTx)
+	item.heapIndex = n
 	pq.txs = append(pq.txs, item)
 }
 
@@ -454,7 +462,8 @@ func (pq *TxPriorityQueue) Pop() interface{} {
 	old := pq.txs
 	n := len(old)
 	item := old[n-1]
-	old[n-1] = nil // avoid memory leak
+	old[n-1] = nil         // avoid memory leak
+	setHeapIndex(item, -1) // for safety
 	pq.txs = old[0 : n-1]
 	return item
 }
@@ -483,4 +492,14 @@ func (pq *TxPriorityQueue) Less(i, j int) bool {
 // Swap implements the Heap interface. It swaps two transactions in the queue.
 func (pq *TxPriorityQueue) Swap(i, j int) {
 	pq.txs[i], pq.txs[j] = pq.txs[j], pq.txs[i]
+	setHeapIndex(pq.txs[i], i)
+	setHeapIndex(pq.txs[j], j)
+}
+
+func setHeapIndex(tx *WrappedTx, i int) {
+	// a removed tx can be nil
+	if tx == nil {
+		return
+	}
+	tx.heapIndex = i
 }