Unbounded queues: add a Close() method.

exp/mpsc/mpsc.go

@@ -46,9 +46,17 @@ func (tx *Tx[T]) Send(v T) {
 	tx.tx.Send(v)
 }
 
+// Close the queue. Calls to Recv on the other end will return io.EOF.
+func (tx *Tx[T]) Close() error {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	return tx.tx.Close()
+}
+
 // Receive a message from the queue. Blocks if the queue is empty.
 // If the context ends before the receive happens, this returns
-// ctx.Err().
+// ctx.Err(). If Close is called on the corresponding Tx, this
+// returns io.EOF
 func (rx *Rx[T]) Recv(ctx context.Context) (T, error) {
 	return rx.rx.Recv(ctx)
 }

exp/spsc/spsc.go

@@ -13,6 +13,7 @@ package spsc
 
 import (
 	"context"
+	"io"
 )
 
 const itemsBuffer = 64
@@ -74,9 +75,17 @@ func (tx *Tx[T]) Send(v T) {
 	}
 }
 
+// Close the queue. Calls to Recv on the other end will return io.EOF.
+func (tx *Tx[T]) Close() error {
+	close(tx.tail.items)
+	close(tx.tail.next)
+	return nil
+}
+
 // Receive a message from the queue. Blocks if the queue is empty.
 // If the context ends before the receive happens, this returns
-// ctx.Err().
+// ctx.Err(). If Close is called on the corresponding Tx, this
+// returns io.EOF
 func (rx *Rx[T]) Recv(ctx context.Context) (T, error) {
 	for {
 		select {
@@ -87,7 +96,11 @@ func (rx *Rx[T]) Recv(ctx context.Context) (T, error) {
 			if ok {
 				return v, nil
 			}
-			rx.head = <-rx.head.next
+			rx.head, ok = <-rx.head.next
+			if !ok {
+				var zero T
+				return zero, io.EOF
+			}
 		}
 	}
 }

exp/spsc/spsc_test.go

@@ -2,6 +2,7 @@ package spsc
 
 import (
 	"context"
+	"io"
 	"testing"
 	"time"
 
@@ -64,3 +65,40 @@ func TestFillThenDrain(t *testing.T) {
 	assert.NotNil(t, err)
 	assert.ErrorIs(t, err, ctx.Err())
 }
+
+// Test that we get io.EOF after closing the tx.
+func TestClose(t *testing.T) {
+	t.Parallel()
+
+	// try it with various numbers of items in the queue:
+	cases := []struct {
+		size int
+		desc string
+	}{
+		{0, "empty"},
+		{1, "one element"},
+		{itemsBuffer, "full buffer"},
+		{itemsBuffer + 1, "overfull buffer"},
+	}
+
+	for _, c := range cases {
+		t.Run("size: "+c.desc, func(t *testing.T) {
+			q := New[int]()
+
+			for i := 0; i < c.size; i++ {
+				q.Send(i)
+			}
+			q.Close()
+
+			ctx := context.Background()
+			for i := 0; i < c.size; i++ {
+				v, err := q.Recv(ctx)
+				assert.NoError(t, err)
+				assert.Equal(t, i, v)
+			}
+
+			_, err := q.Recv(ctx)
+			assert.Equal(t, err, io.EOF)
+		})
+	}
+}