casino/cmd/loadtest/main.go

// Command loadtest measures how many concurrent players an instance holds.
//
// It opens real WebSocket connections and, optionally, places real bets, then
// reports connection success, frame delivery, and latency. The point is to
// produce numbers rather than adjectives: run it against a candidate machine
// and read the ceiling off the output.
//
//	go run ./cmd/loadtest -conns 2000 -addr localhost:8080
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"runtime"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	"github.com/coder/websocket"
)

func main() {
	var (
		addr     = flag.String("addr", "localhost:8080", "instance to load")
		conns    = flag.Int("conns", 500, "concurrent websocket connections")
		duration = flag.Duration("duration", 20*time.Second, "how long to hold them")
		game     = flag.String("game", "rocket", "game room to join")
		ramp     = flag.Duration("ramp", 5*time.Second, "time to open all connections")
	)
	flag.Parse()

	ctx, cancel := context.WithTimeout(context.Background(), *duration+*ramp+30*time.Second)
	defer cancel()

	var (
		connected atomic.Int64
		failed    atomic.Int64
		frames    atomic.Int64
		bytesRecv atomic.Int64
		dialMu    sync.Mutex
		dialTimes []time.Duration
	)

	fmt.Printf("opening %d connections to %s over %v\n", *conns, *addr, *ramp)
	start := time.Now()

	// Stagger dialling: slamming every connection open at once measures the
	// accept backlog rather than the steady state anyone actually runs at.
	gap := *ramp / time.Duration(max(1, *conns))

	var wg sync.WaitGroup
	for i := 0; i < *conns; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			time.Sleep(time.Duration(i) * gap)

			dialStart := time.Now()
			conn, _, err := websocket.Dial(ctx,
				fmt.Sprintf("ws://%s/ws/%s", *addr, *game), nil)
			if err != nil {
				failed.Add(1)
				return
			}
			took := time.Since(dialStart)
			defer conn.CloseNow()

			connected.Add(1)
			dialMu.Lock()
			dialTimes = append(dialTimes, took)
			dialMu.Unlock()

			// Read until the run ends. A client that stops reading is exactly
			// the slow-subscriber case the server has to survive, but here we
			// want the healthy path.
			readCtx, stop := context.WithTimeout(ctx, *duration)
			defer stop()
			for {
				_, data, err := conn.Read(readCtx)
				if err != nil {
					return
				}
				frames.Add(1)
				bytesRecv.Add(int64(len(data)))
			}
		}(i)
	}

	// Report progress while the run is in flight.
	done := make(chan struct{})
	go func() {
		t := time.NewTicker(5 * time.Second)
		defer t.Stop()
		for {
			select {
			case <-done:
				return
			case <-t.C:
				var m runtime.MemStats
				runtime.ReadMemStats(&m)
				fmt.Printf("  t+%-5s connected=%-6d failed=%-5d frames=%-8d client heap=%dMB\n",
					time.Since(start).Round(time.Second),
					connected.Load(), failed.Load(), frames.Load(),
					m.Alloc/1024/1024)
			}
		}
	}()

	wg.Wait()
	close(done)

	elapsed := time.Since(start)
	sort.Slice(dialTimes, func(i, j int) bool { return dialTimes[i] < dialTimes[j] })

	fmt.Println()
	fmt.Println("results")
	fmt.Printf("  connections attempted : %d\n", *conns)
	fmt.Printf("  connected             : %d\n", connected.Load())
	fmt.Printf("  failed                : %d\n", failed.Load())
	if len(dialTimes) > 0 {
		fmt.Printf("  dial p50 / p99 / max  : %v / %v / %v\n",
			dialTimes[len(dialTimes)/2].Round(time.Millisecond),
			dialTimes[len(dialTimes)*99/100].Round(time.Millisecond),
			dialTimes[len(dialTimes)-1].Round(time.Millisecond))
	}
	fmt.Printf("  frames received       : %d\n", frames.Load())
	fmt.Printf("  bytes received        : %.1f MB\n", float64(bytesRecv.Load())/1e6)
	if connected.Load() > 0 {
		fmt.Printf("  frames per connection : %.1f\n",
			float64(frames.Load())/float64(connected.Load()))
		fmt.Printf("  server egress         : %.2f MB/s\n",
			float64(bytesRecv.Load())/1e6/elapsed.Seconds())
	}

	if failed.Load() > 0 {
		fmt.Fprintf(os.Stderr, "\n%d connections were refused: the ceiling is at or below %d\n",
			failed.Load(), *conns)
		os.Exit(1)
	}
	log.Printf("held %d concurrent connections for %v with no failures",
		connected.Load(), duration)
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}