From 8f723623fea9649554ab385912b3e2ae54943683 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <m@mitchellh.com>
Date: Fri, 1 May 2026 13:21:06 -0700
Subject: [PATCH] eventconsumer: add RemoveSource and consolidate per-source
 state

Adds `RemoveSource` to stop streaming from a previously registered
source. Removing a source cancels its connection loop and closes any
in-flight websocket so the loop exits promptly instead of waiting for
the next reconnect tick.

Per-source runtime state (the loop's cancel func and the active
websocket conn) was previously split across two sync.Maps and a
separate map of registered sources, each with its own locking. That
left several races, e.g. the conn was stored in connMap only after a successful
dial,  so a remove during the dial would miss it.

The new design collapses everything into a single
`sources map[Source]*sourceState` guarded by `sourcesMu`. The lock is only
held for short, non-blocking map mutations and is always released before
any side effect possibly-blocking calls.

`runConnection` now re-checks under the lock that the source is still
registered (and the ctx not cancelled) before installing a freshly
dialed conn, and its deferred cleanup only clears the conn slot if the
entry still points at the conn it installed. This makes 'remove
during dial' and 'remove during read' both deterministic: the cancel
cuts the loop, the close breaks ReadMessage, and a concurrent
runConnection cannot resurrect state for a removed source.
---
 eventconsumer/consumer.go | 100 ++++++++++++++++++++++++++++++--------
 1 file changed, 80 insertions(+), 20 deletions(-)

diff --git a/eventconsumer/consumer.go b/eventconsumer/consumer.go
index 25b22734..acce0b0a 100644
--- a/eventconsumer/consumer.go
+++ b/eventconsumer/consumer.go
@@ -54,14 +54,21 @@ type Source interface {
 type Consumer struct {
 	wg         sync.WaitGroup
 	dialer     *websocket.Dialer
-	connMap    sync.Map
 	jobQueue   chan job
 	logger     *slog.Logger
 	randSource *rand.Rand
 
-	// rw lock over edits to ConsumerConfig
-	cfgMu sync.RWMutex
-	cfg   ConsumerConfig
+	// sourcesMu guards sources. It must only be held for short, non-blocking
+	// map operations; never across a blocking call (dial, read, close).
+	sourcesMu sync.Mutex
+	sources   map[Source]*sourceState
+
+	cfg ConsumerConfig
+}
+
+type sourceState struct {
+	cancel context.CancelFunc
+	conn   *websocket.Conn
 }
 
 type job struct {
@@ -97,6 +104,7 @@ func NewConsumer(cfg ConsumerConfig) *Consumer {
 		jobQueue:   make(chan job, cfg.QueueSize), // buffered job queue
 		logger:     cfg.Logger,
 		randSource: rand.New(rand.NewSource(time.Now().UnixNano())),
+		sources:    make(map[Source]*sourceState),
 	}
 }
 
@@ -111,34 +119,64 @@ func (c *Consumer) Start(ctx context.Context) {
 
 	// start streaming
 	for source := range c.cfg.Sources {
-		c.wg.Add(1)
-		go c.startConnectionLoop(ctx, source)
+		c.AddSource(ctx, source)
 	}
 }
 
 func (c *Consumer) Stop() {
-	c.connMap.Range(func(_, val any) bool {
-		if conn, ok := val.(*websocket.Conn); ok {
-			conn.Close()
+	// snapshot conns under lock so we don't hold sourcesMu across Close
+	c.sourcesMu.Lock()
+	conns := make([]*websocket.Conn, 0, len(c.sources))
+	for _, st := range c.sources {
+		if st.conn != nil {
+			conns = append(conns, st.conn)
 		}
-		return true
-	})
+	}
+	c.sourcesMu.Unlock()
+
+	for _, conn := range conns {
+		conn.Close()
+	}
+
 	c.wg.Wait()
 	close(c.jobQueue)
 }
 
 func (c *Consumer) AddSource(ctx context.Context, s Source) {
-	// we are already listening to this source
-	if _, ok := c.cfg.Sources[s]; ok {
+	c.sourcesMu.Lock()
+	if _, ok := c.sources[s]; ok {
+		c.sourcesMu.Unlock()
 		c.logger.Info("source already present", "source", s)
 		return
 	}
+	srcCtx, cancel := context.WithCancel(ctx)
+	c.sources[s] = &sourceState{cancel: cancel}
+	c.sourcesMu.Unlock()
 
-	c.cfgMu.Lock()
-	c.cfg.Sources[s] = struct{}{}
 	c.wg.Add(1)
-	go c.startConnectionLoop(ctx, s)
-	c.cfgMu.Unlock()
+	go c.startConnectionLoop(srcCtx, s)
+}
+
+func (c *Consumer) RemoveSource(s Source) {
+	c.sourcesMu.Lock()
+	st, ok := c.sources[s]
+	if !ok {
+		c.sourcesMu.Unlock()
+		c.logger.Info("source not present", "source", s)
+		return
+	}
+	delete(c.sources, s)
+	cancel := st.cancel
+	conn := st.conn
+	c.sourcesMu.Unlock()
+
+	// release lock before any potentially blocking call
+	if cancel != nil {
+		cancel()
+	}
+	if conn != nil {
+		conn.Close()
+	}
 }
 
 func (c *Consumer) worker(ctx context.Context) {
@@ -238,9 +276,31 @@ func (c *Consumer) runConnection(ctx context.Context, source Source) error {
 		return err
 	}
 
-	c.connMap.Store(source, conn)
-	defer conn.Close()
-	defer c.connMap.Delete(source)
+	// Register the conn. If the source was removed (or our ctx cancelled)
+	// while we were dialing, drop this conn instead of installing it.
+	c.sourcesMu.Lock()
+	st, ok := c.sources[source]
+	if !ok || ctx.Err() != nil {
+		c.sourcesMu.Unlock()
+		conn.Close()
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+		return nil
+	}
+	st.conn = conn
+	c.sourcesMu.Unlock()
+
+	defer func() {
+		// Clear the conn from state, but only if it's still our conn (a
+		// concurrent RemoveSource may have already done it).
+		c.sourcesMu.Lock()
+		if st, ok := c.sources[source]; ok && st.conn == conn {
+			st.conn = nil
+		}
+		c.sourcesMu.Unlock()
+		conn.Close()
+	}()
 
 	c.logger.Info("connected", "source", source)
 
-- 
2.51.2