more features

This commit is contained in:
Para Dox
2025-05-29 01:03:59 +07:00
parent 73f70555d2
commit 596dba9ad9

View File

@@ -1263,22 +1263,21 @@ func main() {
handleWebSocketRequest(w, r, backends, client, &upgrader, statsCollector) handleWebSocketRequest(w, r, backends, client, &upgrader, statsCollector)
} else { } else {
// Handle regular HTTP request // Handle regular HTTP request
stats := handleRequest(w, r, backends, client, enableDetailedLogs == "true", statsCollector) 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
} }
}) })
log.Fatal(http.ListenAndServe(listenAddr, nil)) 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() startTime := time.Now()
// Read the entire request body // Read the entire request body
body, err := io.ReadAll(r.Body) body, err := io.ReadAll(r.Body)
if err != nil { if err != nil {
http.Error(w, "Error reading request body", http.StatusBadRequest) http.Error(w, "Error reading request body", http.StatusBadRequest)
return nil return
} }
defer r.Body.Close() defer r.Body.Close()
@@ -1469,7 +1468,7 @@ func handleRequest(w http.ResponseWriter, r *http.Request, backends []Backend, c
for stat := range statsChan { for stat := range statsChan {
stats = append(stats, stat) stats = append(stats, stat)
} }
return stats return
} }
} }
@@ -1483,57 +1482,65 @@ func handleRequest(w http.ResponseWriter, r *http.Request, backends []Backend, c
} }
w.WriteHeader(response.resp.StatusCode) w.WriteHeader(response.resp.StatusCode)
w.Write(response.body) w.Write(response.body)
} else {
// No valid response received from any backend
http.Error(w, "All backends failed", http.StatusBadGateway)
} }
// Always wait for primary backend to complete before collecting stats // Collect stats asynchronously to avoid blocking the response
// This ensures primary backend stats are always included
go func() { go func() {
// Always wait for primary backend to complete before collecting stats
// This ensures primary backend stats are always included
primaryWg.Wait() // Wait for primary backend to complete first primaryWg.Wait() // Wait for primary backend to complete first
wg.Wait() // Then wait for all other backends wg.Wait() // Then wait for all other backends
close(statsChan) close(statsChan)
}()
// Collect stats // Collect stats
var stats []ResponseStats var stats []ResponseStats
for stat := range statsChan { for stat := range statsChan {
stats = append(stats, stat) stats = append(stats, stat)
} }
// Log response times if enabled // Log response times if enabled
totalDuration := time.Since(startTime) totalDuration := time.Since(startTime)
if enableDetailedLogs { if enableDetailedLogs {
logResponseStats(totalDuration, stats) logResponseStats(totalDuration, stats)
} }
// Add the actual user-experienced duration for the winning response // Add the actual user-experienced duration for the winning response
if response.err == nil && response.backend != "" { if response.err == nil && response.backend != "" {
// Find the stat for the winning backend and update it with the actual user-experienced duration // Find the stat for the winning backend and update it with the actual user-experienced duration
for i := range stats { for i := range stats {
if stats[i].Backend == response.backend && stats[i].Error == nil { if stats[i].Backend == response.backend && stats[i].Error == nil {
// Calculate user latency from when the first backend started processing // Calculate user latency from when the first backend started processing
var userLatency time.Duration var userLatency time.Duration
if firstStart := firstBackendStartTime.Load(); firstStart != nil && !responseReceivedTime.IsZero() { if firstStart := firstBackendStartTime.Load(); firstStart != nil && !responseReceivedTime.IsZero() {
userLatency = responseReceivedTime.Sub(*firstStart) userLatency = responseReceivedTime.Sub(*firstStart)
} else { } else {
// Fallback to original calculation if somehow we don't have the times // Fallback to original calculation if somehow we don't have the times
userLatency = time.Since(startTime) userLatency = time.Since(startTime)
}
// Create a special stat entry for the actual first response time
actualFirstResponseStat := ResponseStats{
Backend: "actual-first-response",
StatusCode: stats[i].StatusCode,
Duration: userLatency,
Error: nil,
Method: stats[i].Method,
}
stats = append(stats, actualFirstResponseStat)
break
} }
// Create a special stat entry for the actual first response time
actualFirstResponseStat := ResponseStats{
Backend: "actual-first-response",
StatusCode: stats[i].StatusCode,
Duration: userLatency,
Error: nil,
Method: stats[i].Method,
}
stats = append(stats, actualFirstResponseStat)
break
} }
} }
}
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) { func logResponseStats(totalDuration time.Duration, stats []ResponseStats) {