diff --git a/internal/server/load_balancer.go b/internal/server/load_balancer.go index 6bec9aec..c175f165 100644 --- a/internal/server/load_balancer.go +++ b/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) { diff --git a/internal/server/target.go b/internal/server/target.go index 09084a05..47c12652 100644 --- a/internal/server/target.go +++ b/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, } } diff --git a/internal/server/target_test.go b/internal/server/target_test.go index bfcaa4b8..161fcc4c 100644 --- a/internal/server/target_test.go +++ b/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{