From 33177e05136a1012b7383ffaf9482145b3810ab3 Mon Sep 17 00:00:00 2001
From: Dave Protasowski <dprotaso@gmail.com>
Date: Wed, 26 Nov 2025 12:11:01 -0500
Subject: [PATCH] Set Connection: close header when draining HTTP/1

This will trigger the http server to close the connection
https://cs.opensource.google/go/go/+/refs/tags/go1.25.4:src/net/http/server.go;l=1389-1391
---
 network/handlers/drain.go      | 22 +++++++++----------
 network/handlers/drain_test.go | 40 +++++++++++++++++++++++++++++++---
 2 files changed, 48 insertions(+), 14 deletions(-)

diff --git a/network/handlers/drain.go b/network/handlers/drain.go
index 3b2f28a077..3089597cef 100644
--- a/network/handlers/drain.go
+++ b/network/handlers/drain.go
@@ -89,9 +89,10 @@ var _ http.Handler = (*Drainer)(nil)
 
 // ServeHTTP implements http.Handler
 func (d *Drainer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	draining := d.draining()
 	// Respond to probes regardless of path.
 	if d.isHealthCheckRequest(r) {
-		if d.draining() {
+		if draining {
 			http.Error(w, "shutting down", http.StatusServiceUnavailable)
 		} else if d.HealthCheck != nil {
 			d.HealthCheck(w, r)
@@ -101,7 +102,7 @@ func (d *Drainer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	if isKProbe(r) {
-		if d.draining() {
+		if draining {
 			http.Error(w, "shutting down", http.StatusServiceUnavailable)
 		} else {
 			serveKProbe(w, r)
@@ -109,7 +110,14 @@ func (d *Drainer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	d.resetTimer()
+	if draining {
+		d.resetTimer()
+		if r.ProtoMajor == 1 {
+			// Setting the connection close header will hint the server to
+			// close the connection after handling the request
+			w.Header().Set("Connection", "close")
+		}
+	}
 	d.Inner.ServeHTTP(w, r)
 }
 
@@ -186,14 +194,6 @@ func (d *Drainer) Reset() {
 }
 
 func (d *Drainer) resetTimer() {
-	if func() bool {
-		d.RLock()
-		defer d.RUnlock()
-		return d.timer == nil
-	}() {
-		return
-	}
-
 	d.Lock()
 	defer d.Unlock()
 	if d.timer != nil && d.timer.Stop() {
diff --git a/network/handlers/drain_test.go b/network/handlers/drain_test.go
index b31218b509..c8b2c726e8 100644
--- a/network/handlers/drain_test.go
+++ b/network/handlers/drain_test.go
@@ -66,7 +66,7 @@ func (mt *mockTimer) tickChan() <-chan time.Time {
 
 func TestDrainMechanics(t *testing.T) {
 	var (
-		w     http.ResponseWriter
+		w     = httptest.NewRecorder()
 		req   = &http.Request{}
 		probe = &http.Request{
 			Header: http.Header{
@@ -239,7 +239,7 @@ func TestDrainMechanics(t *testing.T) {
 
 func TestDrainerKProbe(t *testing.T) {
 	var (
-		w          http.ResponseWriter
+		w          = httptest.NewRecorder()
 		req        = &http.Request{}
 		kprobehash = "hash"
 		kprobe     = &http.Request{
@@ -340,7 +340,7 @@ func TestHealthCheckWithProbeType(t *testing.T) {
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
 			var (
-				w       http.ResponseWriter
+				w       = httptest.NewRecorder()
 				req     = &http.Request{}
 				cnt     = 0
 				inner   = http.HandlerFunc(func(http.ResponseWriter, *http.Request) { cnt++ })
@@ -586,3 +586,37 @@ func TestResetWithActiveRequests(t *testing.T) {
 	// We need requests to be active for a bit
 	time.Sleep(time.Second)
 }
+
+func BenchmarkNoDrain(b *testing.B) {
+	r, _ := http.NewRequest(http.MethodGet, "knative.dev", nil)
+	w := httptest.NewRecorder()
+	handler := func(http.ResponseWriter, *http.Request) {
+	}
+
+	d := Drainer{
+		QuietPeriod: 10 * time.Second,
+		Inner:       http.HandlerFunc(handler),
+	}
+
+	for b.Loop() {
+		d.ServeHTTP(w, r)
+	}
+}
+
+func BenchmarkDrain(b *testing.B) {
+	r, _ := http.NewRequest(http.MethodGet, "knative.dev", nil)
+	w := httptest.NewRecorder()
+	handler := func(http.ResponseWriter, *http.Request) {
+	}
+
+	d := Drainer{
+		QuietPeriod: 10 * time.Second,
+		Inner:       http.HandlerFunc(handler),
+	}
+
+	go d.Drain()
+
+	for b.Loop() {
+		d.ServeHTTP(w, r)
+	}
+}