Close idle transport connections when targets are drained or disposed

internal/server/load_balancer.go

@@ -74,6 +74,12 @@ func (tl TargetList) StopHealthChecks() {
 	}
 }
 
+func (tl TargetList) CloseIdleConnections() {
+	for _, target := range tl {
+		target.transport.CloseIdleConnections()
+	}
+}
+
 func (tl TargetList) targetsMatchingReadonly(readonly bool) TargetList {
 	result := TargetList{}
 	for _, target := range tl {
@@ -155,6 +161,7 @@ func (lb *LoadBalancer) MarkAllHealthy() {
 
 func (lb *LoadBalancer) Dispose() {
 	lb.all.StopHealthChecks()
+	lb.all.CloseIdleConnections()
 }
 
 func (lb *LoadBalancer) DrainAll(timeout time.Duration) {

internal/server/target.go

@@ -93,6 +93,7 @@ type Target struct {
 	targetURL    *url.URL
 	readonly     bool
 	options      TargetOptions
+	transport    *http.Transport
 	proxyHandler http.Handler
 
 	state        TargetState
@@ -214,6 +215,8 @@ WAIT_FOR_REQUESTS_TO_COMPLETE:
 	for _, inflight := range toCancel {
 		inflight.cancel(ErrorDraining)
 	}
+
+	t.transport.CloseIdleConnections()
 }
 
 func (t *Target) BeginHealthChecks(stateConsumer TargetStateConsumer) {
@@ -293,14 +296,16 @@ func (t *Target) buildHealthCheckURL() *url.URL {
 func (t *Target) createProxyHandler() http.Handler {
 	bufferPool := NewBufferPool(ProxyBufferSize)
 
+	t.transport = &http.Transport{
+		MaxIdleConnsPerHost:   MaxIdleConnsPerHost,
+		ResponseHeaderTimeout: t.options.ResponseTimeout,
+	}
+
 	return &httputil.ReverseProxy{
 		BufferPool:   bufferPool,
 		Rewrite:      t.rewrite,
 		ErrorHandler: t.handleProxyError,
-		Transport: &http.Transport{
-			MaxIdleConnsPerHost:   MaxIdleConnsPerHost,
-			ResponseHeaderTimeout: t.options.ResponseTimeout,
-		},
+		Transport:    t.transport,
 	}
 }
 

internal/server/target_test.go

@@ -6,6 +6,7 @@ import (
 	"net"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -398,6 +399,46 @@ func TestTarget_DrainHijackedConnectionsImmediately(t *testing.T) {
 	assert.Less(t, time.Since(startedDraining).Seconds(), 1.0)
 }
 
+func TestTarget_DrainClosesIdleConnections(t *testing.T) {
+	var activeConns atomic.Int32
+
+	backend := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("ok"))
+	}))
+	backend.Config.ConnState = func(conn net.Conn, state http.ConnState) {
+		switch state {
+		case http.StateNew:
+			activeConns.Add(1)
+		case http.StateClosed:
+			activeConns.Add(-1)
+		}
+	}
+	backend.Start()
+	t.Cleanup(backend.Close)
+
+	backendURL, err := url.Parse(backend.URL)
+	require.NoError(t, err)
+
+	target, err := NewTarget(backendURL.Host, defaultTargetOptions)
+	require.NoError(t, err)
+
+	// Make a request to establish a connection in the pool.
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	w := httptest.NewRecorder()
+	testServeRequestWithTarget(t, target, w, req)
+	require.Equal(t, http.StatusOK, w.Result().StatusCode)
+
+	// Connection should be idle in the pool.
+	require.Equal(t, int32(1), activeConns.Load())
+
+	target.Drain(time.Second)
+
+	// After draining, the idle connection should be closed.
+	require.Eventually(t, func() bool {
+		return activeConns.Load() == 0
+	}, time.Second, 10*time.Millisecond)
+}
+
 func TestTarget_EnforceMaxBodySizes(t *testing.T) {
 	sendRequest := func(bufferRequests, bufferResponses bool, maxMemorySize, maxBodySize int64, requestBody, responseBody string) *httptest.ResponseRecorder {
 		targetOptions := TargetOptions{