story/cmd/load-sse/main.go

package main

import (
	"bufio"
	"context"
	"errors"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

type connection struct {
	body   io.ReadCloser
	cursor int64
}

func main() {
	url := flag.String("url", "http://localhost:58080/api/v1/events", "SSE endpoint")
	count := flag.Int("connections", 500, "simultaneous connections")
	rate := flag.Int("open-rate", 40, "new connections per second")
	hold := flag.Duration("hold", 5*time.Second, "hold after all connections are ready")
	reconnects := flag.Int("reconnects", 50, "connections to reopen with Last-Event-ID")
	flag.Parse()
	token := strings.TrimSpace(os.Getenv("ACCESS_TOKEN"))
	if token == "" || *count < 1 || *count > 5000 || *rate < 1 || *rate > 1000 || *reconnects < 0 || *reconnects > *count {
		fatal("ACCESS_TOKEN and safe positive load parameters are required")
	}
	transport := &http.Transport{
		DialContext:         (&net.Dialer{Timeout: 3 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
		MaxIdleConns:        *count + *reconnects,
		MaxIdleConnsPerHost: *count + *reconnects,
		IdleConnTimeout:     30 * time.Second,
	}
	client := &http.Client{Transport: transport}
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(*count/(*rate)+20)*time.Second)
	defer cancel()

	connections, err := openBatch(ctx, client, *url, token, *count, *rate, nil)
	if err != nil {
		fatal(err.Error())
	}
	if err = waitHealthy(connections, *hold); err != nil {
		closeAll(connections)
		fatal(err.Error())
	}
	cursors := make([]int64, 0, *reconnects)
	for i := 0; i < *reconnects; i++ {
		if connections[i].cursor < 1 {
			closeAll(connections)
			fatal("SSE fixture did not provide a positive resumable cursor")
		}
		cursors = append(cursors, connections[i].cursor)
	}
	closeAll(connections)

	reconnectCtx, reconnectCancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer reconnectCancel()
	reopened, err := openBatch(reconnectCtx, client, *url, token, *reconnects, *rate, cursors)
	if err != nil {
		fatal("Last-Event-ID reconnect failed: " + err.Error())
	}
	closeAll(reopened)
	fmt.Printf("SSE connection load passed: simultaneous=%d, open_rate=%d/s, held=%s, Last-Event-ID_reconnects=%d\n", *count, *rate, hold.String(), *reconnects)
}

func openBatch(ctx context.Context, client *http.Client, url, token string, count, rate int, cursors []int64) ([]connection, error) {
	results := make([]connection, count)
	errorsByIndex := make([]error, count)
	var wg sync.WaitGroup
	interval := time.Second / time.Duration(rate)
	started := time.Now()
	for i := 0; i < count; i++ {
		target := started.Add(time.Duration(i) * interval)
		if wait := time.Until(target); wait > 0 {
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(wait):
			}
		}
		wg.Add(1)
		go func(index int) {
			defer wg.Done()
			cursor := int64(0)
			if cursors != nil {
				cursor = cursors[index]
			}
			results[index], errorsByIndex[index] = openOne(ctx, client, url, token, cursor)
		}(i)
	}
	wg.Wait()
	for i, err := range errorsByIndex {
		if err != nil {
			closeAll(results)
			return nil, fmt.Errorf("connection %d: %w", i, err)
		}
	}
	return results, nil
}

func openOne(ctx context.Context, client *http.Client, url, token string, cursor int64) (connection, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return connection{}, err
	}
	req.Header.Set("Authorization", "Bearer "+token)
	req.Header.Set("Accept", "text/event-stream")
	if cursor > 0 {
		req.Header.Set("Last-Event-ID", strconv.FormatInt(cursor, 10))
	}
	response, err := client.Do(req)
	if err != nil {
		return connection{}, err
	}
	if response.StatusCode != http.StatusOK || !strings.HasPrefix(response.Header.Get("Content-Type"), "text/event-stream") {
		body, _ := io.ReadAll(io.LimitReader(response.Body, 512))
		response.Body.Close()
		return connection{}, fmt.Errorf("unexpected response %d: %s", response.StatusCode, strings.TrimSpace(string(body)))
	}
	reader := bufio.NewReader(response.Body)
	readCursor := int64(0)
	for {
		line, readErr := reader.ReadString('\n')
		if readErr != nil {
			response.Body.Close()
			return connection{}, readErr
		}
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, "id:") {
			readCursor, _ = strconv.ParseInt(strings.TrimSpace(strings.TrimPrefix(line, "id:")), 10, 64)
		}
		if line == "" {
			break
		}
	}
	return connection{body: response.Body, cursor: readCursor}, nil
}

func waitHealthy(items []connection, hold time.Duration) error {
	if len(items) == 0 {
		return errors.New("no SSE connections were opened")
	}
	timer := time.NewTimer(hold)
	defer timer.Stop()
	<-timer.C
	for i, item := range items {
		if item.body == nil || item.cursor < 1 {
			return fmt.Errorf("connection %d is not healthy", i)
		}
	}
	return nil
}

func closeAll(items []connection) {
	for _, item := range items {
		if item.body != nil {
			_ = item.body.Close()
		}
	}
}

func fatal(message string) {
	fmt.Fprintln(os.Stderr, message)
	os.Exit(1)
}