uopi/internal/broker/broker_stress_test.go

package broker_test

import (
	"context"
	"fmt"
	"log/slog"
	"runtime"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/uopi/uopi/internal/broker"
	"github.com/uopi/uopi/internal/datasource/stub"
)

// newBrokerN builds a broker backed by a stub with n dynamically generated signals.
func newBrokerN(tb testing.TB, n int) (*broker.Broker, context.CancelFunc) {
	tb.Helper()
	ctx, cancel := context.WithCancel(context.Background())
	log := slog.Default()
	b := broker.New(ctx, log)
	ds := stub.NewN(n)
	if err := ds.Connect(ctx); err != nil {
		cancel()
		tb.Fatal(err)
	}
	b.Register(ds)
	return b, cancel
}

// TestStress_ManySignalsManyClients subscribes 20 clients to 500 signals each,
// runs for 2 seconds, then verifies delivery counts and zero subscription leaks.
func TestStress_ManySignalsManyClients(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping stress test in short mode")
	}

	const (
		nSignals = 500
		nClients = 20
		duration = 2 * time.Second
	)

	goroutinesBefore := runtime.NumGoroutine()

	brk, cancel := newBrokerN(t, nSignals)
	defer cancel()

	// Each client has its own channel and subscribes to all nSignals signals.
	type clientState struct {
		ch    chan broker.Update
		unsub []func()
	}
	clients := make([]clientState, nClients)

	for i := range nClients {
		ch := make(chan broker.Update, 2048)
		unsubs := make([]func(), nSignals)
		for j := range nSignals {
			ref := broker.SignalRef{DS: "stub", Name: fmt.Sprintf("pv_%d", j)}
			unsub, err := brk.Subscribe(ref, ch)
			if err != nil {
				t.Fatalf("client %d subscribe pv_%d: %v", i, j, err)
			}
			unsubs[j] = unsub
		}
		clients[i] = clientState{ch: ch, unsub: unsubs}
	}

	if n := brk.ActiveSubscriptions(); n != nSignals {
		t.Errorf("expected %d active upstream subscriptions, got %d", nSignals, n)
	}

	// Drain updates concurrently for the test duration.
	var totalReceived atomic.Int64
	stop := make(chan struct{})
	var wg sync.WaitGroup

	for i := range nClients {
		wg.Add(1)
		ch := clients[i].ch
		go func() {
			defer wg.Done()
			var n int64
			for {
				select {
				case <-ch:
					n++
				case <-stop:
					totalReceived.Add(n)
					return
				}
			}
		}()
	}

	time.Sleep(duration)
	close(stop)
	wg.Wait()

	total := totalReceived.Load()
	// At 10 Hz for duration seconds: nSignals * nClients * duration.Seconds() * 10
	minExpected := int64(nSignals) * int64(nClients) * int64(duration.Seconds()) * 5 // conservative: 50%
	t.Logf("received %d updates (%d clients × %d signals × %.0fs @ 10Hz, min=%d)",
		total, nClients, nSignals, duration.Seconds(), minExpected)
	if total < minExpected {
		t.Errorf("too few updates: got %d, want >= %d", total, minExpected)
	}

	// Unsubscribe all clients.
	for _, c := range clients {
		for _, unsub := range c.unsub {
			unsub()
		}
	}

	// Give the broker time to tear down upstream subscriptions.
	time.Sleep(200 * time.Millisecond)

	if n := brk.ActiveSubscriptions(); n != 0 {
		t.Errorf("subscription leak: %d active subscriptions remain after all clients unsubscribed", n)
	}

	cancel()
	time.Sleep(300 * time.Millisecond)

	goroutinesAfter := runtime.NumGoroutine()
	leaked := goroutinesAfter - goroutinesBefore
	t.Logf("goroutines: before=%d after=%d delta=%d", goroutinesBefore, goroutinesAfter, leaked)
	if leaked > 20 {
		t.Errorf("goroutine leak: started with %d, ended with %d (%d extra)", goroutinesBefore, goroutinesAfter, leaked)
	}
}

