v0.2.0 (#50)

Co-authored-by: Tasnim Kabir Sadik <tksadik@omukk.dev>

Reviewed-on: wrenn/wrenn#50

pkg/service/build_broker.go

@@ -0,0 +1,143 @@
+package service
+
+import (
+	"context"
+	"encoding/json"
+	"log/slog"
+	"sync"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+)
+
+// buildSubBuffer is the per-subscriber channel buffer. A slow WebSocket
+// consumer that fills the buffer drops live events; it recovers the full
+// build state from the DB log on reconnect.
+const buildSubBuffer = 256
+
+// buildBrokerReconnect is the backoff before re-subscribing to Redis after a
+// subscription error.
+const buildBrokerReconnect = 2 * time.Second
+
+// BuildBroker fans build events out from per-build Redis pub/sub channels to
+// in-process WebSocket subscribers. A Redis subscription is started lazily for
+// a build when its first client connects and torn down when the last leaves.
+//
+// The build worker publishes via publishBuildEvent (Redis only); the broker is
+// purely the read/fan-out side. Decoupling through Redis means the worker and
+// the WebSocket handler need not run in the same process.
+type BuildBroker struct {
+	rdb    *redis.Client
+	mu     sync.Mutex
+	builds map[string]*buildFanout
+}
+
+type buildFanout struct {
+	subs   map[chan BuildStreamEvent]struct{}
+	cancel context.CancelFunc
+}
+
+// NewBuildBroker creates a broker reading from the given Redis client.
+func NewBuildBroker(rdb *redis.Client) *BuildBroker {
+	return &BuildBroker{rdb: rdb, builds: make(map[string]*buildFanout)}
+}
+
+// Subscribe registers an in-process subscriber for buildID's event stream and
+// returns the receive channel plus a release function. The first subscriber
+// for a build starts its Redis subscription; the last to release stops it.
+// The release function is idempotent and closes the channel.
+func (b *BuildBroker) Subscribe(buildID string) (<-chan BuildStreamEvent, func()) {
+	ch := make(chan BuildStreamEvent, buildSubBuffer)
+
+	b.mu.Lock()
+	fan, ok := b.builds[buildID]
+	if !ok {
+		ctx, cancel := context.WithCancel(context.Background())
+		fan = &buildFanout{subs: make(map[chan BuildStreamEvent]struct{}), cancel: cancel}
+		b.builds[buildID] = fan
+		go b.run(ctx, buildID)
+	}
+	fan.subs[ch] = struct{}{}
+	b.mu.Unlock()
+
+	var once sync.Once
+	release := func() {
+		once.Do(func() {
+			b.mu.Lock()
+			defer b.mu.Unlock()
+			fan, ok := b.builds[buildID]
+			if !ok {
+				return
+			}
+			if _, present := fan.subs[ch]; !present {
+				return
+			}
+			delete(fan.subs, ch)
+			close(ch)
+			if len(fan.subs) == 0 {
+				fan.cancel()
+				delete(b.builds, buildID)
+			}
+		})
+	}
+	return ch, release
+}
+
+// run keeps a Redis subscription alive for buildID, reconnecting on error,
+// until the fanout's context is cancelled (last subscriber left).
+func (b *BuildBroker) run(ctx context.Context, buildID string) {
+	for ctx.Err() == nil {
+		b.subscribeOnce(ctx, buildID)
+		if ctx.Err() != nil {
+			return
+		}
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.After(buildBrokerReconnect):
+		}
+	}
+}
+
+func (b *BuildBroker) subscribeOnce(ctx context.Context, buildID string) {
+	sub := b.rdb.Subscribe(ctx, buildStreamChannel(buildID))
+	defer sub.Close()
+
+	msgCh := sub.Channel()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case msg, ok := <-msgCh:
+			if !ok {
+				return
+			}
+			var ev BuildStreamEvent
+			if err := json.Unmarshal([]byte(msg.Payload), &ev); err != nil {
+				slog.Warn("build broker: bad event payload", "build_id", buildID, "error", err)
+				continue
+			}
+			b.dispatch(buildID, ev)
+		}
+	}
+}
+
+// dispatch fans one event to every in-process subscriber. The send is
+// non-blocking; a full subscriber buffer drops the event. The mutex is held
+// for the whole dispatch so a concurrent release cannot close a channel
+// mid-send.
+func (b *BuildBroker) dispatch(buildID string, ev BuildStreamEvent) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	fan, ok := b.builds[buildID]
+	if !ok {
+		return
+	}
+	for ch := range fan.subs {
+		select {
+		case ch <- ev:
+		default:
+			slog.Debug("build broker: dropped event for slow consumer", "build_id", buildID)
+		}
+	}
+}