Added txpool UTs for proposed queue

Ethernal-Tech · Apr 25, 2024 · 7c84654 · 7c84654
1 parent b633fb1
commit 7c84654
Show file tree

Hide file tree

Showing 3 changed files with 217 additions and 2 deletions.
diff --git a/txpool/account.go b/txpool/account.go
@@ -132,7 +132,9 @@ func (m *accountsMap) allTxs(includeEnqueued bool) (
 	return
 }
 
-func (m *accountsMap) reinsertProposed() {
+func (m *accountsMap) reinsertProposed() uint64 {
+	var count uint64
+
 	m.Range(func(key, value interface{}) bool {
 		accountKey, ok := key.(types.Address)
 		if !ok {
@@ -162,6 +164,8 @@ func (m *accountsMap) reinsertProposed() {
 
 				account.promoted.push(tx)
 				account.proposed.pop()
+
+				count++
 			}
 			account.proposed.clear()
 
@@ -171,6 +175,8 @@ func (m *accountsMap) reinsertProposed() {
 
 		return true
 	})
+
+	return count
 }
 
 func (m *accountsMap) clearProposed() {

diff --git a/txpool/txpool.go b/txpool/txpool.go
@@ -545,7 +545,8 @@ func (p *TxPool) ResetWithBlock(block *types.Block) {
 }
 
 func (p *TxPool) ReinsertProposed() {
-	p.accounts.reinsertProposed()
+	count := p.accounts.reinsertProposed()
+	p.gauge.increase(count)
 	p.Prepare()
 }
 

diff --git a/txpool/txpool_test.go b/txpool/txpool_test.go
@@ -109,6 +109,7 @@ func newTestPoolWithSlots(maxSlots uint64, mockStore ...store) (*TxPool, error)
 type accountState struct {
 	enqueued,
 	promoted,
+	proposed,
 	nextNonce uint64
 }
 
@@ -3188,6 +3189,213 @@ func TestRecovery(t *testing.T) {
 	}
 }
 
+func TestProposed(t *testing.T) {
+	commonAssert := func(accounts map[types.Address]accountState, pool *TxPool) {
+		for addr := range accounts {
+			assert.Equal(t, // proposed
+				accounts[addr].proposed,
+				pool.accounts.get(addr).proposed.length())
+
+			assert.Equal(t, // promoted
+				accounts[addr].promoted,
+				pool.accounts.get(addr).promoted.length())
+		}
+	}
+
+	const (
+		REINSERT = 1
+		CLEAN    = 2
+	)
+
+	testCases := []struct {
+		name       string
+		method     int
+		doPop      bool
+		allTxs     map[types.Address][]*types.Transaction
+		beforeCall result
+		afterCall  result
+	}{
+		{
+			name:   "reinsert with pop",
+			method: REINSERT,
+			doPop:  true,
+			allTxs: map[types.Address][]*types.Transaction{
+				addr1: {
+					newTx(addr1, 0, 1, types.LegacyTxType),
+					newTx(addr1, 1, 1, types.LegacyTxType),
+				},
+			},
+			beforeCall: result{
+				slots: 1,
+				accounts: map[types.Address]accountState{
+					addr1: {
+						proposed: 1,
+						promoted: 1,
+					},
+				},
+			},
+			afterCall: result{
+				slots: 2,
+				accounts: map[types.Address]accountState{
+					addr1: {
+						proposed: 0,
+						promoted: 2,
+					},
+				},
+			},
+		},
+		{
+			name:   "reinsert without pop",
+			method: REINSERT,
+			doPop:  false,
+			allTxs: map[types.Address][]*types.Transaction{
+				addr1: {
+					newTx(addr1, 0, 1, types.LegacyTxType),
+					newTx(addr1, 1, 1, types.LegacyTxType),
+				},
+			},
+			beforeCall: result{
+				slots: 2,
+				accounts: map[types.Address]accountState{
+					addr1: {
+						proposed: 0,
+						promoted: 2,
+					},
+				},
+			},
+			afterCall: result{
+				slots: 2,
+				accounts: map[types.Address]accountState{
+					addr1: {
+						proposed: 0,
+						promoted: 2,
+					},
+				},
+			},
+		},
+		{
+			name:   "clean with pop",
+			method: CLEAN,
+			doPop:  true,
+			allTxs: map[types.Address][]*types.Transaction{
+				addr1: {
+					newTx(addr1, 0, 1, types.LegacyTxType),
+					newTx(addr1, 1, 1, types.LegacyTxType),
+				},
+			},
+			beforeCall: result{
+				slots: 1,
+				accounts: map[types.Address]accountState{
+					addr1: {
+						proposed: 1,
+						promoted: 1,
+					},
+				},
+			},
+			afterCall: result{
+				slots: 1,
+				accounts: map[types.Address]accountState{
+					addr1: {
+						proposed: 0,
+						promoted: 1,
+					},
+				},
+			},
+		},
+		{
+			name:   "clean without pop",
+			method: CLEAN,
+			doPop:  false,
+			allTxs: map[types.Address][]*types.Transaction{
+				addr1: {
+					newTx(addr1, 0, 1, types.LegacyTxType),
+					newTx(addr1, 1, 1, types.LegacyTxType),
+				},
+			},
+			beforeCall: result{
+				slots: 2,
+				accounts: map[types.Address]accountState{
+					addr1: {
+						proposed: 0,
+						promoted: 2,
+					},
+				},
+			},
+			afterCall: result{
+				slots: 2,
+				accounts: map[types.Address]accountState{
+					addr1: {
+						proposed: 0,
+						promoted: 2,
+					},
+				},
+			},
+		},
+	}
+
+	for _, test := range testCases {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			// create pool
+			pool, err := newTestPool()
+			assert.NoError(t, err)
+			pool.SetSigner(&mockSigner{})
+
+			pool.Start()
+			defer pool.Close()
+
+			promoteSubscription := pool.eventManager.subscribe(
+				[]proto.EventType{proto.EventType_PROMOTED},
+			)
+
+			// setup prestate
+			totalTx := 0
+			expectedEnqueued := uint64(0)
+
+			for addr, txs := range test.allTxs {
+				// preset nonce so promotions can happen
+				acc := pool.getOrCreateAccount(addr)
+				acc.setNonce(txs[0].Nonce())
+
+				expectedEnqueued += test.afterCall.accounts[addr].enqueued
+
+				// send txs
+				for _, tx := range txs {
+					totalTx++
+
+					assert.NoError(t, pool.addTx(local, tx))
+				}
+			}
+
+			ctx, cancelFn := context.WithTimeout(context.Background(), time.Second*10)
+			defer cancelFn()
+
+			// All txns should get added
+			assert.Len(t, waitForEvents(ctx, promoteSubscription, totalTx), totalTx)
+
+			pool.Prepare()
+			tx := pool.Peek()
+			assert.NotNil(t, tx)
+
+			if test.doPop {
+				pool.Pop(tx)
+			}
+
+			assert.Equal(t, test.beforeCall.slots, pool.gauge.read())
+			commonAssert(test.beforeCall.accounts, pool)
+
+			if test.method == REINSERT {
+				pool.ReinsertProposed()
+			} else {
+				pool.ClearProposed()
+			}
+
+			assert.Equal(t, test.afterCall.slots, pool.gauge.read())
+			commonAssert(test.afterCall.accounts, pool)
+		})
+	}
+}
+
 func TestGetTxs(t *testing.T) {
 	t.Parallel()