fix: unblock Repository producers when a subscriber disconnects mid-replay

interface.go

@@ -5,7 +5,10 @@
 // If the Repository interface is implemented on the server, events can be replayed in case of a network disconnection.
 package eventsource
 
-import "net/http"
+import (
+	"context"
+	"net/http"
+)
 
 // Event is the interface for any event received by the client or sent by the server.
 type Event interface {
@@ -61,6 +64,45 @@ type Repository interface {
 	Replay(channel, id string) chan Event
 }
 
+// RepositoryWithContext is an optional extension of Repository that allows a Repository to be
+// notified when the subscriber that requested a replay has gone away.
+//
+// A plain Repository.Replay receives only the channel and event id, and returns a bare channel;
+// it has no way to learn that the subscriber disconnected. If the subscriber disconnects while
+// the Repository's producer goroutine is still blocked sending an event, that goroutine can be
+// stranded until the server drains the channel on its behalf.
+//
+// If a registered Repository also implements RepositoryWithContext, the Server will call
+// ReplayWithContext instead of Replay, passing the subscribing request's context. That context
+// is cancelled when the subscriber disconnects (or when the connection is otherwise terminated),
+// so the producer can select on ctx.Done() and stop sending immediately, for example:
+//
+//	func (r *myRepo) ReplayWithContext(ctx context.Context, channel, id string) <-chan eventsource.Event {
+//	    out := make(chan eventsource.Event)
+//	    go func() {
+//	        defer close(out)
+//	        for _, event := range r.eventsFor(channel, id) {
+//	            select {
+//	            case out <- event:
+//	            case <-ctx.Done():
+//	                return
+//	            }
+//	        }
+//	    }()
+//	    return out
+//	}
+//
+// Implementing this interface is optional and does not change the behavior of Replay. A Repository
+// that implements only Replay continues to work unchanged; a Repository that implements both must
+// still provide Replay for backward compatibility.
+type RepositoryWithContext interface {
+	Repository
+	// ReplayWithContext behaves like Replay, but additionally receives a context that is cancelled
+	// when the subscriber goes away. It has the same channel-closing responsibilities as Replay,
+	// and may likewise return nil if there are no events to be sent.
+	ReplayWithContext(ctx context.Context, channel, id string) <-chan Event
+}
+
 // Logger is the interface for a custom logging implementation that can handle log output for a Stream.
 type Logger interface {
 	Println(...interface{})

server.go

@@ -1,6 +1,7 @@
 package eventsource
 
 import (
+	"context"
 	"net/http"
 	"strings"
 	"sync"
@@ -11,6 +12,9 @@ type subscription struct {
 	channel     string
 	lastEventID string
 	out         chan<- eventOrComment
+	// ctx is the subscribing request's context. It is cancelled when the subscriber disconnects,
+	// and is passed to a Repository that implements RepositoryWithContext.
+	ctx context.Context
 }
 
 type eventOrComment interface{}
@@ -131,6 +135,7 @@ func (srv *Server) Handler(channel string) http.HandlerFunc {
 			channel:     channel,
 			lastEventID: req.Header.Get("Last-Event-ID"),
 			out:         eventCh,
+			ctx:         req.Context(),
 		}
 		srv.subs <- sub
 		flusher := w.(http.Flusher)
@@ -270,12 +275,33 @@ func (srv *Server) Handler(channel string) http.HandlerFunc {
 				}
 			}
 		}
+		if readBatchCh != nil {
+			// We are exiting the read loop while still in the middle of consuming a batch of replayed
+			// events from a Repository (e.g. the subscriber disconnected, or MaxConnTime elapsed). The
+			// Repository's producer goroutine may be blocked trying to send the remaining events on this
+			// channel. Since we hold the receiving end -- we can neither close it nor keep reading it on
+			// this exiting goroutine -- drain it in the background so the producer can unblock and release
+			// its resources promptly rather than leaking until the process exits.
+			//
+			// A Repository that implements RepositoryWithContext will already have been told to stop via
+			// context cancellation, but draining is harmless in that case and remains the safety net for
+			// repositories that only implement Replay.
+			go drainReplayedEvents(readBatchCh)
+		}
 		if !closedNormally {
 			srv.unsubs <- sub // the server didn't tell us to close, so we must tell it that we're closing
 		}
 	}
 }
 
+// drainReplayedEvents consumes and discards the events from a Repository replay batch channel that
+// no subscriber will read again, so the goroutine producing those events can complete and release
+// its resources instead of blocking forever on a send.
+func drainReplayedEvents(ch <-chan Event) {
+	for range ch { //nolint:revive // draining until the channel is closed
+	}
+}
+
 // Register registers a Repository to be used for the specified channel. The Repository will be used to
 // determine whether new subscribers should receive data that was generated before they subscribed.
 //
@@ -383,9 +409,22 @@ func (srv *Server) run() {
 			if srv.ReplayAll || len(sub.lastEventID) > 0 {
 				repo, ok := repos[sub.channel]
 				if ok {
-					batchCh := repo.Replay(sub.channel, sub.lastEventID)
-					if batchCh != nil {
-						trySend(sub, eventBatch{events: batchCh})
+					// If the repository supports it, pass the subscriber's context so its producer can
+					// stop sending promptly when the subscriber disconnects. Otherwise fall back to the
+					// original context-less Replay; the handler's background drain (see Handler) still
+					// ensures such a producer eventually unblocks.
+					var batchCh <-chan Event
+					if repoCtx, ok := repo.(RepositoryWithContext); ok {
+						batchCh = repoCtx.ReplayWithContext(sub.ctx, sub.channel, sub.lastEventID)
+					} else {
+						batchCh = repo.Replay(sub.channel, sub.lastEventID)
+					}
+					if batchCh != nil && !sub.send(eventBatch{events: batchCh}) {
+						// The subscriber was already closed, so it will never consume this batch and
+						// its producer would otherwise block forever; drain it in the background.
+						sub.close()
+						delete(subs[sub.channel], sub)
+						go drainReplayedEvents(batchCh)
 					}
 				}
 			}