Implement a true priority queue.
* Add a binary max-heap implementation, `binheap`. * Rename `precise` package to `mq`.
This commit is contained in:
164
mq/lib.go
Normal file
164
mq/lib.go
Normal file
@@ -0,0 +1,164 @@
|
||||
package mq
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"gogs.humancabbage.net/sam/priorityq/circ"
|
||||
)
|
||||
|
||||
// Q is a precise, concurrent, prioritized message queue.
|
||||
//
|
||||
// Each queue has two internal buffers, high and low. This implementation
|
||||
// guarantees that when there are items in both buffers, consumers receive
|
||||
// ones from the high priority buffer first.
|
||||
//
|
||||
// Each buffer has the same capacity, set on initial construction. Sending to
|
||||
// a buffer will block if it is full, even if the other buffer has space.
|
||||
type Q[T any] struct {
|
||||
*state[T]
|
||||
}
|
||||
|
||||
// Make a new queue.
|
||||
func Make[T any](cap int) Q[T] {
|
||||
high := circ.Make[T](cap)
|
||||
low := circ.Make[T](cap)
|
||||
s := &state[T]{
|
||||
high: high,
|
||||
low: low,
|
||||
}
|
||||
s.canRecv = sync.NewCond(&s.mu)
|
||||
s.canSendHigh = sync.NewCond(&s.mu)
|
||||
s.canSendLow = sync.NewCond(&s.mu)
|
||||
return Q[T]{s}
|
||||
}
|
||||
|
||||
type state[T any] struct {
|
||||
mu sync.Mutex
|
||||
high circ.B[T]
|
||||
low circ.B[T]
|
||||
canSendHigh *sync.Cond
|
||||
canSendLow *sync.Cond
|
||||
canRecv *sync.Cond
|
||||
closed bool
|
||||
}
|
||||
|
||||
// Close marks the queue as closed.
|
||||
//
|
||||
// Subsequent attempts to send will panic. Subsequent calls to Recv will
|
||||
// continue to return the remaining items in the queue.
|
||||
func (s *state[T]) Close() {
|
||||
s.mu.Lock()
|
||||
s.closed = true
|
||||
s.mu.Unlock()
|
||||
s.canRecv.Broadcast()
|
||||
}
|
||||
|
||||
// Recv returns an item from the prioritized buffers, blocking if empty.
|
||||
//
|
||||
// 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[T]) Recv() (T, bool) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for {
|
||||
for !s.closed && !s.high.CanPop() && !s.low.CanPop() {
|
||||
s.canRecv.Wait()
|
||||
}
|
||||
if s.closed && !s.high.CanPop() && !s.low.CanPop() {
|
||||
var empty T
|
||||
return empty, false
|
||||
}
|
||||
if s.high.CanPop() {
|
||||
value := s.high.PopFront()
|
||||
s.canSendHigh.Broadcast()
|
||||
return value, true
|
||||
}
|
||||
if s.low.CanPop() {
|
||||
value := s.low.PopFront()
|
||||
s.canSendLow.Broadcast()
|
||||
return value, true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Send is an alias for SendLow.
|
||||
func (s *state[T]) Send(value T) {
|
||||
s.SendLow(value)
|
||||
}
|
||||
|
||||
// SendHigh adds an item to the high priority buffer, blocking if full.
|
||||
func (s *state[T]) SendHigh(value T) {
|
||||
s.send(value, &s.high, s.canSendHigh)
|
||||
}
|
||||
|
||||
// SendLow adds an item to the low priority buffer, blocking if full.
|
||||
func (s *state[T]) SendLow(value T) {
|
||||
s.send(value, &s.low, s.canSendLow)
|
||||
}
|
||||
|
||||
// TryRecv attempts to return an item from the prioritized buffers.
|
||||
//
|
||||
// This method does not block. If there is an item in a buffer, it returns
|
||||
// true. If the buffer is empty, it returns false.
|
||||
func (s *state[T]) TryRecv() (value T, ok bool) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.high.CanPop() {
|
||||
value = s.high.PopFront()
|
||||
ok = true
|
||||
s.canSendHigh.Broadcast()
|
||||
return
|
||||
}
|
||||
if s.low.CanPop() {
|
||||
value = s.low.PopFront()
|
||||
ok = true
|
||||
s.canSendLow.Broadcast()
|
||||
return
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// TrySendHigh attempts to add an item to the high priority buffer.
|
||||
//
|
||||
// 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[T]) TrySendHigh(value T) bool {
|
||||
return s.trySend(value, &s.high)
|
||||
}
|
||||
|
||||
// TrySendLow attempts to add an item to the low priority buffer.
|
||||
//
|
||||
// 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[T]) TrySendLow(value T) bool {
|
||||
return s.trySend(value, &s.low)
|
||||
}
|
||||
|
||||
func (s *state[T]) send(value T, buf *circ.B[T], cond *sync.Cond) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for {
|
||||
for !s.closed && !buf.CanPush() {
|
||||
cond.Wait()
|
||||
}
|
||||
if s.closed {
|
||||
panic("send on closed queue")
|
||||
}
|
||||
if buf.CanPush() {
|
||||
buf.PushBack(value)
|
||||
s.canRecv.Broadcast()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *state[T]) trySend(value T, buf *circ.B[T]) bool {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if !buf.CanPush() {
|
||||
return false
|
||||
}
|
||||
buf.PushBack(value)
|
||||
s.canRecv.Broadcast()
|
||||
return true
|
||||
}
|
||||
228
mq/lib_test.go
Normal file
228
mq/lib_test.go
Normal file
@@ -0,0 +1,228 @@
|
||||
package mq_test
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"gogs.humancabbage.net/sam/priorityq/mq"
|
||||
)
|
||||
|
||||
func TestRecvHighFirst(t *testing.T) {
|
||||
t.Parallel()
|
||||
q := mq.Make[int](4)
|
||||
q.Send(1)
|
||||
q.Send(2)
|
||||
q.Send(3)
|
||||
q.Send(4)
|
||||
q.SendHigh(5)
|
||||
q.SendHigh(6)
|
||||
q.SendHigh(7)
|
||||
q.SendHigh(8)
|
||||
checkRecv := func(n int) {
|
||||
if v, _ := q.Recv(); v != n {
|
||||
t.Errorf("popped %d, expected %d", v, n)
|
||||
}
|
||||
}
|
||||
checkRecv(5)
|
||||
checkRecv(6)
|
||||
checkRecv(7)
|
||||
checkRecv(8)
|
||||
checkRecv(1)
|
||||
checkRecv(2)
|
||||
checkRecv(3)
|
||||
checkRecv(4)
|
||||
}
|
||||
|
||||
func TestSendClosedPanic(t *testing.T) {
|
||||
t.Parallel()
|
||||
defer func() {
|
||||
if r := recover(); r == nil {
|
||||
t.Errorf("sending to closed queue did not panic")
|
||||
}
|
||||
}()
|
||||
q := mq.Make[int](4)
|
||||
q.Close()
|
||||
q.Send(1)
|
||||
}
|
||||
|
||||
func TestRecvClosed(t *testing.T) {
|
||||
t.Parallel()
|
||||
q := mq.Make[int](4)
|
||||
q.Send(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 TestTrySendRecv(t *testing.T) {
|
||||
t.Parallel()
|
||||
q := mq.Make[int](4)
|
||||
assumeSendOk := func(n int, f func(int) bool) {
|
||||
ok := f(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, q.TrySendLow)
|
||||
assumeSendOk(2, q.TrySendLow)
|
||||
assumeSendOk(3, q.TrySendLow)
|
||||
assumeSendOk(4, q.TrySendLow)
|
||||
ok := q.TrySendLow(5)
|
||||
if ok {
|
||||
t.Errorf("expected low buffer to be full")
|
||||
}
|
||||
assumeRecvOk(1)
|
||||
assumeRecvOk(2)
|
||||
assumeRecvOk(3)
|
||||
assumeRecvOk(4)
|
||||
|
||||
assumeSendOk(5, q.TrySendHigh)
|
||||
assumeSendOk(6, q.TrySendHigh)
|
||||
assumeSendOk(7, q.TrySendHigh)
|
||||
assumeSendOk(8, q.TrySendHigh)
|
||||
ok = q.TrySendHigh(5)
|
||||
if ok {
|
||||
t.Errorf("expected high buffer to be full")
|
||||
}
|
||||
assumeRecvOk(5)
|
||||
assumeRecvOk(6)
|
||||
assumeRecvOk(7)
|
||||
assumeRecvOk(8)
|
||||
|
||||
_, ok = q.TryRecv()
|
||||
if ok {
|
||||
t.Errorf("expected queue to be empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcProducerConsumer(t *testing.T) {
|
||||
t.Parallel()
|
||||
q := mq.Make[int](4)
|
||||
var wg sync.WaitGroup
|
||||
produceDone := make(chan struct{})
|
||||
wg.Add(2)
|
||||
go func() {
|
||||
for i := 0; i < 10000; i++ {
|
||||
if rand.Intn(2) == 0 {
|
||||
q.Send(i)
|
||||
} else {
|
||||
q.SendHigh(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 := mq.Make[int](b.N)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
q.Send(i)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkSendChan(b *testing.B) {
|
||||
c := make(chan int, b.N)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
c <- i
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkRecv(b *testing.B) {
|
||||
q := mq.Make[int](b.N)
|
||||
for i := 0; i < b.N; i++ {
|
||||
q.Send(i)
|
||||
}
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
q.Recv()
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkRecvChan(b *testing.B) {
|
||||
c := make(chan int, b.N)
|
||||
for i := 0; i < b.N; i++ {
|
||||
c <- i
|
||||
}
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
<-c
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkConcSendRecv(b *testing.B) {
|
||||
q := mq.Make[int](b.N)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
start := make(chan struct{})
|
||||
go func() {
|
||||
<-start
|
||||
for i := 0; i < b.N; i++ {
|
||||
q.Send(i)
|
||||
}
|
||||
wg.Done()
|
||||
}()
|
||||
go func() {
|
||||
<-start
|
||||
for i := 0; i < b.N; i++ {
|
||||
q.Recv()
|
||||
}
|
||||
wg.Done()
|
||||
}()
|
||||
b.ResetTimer()
|
||||
close(start)
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func BenchmarkConcSendRecvChan(b *testing.B) {
|
||||
c := make(chan int, b.N)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
start := make(chan struct{})
|
||||
go func() {
|
||||
<-start
|
||||
for i := 0; i < b.N; i++ {
|
||||
c <- i
|
||||
}
|
||||
wg.Done()
|
||||
}()
|
||||
go func() {
|
||||
<-start
|
||||
for i := 0; i < b.N; i++ {
|
||||
<-c
|
||||
}
|
||||
wg.Done()
|
||||
}()
|
||||
b.ResetTimer()
|
||||
close(start)
|
||||
wg.Wait()
|
||||
}
|
||||
Reference in New Issue
Block a user