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"
	"time"

	"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

	maxInterval time.Duration // 0 = unlimited
	log         *slog.Logger
}

// Option is a functional option for Broker configuration.
type Option func(*Broker)

// WithMaxUpdateRate limits fan-out to at most hz updates per second per signal.
// When upstream delivers faster, intermediate values are coalesced: the most
// recent value in each interval is forwarded once the interval elapses.
// A value ≤ 0 disables rate limiting (the default).
func WithMaxUpdateRate(hz float64) Option {
	return func(b *Broker) {
		if hz > 0 {
			b.maxInterval = time.Duration(float64(time.Second) / hz)
		}
	}
}

// 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, opts ...Option) *Broker {
	b := &Broker{
		ctx:     ctx,
		sources: make(map[string]datasource.DataSource),
		subs:    make(map[SignalRef]*signalSub),
		log:     log,
	}
	for _, opt := range opts {
		opt(b)
	}
	return b
}

// 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.
//
// The broker mutex is NOT held while the upstream ds.Subscribe call is in
// progress — some data sources (e.g. EPICS) block until the channel connects
// or times out. Releasing the lock allows other goroutines to proceed.
func (b *Broker) Subscribe(ref SignalRef, ch chan<- Update) (func(), error) {
	b.mu.Lock()

	sub, exists := b.subs[ref]
	if exists {
		// Fast path: already subscribed upstream; just add this client.
		sub.mu.Lock()
		sub.clients[ch] = struct{}{}
		sub.mu.Unlock()
		b.mu.Unlock()
		return func() { b.unsubscribe(ref, ch) }, nil
	}

	// Slow path: need to start an upstream subscription.
	// Release the lock before the potentially-blocking ds.Subscribe call.
	ds, found := b.sources[ref.DS]
	b.mu.Unlock()

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

	// Re-acquire the lock to register the new subscription.
	// Another goroutine might have registered the same ref while the lock
	// was released; if so, use that subscription and discard ours.
	b.mu.Lock()
	if existing, ok := b.subs[ref]; ok {
		// Race: someone else subscribed first. Discard duplicate upstream.
		dsCancel()
		sub = existing
	} else {
		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()
	b.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.
//
// When b.maxInterval > 0, updates that arrive faster than the interval are
// coalesced: only the most recent value in each interval window is forwarded,
// delivered at the end of the interval by a flush ticker. This ensures the
// latest value is always eventually seen even when the source fires rapidly.
func (b *Broker) fanOut(ref SignalRef, sub *signalSub, rawCh <-chan datasource.Value) {
	maxInterval := b.maxInterval

	var lastSent time.Time
	var pending *Update

	// Only allocate a ticker when rate-limiting is configured.
	var flushC <-chan time.Time
	if maxInterval > 0 {
		t := time.NewTicker(maxInterval)
		defer t.Stop()
		flushC = t.C
	}

	dispatch := func(u Update) {
		sub.mu.RLock()
		for ch := range sub.clients {
			select {
			case ch <- u:
			default: // slow consumer: drop rather than block
			}
		}
		sub.mu.RUnlock()
	}

	for {
		select {
		case v, ok := <-rawCh:
			if !ok {
				return
			}
			update := Update{Ref: ref, Value: v}
			// Rate-limit: if an interval is configured and the last delivery
			// was recent, hold this update as pending (coalesce).
			if maxInterval > 0 && !lastSent.IsZero() && time.Since(lastSent) < maxInterval {
				pending = &update
				continue
			}
			lastSent = time.Now()
			pending = nil
			dispatch(update)

		case <-flushC:
			// Deliver the most recent coalesced value, if any.
			if pending != nil {
				lastSent = time.Now()
				dispatch(*pending)
				pending = nil
			}

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