feat: add blocking pop call

Author: Bychin

kpq/keyed_priority_queue.go

@@ -60,6 +60,12 @@ func newKeyNotFoundError[K comparable](k K) error {
 // CmpFunc is a generic function type used for ordering the priority queue.
 type CmpFunc[V any] func(x, y V) bool
 
+// TODO comment
+type pair[K comparable, V any] struct {
+	k K
+	v V
+}
+
 // KeyedPriorityQueue represents a generic keyed priority queue,
 // where K is the key type and V is the priority value type.
 //
@@ -71,6 +77,8 @@ type KeyedPriorityQueue[K comparable, V any] struct {
 	im   map[K]int // inverse map of pm; note that for a given key k, pm[im[k]] == k
 	vals map[K]V   // generic priority values of key k
 	cmp  CmpFunc[V]
+
+	fastTrack chan pair[K, V] // used for blocking pop calls
 }
 
 // NewKeyedPriorityQueue returns a new keyed priority queue
@@ -82,10 +90,11 @@ func NewKeyedPriorityQueue[K comparable, V any](cmp CmpFunc[V]) *KeyedPriorityQu
 		panic("keyed priority queue: comparison function cannot be nil")
 	}
 	return &KeyedPriorityQueue[K, V]{
-		pm:   make([]K, 0),
-		im:   make(map[K]int),
-		vals: make(map[K]V),
-		cmp:  cmp,
+		pm:        make([]K, 0),
+		im:        make(map[K]int),
+		vals:      make(map[K]V),
+		cmp:       cmp,
+		fastTrack: make(chan pair[K, V]),
 	}
 }
 
@@ -104,6 +113,12 @@ func (pq *KeyedPriorityQueue[K, V]) Push(k K, v V) error {
 }
 
 func (pq *KeyedPriorityQueue[K, V]) push(k K, v V) {
+	select {
+	case pq.fastTrack <- pair[K, V]{k: k, v: v}:
+		return
+	default:
+		// pass
+	}
 	n := len(pq.pm)
 	pq.pm = append(pq.pm, k)
 	pq.im[k] = n
@@ -122,6 +137,29 @@ func (pq *KeyedPriorityQueue[K, V]) Pop() (K, V, bool) {
 		var v V
 		return k, v, false
 	}
+
+	k, v := pq.pop()
+	return k, v, true
+}
+
+// TODO BlockingPop
+// Pop removes and returns the highest priority key and value from the priority queue.
+// It returns false as its last return value if the priority queue is empty; otherwise, true.
+func (pq *KeyedPriorityQueue[K, V]) BlockingPop() (K, V) {
+	pq.mu.Lock()
+
+	if len(pq.pm) == 0 {
+		pq.mu.Unlock()
+
+		pair := <-pq.fastTrack
+		return pair.k, pair.v
+	}
+
+	defer pq.mu.Unlock()
+	return pq.pop()
+}
+
+func (pq *KeyedPriorityQueue[K, V]) pop() (K, V) {
 	n := len(pq.pm) - 1
 	k := pq.pm[0]
 	v := pq.vals[k]
@@ -130,7 +168,7 @@ func (pq *KeyedPriorityQueue[K, V]) Pop() (K, V, bool) {
 	pq.pm = pq.pm[:n]
 	delete(pq.im, k)
 	delete(pq.vals, k)
-	return k, v, true
+	return k, v
 }
 
 // Set inserts a new entry in the priority queue with the given key and value,

kpq/keyed_priority_queue_test.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"testing"
+	"time"
 )
 
 func TestNewKeyedPriorityQueue_NilCmp(t *testing.T) {
@@ -260,6 +261,126 @@ func TestKeyedPriorityQueue_Pop(t *testing.T) {
 	})
 }
 
