Martino Ferrari
344 lines
8.7 KiB
Go
344 lines
8.7 KiB
Go
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)
|
||
}
|
||
}
|