package benchmarks import ( "bytes" "context" "encoding/json" "fmt" "io" "net" "net/http" "net/http/httptest" "strings" "testing" "time" "github.com/gorilla/websocket" "github.com/prometheus/client_golang/prometheus" ) func newBenchmarkHTTPClient() *http.Client { transport := &http.Transport{ Proxy: http.ProxyFromEnvironment, DialContext: (&net.Dialer{ Timeout: 5 * time.Second, KeepAlive: 30 * time.Second, }).DialContext, ForceAttemptHTTP2: false, // Disable HTTP/2 for benchmark stability MaxIdleConns: 500, MaxIdleConnsPerHost: 500, MaxConnsPerHost: 500, IdleConnTimeout: 90 * time.Second, TLSHandshakeTimeout: 5 * time.Second, ExpectContinueTimeout: 1 * time.Second, } return &http.Client{ Timeout: 10 * time.Second, Transport: transport, } } // BenchmarkAPIServerCreateJob tests job creation performance func BenchmarkAPIServerCreateJob(b *testing.B) { server := setupTestAPIServer(b) defer server.Close() client := newBenchmarkHTTPClient() b.ResetTimer() b.ReportAllocs() benchmarkCreateJob(b, server.URL, client) } // BenchmarkAPIServerCreateJobSimple tests job creation with simplified setup func BenchmarkAPIServerCreateJobSimple(b *testing.B) { // Create a simple HTTP server without httptest mux := http.NewServeMux() mux.HandleFunc("/api/v1/jobs", func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusCreated) _ = json.NewEncoder(w).Encode(map[string]string{"id": "test-job-id"}) }) server := &http.Server{ Addr: "127.0.0.1:0", // Use random available port Handler: mux, ReadHeaderTimeout: 5 * time.Second, } // Start server in goroutine go func() { _ = server.ListenAndServe() }() // Get the actual port addr := server.Addr if addr == "" { addr = "127.0.0.1:8080" } client := &http.Client{Timeout: 30 * time.Second} baseURL := "http://" + addr b.ReportAllocs() for i := 0; b.Loop(); i++ { jobData := map[string]any{ "job_name": fmt.Sprintf("benchmark-job-%d", i), "args": map[string]any{ "model": "test-model", "data": generateTestPayload(1024), }, "priority": 0, } jsonData, _ := json.Marshal(jobData) req, err := http.NewRequestWithContext(context.Background(), "POST", baseURL+"/api/v1/jobs", bytes.NewBuffer(jsonData)) if err != nil { b.Fatalf("Failed to create request: %v", err) } req.Header.Set("Content-Type", "application/json") resp, err := client.Do(req) if err != nil { // Skip this iteration if server isn't ready continue } _, _ = io.Copy(io.Discard, resp.Body) _ = resp.Body.Close() } _ = server.Close() } func generateTestPayload(size int) string { data := make([]byte, size) for i := range data { data[i] = byte(i % 256) } return string(data) } // BenchmarkMetrics measures the performance impact of metrics collection func BenchmarkMetricsCollection(b *testing.B) { registry := prometheus.NewRegistry() // Create test metrics counter := prometheus.NewCounter(prometheus.CounterOpts{ Name: "test_operations_total", Help: "Total number of test operations", }) histogram := prometheus.NewHistogram(prometheus.HistogramOpts{ Name: "test_duration_seconds", Help: "Test operation duration", Buckets: prometheus.DefBuckets, }) registry.MustRegister(counter, histogram) b.ReportAllocs() for i := 0; b.Loop(); i++ { counter.Inc() histogram.Observe(float64(i) * 0.001) } } // BenchmarkConcurrentRequests tests concurrent API performance func BenchmarkConcurrentRequests(b *testing.B) { server := setupTestAPIServer(b) defer server.Close() client := newBenchmarkHTTPClient() b.ResetTimer() // Test different concurrency levels for _, concurrency := range []int{1, 5, 10, 25, 50} { b.Run(fmt.Sprintf("Concurrency-%d", concurrency), func(b *testing.B) { benchmarkConcurrentRequests(b, server.URL, client, concurrency) }) } } func benchmarkConcurrentRequests(b *testing.B, baseURL string, client *http.Client, concurrency int) { b.SetParallelism(concurrency) b.RunParallel(func(pb *testing.PB) { i := 0 for pb.Next() { req, _ := http.NewRequestWithContext(context.Background(), "GET", baseURL+"/api/v1/jobs?limit=10", nil) req.Header.Set("Authorization", "Bearer test-token") resp, err := client.Do(req) if err == nil && resp != nil { _, _ = io.Copy(io.Discard, resp.Body) _ = resp.Body.Close() } i++ } }) } // setupTestAPIServer creates a test HTTP server for benchmarking func setupTestAPIServer(_ *testing.B) *httptest.Server { mux := http.NewServeMux() // Add basic API routes for benchmarking mux.HandleFunc("/api/v1/jobs", func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") if r.Method == "POST" { w.WriteHeader(http.StatusCreated) _ = json.NewEncoder(w).Encode(map[string]string{"id": "test-job-id"}) } else { w.WriteHeader(http.StatusOK) _ = json.NewEncoder(w).Encode([]map[string]string{{"id": "test-job-id", "status": "pending"}}) } }) mux.HandleFunc("/api/v1/jobs/", func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _ = json.NewEncoder(w).Encode(map[string]string{"status": "pending"}) }) mux.HandleFunc("/api/v1/metrics", func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _ = json.NewEncoder(w).Encode(map[string]string{"status": "ok"}) }) mux.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) { upgrader := websocket.Upgrader{ CheckOrigin: func(_ *http.Request) bool { return true }, } conn, err := upgrader.Upgrade(w, r, nil) if err != nil { return } defer conn.Close() // Echo back any messages for { mt, message, err := conn.ReadMessage() if err != nil { break } _ = conn.WriteMessage(mt, message) } }) return httptest.NewServer(mux) } // benchmarkCreateJob tests job creation performance func benchmarkCreateJob(b *testing.B, baseURL string, client *http.Client) { // Pre-generate test payload to avoid allocation overhead in hot path testPayload := generateTestPayload(1024) // Use static job name to avoid fmt.Sprintf overhead in benchmark loop jobName := "benchmark-job" for i := 0; b.Loop(); i++ { _ = i // Avoid unused variable warning jobData := map[string]any{ "job_name": jobName, "args": map[string]any{ "model": "test-model", "data": testPayload, }, "priority": 0, } jsonData, _ := json.Marshal(jobData) req, err := http.NewRequestWithContext(context.Background(), "POST", baseURL+"/api/v1/jobs", bytes.NewBuffer(jsonData)) if err != nil { b.Fatalf("Failed to create request: %v", err) } req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", "Bearer test-token") resp, err := client.Do(req) if err != nil { b.Fatalf("Request failed: %v", err) } _, _ = io.Copy(io.Discard, resp.Body) _ = resp.Body.Close() } } // benchmarkListJobs tests job listing performance func benchmarkListJobs(b *testing.B, baseURL string, client *http.Client) { for i := 0; b.Loop(); i++ { req, err := http.NewRequestWithContext(context.Background(), "GET", baseURL+"/api/v1/jobs", nil) if err != nil { b.Fatalf("Failed to create request: %v", err) } req.Header.Set("Authorization", "Bearer test-token") resp, err := client.Do(req) if err != nil { b.Fatalf("Request failed: %v", err) } _, _ = io.Copy(io.Discard, resp.Body) _ = resp.Body.Close() } } // BenchmarkAPIServerListJobs tests job listing performance func BenchmarkAPIServerListJobs(b *testing.B) { server := setupTestAPIServer(b) defer server.Close() client := newBenchmarkHTTPClient() b.ResetTimer() b.ReportAllocs() benchmarkListJobs(b, server.URL, client) } // BenchmarkWebSocketConnection tests WebSocket connection performance func BenchmarkWebSocketConnection(b *testing.B) { server := setupTestAPIServer(b) defer server.Close() b.ReportAllocs() // Convert HTTP URL to WebSocket URL once wsURL := strings.Replace(server.URL, "http://", "ws://", 1) + "/ws" for b.Loop() { conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) if resp != nil && resp.Body != nil { _ = resp.Body.Close() } if err != nil { b.Fatalf("WebSocket dial failed: %v", err) } // Send close message and wait for server to close _ = conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) _ = conn.Close() } }