diff --git a/go.mod b/go.mod index fc8bb4bd6..c1162306f 100644 --- a/go.mod +++ b/go.mod @@ -6,7 +6,7 @@ require ( github.com/bytedance/gopkg v0.1.4 github.com/bytedance/sonic v1.15.0 github.com/cloudwego/gopkg v0.2.0 - github.com/cloudwego/netpoll v0.7.3 + github.com/cloudwego/netpoll v0.7.5-0.20260804042038-d63d249a3ba0 github.com/fsnotify/fsnotify v1.5.4 github.com/stretchr/testify v1.10.0 github.com/tidwall/gjson v1.14.4 diff --git a/go.sum b/go.sum index 54bcdee4c..9d2bcea51 100644 --- a/go.sum +++ b/go.sum @@ -8,8 +8,8 @@ github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU= github.com/cloudwego/gopkg v0.2.0 h1:EU8Ahrj0rCfKZQdah50zKnlrQ1o2AdPYM87UclIqLME= github.com/cloudwego/gopkg v0.2.0/go.mod h1:WjQPYI8PesfQalIVcLzVJBb1EAopioZ+D+3UGJ+dNBs= -github.com/cloudwego/netpoll v0.7.3 h1:E9ImEseXM9BdHS+5aLxcE9Z0c7okFbM11XMwwJ00LxY= -github.com/cloudwego/netpoll v0.7.3/go.mod h1:KiNpLI5MX9vR0xj4gKqyioOrHlp8G0XBMqIV9HsvMCc= +github.com/cloudwego/netpoll v0.7.5-0.20260804042038-d63d249a3ba0 h1:QHIJivU6fuCdTu915Z5jyZ9FpX+OwvzG+8cRUgM8oA4= +github.com/cloudwego/netpoll v0.7.5-0.20260804042038-d63d249a3ba0/go.mod h1:KiNpLI5MX9vR0xj4gKqyioOrHlp8G0XBMqIV9HsvMCc= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= diff --git a/pkg/app/client/client.go b/pkg/app/client/client.go index 7d8fb06d0..503e4c941 100644 --- a/pkg/app/client/client.go +++ b/pkg/app/client/client.go @@ -669,6 +669,7 @@ func newHttp1OptionFromClient(c *Client) *http1.ClientOptions { MaxConns: c.options.MaxConnsPerHost, MaxConnDuration: c.options.MaxConnDuration, MaxIdleConnDuration: c.options.MaxIdleConnDuration, + PooledConnHealthCheck: c.options.PooledConnHealthCheck, ReadTimeout: c.options.ReadTimeout, WriteTimeout: c.options.WriteTimeout, MaxResponseBodySize: c.options.MaxResponseBodySize, diff --git a/pkg/app/client/client_test.go b/pkg/app/client/client_test.go index 21518e9b9..9b3d87289 100644 --- a/pkg/app/client/client_test.go +++ b/pkg/app/client/client_test.go @@ -201,6 +201,127 @@ func TestCloseIdleTLSConnections(t *testing.T) { } } +func TestPooledConnHealthCheckTLS(t *testing.T) { + var accepted, requests int32 + firstIdleClosed := make(chan struct{}) + var closeFirstIdle sync.Once + + httpsServer := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != consts.MethodPost { + http.Error(w, "unexpected method", http.StatusMethodNotAllowed) + return + } + body, err := io.ReadAll(r.Body) + if err != nil || string(body) != "payload" { + http.Error(w, "unexpected body", http.StatusBadRequest) + return + } + atomic.AddInt32(&requests, 1) + _, _ = w.Write([]byte("ok")) + })) + httpsServer.Config.ConnState = func(conn net.Conn, state http.ConnState) { + switch state { + case http.StateNew: + atomic.AddInt32(&accepted, 1) + case http.StateIdle: + if atomic.LoadInt32(&requests) == 1 { + closeFirstIdle.Do(func() { + _ = conn.Close() + close(firstIdleClosed) + }) + } + } + } + httpsServer.StartTLS() + defer httpsServer.Close() + + c, err := NewClient( + WithTLSConfig(httpsServer.Client().Transport.(*http.Transport).TLSClientConfig), + WithPooledConnHealthCheck(true), + ) + if err != nil { + t.Fatal(err) + } + defer c.CloseIdleConnections() + + doPost := func() { + req := protocol.AcquireRequest() + defer protocol.ReleaseRequest(req) + req.Header.SetMethod(consts.MethodPost) + req.SetRequestURI(httpsServer.URL + "/health-check") + req.SetBodyString("payload") + resp := protocol.AcquireResponse() + defer protocol.ReleaseResponse(resp) + + if err := c.Do(context.Background(), req, resp); err != nil { + t.Fatal(err) + } + assert.DeepEqual(t, http.StatusOK, resp.StatusCode()) + assert.DeepEqual(t, "ok", string(resp.Body())) + } + + doPost() + select { + case <-firstIdleClosed: + case <-time.After(time.Second): + t.Fatal("server did not close the first idle TLS connection") + } + doPost() + + assert.DeepEqual(t, int32(2), atomic.LoadInt32(&accepted)) + assert.DeepEqual(t, int32(2), atomic.LoadInt32(&requests)) +} + +func TestPooledConnHealthCheckReusesTLS(t *testing.T) { + var accepted, requests int32 + httpsServer := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != consts.MethodPost { + http.Error(w, "unexpected method", http.StatusMethodNotAllowed) + return + } + body, err := io.ReadAll(r.Body) + if err != nil || string(body) != "payload" { + http.Error(w, "unexpected body", http.StatusBadRequest) + return + } + atomic.AddInt32(&requests, 1) + _, _ = w.Write([]byte("ok")) + })) + httpsServer.Config.ConnState = func(_ net.Conn, state http.ConnState) { + if state == http.StateNew { + atomic.AddInt32(&accepted, 1) + } + } + httpsServer.StartTLS() + defer httpsServer.Close() + + c, err := NewClient( + WithTLSConfig(httpsServer.Client().Transport.(*http.Transport).TLSClientConfig), + WithPooledConnHealthCheck(true), + ) + if err != nil { + t.Fatal(err) + } + defer c.CloseIdleConnections() + + for i := 0; i < 2; i++ { + req := protocol.AcquireRequest() + req.Header.SetMethod(consts.MethodPost) + req.SetRequestURI(httpsServer.URL + "/health-check") + req.SetBodyString("payload") + resp := protocol.AcquireResponse() + err := c.Do(context.Background(), req, resp) + assert.Nil(t, err) + assert.DeepEqual(t, http.StatusOK, resp.StatusCode()) + assert.DeepEqual(t, "ok", string(resp.Body())) + protocol.ReleaseRequest(req) + protocol.ReleaseResponse(resp) + } + + assert.DeepEqual(t, int32(1), atomic.LoadInt32(&accepted)) + assert.DeepEqual(t, int32(2), atomic.LoadInt32(&requests)) +} + func TestClientInvalidURI(t *testing.T) { opt, ln := newTestOptions(t) defer ln.Close() @@ -2290,6 +2411,14 @@ func TestClientHostClientConfigHook(t *testing.T) { assert.DeepEqual(t, "FOO.BAR:443", hcr.Addr) } +func TestPooledConnHealthCheckOptionPropagation(t *testing.T) { + c, err := NewClient(WithPooledConnHealthCheck(true)) + assert.Nil(t, err) + + http1Options := newHttp1OptionFromClient(c) + assert.DeepEqual(t, true, http1Options.PooledConnHealthCheck) +} + func TestClientDialerName(t *testing.T) { client, _ := NewClient() dName, err := client.GetDialerName() diff --git a/pkg/app/client/option.go b/pkg/app/client/option.go index c65cc76c6..6ba562023 100644 --- a/pkg/app/client/option.go +++ b/pkg/app/client/option.go @@ -49,6 +49,17 @@ func WithMaxIdleConnDuration(t time.Duration) config.ClientOption { }} } +// WithPooledConnHealthCheck determines whether idle pooled connections are +// probed before reuse. Unhealthy connections are discarded before writing a +// request. +// +// The check uses a short read probe and is disabled by default. +func WithPooledConnHealthCheck(enabled bool) config.ClientOption { + return config.ClientOption{F: func(o *config.ClientOptions) { + o.PooledConnHealthCheck = enabled + }} +} + // WithMaxConnDuration sets max connection duration, keep-alive connections are closed after this duration. func WithMaxConnDuration(t time.Duration) config.ClientOption { return config.ClientOption{F: func(o *config.ClientOptions) { diff --git a/pkg/app/client/option_test.go b/pkg/app/client/option_test.go index be52a9622..4bdb0de81 100644 --- a/pkg/app/client/option_test.go +++ b/pkg/app/client/option_test.go @@ -33,6 +33,7 @@ func TestClientOptions(t *testing.T) { assert.DeepEqual(t, 0, opt.MaxConnsPerHost) assert.DeepEqual(t, consts.DefaultDialTimeout, opt.DialTimeout) assert.DeepEqual(t, consts.DefaultMaxIdleConnDuration, opt.MaxIdleConnDuration) + assert.DeepEqual(t, false, opt.PooledConnHealthCheck) assert.DeepEqual(t, true, opt.KeepAlive) assert.DeepEqual(t, 5*time.Second, opt.ObservationInterval) @@ -41,6 +42,7 @@ func TestClientOptions(t *testing.T) { WithDialTimeout(100 * time.Millisecond), WithMaxConnsPerHost(128), WithMaxIdleConnDuration(5 * time.Second), + WithPooledConnHealthCheck(true), WithMaxConnDuration(10 * time.Second), WithMaxConnWaitTimeout(5 * time.Second), WithKeepAlive(false), @@ -59,6 +61,7 @@ func TestClientOptions(t *testing.T) { assert.DeepEqual(t, 100*time.Millisecond, opt.DialTimeout) assert.DeepEqual(t, 128, opt.MaxConnsPerHost) assert.DeepEqual(t, 5*time.Second, opt.MaxIdleConnDuration) + assert.DeepEqual(t, true, opt.PooledConnHealthCheck) assert.DeepEqual(t, 10*time.Second, opt.MaxConnDuration) assert.DeepEqual(t, 5*time.Second, opt.MaxConnWaitTimeout) assert.DeepEqual(t, false, opt.KeepAlive) diff --git a/pkg/common/config/client_option.go b/pkg/common/config/client_option.go index 2de181697..f09f414b6 100644 --- a/pkg/common/config/client_option.go +++ b/pkg/common/config/client_option.go @@ -56,12 +56,15 @@ type ClientOptions struct { MaxConnsPerHost int MaxIdleConnDuration time.Duration - MaxConnDuration time.Duration - MaxConnWaitTimeout time.Duration - KeepAlive bool - ReadTimeout time.Duration - TLSConfig *tls.Config - ResponseBodyStream bool + // PooledConnHealthCheck determines whether an idle pooled connection is + // probed before reuse. It is disabled by default. + PooledConnHealthCheck bool + MaxConnDuration time.Duration + MaxConnWaitTimeout time.Duration + KeepAlive bool + ReadTimeout time.Duration + TLSConfig *tls.Config + ResponseBodyStream bool // Client name. Used in User-Agent request header. // diff --git a/pkg/network/netpoll/healthcheck_linux.go b/pkg/network/netpoll/healthcheck_linux.go new file mode 100644 index 000000000..ca2ff69cf --- /dev/null +++ b/pkg/network/netpoll/healthcheck_linux.go @@ -0,0 +1,64 @@ +// Copyright 2022 CloudWeGo Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build linux + +package netpoll + +import ( + "errors" + "net" + "time" + + errs "github.com/cloudwego/hertz/pkg/common/errors" +) + +type reusableConnHealthChecker interface { + IsHealthyForReuse(ownerWaitTimeout time.Duration) bool +} + +// IsHealthy checks whether an idle pooled connection is safe to reuse without +// consuming application data. probeTimeout bounds the fallback timed Peek; +// ownerWaitTimeout bounds waiting for the netpoll operator. The owner-side +// probe itself is non-blocking once acquired. +func (c *Conn) IsHealthy(probeTimeout, ownerWaitTimeout time.Duration) bool { + if c == nil || c.Conn == nil || probeTimeout <= 0 || ownerWaitTimeout <= 0 || c.Len() > 0 { + return false + } + + // Newer netpoll versions synchronize this check with their receive owner. + // Keep the narrow assertion here so Hertz remains compatible with older + // netpoll versions and third-party network.Conn implementations. + if checker, ok := c.Conn.(reusableConnHealthChecker); ok { + return checker.IsHealthyForReuse(ownerWaitTimeout) + } + + // Older netpoll versions have no owner-side probe. Preserve the bounded + // Peek fallback for rolling upgrades instead of inspecting their raw fd. + if err := c.SetReadTimeout(probeTimeout); err != nil { + return false + } + p, err := c.Peek(1) + if resetErr := c.SetReadTimeout(0); resetErr != nil { + return false + } + if len(p) != 0 { + return false + } + if errors.Is(err, errs.ErrTimeout) { + return true + } + var timeoutErr net.Error + return errors.As(err, &timeoutErr) && timeoutErr.Timeout() +} diff --git a/pkg/network/netpoll/healthcheck_linux_test.go b/pkg/network/netpoll/healthcheck_linux_test.go new file mode 100644 index 000000000..67e52b5b4 --- /dev/null +++ b/pkg/network/netpoll/healthcheck_linux_test.go @@ -0,0 +1,144 @@ +// Copyright 2022 CloudWeGo Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build linux + +package netpoll + +import ( + "errors" + "testing" + "time" + + "github.com/cloudwego/hertz/pkg/common/test/mock" + "github.com/cloudwego/hertz/pkg/network" +) + +type reuseHealthProbe struct { + network.Conn + healthy bool + calls int + ownerWaitTimeout time.Duration +} + +func (p *reuseHealthProbe) IsHealthyForReuse(ownerWaitTimeout time.Duration) bool { + p.calls++ + p.ownerWaitTimeout = ownerWaitTimeout + return p.healthy +} + +type fallbackHealthProbe struct { + network.Conn + failTimeoutCall int + timeoutCalls int + peekData []byte +} + +func (p *fallbackHealthProbe) SetReadTimeout(timeout time.Duration) error { + p.timeoutCalls++ + if p.timeoutCalls == p.failTimeoutCall { + return errors.New("set read timeout failed") + } + return p.Conn.SetReadTimeout(timeout) +} + +func (p *fallbackHealthProbe) Peek(int) ([]byte, error) { + if p.peekData != nil { + return p.peekData, nil + } + return p.Conn.Peek(1) +} + +func TestConnIsHealthyDelegatesToOwnerProbe(t *testing.T) { + for _, tt := range []struct { + name string + healthy bool + }{ + {name: "healthy", healthy: true}, + {name: "unhealthy", healthy: false}, + } { + t.Run(tt.name, func(t *testing.T) { + probe := &reuseHealthProbe{Conn: mock.NewConn(""), healthy: tt.healthy} + conn := &Conn{Conn: probe} + if got := conn.IsHealthy(50*time.Microsecond, time.Second); got != tt.healthy { + t.Fatalf("owner probe result: got %v, want %v", got, tt.healthy) + } + if probe.calls != 1 { + t.Fatalf("owner probe calls: got %d, want 1", probe.calls) + } + if probe.ownerWaitTimeout != time.Second { + t.Fatalf("owner wait timeout: got %v, want %v", probe.ownerWaitTimeout, time.Second) + } + }) + } +} + +func TestConnIsHealthyFallsBackToTimedPeek(t *testing.T) { + conn := &Conn{Conn: mock.NewConn("")} + if !conn.IsHealthy(50*time.Microsecond, time.Second) { + t.Fatal("old netpoll without owner probe should use the timed Peek fallback") + } +} + +func TestConnIsHealthyFallbackRejectsStaleConnection(t *testing.T) { + conn := &Conn{Conn: mock.NewBrokenConn("")} + if conn.IsHealthy(50*time.Microsecond, time.Second) { + t.Fatal("timed Peek fallback must reject a stale connection") + } +} + +func TestConnIsHealthyRejectsInvalidInput(t *testing.T) { + for _, tt := range []struct { + name string + conn *Conn + probeTimeout time.Duration + ownerWaitTimeout time.Duration + }{ + {name: "nil receiver", probeTimeout: time.Second, ownerWaitTimeout: time.Second}, + {name: "nil connection", conn: &Conn{}, probeTimeout: time.Second, ownerWaitTimeout: time.Second}, + {name: "zero probe timeout", conn: &Conn{Conn: mock.NewConn("")}, ownerWaitTimeout: time.Second}, + {name: "zero owner wait timeout", conn: &Conn{Conn: mock.NewConn("")}, probeTimeout: time.Second}, + {name: "buffered data", conn: &Conn{Conn: mock.NewConn("x")}, probeTimeout: time.Second, ownerWaitTimeout: time.Second}, + } { + t.Run(tt.name, func(t *testing.T) { + if tt.conn.IsHealthy(tt.probeTimeout, tt.ownerWaitTimeout) { + t.Fatal("invalid health check input must fail closed") + } + }) + } +} + +func TestConnIsHealthyFallbackRejectsReadTimeoutFailures(t *testing.T) { + for _, failCall := range []int{1, 2} { + probe := &fallbackHealthProbe{Conn: mock.NewConn(""), failTimeoutCall: failCall} + conn := &Conn{Conn: probe} + if conn.IsHealthy(50*time.Microsecond, time.Second) { + t.Fatalf("read timeout failure on call %d must fail closed", failCall) + } + if probe.timeoutCalls != failCall { + t.Fatalf("read timeout calls: got %d, want %d", probe.timeoutCalls, failCall) + } + } +} + +func TestConnIsHealthyFallbackRejectsUnreadData(t *testing.T) { + probe := &fallbackHealthProbe{Conn: mock.NewConn(""), peekData: []byte("x")} + conn := &Conn{Conn: probe} + if conn.IsHealthy(50*time.Microsecond, time.Second) { + t.Fatal("timed Peek fallback must reject unread data") + } + if probe.timeoutCalls != 2 { + t.Fatalf("read timeout calls: got %d, want 2", probe.timeoutCalls) + } +} diff --git a/pkg/network/standard/connection.go b/pkg/network/standard/connection.go index 485bca488..230b55b3f 100644 --- a/pkg/network/standard/connection.go +++ b/pkg/network/standard/connection.go @@ -33,11 +33,17 @@ import ( "github.com/cloudwego/hertz/pkg/network" ) +type connHealthChecker interface { + isHealthy() bool +} + type Conn struct { - c net.Conn - br *bufiox.DefaultReader - bw *bufiox.DefaultWriter - stater connstate.ConnStater + c net.Conn + br *bufiox.DefaultReader + bw *bufiox.DefaultWriter + stater connstate.ConnStater + healthChecker connHealthChecker + readerSize int buf [8]byte } @@ -67,6 +73,43 @@ func (c *Conn) SetReadTimeout(t time.Duration) error { return c.c.SetReadDeadline(time.Now().Add(t)) } +// IsHealthy checks whether the peer has closed the connection without +// consuming data from the buffered reader used by the HTTP protocol. +func (c *Conn) IsHealthy(probeTimeout, _ time.Duration) bool { + if c == nil || c.c == nil || probeTimeout <= 0 || c.Len() > 0 { + return false + } + if c.healthChecker != nil { + return c.healthChecker.isHealthy() + } + return c.isHealthyWithTimedPeek(probeTimeout) +} + +func (c *Conn) isHealthyWithTimedPeek(timeout time.Duration) bool { + if err := c.SetReadTimeout(timeout); err != nil { + return false + } + p, err := c.Peek(1) + if resetErr := c.SetReadTimeout(0); resetErr != nil { + return false + } + if len(p) != 0 { + return false + } + var timeoutErr net.Error + if !errors.As(err, &timeoutErr) || !timeoutErr.Timeout() { + return false + } + // DefaultReader intentionally retains read errors. A successful timeout + // probe has no buffered data, so release its empty buffer and replace the + // reader before the connection is returned to the protocol. + if err := c.br.Release(nil); err != nil { + return false + } + c.br = bufiox.NewDefaultReaderSize(c.c, c.readerSize) + return true +} + type TLSConn struct { Conn } @@ -237,18 +280,21 @@ func (c *TLSConn) ConnectionState() tls.ConnectionState { func newConn(c net.Conn, size int) network.Conn { return &Conn{ - c: c, - br: bufiox.NewDefaultReaderSize(c, size), - bw: bufiox.NewDefaultWriter(c), + c: c, + br: bufiox.NewDefaultReaderSize(c, size), + bw: bufiox.NewDefaultWriter(c), + healthChecker: newTCPHealthChecker(c), + readerSize: size, } } func newTLSConn(c net.Conn, size int) network.Conn { return &TLSConn{ Conn{ - c: c, - br: bufiox.NewDefaultReaderSize(c, size), - bw: bufiox.NewDefaultWriter(c), + c: c, + br: bufiox.NewDefaultReaderSize(c, size), + bw: bufiox.NewDefaultWriter(c), + readerSize: size, }, } } diff --git a/pkg/network/standard/connection_test.go b/pkg/network/standard/connection_test.go index 00e652b2d..51c718082 100644 --- a/pkg/network/standard/connection_test.go +++ b/pkg/network/standard/connection_test.go @@ -239,6 +239,71 @@ func TestToHertzError(t *testing.T) { assert.DeepEqual(t, other, conn.ToHertzError(other)) } +func TestConnIsHealthy(t *testing.T) { + serverConn, clientConn := net.Pipe() + defer serverConn.Close() + defer clientConn.Close() + + conn := newConn(clientConn, 0).(*Conn) + assert.True(t, conn.IsHealthy(time.Millisecond, time.Second)) + + writeDone := make(chan error, 1) + go func() { + _, err := serverConn.Write([]byte("x")) + writeDone <- err + }() + b, err := conn.Peek(1) + assert.Nil(t, err) + assert.DeepEqual(t, []byte("x"), b) + assert.Nil(t, <-writeDone) + assert.Nil(t, conn.Skip(1)) + assert.Nil(t, conn.Release()) + + assert.Nil(t, serverConn.Close()) + assert.DeepEqual(t, false, conn.IsHealthy(time.Millisecond, time.Second)) + + t.Run("unexpected data", func(t *testing.T) { + serverConn, clientConn := net.Pipe() + defer serverConn.Close() + defer clientConn.Close() + + conn := newConn(clientConn, 0).(*Conn) + writeDone := make(chan error, 1) + go func() { + _, err := serverConn.Write([]byte("x")) + writeDone <- err + }() + assert.DeepEqual(t, false, conn.IsHealthy(time.Second, time.Second)) + assert.Nil(t, <-writeDone) + b, err := conn.Peek(1) + assert.Nil(t, err) + assert.DeepEqual(t, []byte("x"), b) + }) + + t.Run("buffered data", func(t *testing.T) { + serverConn, clientConn := net.Pipe() + defer serverConn.Close() + defer clientConn.Close() + + conn := newConn(clientConn, 0).(*Conn) + writeDone := make(chan error, 1) + go func() { + _, err := serverConn.Write([]byte("x")) + writeDone <- err + }() + b, err := conn.Peek(1) + assert.Nil(t, err) + assert.DeepEqual(t, []byte("x"), b) + assert.Nil(t, <-writeDone) + + assert.DeepEqual(t, false, conn.IsHealthy(time.Millisecond, time.Second)) + b, err = conn.Peek(1) + assert.Nil(t, err) + assert.DeepEqual(t, []byte("x"), b) + }) + +} + func TestHandleSpecificError(t *testing.T) { conn := &Conn{} assert.DeepEqual(t, false, conn.HandleSpecificError(nil, "")) diff --git a/pkg/network/standard/healthcheck_default.go b/pkg/network/standard/healthcheck_default.go new file mode 100644 index 000000000..686e271d8 --- /dev/null +++ b/pkg/network/standard/healthcheck_default.go @@ -0,0 +1,23 @@ +// Copyright 2022 CloudWeGo Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build !linux + +package standard + +import "net" + +func newTCPHealthChecker(net.Conn) connHealthChecker { + return nil +} diff --git a/pkg/network/standard/healthcheck_linux.go b/pkg/network/standard/healthcheck_linux.go new file mode 100644 index 000000000..c174d02a6 --- /dev/null +++ b/pkg/network/standard/healthcheck_linux.go @@ -0,0 +1,97 @@ +// Copyright 2022 CloudWeGo Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build linux + +package standard + +import ( + "net" + "sync" + "syscall" + "time" +) + +// tcpHealthChecker owns the RawConn.Read callback for a native TCP connection. +// The HTTP pool serializes use of an idle connection; the mutex also serializes +// accidental concurrent probes before the result is returned. +type tcpHealthChecker struct { + mu sync.Mutex + tcpConn *net.TCPConn + rawConn syscall.RawConn + healthy bool + probe func(uintptr) bool +} + +type unhealthyTCPHealthChecker struct{} + +func (unhealthyTCPHealthChecker) isHealthy() bool { return false } + +func newTCPHealthChecker(conn net.Conn) connHealthChecker { + tcpConn, ok := conn.(*net.TCPConn) + if !ok { + return nil + } + + rawConn, err := tcpConn.SyscallConn() + if err != nil { + return unhealthyTCPHealthChecker{} + } + checker := &tcpHealthChecker{tcpConn: tcpConn, rawConn: rawConn} + checker.probe = func(fd uintptr) bool { + checker.healthy = recvHealthyTCP(fd) + // MSG_DONTWAIT makes the probe decisive at this instant. Returning true + // prevents RawConn.Read from registering or waiting in the runtime poller. + return true + } + return checker +} + +func (c *tcpHealthChecker) isHealthy() bool { + c.mu.Lock() + defer c.mu.Unlock() + + // Keep the timed-Peek health check's deadline-reset contract: an HTTP + // response may leave an expired read deadline on a pooled connection. Clear + // it before RawConn.Read so a healthy idle connection is not rejected. This + // does not install a probe deadline: MSG_DONTWAIT keeps the socket + // observation non-blocking. + if err := c.tcpConn.SetReadDeadline(time.Time{}); err != nil { + return false + } + c.healthy = false + if err := c.rawConn.Read(c.probe); err != nil { + return false + } + return c.healthy +} + +func recvHealthyTCP(fd uintptr) bool { + var buffer [1]byte + for { + _, _, err := syscall.Recvfrom(int(fd), buffer[:], syscall.MSG_PEEK|syscall.MSG_DONTWAIT) + switch err { + case nil: + // A positive result is unread data; a zero-length result is the + // peer's FIN. Neither connection is safe to reuse. + return false + case syscall.EAGAIN: // EWOULDBLOCK is the same errno on Linux. + return true + case syscall.EINTR: + continue + default: + return false + } + } +} diff --git a/pkg/network/standard/healthcheck_linux_test.go b/pkg/network/standard/healthcheck_linux_test.go new file mode 100644 index 000000000..9c7d02063 --- /dev/null +++ b/pkg/network/standard/healthcheck_linux_test.go @@ -0,0 +1,281 @@ +// Copyright 2022 CloudWeGo Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build linux + +package standard + +import ( + "crypto/tls" + "net" + "sync/atomic" + "testing" + "time" +) + +func TestConnIsHealthyTCPFastPath(t *testing.T) { + t.Run("clears an existing read deadline", func(t *testing.T) { + conn, peer := newStandardTCPPair(t) + if err := conn.c.SetReadDeadline(time.Now().Add(-time.Second)); err != nil { + t.Fatalf("set expired read deadline: %v", err) + } + if !conn.IsHealthy(5*time.Millisecond, time.Second) { + t.Fatal("idle TCP connection should be healthy after clearing its prior deadline") + } + if _, err := peer.Write([]byte("x")); err != nil { + t.Fatalf("write after health check: %v", err) + } + data, err := conn.Peek(1) + if err != nil { + t.Fatalf("peek after health check: %v", err) + } + if string(data) != "x" { + t.Fatalf("read after health check: got %q, want %q", data, "x") + } + }) + + t.Run("pending data is rejected without consumption", func(t *testing.T) { + conn, peer := newStandardTCPPair(t) + if _, err := peer.Write([]byte("x")); err != nil { + t.Fatalf("write peer data: %v", err) + } + waitForStandardTCPUnhealthy(t, conn) + data, err := conn.Peek(1) + if err != nil { + t.Fatalf("peek after health check: %v", err) + } + if string(data) != "x" { + t.Fatalf("health check consumed or changed data: %q", data) + } + }) + + t.Run("peer FIN is rejected", func(t *testing.T) { + conn, peer := newStandardTCPPair(t) + if err := peer.Close(); err != nil { + t.Fatalf("close peer: %v", err) + } + waitForStandardTCPUnhealthy(t, conn) + }) + + t.Run("peer reset is rejected", func(t *testing.T) { + conn, peer := newStandardTCPPair(t) + if err := peer.SetLinger(0); err != nil { + t.Fatalf("set peer linger: %v", err) + } + if err := peer.Close(); err != nil { + t.Fatalf("reset peer: %v", err) + } + waitForStandardTCPUnhealthy(t, conn) + }) + + t.Run("local close is rejected", func(t *testing.T) { + conn, _ := newStandardTCPPair(t) + if err := conn.Close(); err != nil { + t.Fatalf("close connection: %v", err) + } + if conn.IsHealthy(5*time.Millisecond, time.Second) { + t.Fatal("locally closed TCP connection should not be healthy") + } + }) +} + +func TestConnIsHealthyTLSUsesTimedFallback(t *testing.T) { + certData, keyData, err := generateTestCertificate("") + if err != nil { + t.Fatalf("generate certificate: %v", err) + } + cert, err := tls.X509KeyPair(certData, keyData) + if err != nil { + t.Fatalf("parse certificate: %v", err) + } + + listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1)}) + if err != nil { + t.Fatalf("listen: %v", err) + } + defer listener.Close() + + serverResult := make(chan *tls.Conn, 1) + serverErr := make(chan error, 1) + go func() { + peer, err := listener.AcceptTCP() + if err != nil { + serverErr <- err + return + } + serverTLS := tls.Server(peer, &tls.Config{Certificates: []tls.Certificate{cert}}) + if err := serverTLS.Handshake(); err != nil { + _ = serverTLS.Close() + serverErr <- err + return + } + serverResult <- serverTLS + }() + + clientTCP, err := net.DialTCP("tcp", nil, listener.Addr().(*net.TCPAddr)) + if err != nil { + t.Fatalf("dial: %v", err) + } + trackingConn := &readDeadlineTrackingConn{Conn: clientTCP} + clientTLS := tls.Client(trackingConn, &tls.Config{InsecureSkipVerify: true}) //nolint:gosec // loopback test certificate + if err := clientTLS.Handshake(); err != nil { + _ = clientTLS.Close() + t.Fatalf("client handshake: %v", err) + } + + var serverTLS *tls.Conn + select { + case serverTLS = <-serverResult: + case err := <-serverErr: + _ = clientTLS.Close() + t.Fatalf("server handshake: %v", err) + case <-time.After(time.Second): + _ = clientTLS.Close() + t.Fatal("server handshake timed out") + } + defer serverTLS.Close() + defer clientTLS.Close() + + trackingConn.readDeadlineCalls.Store(0) + conn := newTLSConn(clientTLS, 0).(*TLSConn) + if !conn.IsHealthy(time.Millisecond, time.Second) { + t.Fatal("idle TLS connection should remain healthy through fallback") + } + if calls := trackingConn.readDeadlineCalls.Load(); calls != 2 { + t.Fatalf("TLS read deadline calls: got %d, want 2 for set and reset", calls) + } + if _, err := serverTLS.Write([]byte("x")); err != nil { + t.Fatalf("write after TLS health check: %v", err) + } + data, err := conn.Peek(1) + if err != nil { + t.Fatalf("peek after TLS health check: %v", err) + } + if string(data) != "x" { + t.Fatalf("read after TLS health check: got %q, want %q", data, "x") + } +} + +func TestConnIsHealthyUnsupportedConnUsesTimedFallback(t *testing.T) { + serverConn, clientConn := net.Pipe() + defer serverConn.Close() + defer clientConn.Close() + + trackingConn := &readDeadlineTrackingConn{Conn: clientConn} + conn := newConn(trackingConn, 0).(*Conn) + if !conn.IsHealthy(time.Millisecond, time.Second) { + t.Fatal("idle unsupported connection should remain healthy through fallback") + } + if calls := trackingConn.readDeadlineCalls.Load(); calls != 2 { + t.Fatalf("read deadline calls: got %d, want 2 for set and reset", calls) + } + writeDone := make(chan error, 1) + go func() { + _, err := serverConn.Write([]byte("x")) + writeDone <- err + }() + data, err := conn.Peek(1) + if err != nil { + t.Fatalf("peek after fallback health check: %v", err) + } + if string(data) != "x" { + t.Fatalf("read after fallback health check: got %q, want %q", data, "x") + } + if err := <-writeDone; err != nil { + t.Fatalf("write after fallback health check: %v", err) + } +} + +func BenchmarkConnIsHealthyTCP(b *testing.B) { + conn, _ := newStandardTCPPair(b) + if !conn.IsHealthy(5*time.Millisecond, time.Second) { + b.Fatal("warmup rejected idle TCP connection") + } + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if !conn.IsHealthy(5*time.Millisecond, time.Second) { + b.Fatal("idle TCP connection was rejected") + } + } +} + +type readDeadlineTrackingConn struct { + net.Conn + readDeadlineCalls atomic.Int32 +} + +func (c *readDeadlineTrackingConn) SetReadDeadline(deadline time.Time) error { + c.readDeadlineCalls.Add(1) + return c.Conn.SetReadDeadline(deadline) +} + +func newStandardTCPPair(tb testing.TB) (*Conn, *net.TCPConn) { + tb.Helper() + listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1)}) + if err != nil { + tb.Fatalf("listen: %v", err) + } + + accepted := make(chan *net.TCPConn, 1) + acceptErr := make(chan error, 1) + go func() { + peer, err := listener.AcceptTCP() + if err != nil { + acceptErr <- err + return + } + accepted <- peer + }() + + clientConn, err := net.DialTCP("tcp", nil, listener.Addr().(*net.TCPAddr)) + if err != nil { + _ = listener.Close() + tb.Fatalf("dial: %v", err) + } + + var peer *net.TCPConn + select { + case peer = <-accepted: + case err := <-acceptErr: + _ = clientConn.Close() + _ = listener.Close() + tb.Fatalf("accept: %v", err) + case <-time.After(time.Second): + _ = clientConn.Close() + _ = listener.Close() + tb.Fatal("accept timed out") + } + + conn := newConn(clientConn, 0).(*Conn) + tb.Cleanup(func() { + _ = conn.Close() + _ = peer.Close() + _ = listener.Close() + }) + return conn, peer +} + +func waitForStandardTCPUnhealthy(tb testing.TB, conn *Conn) { + tb.Helper() + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if !conn.IsHealthy(5*time.Millisecond, time.Second) { + return + } + time.Sleep(time.Millisecond) + } + tb.Fatal("TCP connection remained healthy") +} diff --git a/pkg/protocol/http1/client.go b/pkg/protocol/http1/client.go index 17caf469d..7f10b4981 100644 --- a/pkg/protocol/http1/client.go +++ b/pkg/protocol/http1/client.go @@ -551,12 +551,22 @@ func (c *HostClient) doNonNilReqResp(req *protocol.Request, resp *protocol.Respo if timeout < 0 { return false, errTimeout } - cc, inPool, err := c.acquireConn(timeout) + cc, inPool, err := c.acquireConn(timeout, deadline) if err != nil { return false, err } conn := cc.c resp.ParseNetAddr(conn) + if c.PooledConnHealthCheck { + // Health checks may discard stale pooled connections and acquire a + // replacement. Recompute the remaining dial/handshake budget after that + // work instead of reusing the timeout calculated before acquireConn. + timeout = calcTimeout(deadline, dtimeout) + if timeout < 0 { + c.closeConn(cc) + return false, errTimeout + } + } if c.IsTLS && timeout > 0 { // force handshake using dial timeout // NOTE: Handshake() here is optional as Write would tirigger handshake @@ -816,7 +826,58 @@ func (c *HostClient) SetMaxConns(newMaxConns int) { c.connsLock.Unlock() } -func (c *HostClient) acquireConn(dialTimeout time.Duration) (cc *clientConn, inPool bool, err error) { +// pooledConnHealthCheckTimeout bounds the read probe so healthy connections +// are not delayed for long before a request is written. +const pooledConnHealthCheckTimeout = 50 * time.Microsecond + +type pooledConnHealthChecker interface { + IsHealthy(probeTimeout, ownerWaitTimeout time.Duration) bool +} + +func (c *HostClient) acquireConn(dialTimeout time.Duration, deadline time.Time) (cc *clientConn, inPool bool, err error) { + if !c.PooledConnHealthCheck { + return c.acquireConnOnce(dialTimeout) + } + + // Probe reused connections before any request bytes are written. A stale + // connection is closed and the loop acquires another one; this is connection + // selection, not a request retry. + for { + // Previous probes and replacement attempts consume the request deadline, + // so every acquisition must use the remaining budget. + currentDialTimeout := calcTimeout(deadline, dialTimeout) + if currentDialTimeout < 0 { + return nil, false, errTimeout + } + cc, inPool, err = c.acquireConnOnce(currentDialTimeout) + // Newly dialed connections have no last-use time and do not need a + // preflight probe. Only connections returned by the pool are checked. + if err != nil || cc == nil || cc.lastUseTime.IsZero() { + return cc, inPool, err + } + // Bound the socket probe by the short fixed window and bound the netpoll + // owner acquisition by the smaller of that window and the remaining + // request/dial budget. Older transports ignore the owner budget and use + // the probe timeout in their timed-Peek fallback. + healthCheckTimeout := calcTimeout(deadline, pooledConnHealthCheckTimeout) + if healthCheckTimeout < 0 { + c.closeConn(cc) + return nil, false, errTimeout + } + ownerWaitTimeout := healthCheckTimeout + if currentDialTimeout > 0 && currentDialTimeout < ownerWaitTimeout { + ownerWaitTimeout = currentDialTimeout + } + if c.isPooledConnHealthy(cc.c, healthCheckTimeout, ownerWaitTimeout) { + return cc, inPool, nil + } + // No request bytes have been written yet, so discard the stale connection + // and safely continue looking for a usable one. + c.closeConn(cc) + } +} + +func (c *HostClient) acquireConnOnce(dialTimeout time.Duration) (cc *clientConn, inPool bool, err error) { createConn := false startCleaner := false @@ -892,6 +953,44 @@ func (c *HostClient) acquireConn(dialTimeout time.Duration) (cc *clientConn, inP return cc, false, nil } +// isPooledConnHealthy performs a non-consuming pre-write read probe. A timeout +// with no buffered data means the connection is still usable; readable data or +// any other error makes the connection ineligible for reuse. The owner wait +// budget is kept separate because transports with an owner-side probe must +// bound operator acquisition without changing the fallback read timeout. +func (c *HostClient) isPooledConnHealthy(conn network.Conn, probeTimeout, ownerWaitTimeout time.Duration) bool { + // Prefer a transport-specific implementation because its buffered reader may + // need special handling to avoid consuming data or retaining a probe timeout. + if checker, ok := conn.(pooledConnHealthChecker); ok { + return checker.IsHealthy(probeTimeout, ownerWaitTimeout) + } + // Keep the probe short so an idle but healthy connection does not delay the + // request while waiting for a byte that should not arrive. + if err := conn.SetReadTimeout(probeTimeout); err != nil { + return false + } + // Peek preserves any byte for the protocol reader; the health check must not + // consume application data. + p, err := conn.Peek(1) + // Restore the normal read deadline before the connection is reused. If the + // deadline cannot be restored, the connection is unsafe to return. + if resetErr := conn.SetReadTimeout(0); resetErr != nil { + return false + } + // Any readable byte is unexpected on an idle connection. Check the data + // independently because some implementations may return data with an error. + if len(p) != 0 { + return false + } + // No data until the probe deadline is the expected healthy result. Accept + // both Hertz's timeout sentinel and transport-specific net.Error values. + if errors.Is(err, errs.ErrTimeout) { + return true + } + var timeoutErr net.Error + return errors.As(err, &timeoutErr) && timeoutErr.Timeout() +} + func (c *HostClient) queueForIdle(w *wantConn) { c.connsLock.Lock() defer c.connsLock.Unlock() @@ -1399,6 +1498,10 @@ type ClientOptions struct { // after DefaultMaxIdleConnDuration. MaxIdleConnDuration time.Duration + // PooledConnHealthCheck determines whether an idle pooled connection is + // probed before reuse. It is disabled by default. + PooledConnHealthCheck bool + // Maximum duration for full response reading (including body). // // By default response read timeout is unlimited. diff --git a/pkg/protocol/http1/client_healthcheck_dialers_nonwindows_test.go b/pkg/protocol/http1/client_healthcheck_dialers_nonwindows_test.go new file mode 100644 index 000000000..8f91ae5cf --- /dev/null +++ b/pkg/protocol/http1/client_healthcheck_dialers_nonwindows_test.go @@ -0,0 +1,29 @@ +// Copyright 2022 CloudWeGo Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build !windows + +package http1 + +import ( + hznetpoll "github.com/cloudwego/hertz/pkg/network/netpoll" + "github.com/cloudwego/hertz/pkg/network/standard" +) + +func pooledConnHealthCheckDialers() []pooledConnHealthCheckDialer { + return []pooledConnHealthCheckDialer{ + {name: "standard", dialer: standard.NewDialer()}, + {name: "netpoll", dialer: hznetpoll.NewDialer()}, + } +} diff --git a/pkg/protocol/http1/client_healthcheck_dialers_windows_test.go b/pkg/protocol/http1/client_healthcheck_dialers_windows_test.go new file mode 100644 index 000000000..400e68a80 --- /dev/null +++ b/pkg/protocol/http1/client_healthcheck_dialers_windows_test.go @@ -0,0 +1,25 @@ +// Copyright 2022 CloudWeGo Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build windows + +package http1 + +import "github.com/cloudwego/hertz/pkg/network/standard" + +func pooledConnHealthCheckDialers() []pooledConnHealthCheckDialer { + return []pooledConnHealthCheckDialer{ + {name: "standard", dialer: standard.NewDialer()}, + } +} diff --git a/pkg/protocol/http1/client_test.go b/pkg/protocol/http1/client_test.go index d4cf07f1a..ffd6a66da 100644 --- a/pkg/protocol/http1/client_test.go +++ b/pkg/protocol/http1/client_test.go @@ -42,14 +42,17 @@ package http1 import ( + "bufio" "bytes" "context" "crypto/tls" "errors" "fmt" + "io" "io/ioutil" "net" "net/http" + "net/http/httptest" "strings" "sync" "sync/atomic" @@ -67,6 +70,7 @@ import ( "github.com/cloudwego/hertz/pkg/common/utils" "github.com/cloudwego/hertz/pkg/network" "github.com/cloudwego/hertz/pkg/network/dialer" + "github.com/cloudwego/hertz/pkg/network/standard" "github.com/cloudwego/hertz/pkg/protocol" "github.com/cloudwego/hertz/pkg/protocol/client" "github.com/cloudwego/hertz/pkg/protocol/consts" @@ -75,6 +79,11 @@ import ( var errDialTimeout = errs.New(errs.ErrTimeout, errs.ErrorTypePublic, "dial timeout") +type pooledConnHealthCheckDialer struct { + name string + dialer network.Dialer +} + func TestHostClientMaxConnWaitTimeoutWithEarlierDeadline(t *testing.T) { var ( emptyBodyCount uint8 @@ -738,6 +747,511 @@ func (w retryConn) SetWriteTimeout(t time.Duration) error { return errors.New("should retry") } +type pooledConnHealthProbe struct { + network.Conn + firstPeekErr error + healthCheckPeeks int + closeCalls int + readTimeouts []time.Duration + currentReadTimeout time.Duration +} + +type pooledConnOwnerHealthProbe struct { + network.Conn + healthy bool + calls int + probeTimeout time.Duration + ownerWaitTimeout time.Duration +} + +type pooledConnFallbackHealthProbe struct { + network.Conn + failReadTimeoutCall int + readTimeoutCalls int + peekData []byte +} + +func (c *pooledConnOwnerHealthProbe) IsHealthy(probeTimeout, ownerWaitTimeout time.Duration) bool { + c.calls++ + c.probeTimeout = probeTimeout + c.ownerWaitTimeout = ownerWaitTimeout + return c.healthy +} + +func (c *pooledConnFallbackHealthProbe) SetReadTimeout(timeout time.Duration) error { + c.readTimeoutCalls++ + if c.readTimeoutCalls == c.failReadTimeoutCall { + return errors.New("set read timeout failed") + } + return c.Conn.SetReadTimeout(timeout) +} + +func (c *pooledConnFallbackHealthProbe) Peek(int) ([]byte, error) { + if c.peekData != nil { + return c.peekData, nil + } + return c.Conn.Peek(1) +} + +func (c *pooledConnHealthProbe) Peek(n int) ([]byte, error) { + if c.currentReadTimeout > 0 && c.healthCheckPeeks == 0 { + c.healthCheckPeeks++ + return nil, c.firstPeekErr + } + return c.Conn.Peek(n) +} + +func (c *pooledConnHealthProbe) SetReadTimeout(timeout time.Duration) error { + c.readTimeouts = append(c.readTimeouts, timeout) + c.currentReadTimeout = timeout + return c.Conn.SetReadTimeout(timeout) +} + +func (c *pooledConnHealthProbe) Close() error { + c.closeCalls++ + return c.Conn.Close() +} + +func addPooledConn(c *HostClient, conn network.Conn) *clientConn { + cc := acquireClientConn(conn) + cc.lastUseTime = time.Now() + c.conns = append(c.conns, cc) + c.connsCount++ + return cc +} + +func TestPooledConnHealthCheckFallbackRejectsProbeFailures(t *testing.T) { + c := &HostClient{} + for _, failCall := range []int{1, 2} { + probe := &pooledConnFallbackHealthProbe{ + Conn: mock.NewConn(""), + failReadTimeoutCall: failCall, + } + if c.isPooledConnHealthy(probe, time.Millisecond, time.Millisecond) { + t.Fatalf("read timeout failure on call %d must fail closed", failCall) + } + if probe.readTimeoutCalls != failCall { + t.Fatalf("read timeout calls: got %d, want %d", probe.readTimeoutCalls, failCall) + } + } + + probe := &pooledConnFallbackHealthProbe{ + Conn: mock.NewConn(""), + peekData: []byte("x"), + } + if c.isPooledConnHealthy(probe, time.Second, time.Second) { + t.Fatal("fallback probe must reject unread data") + } + if probe.readTimeoutCalls != 2 { + t.Fatalf("read timeout calls: got %d, want 2", probe.readTimeoutCalls) + } +} + +func TestPooledConnHealthCheckDisabledByDefault(t *testing.T) { + stale := &pooledConnHealthProbe{ + Conn: mock.NewBrokenConn(""), + firstPeekErr: io.EOF, + } + c := &HostClient{ClientOptions: &ClientOptions{}} + want := addPooledConn(c, stale) + + got, inPool, err := c.acquireConn(time.Second, time.Time{}) + assert.Nil(t, err) + assert.True(t, inPool) + assert.True(t, got == want) + assert.DeepEqual(t, 0, stale.healthCheckPeeks) + + c.closeConn(got) +} + +func TestPooledConnHealthCheckReusesHealthyConnection(t *testing.T) { + response := "HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok" + healthy := &pooledConnHealthProbe{ + Conn: mock.NewConn(response), + firstPeekErr: errs.ErrTimeout, + } + var dialCalls int + c := &HostClient{ + ClientOptions: &ClientOptions{ + PooledConnHealthCheck: true, + Dialer: dialerFunc(func(network, addr string, timeout time.Duration) (network.Conn, error) { + dialCalls++ + return nil, errors.New("unexpected dial") + }), + }, + Addr: "foobar", + } + addPooledConn(c, healthy) + + req := protocol.AcquireRequest() + defer protocol.ReleaseRequest(req) + req.Header.SetMethod(consts.MethodPost) + req.SetRequestURI("http://foobar/health-check") + req.SetBodyString("payload") + resp := protocol.AcquireResponse() + defer protocol.ReleaseResponse(resp) + + err := c.Do(context.Background(), req, resp) + assert.Nil(t, err) + assert.DeepEqual(t, 200, resp.StatusCode()) + assert.DeepEqual(t, "ok", string(resp.Body())) + assert.DeepEqual(t, 0, dialCalls) + assert.DeepEqual(t, 1, healthy.healthCheckPeeks) + assert.DeepEqual(t, []time.Duration{pooledConnHealthCheckTimeout, 0}, healthy.readTimeouts[:2]) +} + +func TestPooledConnHealthCheckUsesShortProbeTimeout(t *testing.T) { + healthy := &pooledConnHealthProbe{ + Conn: mock.NewConn(""), + firstPeekErr: errs.ErrTimeout, + } + c := &HostClient{ClientOptions: &ClientOptions{ + PooledConnHealthCheck: true, + }} + addPooledConn(c, healthy) + + got, inPool, err := c.acquireConn(time.Second, time.Time{}) + assert.Nil(t, err) + assert.True(t, inPool) + assert.DeepEqual(t, 1, healthy.healthCheckPeeks) + assert.DeepEqual(t, pooledConnHealthCheckTimeout, healthy.readTimeouts[0]) + assert.DeepEqual(t, time.Duration(0), healthy.readTimeouts[1]) + + c.closeConn(got) +} + +func TestPooledConnHealthCheckBoundsOwnerWaitBudget(t *testing.T) { + const requestBudget = 200 * time.Millisecond + probe := &pooledConnOwnerHealthProbe{ + Conn: mock.NewConn(""), + healthy: true, + } + c := &HostClient{ClientOptions: &ClientOptions{ + PooledConnHealthCheck: true, + }} + addPooledConn(c, probe) + + got, inPool, err := c.acquireConn(requestBudget, time.Time{}) + assert.Nil(t, err) + assert.True(t, inPool) + assert.DeepEqual(t, 1, probe.calls) + assert.DeepEqual(t, pooledConnHealthCheckTimeout, probe.probeTimeout) + assert.DeepEqual(t, pooledConnHealthCheckTimeout, probe.ownerWaitTimeout) + + c.closeConn(got) + + probe = &pooledConnOwnerHealthProbe{Conn: mock.NewConn(""), healthy: true} + c = &HostClient{ClientOptions: &ClientOptions{PooledConnHealthCheck: true}} + addPooledConn(c, probe) + got, inPool, err = c.acquireConn(10*time.Microsecond, time.Time{}) + assert.Nil(t, err) + assert.True(t, inPool) + assert.DeepEqual(t, 10*time.Microsecond, probe.ownerWaitTimeout) + + c.closeConn(got) +} + +func TestPooledConnHealthCheckDiscardsStaleConnectionBeforePost(t *testing.T) { + staleBase := mock.NewBrokenConn("") + stale := &pooledConnHealthProbe{ + Conn: staleBase, + firstPeekErr: io.EOF, + } + fresh := mock.NewConn("HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok") + var dialCalls int + c := &HostClient{ + ClientOptions: &ClientOptions{ + PooledConnHealthCheck: true, + Dialer: dialerFunc(func(network, addr string, timeout time.Duration) (network.Conn, error) { + dialCalls++ + return fresh, nil + }), + }, + Addr: "foobar", + } + addPooledConn(c, stale) + + req := protocol.AcquireRequest() + defer protocol.ReleaseRequest(req) + req.Header.SetMethod(consts.MethodPost) + req.SetRequestURI("http://foobar/health-check") + req.SetBodyString("payload") + resp := protocol.AcquireResponse() + defer protocol.ReleaseResponse(resp) + + err := c.Do(context.Background(), req, resp) + assert.Nil(t, err) + assert.DeepEqual(t, 200, resp.StatusCode()) + assert.DeepEqual(t, "ok", string(resp.Body())) + assert.DeepEqual(t, 1, stale.healthCheckPeeks) + assert.DeepEqual(t, 1, stale.closeCalls) + assert.DeepEqual(t, 0, staleBase.WriterRecorder().WroteLen()) + assert.DeepEqual(t, 1, dialCalls) +} + +func TestPooledConnHealthCheckChecksConnectionDeliveredToWaiter(t *testing.T) { + stale := &pooledConnHealthProbe{ + Conn: mock.NewBrokenConn(""), + firstPeekErr: io.EOF, + } + fresh := mock.NewConn("") + var dialCalls int + c := &HostClient{ + ClientOptions: &ClientOptions{ + Dialer: dialerFunc(func(network, addr string, timeout time.Duration) (network.Conn, error) { + dialCalls++ + if dialCalls == 1 { + return stale, nil + } + return fresh, nil + }), + MaxConns: 1, + MaxConnWaitTimeout: time.Second, + PooledConnHealthCheck: true, + }, + Addr: "foobar", + } + + cc, inPool, err := c.acquireConn(time.Second, time.Time{}) + assert.Nil(t, err) + assert.DeepEqual(t, false, inPool) + + type acquireResult struct { + cc *clientConn + inPool bool + err error + } + resultCh := make(chan acquireResult, 1) + go func() { + cc, inPool, err := c.acquireConn(time.Second, time.Time{}) + resultCh <- acquireResult{cc: cc, inPool: inPool, err: err} + }() + + deadline := time.After(time.Second) + for { + c.connsLock.Lock() + waiting := c.connsWait != nil && c.connsWait.len() == 1 + c.connsLock.Unlock() + if waiting { + break + } + select { + case <-deadline: + t.Fatal("waiter was not queued") + default: + time.Sleep(time.Millisecond) + } + } + + c.releaseConn(cc) + select { + case got := <-resultCh: + assert.Nil(t, got.err) + assert.DeepEqual(t, false, got.inPool) + assert.True(t, got.cc.c == fresh) + c.closeConn(got.cc) + case <-time.After(time.Second): + t.Fatal("waiter did not receive a replacement connection") + } + + assert.DeepEqual(t, 1, stale.healthCheckPeeks) + assert.DeepEqual(t, 1, stale.closeCalls) + assert.DeepEqual(t, 2, dialCalls) +} + +func TestPooledConnHealthCheckRealConnection(t *testing.T) { + for _, tt := range pooledConnHealthCheckDialers() { + t.Run(tt.name, func(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + assert.Nil(t, err) + defer ln.Close() + + serverResult := make(chan error, 1) + var accepted, requests int + go func() { + for i := 0; i < 2; i++ { + conn, err := ln.Accept() + if err != nil { + serverResult <- err + return + } + accepted++ + req, err := http.ReadRequest(bufio.NewReader(conn)) + if err != nil { + conn.Close() + serverResult <- err + return + } + body, err := io.ReadAll(req.Body) + req.Body.Close() + if err != nil || req.Method != consts.MethodPost || string(body) != "payload" { + conn.Close() + serverResult <- fmt.Errorf("unexpected request method=%q body=%q err=%v", req.Method, body, err) + return + } + requests++ + if _, err := io.WriteString(conn, "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: keep-alive\r\n\r\nok"); err != nil { + conn.Close() + serverResult <- err + return + } + conn.Close() + } + serverResult <- nil + }() + + c := &HostClient{ + ClientOptions: &ClientOptions{ + Dialer: tt.dialer, + PooledConnHealthCheck: true, + }, + Addr: ln.Addr().String(), + } + for i := 0; i < 2; i++ { + req := protocol.AcquireRequest() + req.Header.SetMethod(consts.MethodPost) + req.SetRequestURI("http://" + ln.Addr().String() + "/health-check") + req.SetBodyString("payload") + resp := protocol.AcquireResponse() + err := c.Do(context.Background(), req, resp) + assert.Nil(t, err) + assert.DeepEqual(t, 200, resp.StatusCode()) + assert.DeepEqual(t, "ok", string(resp.Body())) + protocol.ReleaseRequest(req) + protocol.ReleaseResponse(resp) + } + + select { + case err := <-serverResult: + assert.Nil(t, err) + case <-time.After(time.Second): + t.Fatal("server did not receive both requests") + } + assert.DeepEqual(t, 2, accepted) + assert.DeepEqual(t, 2, requests) + }) + } +} + +func TestPooledConnHealthCheckReusesRealConnection(t *testing.T) { + for _, tt := range pooledConnHealthCheckDialers() { + t.Run(tt.name, func(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + assert.Nil(t, err) + defer ln.Close() + + serverResult := make(chan error, 1) + var accepted, requests int + go func() { + conn, err := ln.Accept() + if err != nil { + serverResult <- err + return + } + defer conn.Close() + accepted++ + reader := bufio.NewReader(conn) + for i := 0; i < 2; i++ { + req, err := http.ReadRequest(reader) + if err != nil { + serverResult <- err + return + } + body, err := io.ReadAll(req.Body) + req.Body.Close() + if err != nil || req.Method != consts.MethodPost || string(body) != "payload" { + serverResult <- fmt.Errorf("unexpected request method=%q body=%q err=%v", req.Method, body, err) + return + } + requests++ + if _, err := io.WriteString(conn, "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: keep-alive\r\n\r\nok"); err != nil { + serverResult <- err + return + } + } + serverResult <- nil + }() + + c := &HostClient{ + ClientOptions: &ClientOptions{ + Dialer: tt.dialer, + PooledConnHealthCheck: true, + }, + Addr: ln.Addr().String(), + } + defer c.CloseIdleConnections() + for i := 0; i < 2; i++ { + req := protocol.AcquireRequest() + req.Header.SetMethod(consts.MethodPost) + req.SetRequestURI("http://" + ln.Addr().String() + "/health-check") + req.SetBodyString("payload") + resp := protocol.AcquireResponse() + err := c.Do(context.Background(), req, resp) + assert.Nil(t, err) + assert.DeepEqual(t, 200, resp.StatusCode()) + assert.DeepEqual(t, "ok", string(resp.Body())) + protocol.ReleaseRequest(req) + protocol.ReleaseResponse(resp) + } + + select { + case err := <-serverResult: + assert.Nil(t, err) + case <-time.After(time.Second): + t.Fatal("server did not receive both requests") + } + assert.DeepEqual(t, 1, accepted) + assert.DeepEqual(t, 2, requests) + }) + } +} + +func TestPooledConnHealthCheckPreservesHealthyStandardReuseAfterReadTimeout(t *testing.T) { + var accepted atomic.Int32 + server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok")) + })) + server.Config.ConnState = func(_ net.Conn, state http.ConnState) { + if state == http.StateNew { + accepted.Add(1) + } + } + server.Start() + defer server.Close() + + c := &HostClient{ + ClientOptions: &ClientOptions{ + Dialer: standard.NewDialer(), + PooledConnHealthCheck: true, + ReadTimeout: 50 * time.Millisecond, + }, + Addr: strings.TrimPrefix(server.URL, "http://"), + } + defer c.CloseIdleConnections() + + for i := 0; i < 2; i++ { + req := protocol.AcquireRequest() + req.Header.SetMethod(consts.MethodGet) + req.SetRequestURI(server.URL) + resp := protocol.AcquireResponse() + if err := c.Do(context.Background(), req, resp); err != nil { + t.Fatalf("request %d: %v", i+1, err) + } + if got := resp.StatusCode(); got != http.StatusOK { + t.Fatalf("request %d status: got %d, want %d", i+1, got, http.StatusOK) + } + protocol.ReleaseRequest(req) + protocol.ReleaseResponse(resp) + if i == 0 { + time.Sleep(75 * time.Millisecond) + } + } + if got := accepted.Load(); got != 1 { + t.Fatalf("healthy standard connection was redialed: got %d accepts, want 1", got) + } +} + func TestConnInPoolRetry(t *testing.T) { c := &HostClient{ ClientOptions: &ClientOptions{