Move priority queue to pq package; improve docs.

pq/lib.go

@@ -0,0 +1,145 @@
+// Package pq implements a concurrent priority queue.
+//
+// [Q] is similar to a buffered channel, except that senders attach to each
+// item a priority, and receivers always get the highest-priority item.
+//
+// For example:
+//
+//	import "gogs.humancabbage.net/sam/priorityq/pq"
+//	q := pq.Make[int, string](8)
+//	q.Send(1, "world")
+//	q.Send(2, "hello")
+//	_, word1, _ := pq.Recv()
+//	_, word2, _ := pq.Recv()
+//	fmt.Println(word1, word2)
+//	pq.Close()
+//	// Output: hello world
+//
+// # Implementation
+//
+// Each queue has a [binary max-heap]. Sending and receiving items require
+// heap-up and heap-down operations, respectively.
+//
+// [binary max-heap]: https://en.wikipedia.org/wiki/Binary_heap
+package pq
+
+import (
+	"sync"
+
+	"gogs.humancabbage.net/sam/priorityq/binheap"
+	"golang.org/x/exp/constraints"
+)
+
+// Q is a generic, concurrent priority queue.
+type Q[P constraints.Ordered, T any] struct {
+	*state[P, T]
+}
+
+// Make a new queue.
+func Make[P constraints.Ordered, T any](cap int) Q[P, T] {
+	heap := binheap.Make[P, T](cap)
+	s := &state[P, T]{
+		heap: heap,
+	}
+	s.canRecv = sync.NewCond(&s.mu)
+	s.canSend = sync.NewCond(&s.mu)
+	return Q[P, T]{s}
+}
+
+type state[P constraints.Ordered, T any] struct {
+	mu      sync.Mutex
+	heap    binheap.H[P, T]
+	canSend *sync.Cond
+	canRecv *sync.Cond
+	closed  bool
+}
+
+// Close marks the queue as closed.
+//
+// Attempting to close an already-closed queue results in a panic.
+func (s *state[P, T]) Close() {
+	s.mu.Lock()
+	if s.closed {
+		panic("close of closed queue")
+	}
+	s.closed = true
+	s.mu.Unlock()
+	s.canRecv.Broadcast()
+}
+
+// Recv gets an item, blocking when empty until one is available.
+//
+// This returns both the item itself and the its assigned priority.
+//
+// The returned bool will be true if the queue still has items or is open.
+// It will be false if the queue is empty and closed.
+func (s *state[P, T]) Recv() (P, T, bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	for {
+		for !s.closed && !s.heap.CanExtract() {
+			s.canRecv.Wait()
+		}
+		if s.closed && !s.heap.CanExtract() {
+			var emptyP P
+			var emptyT T
+			return emptyP, emptyT, false
+		}
+		if s.heap.CanExtract() {
+			priority, value := s.heap.Extract()
+			s.canSend.Broadcast()
+			return priority, value, true
+		}
+	}
+}
+
+// Send adds an item with some priority, blocking if full.
+func (s *state[P, T]) Send(priority P, value T) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	for {
+		for !s.closed && !s.heap.CanInsert() {
+			s.canSend.Wait()
+		}
+		if s.closed {
+			panic("send on closed queue")
+		}
+		if s.heap.CanInsert() {
+			s.heap.Insert(priority, value)
+			s.canRecv.Broadcast()
+			return
+		}
+	}
+}
+
+// TryRecv attempts to get an item without blocking.
+//
+// This returns both the item itself and the its assigned priority.
+//
+// If the attempt succeeds, the returned bool is true. Otherwise, it is false.
+func (s *state[P, T]) TryRecv() (priority P, value T, ok bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.heap.CanExtract() {
+		priority, value = s.heap.Extract()
+		ok = true
+		s.canSend.Broadcast()
+		return
+	}
+	return
+}
+
+// TrySend attempts to add an item with some priority, without blocking.
+//
+// This method does not block. If there is space in the buffer, it returns
+// true. If the buffer is full, it returns false.
+func (s *state[P, T]) TrySend(priority P, value T) bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if !s.heap.CanInsert() {
+		return false
+	}
+	s.heap.Insert(priority, value)
+	s.canRecv.Broadcast()
+	return true
+}

pq/lib_test.go

