ai-gateway-go/cmd/gateway-loadtest/main.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type sample struct {
	latency time.Duration
	status  int
	err     string
}
type report struct {
	Requests          int         `json:"requests"`
	Success           int         `json:"success"`
	Errors            int         `json:"errors"`
	ErrorRate         float64     `json:"error_rate"`
	RequestsPerSecond float64     `json:"requests_per_second"`
	P50MS             float64     `json:"p50_ms"`
	P95MS             float64     `json:"p95_ms"`
	P99MS             float64     `json:"p99_ms"`
	DurationSeconds   float64     `json:"duration_seconds"`
	StatusCounts      map[int]int `json:"status_counts"`
}

func main() {
	base := flag.String("url", "http://127.0.0.1:8080", "gateway base URL")
	path := flag.String("path", "/v1/chat/completions", "request path")
	key := flag.String("api-key", os.Getenv("GATEWAY_LOADTEST_API_KEY"), "dedicated load-test API key")
	model := flag.String("model", "test-model", "model or model alias")
	concurrency := flag.Int("concurrency", 16, "concurrent workers")
	duration := flag.Duration("duration", 30*time.Second, "test duration")
	timeout := flag.Duration("timeout", 60*time.Second, "per-request timeout")
	maxError := flag.Float64("max-error-rate", 0.01, "failure threshold")
	maxP95 := flag.Duration("max-p95", 2*time.Second, "p95 latency threshold")
	flag.Parse()
	if *concurrency < 1 || *concurrency > 2000 || *duration < time.Second || *key == "" {
		fmt.Fprintln(os.Stderr, "invalid arguments: api-key, positive duration and concurrency 1..2000 are required")
		os.Exit(2)
	}
	target := strings.TrimRight(*base, "/") + *path
	payload, _ := json.Marshal(map[string]any{"model": *model, "messages": []map[string]string{{"role": "user", "content": "Reply with OK."}}, "stream": false, "max_tokens": 8})
	client := &http.Client{Timeout: *timeout, Transport: &http.Transport{MaxIdleConns: *concurrency * 2, MaxIdleConnsPerHost: *concurrency, IdleConnTimeout: 90 * time.Second, TLSHandshakeTimeout: 5 * time.Second}}
	ctx, cancel := context.WithTimeout(context.Background(), *duration)
	defer cancel()
	started := time.Now()
	results := make(chan sample, *concurrency*4)
	var sequence atomic.Uint64
	var workers sync.WaitGroup
	for worker := 0; worker < *concurrency; worker++ {
		workers.Add(1)
		go func() {
			defer workers.Done()
			for ctx.Err() == nil {
				number := sequence.Add(1)
				request, err := http.NewRequestWithContext(ctx, http.MethodPost, target, bytes.NewReader(payload))
				if err != nil {
					results <- sample{err: err.Error()}
					continue
				}
				request.Header.Set("Authorization", "Bearer "+*key)
				request.Header.Set("Content-Type", "application/json")
				request.Header.Set("X-Request-ID", fmt.Sprintf("load-%d", number))
				request.Header.Set("Idempotency-Key", fmt.Sprintf("load-%d", number))
				begin := time.Now()
				response, err := client.Do(request)
				elapsed := time.Since(begin)
				if err != nil {
					if ctx.Err() == nil {
						results <- sample{latency: elapsed, err: err.Error()}
					}
					continue
				}
				_, readErr := io.Copy(io.Discard, io.LimitReader(response.Body, 2<<20))
				response.Body.Close()
				entry := sample{latency: elapsed, status: response.StatusCode}
				if readErr != nil {
					entry.err = readErr.Error()
				}
				results <- entry
			}
		}()
	}
	go func() { workers.Wait(); close(results) }()
	samples := []sample{}
	for result := range results {
		samples = append(samples, result)
	}
	elapsed := time.Since(started)
	summary := summarize(samples, elapsed)
	encoded, _ := json.MarshalIndent(summary, "", "  ")
	fmt.Println(string(encoded))
	if summary.ErrorRate > *maxError || time.Duration(summary.P95MS*float64(time.Millisecond)) > *maxP95 {
		os.Exit(1)
	}
}

func summarize(samples []sample, elapsed time.Duration) report {
	summary := report{Requests: len(samples), StatusCounts: map[int]int{}, DurationSeconds: elapsed.Seconds()}
	latencies := make([]time.Duration, 0, len(samples))
	for _, item := range samples {
		summary.StatusCounts[item.status]++
		latencies = append(latencies, item.latency)
		if item.err == "" && item.status >= 200 && item.status < 300 {
			summary.Success++
		} else {
			summary.Errors++
		}
	}
	sort.Slice(latencies, func(i, j int) bool { return latencies[i] < latencies[j] })
	if summary.Requests > 0 {
		summary.ErrorRate = float64(summary.Errors) / float64(summary.Requests)
		summary.RequestsPerSecond = float64(summary.Requests) / elapsed.Seconds()
		summary.P50MS = percentile(latencies, .50)
		summary.P95MS = percentile(latencies, .95)
		summary.P99MS = percentile(latencies, .99)
	}
	return summary
}
func percentile(values []time.Duration, p float64) float64 {
	if len(values) == 0 {
		return 0
	}
	index := int(float64(len(values)-1) * p)
	return float64(values[index]) / float64(time.Millisecond)
}