+func TestKeyedPriorityQueue_BlockingPop(t *testing.T) {
+	t.Run("Keys", func(t *testing.T) {
+		pq := NewKeyedPriorityQueue[string](func(x, y int) bool {
+			return x < y
+		})
+
+		items := []struct {
+			key string
+			val int
+		}{
+			{key: "fourth", val: 10},
+			{key: "second", val: 8},
+			{key: "third", val: 9},
+			{key: "first", val: 6},
+			{key: "last", val: 20},
+		}
+
+		go func() {
+			time.Sleep(10 * time.Millisecond)
+			for _, item := range items {
+				err := pq.Push(item.key, item.val)
+				if err != nil {
+					panic(fmt.Sprintf("Push(%v, %v): got unexpected error %v", item.key, item.val, err))
+				}
+			}
+		}()
+
+		testCases := []struct {
+			wantKey       string
+			wantValue     int
+			wantPeekKey   string
+			wantPeekValue int
+			wantLen       int
+		}{
+			{
+				wantKey:       "fourth", //"first",
+				wantValue:     10,       // 6,
+				wantPeekKey:   "first",  //"second",
+				wantPeekValue: 6,        // 8,
+				wantLen:       4,
+			},
+			{
+				wantKey:       "first",  //"second",
+				wantValue:     6,        // 8,
+				wantPeekKey:   "second", //"third",
+				wantPeekValue: 8,        // 9,
+				wantLen:       3,
+			},
+		}
+
+		for _, tc := range testCases {
+			t.Run(fmt.Sprintf("%s_%d", tc.wantKey, tc.wantValue), func(t *testing.T) {
+				gotKey, gotValue := pq.BlockingPop()
+
+				if gotKey != tc.wantKey {
+					t.Errorf("pq.BlockingPop(): got key %q; want %q", gotKey, tc.wantKey)
+				}
+
+				if gotValue != tc.wantValue {
+					t.Errorf("pq.BlockingPop(): got value %d; want %d", gotValue, tc.wantValue)
+				}
+
+				time.Sleep(10 * time.Millisecond)
+
+				gotPeekKey, gotPeekValue, ok := pq.Peek()
+				if !ok {
+					t.Fatal("got no min key and value in the priority queue")
+				}
+
+				if gotPeekKey != tc.wantPeekKey {
+					t.Errorf("pq.Peek(): got key %q; want %q", gotPeekKey, tc.wantPeekKey)
+				}
+
+				if gotPeekValue != tc.wantPeekValue {
+					t.Errorf("pq.Peek(): got value %d; want %d", gotPeekValue, tc.wantPeekValue)
+				}
+
+				if got := pq.Len(); got != tc.wantLen {
+					t.Errorf("pq.Len(): got %d; want %d", got, tc.wantLen)
+				}
+			})
+		}
+	})
+
+	/*t.Run("EmptyPQ", func(t *testing.T) {
+		pq := NewKeyedPriorityQueue[string](func(x, y int) bool {
+			return x < y
+		})
+
+		items := []struct {
+			key string
+			val int
+		}{
+			{key: "fourth", val: 10},
+			{key: "second", val: 8},
+			{key: "third", val: 9},
+			{key: "first", val: 6},
+			{key: "last", val: 20},
+		}
+
+		go func() {
+			time.Sleep(10 * time.Millisecond)
+			for _, item := range items {
+				err := pq.Push(item.key, item.val)
+				if err != nil {
+					panic(fmt.Sprintf("Push(%v, %v): got unexpected error %v", item.key, item.val, err))
+				}
+			}
+		}()
+
+		gotKey, gotValue := pq.BlockingPop()
+		if gotKey != items[0].key {
+			t.Errorf("pq.BlockingPop(): got key %q; want %q", gotKey, items[0].key)
+		}
+		if gotValue != items[0].val {
+			t.Errorf("pq.BlockingPop(): got value %d; want %d", gotValue, items[0].val)
+		}
+	})*/
+}
+
 func TestKeyedPriorityQueue_Remove(t *testing.T) {
 	pq := NewKeyedPriorityQueue[string](func(x, y int) bool {
 		return x < y