Added a priority semaphore

runtime/pkg/priorityqueue/priorityqueue.go

@@ -0,0 +1,87 @@
+package priorityqueue
+
+import (
+	"container/heap"
+)
+
+// Item is a value in PriorityQueue
+type Item[V any] struct {
+	Value    V
+	priority int
+	index    int
+}
+
+// PriorityQueue is a generic priority queue.
+// It returns items with a higher priority first. It is not concurrency safe.
+type PriorityQueue[V any] struct {
+	heap priorityHeap[V]
+}
+
+func New[V any]() *PriorityQueue[V] {
+	pq := &PriorityQueue[V]{}
+	heap.Init(&pq.heap)
+	return pq
+}
+
+func (pq *PriorityQueue[V]) Push(val V, priority int) *Item[V] {
+	itm := &Item[V]{
+		Value:    val,
+		priority: priority,
+		index:    -1,
+	}
+	heap.Push(&pq.heap, itm)
+	return itm
+}
+
+func (pq *PriorityQueue[V]) Pop() V {
+	itm := heap.Pop(&pq.heap).(*Item[V])
+	return itm.Value
+}
+
+func (pq *PriorityQueue[V]) Remove(itm *Item[V]) {
+	if itm.index >= 0 {
+		heap.Remove(&pq.heap, itm.index)
+	}
+}
+
+func (pq *PriorityQueue[V]) Contains(itm *Item[V]) bool {
+	return itm.index >= 0
+}
+
+func (pq *PriorityQueue[V]) Len() int {
+	return pq.heap.Len()
+}
+
+// priorityHeap implements heap.Interface to serve as a priority queue.
+// See the heap docs for usage details: https://pkg.go.dev/container/heap#example-package-PriorityQueue
+type priorityHeap[V any] []*Item[V]
+
+func (pq priorityHeap[V]) Len() int { return len(pq) }
+
+func (pq priorityHeap[V]) Less(i, j int) bool {
+	// We use greater than here so that Pop gives us the highest priority item (not lowest)
+	return pq[i].priority > pq[j].priority
+}
+
+func (pq priorityHeap[V]) Swap(i, j int) {
+	pq[i], pq[j] = pq[j], pq[i]
+	pq[i].index = i
+	pq[j].index = j
+}
+
+func (pq *priorityHeap[V]) Push(x any) {
+	n := len(*pq)
+	itm := x.(*Item[V])
+	itm.index = n
+	*pq = append(*pq, itm)
+}
+
+func (pq *priorityHeap[V]) Pop() any {
+	old := *pq
+	n := len(old)
+	itm := old[n-1]
+	old[n-1] = nil // avoid memory leak
+	itm.index = -1 // for safety
+	*pq = old[0 : n-1]
+	return itm
+}

runtime/pkg/priorityqueue/priorityqueue_test.go

@@ -0,0 +1,23 @@
+package priorityqueue
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestPriorityQueue(t *testing.T) {
+	pq := New[int]()
+	require.Equal(t, 0, pq.Len())
+	pq.Push(1, 1)
+	pq.Push(2, 2)
+	itm := pq.Push(3, 3)
+	require.Equal(t, 3, pq.Len())
+	require.True(t, pq.Contains(itm))
+	pq.Remove(itm)
+	require.False(t, pq.Contains(itm))
+	require.Equal(t, 2, pq.Pop())
+	require.Equal(t, 1, pq.Len())
+	require.Equal(t, 1, pq.Pop())
+	require.Equal(t, 0, pq.Len())
+}

runtime/pkg/priorityqueue/semaphore.go

