Drain HTTP responses before scale-to-zero

Author: tnsardesai

server/cmd/api/main.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log/slog"
+	"net"
 	"net/http"
 	"net/url"
 	"os"
@@ -377,6 +378,7 @@ func main() {
 		metrics.NewChromeCollector(upstreamMgr),
 		metrics.NewGPUCollector(),
 		metrics.NewSystemCollector(),
+		metrics.NewResponseDrainCollector(scaletozero.ResponseDrainOutcomeCounts, scaletozero.ActiveResponseHolds, scaletozero.FailClosedResponseHolds),
 	}
 	if otlpMetrics != nil {
 		metricsCollectors = append(metricsCollectors, metrics.NewOTLPCollector(otlpMetrics))
@@ -387,29 +389,22 @@ func main() {
 		Handler: rMetrics,
 	}
 
-	go func() {
-		slogger.Info("http server starting", "addr", srv.Addr)
-		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
-			slogger.Error("http server failed", "err", err)
-			stop()
-		}
-	}()
-
-	go func() {
-		slogger.Info("devtools websocket proxy starting", "addr", srvDevtools.Addr)
-		if err := srvDevtools.ListenAndServe(); err != nil && err != http.ErrServerClosed {
-			slogger.Error("devtools websocket proxy failed", "err", err)
-			stop()
-		}
-	}()
-
-	go func() {
-		slogger.Info("chromedriver proxy starting", "addr", srvChromeDriver.Addr)
-		if err := srvChromeDriver.ListenAndServe(); err != nil && err != http.ErrServerClosed {
-			slogger.Error("chromedriver proxy failed", "err", err)
-			stop()
-		}
-	}()
+	serveHTTP := func(name string, server *http.Server) {
+		go func() {
+			slogger.Info(name+" starting", "addr", server.Addr)
+			listener, err := net.Listen("tcp", server.Addr)
+			if err == nil {
+				err = scaletozero.Serve(server, listener)
+			}
+			if err != nil && err != http.ErrServerClosed {
+				slogger.Error(name+" failed", "err", err)
+				stop()
+			}
+		}()
+	}
+	serveHTTP("http server", srv)
+	serveHTTP("devtools websocket proxy", srvDevtools)
+	serveHTTP("chromedriver proxy", srvChromeDriver)
 
 	go func() {
 		slogger.Info("metrics server starting", "addr", srvMetrics.Addr)

server/lib/metrics/scaletozero.go

@@ -0,0 +1,40 @@
+package metrics
+
+import (
+	"context"
+	"sort"
+)
+
+type ResponseDrainSource func() map[string]uint64
+type ResponseDrainGauge func() int64
+
+type ResponseDrainCollector struct {
+	snapshot   ResponseDrainSource
+	active     ResponseDrainGauge
+	failClosed ResponseDrainGauge
+}
+
+func NewResponseDrainCollector(snapshot ResponseDrainSource, active, failClosed ResponseDrainGauge) *ResponseDrainCollector {
+	return &ResponseDrainCollector{snapshot: snapshot, active: active, failClosed: failClosed}
+}
+
+func (c *ResponseDrainCollector) Name() string { return "scale-to-zero response drain" }
+
+func (c *ResponseDrainCollector) Collect(_ context.Context, w *Writer) error {
+	w.Metric("kernel_scale_to_zero_response_drain_total", "HTTP response drain events before scale-to-zero is re-enabled.", "counter")
+	counts := c.snapshot()
+	outcomes := make([]string, 0, len(counts))
+	for outcome := range counts {
+		outcomes = append(outcomes, outcome)
+	}
+	sort.Strings(outcomes)
+	for _, outcome := range outcomes {
+		w.Sample("kernel_scale_to_zero_response_drain_total", []Label{{Name: "outcome", Value: outcome}}, float64(counts[outcome]))
+	}
+
+	w.Metric("kernel_scale_to_zero_response_holds", "HTTP response scale-to-zero holds currently active.", "gauge")
+	w.Sample("kernel_scale_to_zero_response_holds", nil, float64(c.active()))
+	w.Metric("kernel_scale_to_zero_response_fail_closed_holds", "HTTP response holds awaiting terminal connection recovery or guest termination.", "gauge")
+	w.Sample("kernel_scale_to_zero_response_fail_closed_holds", nil, float64(c.failClosed()))
+	return nil
+}

server/lib/metrics/scaletozero_test.go

@@ -0,0 +1,33 @@
+package metrics
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResponseDrainCollector(t *testing.T) {
+	collector := NewResponseDrainCollector(func() map[string]uint64 {
+		return map[string]uint64{
+			"timeout": 2,
+			"drained": 7,
+		}
+	}, func() int64 { return 3 }, func() int64 { return 1 })
+	writer := &Writer{}
+
+	require.NoError(t, collector.Collect(context.Background(), writer))
+
+	assert.Equal(t, `# HELP kernel_scale_to_zero_response_drain_total HTTP response drain events before scale-to-zero is re-enabled.
+# TYPE kernel_scale_to_zero_response_drain_total counter
+kernel_scale_to_zero_response_drain_total{outcome="drained"} 7
+kernel_scale_to_zero_response_drain_total{outcome="timeout"} 2
+# HELP kernel_scale_to_zero_response_holds HTTP response scale-to-zero holds currently active.
+# TYPE kernel_scale_to_zero_response_holds gauge
+kernel_scale_to_zero_response_holds 3
+# HELP kernel_scale_to_zero_response_fail_closed_holds HTTP response holds awaiting terminal connection recovery or guest termination.
+# TYPE kernel_scale_to_zero_response_fail_closed_holds gauge
+kernel_scale_to_zero_response_fail_closed_holds 1
+`, string(writer.Bytes()))
+}