Add queue util

code-asher · code-asher · commit f696ff4e9b4e · 2025-06-11T01:20:17.000-08:00
diff --git a/cli/cliutil/queue.go b/cli/cliutil/queue.go
@@ -0,0 +1,72 @@
+package cliutil
+
+import (
+	"sync"
+
+	"golang.org/x/xerrors"
+)
+
+// Queue is a FIFO queue with a fixed size.  If the size is exceeded, the first
+// item is dropped.
+type Queue[T any] struct {
+	cond   *sync.Cond
+	items  []T
+	mu     sync.Mutex
+	size   int
+	closed bool
+}
+
+// NewQueue creates a queue with the given size.
+func NewQueue[T any](size int) *Queue[T] {
+	q := &Queue[T]{
+		items: make([]T, 0, size),
+		size:  size,
+	}
+	q.cond = sync.NewCond(&q.mu)
+	return q
+}
+
+// Close aborts any pending pops and makes future pushes error.
+func (q *Queue[T]) Close() {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	q.closed = true
+	q.cond.Broadcast()
+}
+
+// Push adds an item to the queue.  If closed, returns an error.
+func (q *Queue[T]) Push(x T) error {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	if q.closed {
+		return xerrors.New("queue has been closed")
+	}
+	if len(q.items) >= q.size {
+		q.items = q.items[1:]
+	}
+	q.items = append(q.items, x)
+	q.cond.Broadcast()
+	return nil
+}
+
+// Pop removes and returns the first item from the queue, waiting until there is
+// something to pop if necessary.  If closed, returns false.
+func (q *Queue[T]) Pop() (T, bool) {
+	var head T
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	for len(q.items) == 0 && !q.closed {
+		q.cond.Wait()
+	}
+	if q.closed {
+		return head, false
+	}
+	head, q.items = q.items[0], q.items[1:]
+	return head, true
+}
+
+func (q *Queue[T]) Len() int {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	return len(q.items)
+}
diff --git a/cli/cliutil/queue_test.go b/cli/cliutil/queue_test.go
@@ -0,0 +1,85 @@
+package cliutil_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/cli/cliutil"
+)
+
+func TestQueue(t *testing.T) {
+	t.Parallel()
+
+	t.Run("DropsFirst", func(t *testing.T) {
+		t.Parallel()
+
+		q := cliutil.NewQueue[int](10)
+		require.Equal(t, 0, q.Len())
+
+		for i := 0; i < 20; i++ {
+			err := q.Push(i)
+			require.NoError(t, err)
+			if i < 10 {
+				require.Equal(t, i+1, q.Len())
+			} else {
+				require.Equal(t, 10, q.Len())
+			}
+		}
+
+		val, ok := q.Pop()
+		require.True(t, ok)
+		require.Equal(t, 10, val)
+		require.Equal(t, 9, q.Len())
+	})
+
+	t.Run("Pop", func(t *testing.T) {
+		t.Parallel()
+
+		q := cliutil.NewQueue[int](10)
+		for i := 0; i < 5; i++ {
+			err := q.Push(i)
+			require.NoError(t, err)
+		}
+
+		// No blocking, should pop immediately.
+		for i := 0; i < 5; i++ {
+			val, ok := q.Pop()
+			require.True(t, ok)
+			require.Equal(t, i, val)
+		}
+
+		// Pop should block until the next push.
+		go func() {
+			err := q.Push(55)
+			assert.NoError(t, err)
+		}()
+
+		item, ok := q.Pop()
+		require.True(t, ok)
+		require.Equal(t, 55, item)
+	})
+
+	t.Run("Close", func(t *testing.T) {
+		t.Parallel()
+
+		q := cliutil.NewQueue[int](10)
+
+		done := make(chan bool)
+		go func() {
+			_, ok := q.Pop()
+			done <- ok
+		}()
+
+		q.Close()
+
+		require.False(t, <-done)
+
+		_, ok := q.Pop()
+		require.False(t, ok)
+
+		err := q.Push(10)
+		require.Error(t, err)
+	})
+}
