uopi/internal/broker/broker.go

// 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)
}