streaming: prevent Close from deadlocking on a stalled subscriber (#77)

The Reader/Sink read loop fans events out to subscriber channels with a
blocking send while holding the lock. Close acquired that same lock as its
first action before closing donechan, so when a subscriber stopped draining
its channel the read loop parked on the send while holding the lock, Close
deadlocked acquiring it, and the read goroutine plus its Redis connection
leaked forever.

Close now closes donechan first, without holding the lock, so the signal
reaches a parked read loop; the fan-out send selects on the done channel and
abandons delivery during teardown. Close is wrapped in a sync.Once so the
full shutdown runs exactly once and concurrent callers block until it
completes.

Adds regression tests covering a stalled subscriber for both Reader and Sink.

Author: dgunther

streaming/reader.go

@@ -44,6 +44,8 @@ type (
 		chans []chan *Event
 		// startOnce is used to ensure the reader is started only once.
 		startOnce sync.Once
+		// closeOnce is used to ensure the reader is closed only once.
+		closeOnce sync.Once
 		// donechan is the reader donechan channel.
 		donechan chan struct{}
 		// streamschan notifies the reader when streams are added or
@@ -185,21 +187,26 @@ func (r *Reader) RemoveStream(ctx context.Context, stream *Stream) error {
 }
 
 // Close stops event polling and closes the reader channel. It is safe to call
-// Close multiple times.
+// Close multiple times; concurrent callers block until the first Close
+// completes. Close returns only once the read goroutine has stopped and its
+// resources are released, which may take up to one block duration.
 func (r *Reader) Close() {
-	r.lock.Lock()
-	if r.closing {
-		return
-	}
-	r.closing = true
-	close(r.donechan)
-	close(r.streamschan)
-	r.lock.Unlock()
-	r.wait.Wait()
-	r.lock.Lock()
-	defer r.lock.Unlock()
-	r.closed = true
-	r.logger.Info("stopped")
+	r.closeOnce.Do(func() {
+		// Close donechan first, without holding the lock, so the signal
+		// reaches the read loop even when it is parked on a fan-out send
+		// to a stalled subscriber (which holds the lock). Otherwise Close
+		// would deadlock acquiring the lock the read loop never releases.
+		close(r.donechan)
+		r.lock.Lock()
+		r.closing = true
+		close(r.streamschan)
+		r.lock.Unlock()
+		r.wait.Wait()
+		r.lock.Lock()
+		defer r.lock.Unlock()
+		r.closed = true
+		r.logger.Info("stopped")
+	})
 }
 
 // IsClosed returns true if the reader is stopped.
@@ -238,7 +245,7 @@ func (r *Reader) read() {
 		r.lock.Lock()
 		for _, events := range streamsEvents {
 			streamName := events.Stream[len(streamKeyPrefix):]
-			streamEvents(streamName, events.Stream, "", events.Messages, r.eventFilter, r.chans, r.rdb, r.logger)
+			streamEvents(streamName, events.Stream, "", events.Messages, r.eventFilter, r.chans, r.donechan, r.rdb, r.logger)
 			for i := range r.streamKeys {
 				if r.streamKeys[i] == events.Stream {
 					r.streamCursors[i] = events.Messages[len(events.Messages)-1].ID
@@ -312,6 +319,7 @@ func streamEvents(
 	msgs []redis.XMessage,
 	eventFilter eventFilterFunc,
 	chans []chan *Event,
+	done <-chan struct{},
 	rdb *redis.Client,
 	logger pulse.Logger,
 ) {
@@ -339,7 +347,17 @@ func streamEvents(
 		}
 		logger.Debug("event", "stream", streamName, "event", ev.EventName, "id", ev.ID, "channels", len(chans))
 		for _, c := range chans {
-			c <- ev
+			select {
+			case c <- ev:
+			case <-done:
+				// The reader/sink is closing; stop fanning out so the
+				// read loop can return and release its resources instead
+				// of blocking forever on a stalled subscriber. Any
+				// remaining subscribers and messages in this batch are
+				// abandoned: delivery is at-most-once and the reader/sink
+				// is being torn down, so partial fan-out is acceptable.
+				return
+			}
 		}
 	}
 }

streaming/reader_test.go

@@ -193,6 +193,53 @@ func TestRemoveReaderStream(t *testing.T) {
 	assert.Equal(t, []byte("payload3"), read.Payload)
 }
 
+func TestReaderCloseWithStalledSubscriber(t *testing.T) {
+	testName := strings.Replace(t.Name(), "/", "_", -1)
+	rdb := ptesting.NewRedisClient(t)
+	defer ptesting.CleanupRedis(t, rdb, true, testName)
+	ctx := ptesting.NewTestContext(t)
+	s, err := NewStream(testName, rdb, options.WithStreamLogger(pulse.ClueLogger(ctx)))
+	require.NoError(t, err)
+
+	// Tiny buffer so the read loop's fan-out send blocks after a couple of
+	// events when the subscriber never drains its channel.
+	reader, err := s.NewReader(ctx,
+		options.WithReaderStartAtOldest(),
+		options.WithReaderBlockDuration(testBlockDuration),
+		options.WithReaderBufferSize(1))
+	require.NoError(t, err)
+
+	// Subscribe but deliberately never read from the channel so the read
+	// loop fills the buffer and then parks on the next fan-out send.
+	c := reader.Subscribe()
+
+	// Add more events than the buffer can hold so the read loop parks on
+	// `c <- ev` inside streamEvents.
+	for range 5 {
+		_, err = s.Add(ctx, "event", []byte("payload"))
+		require.NoError(t, err)
+	}
+
+	// Wait until the buffer is full, which means the read loop has consumed
+	// events and is now parked on the fan-out send to the stalled subscriber.
+	require.Eventually(t, func() bool { return len(c) == cap(c) }, max, delay)
+
+	// Close must return even though the subscriber stalled the read loop.
+	done := make(chan struct{})
+	go func() {
+		reader.Close()
+		close(done)
+	}()
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("reader.Close() hung with a stalled subscriber")
+	}
+	assert.True(t, reader.IsClosed())
+
+	require.NoError(t, s.Destroy(ctx))
+}
+
 func TestEventCreatedAt(t *testing.T) {
 	rdb := ptesting.NewRedisClient(t)
 	defer ptesting.CleanupRedis(t, rdb, false, "")

streaming/sink.go

@@ -75,6 +75,8 @@ type (
 		donechan chan struct{}
 		// wait is the sink cleanup wait group.
 		wait sync.WaitGroup
+		// closeOnce is used to ensure the sink is closed only once.
+		closeOnce sync.Once
 		// closing is true if Close was called.
 		closing bool
 		// eventFilter is the event filter if any.
@@ -291,30 +293,33 @@ func (s *Sink) RemoveStream(ctx context.Context, stream *Stream) error {
 }
 
 // Close stops event polling, waits for all events to be processed, and closes the sink channel.
-// It is safe to call Close multiple times.
+// It is safe to call Close multiple times; concurrent callers block until the
+// first Close completes.
 func (s *Sink) Close(ctx context.Context) {
-	s.lock.Lock()
-	if s.closing {
+	s.closeOnce.Do(func() {
+		// Close donechan first, without holding the lock, so the signal
+		// reaches the read loop even when it is parked on a fan-out send
+		// to a stalled subscriber (which holds the lock). Otherwise Close
+		// would deadlock acquiring the lock the read loop never releases.
+		close(s.donechan)
+		s.lock.Lock()
+		s.closing = true
 		s.lock.Unlock()
-		return
-	}
-	s.closing = true
-	close(s.donechan)
-	s.lock.Unlock()
-	s.wait.Wait()
-	s.lock.Lock()
-	defer s.lock.Unlock()
-	for _, c := range s.chans {
-		close(c)
-	}
-	// Note: we do not delete the consumer from the keep-alive and consumer maps
-	// so that another instance may claim any pending messages.
-	s.consumersKeepAliveMap.Close()
-	for _, cm := range s.consumersMap {
-		cm.Close()
-	}
-	s.closed = true
-	s.logger.Info("closed")
+		s.wait.Wait()
+		s.lock.Lock()
+		defer s.lock.Unlock()
+		for _, c := range s.chans {
+			close(c)
+		}
+		// Note: we do not delete the consumer from the keep-alive and consumer maps
+		// so that another instance may claim any pending messages.
+		s.consumersKeepAliveMap.Close()
+		for _, cm := range s.consumersMap {
+			cm.Close()
+		}
+		s.closed = true
+		s.logger.Info("closed")
+	})
 }
 
 // IsClosed returns true if the sink was closed.
@@ -457,7 +462,7 @@ func (s *Sink) read(ctx context.Context) {
 		}
 		for _, events := range streams {
 			streamName := events.Stream[len(streamKeyPrefix):]
-			streamEvents(streamName, events.Stream, s.Name, events.Messages, s.eventFilter, s.chans, s.rdb, s.logger)
+			streamEvents(streamName, events.Stream, s.Name, events.Messages, s.eventFilter, s.chans, s.donechan, s.rdb, s.logger)
 		}
 		s.lock.Unlock()
 	}
@@ -574,7 +579,7 @@ func (s *Sink) claim(ctx context.Context, streamName string, args redis.XAutoCla
 	messages, start, err := s.rdb.XAutoClaim(ctx, &args).Result()
 	if len(messages) > 0 {
 		s.logger.Info("claimed", "stream", streamName, "messages", len(messages))
-		streamEvents(streamName, args.Stream, s.Name, messages, s.eventFilter, s.chans, s.rdb, s.logger)
+		streamEvents(streamName, args.Stream, s.Name, messages, s.eventFilter, s.chans, s.donechan, s.rdb, s.logger)
 	}
 	return start, err
 }

streaming/sink_test.go

@@ -19,6 +19,53 @@ var (
 	testAckDuration     = 50 * time.Millisecond
 )
 
+func TestSinkCloseWithStalledSubscriber(t *testing.T) {
+	testName := strings.Replace(t.Name(), "/", "_", -1)
+	rdb := ptesting.NewRedisClient(t)
+	defer ptesting.CleanupRedis(t, rdb, true, testName)
+	ctx := ptesting.NewTestContext(t)
+	s, err := NewStream(testName, rdb, options.WithStreamLogger(pulse.ClueLogger(ctx)))
+	require.NoError(t, err)
+
+	// Tiny buffer so the read loop's fan-out send blocks after a couple of
+	// events when the subscriber never drains its channel.
+	sink, err := s.NewSink(ctx, "sink",
+		options.WithSinkStartAtOldest(),
+		options.WithSinkBlockDuration(testBlockDuration),
+		options.WithSinkBufferSize(1))
+	require.NoError(t, err)
+
+	// Subscribe but deliberately never read from the channel so the read
+	// loop fills the buffer and then parks on the next fan-out send.
+	c := sink.Subscribe()
+
+	// Add more events than the buffer can hold so the read loop parks on
+	// `c <- ev` inside streamEvents.
+	for range 5 {
+		_, err = s.Add(ctx, "event", []byte("payload"))
+		require.NoError(t, err)
+	}
+
+	// Wait until the buffer is full, which means the read loop has consumed
+	// events and is now parked on the fan-out send to the stalled subscriber.
+	require.Eventually(t, func() bool { return len(c) == cap(c) }, max, delay)
+
+	// Close must return even though the subscriber stalled the read loop.
+	done := make(chan struct{})
+	go func() {
+		sink.Close(ctx)
+		close(done)
+	}()
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("sink.Close() hung with a stalled subscriber")
+	}
+	assert.True(t, sink.IsClosed())
+
+	require.NoError(t, s.Destroy(ctx))
+}
+
 func TestNewSink(t *testing.T) {
 	testName := strings.Replace(t.Name(), "/", "_", -1)
 	rdb := ptesting.NewRedisClient(t)