@@ -0,0 +1,87 @@
+package priorityqueue
+
+import (
+	"context"
+	"sync"
+)
+
+type Semaphore struct {
+	mu   sync.Mutex
+	pq   *PriorityQueue[chan struct{}]
+	size int
+	cur  int
+}
+
+func NewSemaphore(size int) *Semaphore {
+	return &Semaphore{
+		mu:   sync.Mutex{},
+		pq:   New[chan struct{}](),
+		size: size,
+		cur:  0,
+	}
+}
+
+func (s *Semaphore) Acquire(ctx context.Context, priority int) error {
+	s.mu.Lock()
+	if s.size-s.cur >= 1 && s.pq.Len() == 0 {
+		s.cur += 1
+		s.mu.Unlock()
+		return nil
+	}
+
+	readyCh := make(chan struct{})
+	itm := s.pq.Push(readyCh, priority)
+	s.mu.Unlock()
+
+	select {
+	case <-readyCh:
+		return nil
+	case <-ctx.Done():
+		s.mu.Lock()
+		if !s.pq.Contains(itm) {
+			// Cancelled and acquired at the same time. Easiest to pretend it was acquired first.
+			s.mu.Unlock()
+			return nil
+		}
+		s.pq.Remove(itm)
+		s.mu.Unlock()
+		return ctx.Err()
+	}
+}
+
+func (s *Semaphore) TryAcquire() bool {
+	s.mu.Lock()
+	ok := s.size-s.cur >= 1 && s.pq.Len() == 0
+	if ok {
+		s.cur += 1
+	}
+	s.mu.Unlock()
+	return ok
+}
+
+func (s *Semaphore) Release() {
+	s.mu.Lock()
+	s.cur -= 1
+	if s.cur < 0 {
+		s.mu.Unlock()
+		panic("semaphore released more times than acquired")
+	}
+	s.notifyWaiters()
+	s.mu.Unlock()
+}
+
+func (s *Semaphore) notifyWaiters() {
+	for {
+		if s.pq.Len() == 0 {
+			break
+		}
+
+		if s.cur == s.size {
+			break
+		}
+
+		readyCh := s.pq.Pop()
+		s.cur++
+		close(readyCh)
+	}
+}

runtime/pkg/priorityqueue/semaphore_test.go

@@ -0,0 +1,125 @@
+package priorityqueue
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
+)
+
+func TestSemaphoreSimple(t *testing.T) {
+	s := NewSemaphore(1)
+	err := s.Acquire(context.Background(), 1)
+	require.NoError(t, err)
+	s.Release()
+	err = s.Acquire(context.Background(), 1)
+	require.NoError(t, err)
+	s.Release()
+	require.True(t, s.TryAcquire())
+	require.False(t, s.TryAcquire())
+	s.Release()
+
+	s = NewSemaphore(2)
+	err = s.Acquire(context.Background(), 1)
+	require.NoError(t, err)
+	err = s.Acquire(context.Background(), 1)
+	require.NoError(t, err)
+	s.Release()
+	s.Release()
+	require.Panics(t, func() {
+		s.Release()
+	})
+}
+
+func TestSemaphorePriority(t *testing.T) {
+	// Prepare
+	n := 10
+	var results []int
+	var g errgroup.Group
+	s := NewSemaphore(1)
+
+	// Acquire to block
+	err := s.Acquire(context.Background(), 1)
+	require.NoError(t, err)
+	require.False(t, s.TryAcquire())
+
+	// Fill up queue
+	for i := 0; i <= n; i++ {
+		priority := i
+		g.Go(func() error {
+			err := s.Acquire(context.Background(), priority)
+			require.NoError(t, err)
+			results = append(results, priority)
+			s.Release()
+			return nil
+		})
+	}
+
+	// Wait a bit to ensure queue fills up
+	time.Sleep(time.Second)
+	s.Release()
+
+	// Wait for processing
+	err = g.Wait()
+	require.NoError(t, err)
+
+	// Check results evaluated in priority order
+	for i := 0; i <= n; i++ {
+		require.Equal(t, i, results[n-i])
+	}
+}
+
+func TestSemaphoreCancel(t *testing.T) {
+	// Prepare
+	n := 100
+	size := 4
+	cancelIdx := 50
+	results := make(chan int, n-1)
+	var g errgroup.Group
+	s := NewSemaphore(size)
+
+	// Acquire up to size
+	for i := 0; i < size; i++ {
+		err := s.Acquire(context.Background(), 1)
+		require.NoError(t, err)
+	}
+	require.False(t, s.TryAcquire())
+
+	// Fill up queue
+	for i := 0; i < n; i++ {
+		priority := i
+		g.Go(func() error {
+			ctx := context.Background()
+			if priority == cancelIdx {
+				cctx, cancel := context.WithCancel(ctx)
+				ctx = cctx
+				cancel()
+			}
+
+			err := s.Acquire(ctx, priority)
+			if priority == cancelIdx {
+				require.Error(t, err)
+				return nil
+			}
+
+			results <- priority
+			s.Release()
+			return nil
+		})
+	}
+
+	// Release to unblock for processing
+	for i := 0; i < size; i++ {
+		s.Release()
+	}
+
+	// Wait for processing
+	err := g.Wait()
+	require.NoError(t, err)
+
+	for i := 0; i < n-1; i++ {
+		require.NotEqual(t, cancelIdx, <-results)
+	}
+}