@@ -0,0 +1,227 @@
+package pq_test
+
+import (
+	"math/rand"
+	"runtime"
+	"sync"
+	"testing"
+
+	"gogs.humancabbage.net/sam/priorityq/pq"
+)
+
+func TestRecvHighestFirst(t *testing.T) {
+	t.Parallel()
+	q := pq.Make[int, int](8)
+	q.Send(4, 4)
+	q.Send(2, 2)
+	q.Send(1, 1)
+	q.Send(5, 5)
+	q.Send(7, 7)
+	q.Send(8, 8)
+	q.Send(3, 3)
+	q.Send(6, 6)
+	checkRecv := func(n int) {
+		if _, v, _ := q.Recv(); v != n {
+			t.Errorf("popped %d, expected %d", v, n)
+		}
+	}
+	checkRecv(8)
+	checkRecv(7)
+	checkRecv(6)
+	checkRecv(5)
+	checkRecv(4)
+	checkRecv(3)
+	checkRecv(2)
+	checkRecv(1)
+}
+
+func TestSendClosedPanic(t *testing.T) {
+	t.Parallel()
+	defer func() {
+		if r := recover(); r == nil {
+			t.Errorf("sending to closed queue did not panic")
+		}
+	}()
+	q := pq.Make[int, int](4)
+	q.Close()
+	q.Send(1, 1)
+}
+
+func TestRecvClosed(t *testing.T) {
+	t.Parallel()
+	q := pq.Make[int, int](4)
+	q.Send(1, 1)
+	q.Close()
+	_, _, ok := q.Recv()
+	if !ok {
+		t.Errorf("queue should have item to receive")
+	}
+	_, _, ok = q.Recv()
+	if ok {
+		t.Errorf("queue should be closed")
+	}
+}
+
+func TestDoubleClose(t *testing.T) {
+	t.Parallel()
+	q := pq.Make[int, int](4)
+	defer func() {
+		if r := recover(); r == nil {
+			t.Errorf("closing a closed queue did not panic")
+		}
+	}()
+	q.Close()
+	q.Close()
+}
+
+func TestTrySendRecv(t *testing.T) {
+	t.Parallel()
+	q := pq.Make[int, int](4)
+	assumeSendOk := func(n int) {
+		ok := q.TrySend(n, n)
+		if !ok {
+			t.Errorf("expected to be able to send")
+		}
+	}
+	assumeRecvOk := func(expected int) {
+		_, actual, ok := q.TryRecv()
+		if !ok {
+			t.Errorf("expected to be able to receive")
+		}
+		if actual != expected {
+			t.Errorf("expected %d, got %d", expected, actual)
+		}
+	}
+	assumeSendOk(1)
+	assumeSendOk(2)
+	assumeSendOk(3)
+	assumeSendOk(4)
+	ok := q.TrySend(5, 5)
+	if ok {
+		t.Errorf("expected queue to be full")
+	}
+	assumeRecvOk(4)
+	assumeRecvOk(3)
+	assumeRecvOk(2)
+	assumeRecvOk(1)
+
+	_, _, ok = q.TryRecv()
+	if ok {
+		t.Errorf("expected queue to be empty")
+	}
+}
+
+func TestConcProducerConsumer(t *testing.T) {
+	t.Parallel()
+	q := pq.Make[int, int](4)
+	var wg sync.WaitGroup
+	produceDone := make(chan struct{})
+	wg.Add(2)
+	go func() {
+		for i := 0; i < 10000; i++ {
+			q.Send(rand.Int(), i)
+		}
+		close(produceDone)
+		wg.Done()
+	}()
+	go func() {
+		ok := true
+		for ok {
+			_, _, ok = q.Recv()
+		}
+		wg.Done()
+	}()
+	<-produceDone
+	t.Logf("producer done, closing channel")
+	q.Close()
+	wg.Wait()
+}
+
+func BenchmarkSend(b *testing.B) {
+	q := pq.Make[int, int](b.N)
+	// randomize priorities to get amortized cost per op
+	ps := make([]int, b.N)
+	for i := 0; i < b.N; i++ {
+		ps[i] = rand.Int()
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		q.Send(ps[i], i)
+	}
+}
+
+func BenchmarkRecv(b *testing.B) {
+	q := pq.Make[int, int](b.N)
+	// randomize priorities to get amortized cost per op
+	for i := 0; i < b.N; i++ {
+		q.Send(rand.Int(), i)
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		q.Recv()
+	}
+}
+
+func BenchmarkConcSendRecv(b *testing.B) {
+	q := pq.Make[int, int](b.N)
+	// randomize priorities to get amortized cost per op
+	ps := make([]int, b.N)
+	for i := 0; i < b.N; i++ {
+		ps[i] = rand.Int()
+	}
+	var wg sync.WaitGroup
+	wg.Add(2)
+	start := make(chan struct{})
+	go func() {
+		<-start
+		for i := 0; i < b.N; i++ {
+			q.Send(ps[i], i)
+		}
+		wg.Done()
+	}()
+	go func() {
+		<-start
+		for i := 0; i < b.N; i++ {
+			q.Recv()
+		}
+		wg.Done()
+	}()
+	b.ResetTimer()
+	close(start)
+	wg.Wait()
+}
+
+func BenchmarkHighContention(b *testing.B) {
+	q := pq.Make[int, int](b.N)
+	var wg sync.WaitGroup
+	start := make(chan struct{})
+	done := make(chan struct{})
+	numProducers := runtime.NumCPU()
+	sendsPerProducer := b.N / numProducers
+	wg.Add(numProducers)
+	for i := 0; i < numProducers; i++ {
+		go func() {
+			ps := make([]int, sendsPerProducer)
+			for i := 0; i < sendsPerProducer; i++ {
+				ps[i] = rand.Int()
+			}
+			<-start
+			for i := 0; i < sendsPerProducer; i++ {
+				q.Send(ps[i], 1)
+			}
+			wg.Done()
+		}()
+	}
+	go func() {
+		ok := true
+		for ok {
+			_, _, ok = q.Recv()
+		}
+		close(done)
+	}()
+	b.ResetTimer()
+	close(start)
+	wg.Wait()
+	q.Close()
+	<-done
+}