priorityq/pq/lib.go

156 lines
3.6 KiB
Go

// 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)
// q.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"
"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.Cond{L: &s.mu}
s.canSend = sync.Cond{L: &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 {
s.mu.Unlock()
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.
//
// The error indicates whether the attempt succeeded, the queue is empty, or
// the queue is closed.
func (s *state[P, T]) TryRecv() (priority P, value T, err error) {
s.mu.Lock()
defer s.mu.Unlock()
if s.heap.CanExtract() {
priority, value = s.heap.Extract()
s.canSend.Broadcast()
return
}
if s.closed {
err = priorityq.ErrClosed
} else {
err = priorityq.ErrEmpty
}
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) error {
s.mu.Lock()
defer s.mu.Unlock()
if s.closed {
return priorityq.ErrClosed
}
if !s.heap.CanInsert() {
return priorityq.ErrFull
}
s.heap.Insert(priority, value)
s.canRecv.Broadcast()
return nil
}