// TestStress_RapidSubscribeUnsubscribe hammers subscribe/unsubscribe on a small
// signal set from many goroutines concurrently to expose races and deadlocks.
func TestStress_RapidSubscribeUnsubscribe(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping stress test in short mode")
	}

	const (
		nSignals    = 20
		nGoroutines = 50
		duration    = 2 * time.Second
	)

	brk, cancel := newBrokerN(t, nSignals)
	defer cancel()

	var wg sync.WaitGroup
	stop := make(chan struct{})
	var ops atomic.Int64

	for range nGoroutines {
		wg.Add(1)
		go func() {
			defer wg.Done()
			ch := make(chan broker.Update, 64)
			for {
				select {
				case <-stop:
					return
				default:
				}
				// Pick a signal and rapidly subscribe/unsubscribe.
				idx := int(ops.Load()) % nSignals
				ref := broker.SignalRef{DS: "stub", Name: fmt.Sprintf("pv_%d", idx)}
				unsub, err := brk.Subscribe(ref, ch)
				if err != nil {
					t.Errorf("subscribe pv_%d: %v", idx, err)
					return
				}
				ops.Add(1)
				// Drain a bit before unsubscribing.
				select {
				case <-ch:
				case <-time.After(50 * time.Millisecond):
				}
				unsub()
			}
		}()
	}

	time.Sleep(duration)
	close(stop)
	wg.Wait()

	t.Logf("completed %d subscribe/unsubscribe cycles in %s", ops.Load(), duration)

	time.Sleep(200 * time.Millisecond)
	if n := brk.ActiveSubscriptions(); n != 0 {
		t.Errorf("subscription leak after rapid churn: %d remain", n)
	}
}

// TestStress_SharedSignalManyClients verifies that a single high-rate signal
// fans out correctly to many concurrent clients with no drops for fast consumers.
func TestStress_SharedSignalManyClients(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping stress test in short mode")
	}

	const (
		nClients = 50
		duration = 2 * time.Second
	)

	brk, cancel := newBroker(t) // uses stub.New() which has counter_fast at 1 ms
	defer cancel()

	ref := broker.SignalRef{DS: "stub", Name: "counter_fast"}
	channels := make([]chan broker.Update, nClients)
	unsubs := make([]func(), nClients)

	for i := range nClients {
		ch := make(chan broker.Update, 4096)
		channels[i] = ch
		unsub, err := brk.Subscribe(ref, ch)
		if err != nil {
			t.Fatalf("subscribe client %d: %v", i, err)
		}
		unsubs[i] = unsub
	}
	defer func() {
		for _, u := range unsubs {
			u()
		}
	}()

	if n := brk.ActiveSubscriptions(); n != 1 {
		t.Errorf("expected exactly 1 upstream subscription for shared signal, got %d", n)
	}

	var totalReceived atomic.Int64
	stop := make(chan struct{})
	var wg sync.WaitGroup

	for i := range nClients {
		wg.Add(1)
		ch := channels[i]
		go func() {
			defer wg.Done()
			var n int64
			for {
				select {
				case <-ch:
					n++
				case <-stop:
					totalReceived.Add(n)
					return
				}
			}
		}()
	}

	time.Sleep(duration)
	close(stop)
	wg.Wait()

	total := totalReceived.Load()
	// counter_fast ticks at 1ms → ~1000 Hz → 2000 updates × 50 clients = 100000 expected minimum
	minExpected := int64(nClients) * int64(duration.Seconds()) * 500 // 50% delivery ok due to buffering
	t.Logf("shared signal: received %d updates across %d clients (min=%d)", total, nClients, minExpected)
	if total < minExpected {
		t.Errorf("too few updates: got %d, want >= %d", total, minExpected)
	}
}