queue: detect close of incoming channel

Author: joostjager

queue/queue.go

@@ -58,14 +58,18 @@ func (cq *ConcurrentQueue) start() {
 	go func() {
 		defer cq.wg.Done()
 
+	readLoop:
 		for {
 			nextElement := cq.overflow.Front()
 			if nextElement == nil {
 				// Overflow queue is empty so incoming items can be pushed
 				// directly to the output channel. If output channel is full
 				// though, push to overflow.
 				select {
-				case item := <-cq.chanIn:
+				case item, ok := <-cq.chanIn:
+					if !ok {
+						break readLoop
+					}
 					select {
 					case cq.chanOut <- item:
 						// Optimistically push directly to chanOut
@@ -79,7 +83,10 @@ func (cq *ConcurrentQueue) start() {
 				// Overflow queue is not empty, so any new items get pushed to
 				// the back to preserve order.
 				select {
-				case item := <-cq.chanIn:
+				case item, ok := <-cq.chanIn:
+					if !ok {
+						break readLoop
+					}
 					cq.overflow.PushBack(item)
 				case cq.chanOut <- nextElement.Value:
 					cq.overflow.Remove(nextElement)
@@ -88,6 +95,22 @@ func (cq *ConcurrentQueue) start() {
 				}
 			}
 		}
+
+		// Incoming channel has been closed. Empty overflow queue into
+		// the outgoing channel.
+		nextElement := cq.overflow.Front()
+		for nextElement != nil {
+			select {
+			case cq.chanOut <- nextElement.Value:
+				cq.overflow.Remove(nextElement)
+			case <-cq.quit:
+				return
+			}
+			nextElement = cq.overflow.Front()
+		}
+
+		// Close outgoing channel.
+		close(cq.chanOut)
 	}()
 }
 

queue/queue_test.go

@@ -63,3 +63,25 @@ func TestConcurrentQueueIdempotentStop(t *testing.T) {
 
 	testQueueAddDrain(t, 100, 1, 10, 1000, 1000)
 }
+
+// TestQueueCloseIncoming tests that the queue properly handles an incoming
+// channel that is closed.
+func TestQueueCloseIncoming(t *testing.T) {
+	t.Parallel()
+
+	queue := queue.NewConcurrentQueue(10)
+	queue.Start()
+
+	queue.ChanIn() <- 1
+	close(queue.ChanIn())
+
+	item := <-queue.ChanOut()
+	if item.(int) != 1 {
+		t.Fatalf("unexpected item")
+	}
+
+	_, ok := <-queue.ChanOut()
+	if ok {
+		t.Fatalf("expected outgoing channel being closed")
+	}
+}