You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

176 lines
3.5 KiB

  1. package v1
  2. import (
  3. "math/rand"
  4. "sort"
  5. "sync"
  6. "testing"
  7. "time"
  8. "github.com/stretchr/testify/require"
  9. )
  10. func TestTxPriorityQueue(t *testing.T) {
  11. pq := NewTxPriorityQueue()
  12. numTxs := 1000
  13. priorities := make([]int, numTxs)
  14. var wg sync.WaitGroup
  15. for i := 1; i <= numTxs; i++ {
  16. priorities[i-1] = i
  17. wg.Add(1)
  18. go func(i int) {
  19. pq.PushTx(&WrappedTx{
  20. priority: int64(i),
  21. timestamp: time.Now(),
  22. })
  23. wg.Done()
  24. }(i)
  25. }
  26. sort.Sort(sort.Reverse(sort.IntSlice(priorities)))
  27. wg.Wait()
  28. require.Equal(t, numTxs, pq.NumTxs())
  29. // Wait a second and push a tx with a duplicate priority
  30. time.Sleep(time.Second)
  31. now := time.Now()
  32. pq.PushTx(&WrappedTx{
  33. priority: 1000,
  34. timestamp: now,
  35. })
  36. require.Equal(t, 1001, pq.NumTxs())
  37. tx := pq.PopTx()
  38. require.Equal(t, 1000, pq.NumTxs())
  39. require.Equal(t, int64(1000), tx.priority)
  40. require.NotEqual(t, now, tx.timestamp)
  41. gotPriorities := make([]int, 0)
  42. for pq.NumTxs() > 0 {
  43. gotPriorities = append(gotPriorities, int(pq.PopTx().priority))
  44. }
  45. require.Equal(t, priorities, gotPriorities)
  46. }
  47. func TestTxPriorityQueue_GetEvictableTxs(t *testing.T) {
  48. pq := NewTxPriorityQueue()
  49. rng := rand.New(rand.NewSource(time.Now().UnixNano()))
  50. values := make([]int, 1000)
  51. for i := 0; i < 1000; i++ {
  52. tx := make([]byte, 5) // each tx is 5 bytes
  53. _, err := rng.Read(tx)
  54. require.NoError(t, err)
  55. x := rng.Intn(100000)
  56. pq.PushTx(&WrappedTx{
  57. tx: tx,
  58. priority: int64(x),
  59. })
  60. values[i] = x
  61. }
  62. sort.Ints(values)
  63. max := values[len(values)-1]
  64. min := values[0]
  65. totalSize := int64(len(values) * 5)
  66. testCases := []struct {
  67. name string
  68. priority, txSize, totalSize, cap int64
  69. expectedLen int
  70. }{
  71. {
  72. name: "larest priority; single tx",
  73. priority: int64(max + 1),
  74. txSize: 5,
  75. totalSize: totalSize,
  76. cap: totalSize,
  77. expectedLen: 1,
  78. },
  79. {
  80. name: "larest priority; multi tx",
  81. priority: int64(max + 1),
  82. txSize: 17,
  83. totalSize: totalSize,
  84. cap: totalSize,
  85. expectedLen: 4,
  86. },
  87. {
  88. name: "larest priority; out of capacity",
  89. priority: int64(max + 1),
  90. txSize: totalSize + 1,
  91. totalSize: totalSize,
  92. cap: totalSize,
  93. expectedLen: 0,
  94. },
  95. {
  96. name: "smallest priority; no tx",
  97. priority: int64(min - 1),
  98. txSize: 5,
  99. totalSize: totalSize,
  100. cap: totalSize,
  101. expectedLen: 0,
  102. },
  103. {
  104. name: "small priority; no tx",
  105. priority: int64(min),
  106. txSize: 5,
  107. totalSize: totalSize,
  108. cap: totalSize,
  109. expectedLen: 0,
  110. },
  111. }
  112. for _, tc := range testCases {
  113. tc := tc
  114. t.Run(tc.name, func(t *testing.T) {
  115. evictTxs := pq.GetEvictableTxs(tc.priority, tc.txSize, tc.totalSize, tc.cap)
  116. require.Len(t, evictTxs, tc.expectedLen)
  117. })
  118. }
  119. }
  120. func TestTxPriorityQueue_RemoveTx(t *testing.T) {
  121. pq := NewTxPriorityQueue()
  122. rng := rand.New(rand.NewSource(time.Now().UnixNano()))
  123. numTxs := 1000
  124. values := make([]int, numTxs)
  125. for i := 0; i < numTxs; i++ {
  126. x := rng.Intn(100000)
  127. pq.PushTx(&WrappedTx{
  128. priority: int64(x),
  129. })
  130. values[i] = x
  131. }
  132. require.Equal(t, numTxs, pq.NumTxs())
  133. sort.Ints(values)
  134. max := values[len(values)-1]
  135. wtx := pq.txs[pq.NumTxs()/2]
  136. pq.RemoveTx(wtx)
  137. require.Equal(t, numTxs-1, pq.NumTxs())
  138. require.Equal(t, int64(max), pq.PopTx().priority)
  139. require.Equal(t, numTxs-2, pq.NumTxs())
  140. require.NotPanics(t, func() {
  141. pq.RemoveTx(&WrappedTx{heapIndex: numTxs})
  142. pq.RemoveTx(&WrappedTx{heapIndex: numTxs + 1})
  143. })
  144. require.Equal(t, numTxs-2, pq.NumTxs())
  145. }