Phase 2/4 done, working on phase 3/5

internal/broker/broker.go

@@ -0,0 +1,177 @@
+// Package broker multiplexes signal subscriptions from data sources to multiple
+// WebSocket clients. For each unique signal only one upstream subscription is
+// maintained regardless of how many clients are watching it.
+package broker
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"sync"
+
+	"github.com/uopi/uopi/internal/datasource"
+)
+
+// SignalRef identifies a signal within a named data source.
+type SignalRef struct {
+	DS   string
+	Name string
+}
+
+// Update is delivered to every subscriber when a signal value changes.
+type Update struct {
+	Ref   SignalRef
+	Value datasource.Value
+}
+
+// signalSub holds the state for one active upstream signal subscription and its
+// set of downstream client channels.
+type signalSub struct {
+	mu      sync.RWMutex
+	clients map[chan<- Update]struct{}
+	done    chan struct{} // closed to stop the fanOut goroutine
+	dsStop  datasource.CancelFunc
+}
+
+// Broker manages data-source registrations and fan-out of signal updates.
+type Broker struct {
+	ctx context.Context
+
+	mu      sync.Mutex
+	sources map[string]datasource.DataSource
+	subs    map[SignalRef]*signalSub
+
+	log *slog.Logger
+}
+
+// New creates a Broker whose upstream subscriptions are bound to ctx.
+// Cancel ctx (or the parent context passed to main) to shut everything down.
+func New(ctx context.Context, log *slog.Logger) *Broker {
+	return &Broker{
+		ctx:     ctx,
+		sources: make(map[string]datasource.DataSource),
+		subs:    make(map[SignalRef]*signalSub),
+		log:     log,
+	}
+}
+
+// Register adds a DataSource to the broker. Must be called before Subscribe.
+func (b *Broker) Register(ds datasource.DataSource) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	b.sources[ds.Name()] = ds
+	b.log.Info("data source registered", "name", ds.Name())
+}
+
+// DataSources returns all registered sources.
+func (b *Broker) DataSources() []datasource.DataSource {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	out := make([]datasource.DataSource, 0, len(b.sources))
+	for _, ds := range b.sources {
+		out = append(out, ds)
+	}
+	return out
+}
+
+// Source returns the named data source, if registered.
+func (b *Broker) Source(name string) (datasource.DataSource, bool) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	ds, ok := b.sources[name]
+	return ds, ok
+}
+
+// Subscribe registers ch to receive updates for ref.
+// If this is the first subscriber for the signal, the upstream DS subscription
+// is started. Subsequent calls for the same signal share the existing one.
+// The returned func must be called when the client no longer needs the signal.
+func (b *Broker) Subscribe(ref SignalRef, ch chan<- Update) (func(), error) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	sub, exists := b.subs[ref]
+	if !exists {
+		ds, found := b.sources[ref.DS]
+		if !found {
+			return nil, fmt.Errorf("unknown data source %q", ref.DS)
+		}
+
+		rawCh := make(chan datasource.Value, 32)
+		dsCancel, err := ds.Subscribe(b.ctx, ref.Name, rawCh)
+		if err != nil {
+			return nil, fmt.Errorf("subscribe %s/%s: %w", ref.DS, ref.Name, err)
+		}
+
+		sub = &signalSub{
+			clients: make(map[chan<- Update]struct{}),
+			done:    make(chan struct{}),
+			dsStop:  dsCancel,
+		}
+		b.subs[ref] = sub
+		go b.fanOut(ref, sub, rawCh)
+		b.log.Debug("upstream subscription started", "ds", ref.DS, "signal", ref.Name)
+	}
+
+	sub.mu.Lock()
+	sub.clients[ch] = struct{}{}
+	sub.mu.Unlock()
+
+	return func() { b.unsubscribe(ref, ch) }, nil
+}
+
+func (b *Broker) unsubscribe(ref SignalRef, ch chan<- Update) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	sub, ok := b.subs[ref]
+	if !ok {
+		return
+	}
+
+	sub.mu.Lock()
+	delete(sub.clients, ch)
+	remaining := len(sub.clients)
+	sub.mu.Unlock()
+
+	if remaining == 0 {
+		close(sub.done) // terminates fanOut goroutine
+		sub.dsStop()
+		delete(b.subs, ref)
+		b.log.Debug("upstream subscription torn down", "ds", ref.DS, "signal", ref.Name)
+	}
+}
+
+// fanOut reads values from rawCh and dispatches them to all registered clients.
+// It exits when sub.done is closed or rawCh is closed.
+func (b *Broker) fanOut(ref SignalRef, sub *signalSub, rawCh <-chan datasource.Value) {
+	for {
+		select {
+		case v, ok := <-rawCh:
+			if !ok {
+				return
+			}
+			update := Update{Ref: ref, Value: v}
+			sub.mu.RLock()
+			for ch := range sub.clients {
+				select {
+				case ch <- update:
+				default:
+					// slow consumer: drop rather than block
+				}
+			}
+			sub.mu.RUnlock()
+
+		case <-sub.done:
+			return
+		}
+	}
+}
+
+// ActiveSubscriptions returns the number of currently active upstream signal
+// subscriptions. Useful for diagnostics and tests.
+func (b *Broker) ActiveSubscriptions() int {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return len(b.subs)
+}