common/prque: generic priority queue (#26290)

* common, core, eth, les, trie: make prque generic * les/vflux/server: fixed issues in priorityPool * common, core, eth, les, trie: make priority also generic in prque * les/flowcontrol: add test case for priority accumulator overflow * les/flowcontrol: avoid priority value overflow * common/prque: use int priority in some tests No need to convert to int64 when we can just change the type used by the queue. * common/prque: remove comment about int64 range --------- Co-authored-by: Zsolt Felfoldi <zsfelfoldi@gmail.com> Co-authored-by: Felix Lange <fjl@twurst.com>
2023-02-09 13:03:54 +02:00
parent 6a148dd5c3
commit bf1798e04e
20 changed files with 277 additions and 264 deletions
@@ -21,6 +21,7 @@ import (
 	"time"

 	"github.com/ethereum/go-ethereum/common/mclock"
+	"golang.org/x/exp/constraints"
 )

 // LazyQueue is a priority queue data structure where priorities can change over
@@ -32,31 +33,31 @@ import (
 //
 // If the upper estimate is exceeded then Update should be called for that item.
 // A global Refresh function should also be called periodically.
-type LazyQueue struct {
+type LazyQueue[P constraints.Ordered, V any] struct {
 	clock mclock.Clock
 	// Items are stored in one of two internal queues ordered by estimated max
 	// priority until the next and the next-after-next refresh. Update and Refresh
 	// always places items in queue[1].
-	queue                      [2]*sstack
-	popQueue                   *sstack
+	queue                      [2]*sstack[P, V]
+	popQueue                   *sstack[P, V]
 	period                     time.Duration
 	maxUntil                   mclock.AbsTime
 	indexOffset                int
-	setIndex                   SetIndexCallback
-	priority                   PriorityCallback
-	maxPriority                MaxPriorityCallback
+	setIndex                   SetIndexCallback[V]
+	priority                   PriorityCallback[P, V]
+	maxPriority                MaxPriorityCallback[P, V]
 	lastRefresh1, lastRefresh2 mclock.AbsTime
 }

 type (
-	PriorityCallback    func(data interface{}) int64                       // actual priority callback
-	MaxPriorityCallback func(data interface{}, until mclock.AbsTime) int64 // estimated maximum priority callback
+	PriorityCallback[P constraints.Ordered, V any]    func(data V) P                       // actual priority callback
+	MaxPriorityCallback[P constraints.Ordered, V any] func(data V, until mclock.AbsTime) P // estimated maximum priority callback
 )

 // NewLazyQueue creates a new lazy queue
-func NewLazyQueue(setIndex SetIndexCallback, priority PriorityCallback, maxPriority MaxPriorityCallback, clock mclock.Clock, refreshPeriod time.Duration) *LazyQueue {
-	q := &LazyQueue{
-		popQueue:     newSstack(nil, false),
+func NewLazyQueue[P constraints.Ordered, V any](setIndex SetIndexCallback[V], priority PriorityCallback[P, V], maxPriority MaxPriorityCallback[P, V], clock mclock.Clock, refreshPeriod time.Duration) *LazyQueue[P, V] {
+	q := &LazyQueue[P, V]{
+		popQueue:     newSstack[P, V](nil),
 		setIndex:     setIndex,
 		priority:     priority,
 		maxPriority:  maxPriority,
@@ -71,13 +72,13 @@ func NewLazyQueue(setIndex SetIndexCallback, priority PriorityCallback, maxPrior
 }

 // Reset clears the contents of the queue
-func (q *LazyQueue) Reset() {
-	q.queue[0] = newSstack(q.setIndex0, false)
-	q.queue[1] = newSstack(q.setIndex1, false)
+func (q *LazyQueue[P, V]) Reset() {
+	q.queue[0] = newSstack[P, V](q.setIndex0)
+	q.queue[1] = newSstack[P, V](q.setIndex1)
 }

 // Refresh performs queue re-evaluation if necessary
-func (q *LazyQueue) Refresh() {
+func (q *LazyQueue[P, V]) Refresh() {
 	now := q.clock.Now()
 	for time.Duration(now-q.lastRefresh2) >= q.period*2 {
 		q.refresh(now)
@@ -87,10 +88,10 @@ func (q *LazyQueue) Refresh() {
 }

 // refresh re-evaluates items in the older queue and swaps the two queues
-func (q *LazyQueue) refresh(now mclock.AbsTime) {
+func (q *LazyQueue[P, V]) refresh(now mclock.AbsTime) {
 	q.maxUntil = now.Add(q.period)
 	for q.queue[0].Len() != 0 {
-		q.Push(heap.Pop(q.queue[0]).(*item).value)
+		q.Push(heap.Pop(q.queue[0]).(*item[P, V]).value)
 	}
 	q.queue[0], q.queue[1] = q.queue[1], q.queue[0]
 	q.indexOffset = 1 - q.indexOffset
@@ -98,22 +99,22 @@ func (q *LazyQueue) refresh(now mclock.AbsTime) {
 }

 // Push adds an item to the queue
-func (q *LazyQueue) Push(data interface{}) {
-	heap.Push(q.queue[1], &item{data, q.maxPriority(data, q.maxUntil)})
+func (q *LazyQueue[P, V]) Push(data V) {
+	heap.Push(q.queue[1], &item[P, V]{data, q.maxPriority(data, q.maxUntil)})
 }

 // Update updates the upper priority estimate for the item with the given queue index
-func (q *LazyQueue) Update(index int) {
+func (q *LazyQueue[P, V]) Update(index int) {
 	q.Push(q.Remove(index))
 }

 // Pop removes and returns the item with the greatest actual priority
-func (q *LazyQueue) Pop() (interface{}, int64) {
+func (q *LazyQueue[P, V]) Pop() (V, P) {
 	var (
-		resData interface{}
-		resPri  int64
+		resData V
+		resPri  P
 	)
-	q.MultiPop(func(data interface{}, priority int64) bool {
+	q.MultiPop(func(data V, priority P) bool {
 		resData = data
 		resPri = priority
 		return false
@@ -123,7 +124,7 @@ func (q *LazyQueue) Pop() (interface{}, int64) {

 // peekIndex returns the index of the internal queue where the item with the
 // highest estimated priority is or -1 if both are empty
-func (q *LazyQueue) peekIndex() int {
+func (q *LazyQueue[P, V]) peekIndex() int {
 	if q.queue[0].Len() != 0 {
 		if q.queue[1].Len() != 0 && q.queue[1].blocks[0][0].priority > q.queue[0].blocks[0][0].priority {
 			return 1
@@ -139,17 +140,17 @@ func (q *LazyQueue) peekIndex() int {
 // MultiPop pops multiple items from the queue and is more efficient than calling
 // Pop multiple times. Popped items are passed to the callback. MultiPop returns
 // when the callback returns false or there are no more items to pop.
-func (q *LazyQueue) MultiPop(callback func(data interface{}, priority int64) bool) {
+func (q *LazyQueue[P, V]) MultiPop(callback func(data V, priority P) bool) {
 	nextIndex := q.peekIndex()
 	for nextIndex != -1 {
-		data := heap.Pop(q.queue[nextIndex]).(*item).value
-		heap.Push(q.popQueue, &item{data, q.priority(data)})
+		data := heap.Pop(q.queue[nextIndex]).(*item[P, V]).value
+		heap.Push(q.popQueue, &item[P, V]{data, q.priority(data)})
 		nextIndex = q.peekIndex()
 		for q.popQueue.Len() != 0 && (nextIndex == -1 || q.queue[nextIndex].blocks[0][0].priority < q.popQueue.blocks[0][0].priority) {
-			i := heap.Pop(q.popQueue).(*item)
+			i := heap.Pop(q.popQueue).(*item[P, V])
 			if !callback(i.value, i.priority) {
 				for q.popQueue.Len() != 0 {
-					q.Push(heap.Pop(q.popQueue).(*item).value)
+					q.Push(heap.Pop(q.popQueue).(*item[P, V]).value)
 				}
 				return
 			}
@@ -159,31 +160,28 @@ func (q *LazyQueue) MultiPop(callback func(data interface{}, priority int64) boo
 }

 // PopItem pops the item from the queue only, dropping the associated priority value.
-func (q *LazyQueue) PopItem() interface{} {
+func (q *LazyQueue[P, V]) PopItem() V {
 	i, _ := q.Pop()
 	return i
 }

 // Remove removes the item with the given index.
-func (q *LazyQueue) Remove(index int) interface{} {
-	if index < 0 {
-		return nil
-	}
-	return heap.Remove(q.queue[index&1^q.indexOffset], index>>1).(*item).value
+func (q *LazyQueue[P, V]) Remove(index int) V {
+	return heap.Remove(q.queue[index&1^q.indexOffset], index>>1).(*item[P, V]).value
 }

 // Empty checks whether the priority queue is empty.
-func (q *LazyQueue) Empty() bool {
+func (q *LazyQueue[P, V]) Empty() bool {
 	return q.queue[0].Len() == 0 && q.queue[1].Len() == 0
 }

 // Size returns the number of items in the priority queue.
-func (q *LazyQueue) Size() int {
+func (q *LazyQueue[P, V]) Size() int {
 	return q.queue[0].Len() + q.queue[1].Len()
 }

 // setIndex0 translates internal queue item index to the virtual index space of LazyQueue
-func (q *LazyQueue) setIndex0(data interface{}, index int) {
+func (q *LazyQueue[P, V]) setIndex0(data V, index int) {
 	if index == -1 {
 		q.setIndex(data, -1)
 	} else {
@@ -192,6 +190,6 @@ func (q *LazyQueue) setIndex0(data interface{}, index int) {
 }

 // setIndex1 translates internal queue item index to the virtual index space of LazyQueue
-func (q *LazyQueue) setIndex1(data interface{}, index int) {
+func (q *LazyQueue[P, V]) setIndex1(data V, index int) {
 	q.setIndex(data, index+index+1)
 }
@@ -19,65 +19,59 @@ package prque

 import (
 	"container/heap"
+
+	"golang.org/x/exp/constraints"
 )

 // Priority queue data structure.
-type Prque struct {
-	cont *sstack
+type Prque[P constraints.Ordered, V any] struct {
+	cont *sstack[P, V]
 }

 // New creates a new priority queue.
-func New(setIndex SetIndexCallback) *Prque {
-	return &Prque{newSstack(setIndex, false)}
-}
-
-// NewWrapAround creates a new priority queue with wrap-around priority handling.
-func NewWrapAround(setIndex SetIndexCallback) *Prque {
-	return &Prque{newSstack(setIndex, true)}
+func New[P constraints.Ordered, V any](setIndex SetIndexCallback[V]) *Prque[P, V] {
+	return &Prque[P, V]{newSstack[P, V](setIndex)}
 }

 // Pushes a value with a given priority into the queue, expanding if necessary.
-func (p *Prque) Push(data interface{}, priority int64) {
-	heap.Push(p.cont, &item{data, priority})
+func (p *Prque[P, V]) Push(data V, priority P) {
+	heap.Push(p.cont, &item[P, V]{data, priority})
 }

 // Peek returns the value with the greatest priority but does not pop it off.
-func (p *Prque) Peek() (interface{}, int64) {
+func (p *Prque[P, V]) Peek() (V, P) {
 	item := p.cont.blocks[0][0]
 	return item.value, item.priority
 }

 // Pops the value with the greatest priority off the stack and returns it.
 // Currently no shrinking is done.
-func (p *Prque) Pop() (interface{}, int64) {
-	item := heap.Pop(p.cont).(*item)
+func (p *Prque[P, V]) Pop() (V, P) {
+	item := heap.Pop(p.cont).(*item[P, V])
 	return item.value, item.priority
 }

 // Pops only the item from the queue, dropping the associated priority value.
-func (p *Prque) PopItem() interface{} {
-	return heap.Pop(p.cont).(*item).value
+func (p *Prque[P, V]) PopItem() V {
+	return heap.Pop(p.cont).(*item[P, V]).value
 }

 // Remove removes the element with the given index.
-func (p *Prque) Remove(i int) interface{} {
-	if i < 0 {
-		return nil
-	}
-	return heap.Remove(p.cont, i)
+func (p *Prque[P, V]) Remove(i int) V {
+	return heap.Remove(p.cont, i).(*item[P, V]).value
 }

 // Checks whether the priority queue is empty.
-func (p *Prque) Empty() bool {
+func (p *Prque[P, V]) Empty() bool {
 	return p.cont.Len() == 0
 }

 // Returns the number of element in the priority queue.
-func (p *Prque) Size() int {
+func (p *Prque[P, V]) Size() int {
 	return p.cont.Len()
 }

 // Clears the contents of the priority queue.
-func (p *Prque) Reset() {
-	*p = *New(p.cont.setIndex)
+func (p *Prque[P, V]) Reset() {
+	*p = *New[P, V](p.cont.setIndex)
 }
@@ -21,22 +21,24 @@ func TestPrque(t *testing.T) {
 	for i := 0; i < size; i++ {
 		data[i] = rand.Int()
 	}
-	queue := New(nil)
+	queue := New[int, int](nil)
+
 	for rep := 0; rep < 2; rep++ {
 		// Fill a priority queue with the above data
 		for i := 0; i < size; i++ {
-			queue.Push(data[i], int64(prio[i]))
+			queue.Push(data[i], prio[i])
 			if queue.Size() != i+1 {
 				t.Errorf("queue size mismatch: have %v, want %v.", queue.Size(), i+1)
 			}
 		}
 		// Create a map the values to the priorities for easier verification
-		dict := make(map[int64]int)
+		dict := make(map[int]int)
 		for i := 0; i < size; i++ {
-			dict[int64(prio[i])] = data[i]
+			dict[prio[i]] = data[i]
 		}
+
 		// Pop out the elements in priority order and verify them
-		prevPrio := int64(size + 1)
+		prevPrio := size + 1
 		for !queue.Empty() {
 			val, prio := queue.Pop()
 			if prio > prevPrio {
@@ -59,22 +61,23 @@ func TestReset(t *testing.T) {
 	for i := 0; i < size; i++ {
 		data[i] = rand.Int()
 	}
-	queue := New(nil)
+	queue := New[int, int](nil)
+
 	for rep := 0; rep < 2; rep++ {
 		// Fill a priority queue with the above data
 		for i := 0; i < size; i++ {
-			queue.Push(data[i], int64(prio[i]))
+			queue.Push(data[i], prio[i])
 			if queue.Size() != i+1 {
 				t.Errorf("queue size mismatch: have %v, want %v.", queue.Size(), i+1)
 			}
 		}
 		// Create a map the values to the priorities for easier verification
-		dict := make(map[int64]int)
+		dict := make(map[int]int)
 		for i := 0; i < size; i++ {
-			dict[int64(prio[i])] = data[i]
+			dict[prio[i]] = data[i]
 		}
 		// Pop out half the elements in priority order and verify them
-		prevPrio := int64(size + 1)
+		prevPrio := size + 1
 		for i := 0; i < size/2; i++ {
 			val, prio := queue.Pop()
 			if prio > prevPrio {
@@ -104,7 +107,7 @@ func BenchmarkPush(b *testing.B) {
 	}
 	// Execute the benchmark
 	b.ResetTimer()
-	queue := New(nil)
+	queue := New[int64, int](nil)
 	for i := 0; i < len(data); i++ {
 		queue.Push(data[i], prio[i])
 	}
@@ -118,7 +121,7 @@ func BenchmarkPop(b *testing.B) {
 		data[i] = rand.Int()
 		prio[i] = rand.Int63()
 	}
-	queue := New(nil)
+	queue := New[int64, int](nil)
 	for i := 0; i < len(data); i++ {
 		queue.Push(data[i], prio[i])
 	}
@@ -10,53 +10,50 @@

 package prque

+import "golang.org/x/exp/constraints"
+
 // The size of a block of data
 const blockSize = 4096

 // A prioritized item in the sorted stack.
-//
-// Note: priorities can "wrap around" the int64 range, a comes before b if (a.priority - b.priority) > 0.
-// The difference between the lowest and highest priorities in the queue at any point should be less than 2^63.
-type item struct {
-	value    interface{}
-	priority int64
+type item[P constraints.Ordered, V any] struct {
+	value    V
+	priority P
 }

 // SetIndexCallback is called when the element is moved to a new index.
 // Providing SetIndexCallback is optional, it is needed only if the application needs
 // to delete elements other than the top one.
-type SetIndexCallback func(data interface{}, index int)
+type SetIndexCallback[V any] func(data V, index int)

 // Internal sortable stack data structure. Implements the Push and Pop ops for
 // the stack (heap) functionality and the Len, Less and Swap methods for the
 // sortability requirements of the heaps.
-type sstack struct {
-	setIndex   SetIndexCallback
-	size       int
-	capacity   int
-	offset     int
-	wrapAround bool
+type sstack[P constraints.Ordered, V any] struct {
+	setIndex SetIndexCallback[V]
+	size     int
+	capacity int
+	offset   int

-	blocks [][]*item
-	active []*item
+	blocks [][]*item[P, V]
+	active []*item[P, V]
 }

 // Creates a new, empty stack.
-func newSstack(setIndex SetIndexCallback, wrapAround bool) *sstack {
-	result := new(sstack)
+func newSstack[P constraints.Ordered, V any](setIndex SetIndexCallback[V]) *sstack[P, V] {
+	result := new(sstack[P, V])
 	result.setIndex = setIndex
-	result.active = make([]*item, blockSize)
-	result.blocks = [][]*item{result.active}
+	result.active = make([]*item[P, V], blockSize)
+	result.blocks = [][]*item[P, V]{result.active}
 	result.capacity = blockSize
-	result.wrapAround = wrapAround
 	return result
 }

 // Pushes a value onto the stack, expanding it if necessary. Required by
 // heap.Interface.
-func (s *sstack) Push(data interface{}) {
+func (s *sstack[P, V]) Push(data any) {
 	if s.size == s.capacity {
-		s.active = make([]*item, blockSize)
+		s.active = make([]*item[P, V], blockSize)
 		s.blocks = append(s.blocks, s.active)
 		s.capacity += blockSize
 		s.offset = 0
@@ -65,16 +62,16 @@ func (s *sstack) Push(data interface{}) {
 		s.offset = 0
 	}
 	if s.setIndex != nil {
-		s.setIndex(data.(*item).value, s.size)
+		s.setIndex(data.(*item[P, V]).value, s.size)
 	}
-	s.active[s.offset] = data.(*item)
+	s.active[s.offset] = data.(*item[P, V])
 	s.offset++
 	s.size++
 }

 // Pops a value off the stack and returns it. Currently no shrinking is done.
 // Required by heap.Interface.
-func (s *sstack) Pop() (res interface{}) {
+func (s *sstack[P, V]) Pop() (res any) {
 	s.size--
 	s.offset--
 	if s.offset < 0 {
@@ -83,28 +80,24 @@ func (s *sstack) Pop() (res interface{}) {
 	}
 	res, s.active[s.offset] = s.active[s.offset], nil
 	if s.setIndex != nil {
-		s.setIndex(res.(*item).value, -1)
+		s.setIndex(res.(*item[P, V]).value, -1)
 	}
 	return
 }

 // Returns the length of the stack. Required by sort.Interface.
-func (s *sstack) Len() int {
+func (s *sstack[P, V]) Len() int {
 	return s.size
 }

 // Compares the priority of two elements of the stack (higher is first).
 // Required by sort.Interface.
-func (s *sstack) Less(i, j int) bool {
-	a, b := s.blocks[i/blockSize][i%blockSize].priority, s.blocks[j/blockSize][j%blockSize].priority
-	if s.wrapAround {
-		return a-b > 0
-	}
-	return a > b
+func (s *sstack[P, V]) Less(i, j int) bool {
+	return s.blocks[i/blockSize][i%blockSize].priority > s.blocks[j/blockSize][j%blockSize].priority
 }

 // Swaps two elements in the stack. Required by sort.Interface.
-func (s *sstack) Swap(i, j int) {
+func (s *sstack[P, V]) Swap(i, j int) {
 	ib, io, jb, jo := i/blockSize, i%blockSize, j/blockSize, j%blockSize
 	a, b := s.blocks[jb][jo], s.blocks[ib][io]
 	if s.setIndex != nil {
@@ -115,6 +108,6 @@ func (s *sstack) Swap(i, j int) {
 }

 // Resets the stack, effectively clearing its contents.
-func (s *sstack) Reset() {
-	*s = *newSstack(s.setIndex, false)
+func (s *sstack[P, V]) Reset() {
+	*s = *newSstack[P, V](s.setIndex)
 }
@@ -17,23 +17,23 @@ import (
 func TestSstack(t *testing.T) {
 	// Create some initial data
 	size := 16 * blockSize
-	data := make([]*item, size)
+	data := make([]*item[int64, int], size)
 	for i := 0; i < size; i++ {
-		data[i] = &item{rand.Int(), rand.Int63()}
+		data[i] = &item[int64, int]{rand.Int(), rand.Int63()}
 	}
-	stack := newSstack(nil, false)
+	stack := newSstack[int64, int](nil)
 	for rep := 0; rep < 2; rep++ {
 		// Push all the data into the stack, pop out every second
-		secs := []*item{}
+		secs := []*item[int64, int]{}
 		for i := 0; i < size; i++ {
 			stack.Push(data[i])
 			if i%2 == 0 {
-				secs = append(secs, stack.Pop().(*item))
+				secs = append(secs, stack.Pop().(*item[int64, int]))
 			}
 		}
-		rest := []*item{}
+		rest := []*item[int64, int]{}
 		for stack.Len() > 0 {
-			rest = append(rest, stack.Pop().(*item))
+			rest = append(rest, stack.Pop().(*item[int64, int]))
 		}
 		// Make sure the contents of the resulting slices are ok
 		for i := 0; i < size; i++ {
@@ -50,12 +50,12 @@ func TestSstack(t *testing.T) {
 func TestSstackSort(t *testing.T) {
 	// Create some initial data
 	size := 16 * blockSize
-	data := make([]*item, size)
+	data := make([]*item[int64, int], size)
 	for i := 0; i < size; i++ {
-		data[i] = &item{rand.Int(), int64(i)}
+		data[i] = &item[int64, int]{rand.Int(), int64(i)}
 	}
 	// Push all the data into the stack
-	stack := newSstack(nil, false)
+	stack := newSstack[int64, int](nil)
 	for _, val := range data {
 		stack.Push(val)
 	}
@@ -72,18 +72,18 @@ func TestSstackSort(t *testing.T) {
 func TestSstackReset(t *testing.T) {
 	// Create some initial data
 	size := 16 * blockSize
-	data := make([]*item, size)
+	data := make([]*item[int64, int], size)
 	for i := 0; i < size; i++ {
-		data[i] = &item{rand.Int(), rand.Int63()}
+		data[i] = &item[int64, int]{rand.Int(), rand.Int63()}
 	}
-	stack := newSstack(nil, false)
+	stack := newSstack[int64, int](nil)
 	for rep := 0; rep < 2; rep++ {
 		// Push all the data into the stack, pop out every second
-		secs := []*item{}
+		secs := []*item[int64, int]{}
 		for i := 0; i < size; i++ {
 			stack.Push(data[i])
 			if i%2 == 0 {
-				secs = append(secs, stack.Pop().(*item))
+				secs = append(secs, stack.Pop().(*item[int64, int]))
 			}
 		}
 		// Reset and verify both pulled and stack contents