more features

benchmark-proxy/main.go

@@ -1263,22 +1263,21 @@ func main() {
 			handleWebSocketRequest(w, r, backends, client, &upgrader, statsCollector)
 		} else {
 			// Handle regular HTTP request
-			stats := handleRequest(w, r, backends, client, enableDetailedLogs == "true", statsCollector)
-			statsCollector.AddStats(stats, 0) // The 0 is a placeholder, we're not using totalDuration in the collector
+			handleRequest(w, r, backends, client, enableDetailedLogs == "true", statsCollector)
 		}
 	})
 
 	log.Fatal(http.ListenAndServe(listenAddr, nil))
 }
 
-func handleRequest(w http.ResponseWriter, r *http.Request, backends []Backend, client *http.Client, enableDetailedLogs bool, statsCollector *StatsCollector) []ResponseStats {
+func handleRequest(w http.ResponseWriter, r *http.Request, backends []Backend, client *http.Client, enableDetailedLogs bool, statsCollector *StatsCollector) {
 	startTime := time.Now()
 
 	// Read the entire request body
 	body, err := io.ReadAll(r.Body)
 	if err != nil {
 		http.Error(w, "Error reading request body", http.StatusBadRequest)
-		return nil
+		return
 	}
 	defer r.Body.Close()
 
@@ -1469,7 +1468,7 @@ func handleRequest(w http.ResponseWriter, r *http.Request, backends []Backend, c
 			for stat := range statsChan {
 				stats = append(stats, stat)
 			}
-			return stats
+			return
 		}
 	}
 
@@ -1483,15 +1482,18 @@ func handleRequest(w http.ResponseWriter, r *http.Request, backends []Backend, c
 		}
 		w.WriteHeader(response.resp.StatusCode)
 		w.Write(response.body)
+	} else {
+		// No valid response received from any backend
+		http.Error(w, "All backends failed", http.StatusBadGateway)
 	}
 
+	// Collect stats asynchronously to avoid blocking the response
+	go func() {
 		// Always wait for primary backend to complete before collecting stats
 		// This ensures primary backend stats are always included
-	go func() {
 		primaryWg.Wait() // Wait for primary backend to complete first
 		wg.Wait()        // Then wait for all other backends
 		close(statsChan)
-	}()
 
 		// Collect stats
 		var stats []ResponseStats
@@ -1533,7 +1535,12 @@ func handleRequest(w http.ResponseWriter, r *http.Request, backends []Backend, c
 			}
 		}
 
-	return stats
+		// Send stats to collector
+		statsCollector.AddStats(stats, 0)
+	}()
+
+	// Return immediately after sending response to client
+	return
 }
 
 func logResponseStats(totalDuration time.Duration, stats []ResponseStats) {