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: 30 * time.Second, KeepAlive: 30 * time.Second, }).DialContext, ForceAttemptHTTP2: true, MaxIdleConns: 256, MaxIdleConnsPerHost: 256, IdleConnTimeout: 90 * time.Second, TLSHandshakeTimeout: 10 * time.Second, ExpectContinueTimeout: 1 * time.Second, } return &http.Client{ Timeout: 30 * 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.ResetTimer() b.ReportAllocs() for i := 0; i < b.N; i++ { jobData := map[string]interface{}{ "job_name": fmt.Sprintf("benchmark-job-%d", i), "args": map[string]interface{}{ "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.ResetTimer() b.ReportAllocs() for i := 0; i < b.N; 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, _ *http.Request) { w.WriteHeader(http.StatusOK) }) return httptest.NewServer(mux) } // benchmarkCreateJob tests job creation performance func benchmarkCreateJob(b *testing.B, baseURL string, client *http.Client) { for i := 0; i < b.N; i++ { jobData := map[string]interface{}{ "job_name": fmt.Sprintf("benchmark-job-%d", i), "args": map[string]interface{}{ "model": "test-model", "data": generateTestPayload(1024), // 1KB payload }, "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; i < b.N; 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() for i := 0; i < b.N; i++ { // Convert HTTP URL to WebSocket URL wsURL := strings.Replace(server.URL, "http://", "ws://", 1) wsURL += "/ws" conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) if resp != nil && resp.Body != nil { _ = resp.Body.Close() } if err != nil { // Skip iteration if WebSocket server isn't available continue } _ = conn.Close() } }