Improved perfs

2026-05-12 10:16:48 +02:00
parent 912ecdd9ed
commit 6ff8fb5c25
14 changed files with 1357 additions and 25 deletions
@@ -0,0 +1,343 @@
+package server
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log/slog"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/coder/websocket"
+
+	"github.com/uopi/uopi/internal/broker"
+	"github.com/uopi/uopi/internal/datasource/stub"
+)
+
+// newStressServer creates an httptest.Server backed by a stub with n signals.
+// The server and broker contexts are tied to t.Context() and cleaned up automatically.
+func newStressServer(t *testing.T, nSignals int) (*httptest.Server, *broker.Broker) {
+	t.Helper()
+	ctx := t.Context()
+	log := slog.New(slog.NewTextHandler(io.Discard, nil))
+
+	var ds *stub.Stub
+	if nSignals > 0 {
+		ds = stub.NewN(nSignals)
+	} else {
+		ds = stub.New()
+	}
+	if err := ds.Connect(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	brk := broker.New(ctx, log)
+	brk.Register(ds)
+
+	mux := http.NewServeMux()
+	mux.Handle("/ws", &wsHandler{broker: brk, log: log})
+
+	srv := httptest.NewServer(mux)
+	t.Cleanup(srv.Close)
+	return srv, brk
+}
+
+// wsURL converts an http:// test server URL to ws://.
+func wsURL(srv *httptest.Server) string {
+	return strings.ReplaceAll(srv.URL, "http://", "ws://") + "/ws"
+}
+
+// buildSubMsg encodes a subscribe message for n stub PVs named "pv_0"…"pv_{n-1}".
+func buildSubMsg(n int) []byte {
+	sigs := make([]sigRef, n)
+	for i := range n {
+		sigs[i] = sigRef{DS: "stub", Name: fmt.Sprintf("pv_%d", i)}
+	}
+	data, _ := json.Marshal(inMsg{Type: "subscribe", Signals: sigs})
+	return data
+}
+
+// connectAndCount dials the WebSocket, sends subMsg, counts incoming messages until
+// ctx is cancelled, adds the count to *total, and signals wg.Done.
+func connectAndCount(ctx context.Context, t *testing.T, url string, subMsg []byte, total *atomic.Int64, wg *sync.WaitGroup) {
+	t.Helper()
+	defer wg.Done()
+
+	conn, _, err := websocket.Dial(ctx, url, nil)
+	if err != nil {
+		t.Errorf("dial: %v", err)
+		return
+	}
+	defer conn.CloseNow()
+
+	if err := conn.Write(ctx, websocket.MessageText, subMsg); err != nil {
+		t.Errorf("subscribe write: %v", err)
+		return
+	}
+
+	var n int64
+	for {
+		_, _, err := conn.Read(ctx)
+		if err != nil {
+			break
+		}
+		n++
+	}
+	total.Add(n)
+}
+
+// TestStress_100PVs_10Clients verifies the full stack (broker + WS handler) can
+// deliver updates from 100 signals to 10 concurrent clients without loss or deadlock.
+func TestStress_100PVs_10Clients(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping stress test in short mode")
+	}
+
+	const (
+		nSignals = 100
+		nClients = 10
+		duration = 3 * time.Second
+	)
+
+	srv, brk := newStressServer(t, nSignals)
+	url := wsURL(srv)
+	subMsg := buildSubMsg(nSignals)
+
+	ctx, cancel := context.WithTimeout(t.Context(), duration+5*time.Second)
+	defer cancel()
+
+	var total atomic.Int64
+	var wg sync.WaitGroup
+
+	for range nClients {
+		wg.Add(1)
+		// Each client runs for exactly `duration`, then its context expires and
+		// conn.Read returns an error, causing the goroutine to exit cleanly.
+		clientCtx, clientCancel := context.WithTimeout(ctx, duration)
+		go func() {
+			defer clientCancel()
+			connectAndCount(clientCtx, t, url, subMsg, &total, &wg)
+		}()
+	}
+
+	wg.Wait()
+
+	got := total.Load()
+	// At 10 Hz, each PV fires ≥ duration*10 times; each client receives all of them.
+	// min = nClients × nSignals × duration_s × 10 × 0.5 (conservative — meta msgs too)
+	minExpected := int64(nClients) * int64(nSignals) * int64(duration.Seconds()) * 5
+	t.Logf("100PVs/10Clients: received %d updates (min expected %d)", got, minExpected)
+	if got < minExpected {
+		t.Errorf("too few updates: got %d, want >= %d", got, minExpected)
+	}
+
+	// Broker must release all upstream subscriptions after clients disconnect.
+	time.Sleep(300 * time.Millisecond)
+	if n := brk.ActiveSubscriptions(); n != 0 {
+		t.Errorf("subscription leak: %d remain after all clients disconnected", n)
+	}
+}
+
+// TestStress_500PVs_20Clients is the full-scale scenario: hundreds of signals,
+// tens of clients. Verifies no crash, no deadlock, and clean teardown.
+func TestStress_500PVs_20Clients(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping stress test in short mode")
+	}
+
+	const (
+		nSignals = 500
+		nClients = 20
+		duration = 3 * time.Second
+	)
+
+	srv, brk := newStressServer(t, nSignals)
+	url := wsURL(srv)
+	subMsg := buildSubMsg(nSignals)
+
+	ctx, cancel := context.WithTimeout(t.Context(), duration+10*time.Second)
+	defer cancel()
+
+	var total atomic.Int64
+	var wg sync.WaitGroup
+
+	for range nClients {
+		wg.Add(1)
+		clientCtx, clientCancel := context.WithTimeout(ctx, duration)
+		go func() {
+			defer clientCancel()
+			connectAndCount(clientCtx, t, url, subMsg, &total, &wg)
+		}()
+	}
+
+	wg.Wait()
+
+	got := total.Load()
+	// Conservative: expect at least 20% of the theoretical maximum
+	minExpected := int64(nClients) * int64(nSignals) * int64(duration.Seconds()) * 2
+	t.Logf("500PVs/20Clients: received %d updates (min expected %d)", got, minExpected)
+	if got < minExpected {
+		t.Errorf("too few updates: got %d, want >= %d", got, minExpected)
+	}
+
+	time.Sleep(500 * time.Millisecond)
+	if n := brk.ActiveSubscriptions(); n != 0 {
+		t.Errorf("subscription leak: %d remain after all clients disconnected", n)
+	}
+}
+
+// TestStress_ClientChurn repeatedly connects and disconnects clients while a
+// stable set of signals is being delivered to ensure subscribe/unsubscribe under
+// load doesn't cause races, panics, or subscription leaks.
+func TestStress_ClientChurn(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping stress test in short mode")
+	}
+
+	const (
+		nSignals       = 50
+		nParallelConns = 30
+		duration       = 3 * time.Second
+	)
+
+	srv, brk := newStressServer(t, nSignals)
+	u := wsURL(srv)
+	subMsg := buildSubMsg(nSignals)
+
+	deadline := time.Now().Add(duration)
+	var wg sync.WaitGroup
+	var total atomic.Int64
+
+	for range nParallelConns {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for time.Now().Before(deadline) {
+				// Each iteration: one short-lived connection (~500 ms).
+				func() {
+					connCtx, connCancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+					defer connCancel()
+
+					conn, _, err := websocket.Dial(connCtx, u, nil)
+					if err != nil {
+						return // dial failed (server busy) — try next iteration
+					}
+					defer conn.CloseNow()
+
+					if err := conn.Write(connCtx, websocket.MessageText, subMsg); err != nil {
+						return
+					}
+					var n int64
+					for {
+						_, _, err := conn.Read(connCtx)
+						if err != nil {
+							break
+						}
+						n++
+					}
+					total.Add(n)
+				}()
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	t.Logf("ClientChurn: %d total messages during %s of churn with %d concurrent connections",
+		total.Load(), duration, nParallelConns)
+
+	time.Sleep(500 * time.Millisecond)
+	if n := brk.ActiveSubscriptions(); n != 0 {
+		t.Errorf("subscription leak after churn: %d remain", n)
+	}
+}
+
+// TestStress_SlowConsumer verifies that a slow-reading client doesn't block
+// fast clients or the broker fan-out (drop-on-full behaviour).
+func TestStress_SlowConsumer(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping stress test in short mode")
+	}
+
+	const (
+		nSignals = 100
+		duration = 3 * time.Second
+	)
+
+	srv, brk := newStressServer(t, nSignals)
+	url := wsURL(srv)
+	subMsg := buildSubMsg(nSignals)
+
+	ctx, cancel := context.WithTimeout(t.Context(), duration+5*time.Second)
+	defer cancel()
+
+	var fastTotal, slowTotal atomic.Int64
+	var wg sync.WaitGroup
+
+	// Fast client: reads as quickly as possible.
+	wg.Add(1)
+	fastCtx, fastCancel := context.WithTimeout(ctx, duration)
+	go func() {
+		defer fastCancel()
+		connectAndCount(fastCtx, t, url, subMsg, &fastTotal, &wg)
+	}()
+
+	// Slow client: sleeps 100 ms between each read.
+	wg.Add(1)
+	slowCtx, slowCancel := context.WithTimeout(ctx, duration)
+	go func() {
+		defer wg.Done()
+		defer slowCancel()
+
+		conn, _, err := websocket.Dial(slowCtx, url, nil)
+		if err != nil {
+			t.Errorf("slow client dial: %v", err)
+			return
+		}
+		defer conn.CloseNow()
+
+		if err := conn.Write(slowCtx, websocket.MessageText, subMsg); err != nil {
+			t.Errorf("slow client subscribe: %v", err)
+			return
+		}
+
+		var n int64
+		for {
+			time.Sleep(100 * time.Millisecond)
+			_, _, err := conn.Read(slowCtx)
+			if err != nil {
+				break
+			}
+			n++
+		}
+		slowTotal.Add(n)
+	}()
+
+	wg.Wait()
+
+	fast := fastTotal.Load()
+	slow := slowTotal.Load()
+	t.Logf("SlowConsumer: fast=%d updates, slow=%d updates in %s", fast, slow, duration)
+
+	// Fast client must receive substantially more than slow client.
+	if fast < slow*3 {
+		t.Errorf("fast client should receive many more updates than slow client (fast=%d slow=%d)", fast, slow)
+	}
+
+	// Fast client must have received a meaningful number of updates.
+	minFast := int64(nSignals) * int64(duration.Seconds()) * 5
+	if fast < minFast {
+		t.Errorf("fast client received too few updates: got %d, want >= %d", fast, minFast)
+	}
+
+	time.Sleep(300 * time.Millisecond)
+	if n := brk.ActiveSubscriptions(); n != 0 {
+		t.Errorf("subscription leak: %d remain", n)
+	}
+}