Skip to content

Commit

Permalink
Add a generic priority queue (used in followup)
Browse files Browse the repository at this point in the history
  • Loading branch information
tylerwilliams committed Dec 20, 2024
1 parent 3ab0efe commit 0ad4507
Show file tree
Hide file tree
Showing 3 changed files with 173 additions and 0 deletions.
17 changes: 17 additions & 0 deletions server/util/priority_queue/BUILD
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")

go_library(
name = "priority_queue",
srcs = ["priority_queue.go"],
importpath = "github.com/buildbuddy-io/buildbuddy/server/util/priority_queue",
visibility = ["//visibility:public"],
)

go_test(
name = "priority_queue_test",
srcs = ["priority_queue_test.go"],
deps = [
":priority_queue",
"@com_github_stretchr_testify//assert",
],
)
134 changes: 134 additions & 0 deletions server/util/priority_queue/priority_queue.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@
package priority_queue

import (
"container/heap"
"sync"
"time"
)

// A pqItem is the element managed by a priority queue.
type pqItem[V any] struct {
value V
priority int
insertTime time.Time
}

// innerPQ implements heap.Interface and holds pqItems.
type innerPQ[V any] []*pqItem[V]

func (pq innerPQ[V]) Len() int { return len(pq) }
func (pq innerPQ[V]) Less(i, j int) bool {
return pq[i].priority > pq[j].priority ||
(pq[i].priority == pq[j].priority && pq[i].insertTime.Before(pq[j].insertTime))
}
func (pq innerPQ[V]) Swap(i, j int) {
pq[i], pq[j] = pq[j], pq[i]
}
func (pq *innerPQ[V]) Push(x any) {
item := x.(*pqItem[V])
*pq = append(*pq, item)
}
func (pq *innerPQ[V]) Pop() any {
old := *pq
n := len(old)
item := old[n-1]
old[n-1] = nil // avoid memory leak
*pq = old[0 : n-1]
return item
}

// PriorityQueue implements a thread safe priority queue for type V.
// If the queue is empty, calling Pop() or Peek() will return a zero value of
// type V, or a specific empty value configured via options.
type PriorityQueue[V any] struct {
inner *innerPQ[V]
mu sync.Mutex
noValueFunc func() V
}

type Options struct {
noValueGetter func() any
}

func defaultOptions() *Options {
return &Options{}
}

type Option func(*Options)

func WithEmptyValue(v any) Option {
return func(o *Options) {
o.noValueGetter = func() any {
return v
}
}
}

func New[V any](opts ...Option) *PriorityQueue[V] {
inner := make(innerPQ[V], 0)
pq := &PriorityQueue[V]{
inner: &inner,
}
options := defaultOptions()
for _, opt := range opts {
opt(options)
}
if options.noValueGetter != nil {
pq.noValueFunc = func() V {
return options.noValueGetter().(V)
}
} else {
pq.noValueFunc = pq.zeroValue
}
return pq
}

func (pq *PriorityQueue[V]) zeroValue() V {
var zero V
return zero
}

func (pq *PriorityQueue[V]) Push(v V, priority int) {
pq.mu.Lock()
defer pq.mu.Unlock()
heap.Push(pq.inner, &pqItem[V]{
value: v,
insertTime: time.Now(),
priority: priority,
})
}

func (pq *PriorityQueue[V]) Pop() V {
pq.mu.Lock()
defer pq.mu.Unlock()
if len(*pq.inner) == 0 {
return pq.noValueFunc()
}
item := heap.Pop(pq.inner).(*pqItem[V])
return item.value
}

func (pq *PriorityQueue[V]) Peek() V {
pq.mu.Lock()
defer pq.mu.Unlock()
if len(*pq.inner) == 0 {
return pq.noValueFunc()
}
return (*pq.inner)[0].value
}

func (pq *PriorityQueue[V]) GetAll() []V {
pq.mu.Lock()
defer pq.mu.Unlock()
var allValues []V
for _, i := range *pq.inner {
allValues = append(allValues, i.value)
}
return allValues
}

func (pq *PriorityQueue[V]) Len() int {
pq.mu.Lock()
defer pq.mu.Unlock()
return len(*pq.inner)
}
22 changes: 22 additions & 0 deletions server/util/priority_queue/priority_queue_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
package priority_queue_test

import (
"testing"

"github.com/buildbuddy-io/buildbuddy/server/util/priority_queue"
"github.com/stretchr/testify/assert"
)

func TestPushPop(t *testing.T) {
q := priority_queue.New[string]()
q.Push("A", 1)
q.Push("E", 5)
q.Push("D", 4)
q.Push("B", 2)

assert.Equal(t, "E", q.Pop())
assert.Equal(t, "D", q.Pop())
assert.Equal(t, "B", q.Pop())
assert.Equal(t, "A", q.Pop())
assert.Equal(t, "", q.Pop())
}

0 comments on commit 0ad4507

Please sign in to comment.