feat: client support context cancel to control the lifecycle of request

Author: DMwangnima

pkg/app/client/option.go

@@ -181,6 +181,14 @@ func WithDialFunc(f network.DialFunc, dialers ...network.Dialer) config.ClientOp
 	}}
 }
 
+// WithSenseContextCancel is used to set whether the ctx passed into Do control the lifecycle of the request,
+// so that when the ctx is canceled, the corresponding request will also be finished
+func WithSenseContextCancel(b bool) config.ClientOption {
+	return config.ClientOption{F: func(o *config.ClientOptions) {
+		o.SenseContextCancel = b
+	}}
+}
+
 // customDialer set customDialerFunc and params to set dailFunc
 type customDialer struct {
 	network.Dialer

pkg/app/client/option_test.go

@@ -44,6 +44,7 @@ func TestClientOptions(t *testing.T) {
 		),
 		WithWriteTimeout(time.Second),
 		WithConnStateObserve(nil, time.Second),
+		WithSenseContextCancel(true),
 	})
 	assert.DeepEqual(t, 100*time.Millisecond, opt.DialTimeout)
 	assert.DeepEqual(t, 128, opt.MaxConnsPerHost)
@@ -59,6 +60,7 @@ func TestClientOptions(t *testing.T) {
 	assert.DeepEqual(t, 5*time.Second, opt.RetryConfig.MaxDelay)
 	assert.DeepEqual(t, 1*time.Second, opt.RetryConfig.MaxJitter)
 	assert.DeepEqual(t, 1*time.Second, opt.ObservationInterval)
+	assert.DeepEqual(t, true, opt.SenseContextCancel)
 	for i := 0; i < 100; i++ {
 		assert.DeepEqual(t, opt.RetryConfig.DelayPolicy(uint(i), nil, opt.RetryConfig), retry.CombineDelay(retry.DefaultDelayPolicy, retry.FixedDelayPolicy, retry.BackOffDelayPolicy)(uint(i), nil, opt.RetryConfig))
 	}

pkg/common/config/client_option.go

@@ -134,6 +134,10 @@ type ClientOptions struct {
 	// Callback hook for re-configuring host client
 	// If an error is returned, the request will be terminated.
 	HostClientConfigHook func(hc interface{}) error
+
+	// Whether to let the ctx passed into Do control the lifecycle of the request,
+	// so that when the ctx is canceled, the corresponding request will also be finished
+	SenseContextCancel bool
 }
 
 func NewClientOptions(opts []ClientOption) *ClientOptions {

pkg/common/config/request_option.go

@@ -29,8 +29,9 @@ type RequestOptions struct {
 	writeTimeout time.Duration
 	// Request timeout. Usually set by DoDeadline or DoTimeout
 	// if <= 0, means not set
-	requestTimeout time.Duration
-	start          time.Time
+	requestTimeout     time.Duration
+	start              time.Time
+	senseContextCancel bool
 }
 
 // RequestOption is the only struct to set request-level options.
@@ -109,6 +110,16 @@ func WithRequestTimeout(t time.Duration) RequestOption {
 	}}
 }
 
+// WithSenseContextCancel is used to set whether the ctx passed into Do control the lifecycle of the request,
+// so that when the ctx is canceled, the corresponding request will also be finished.
+//
+// This is the request level configuration.
+func WithSenseContextCancel(b bool) RequestOption {
+	return RequestOption{F: func(o *RequestOptions) {
+		o.senseContextCancel = b
+	}}
+}
+
 func (o *RequestOptions) Apply(opts []RequestOption) {
 	for _, op := range opts {
 		op.F(o)
@@ -143,6 +154,10 @@ func (o *RequestOptions) RequestTimeout() time.Duration {
 	return o.requestTimeout
 }
 
+func (o *RequestOptions) SenseContextCancel() bool {
+	return o.senseContextCancel
+}
+
 // StartRequest records the start time of the request.
 //
 // Note: Users should not call this method.
@@ -171,6 +186,7 @@ func (o *RequestOptions) CopyTo(dst *RequestOptions) {
 	dst.dialTimeout = o.dialTimeout
 	dst.requestTimeout = o.requestTimeout
 	dst.start = o.start
+	dst.senseContextCancel = o.senseContextCancel
 }
 
 // SetPreDefinedOpts Pre define some RequestOption here

pkg/common/config/request_option_test.go

@@ -33,13 +33,15 @@ func TestRequestOptions(t *testing.T) {
 		WithDialTimeout(time.Second),
 		WithReadTimeout(time.Second),
 		WithWriteTimeout(time.Second),
+		WithSenseContextCancel(true),
 	})
 	assert.DeepEqual(t, "b", opt.Tag("a"))
 	assert.DeepEqual(t, "d", opt.Tag("c"))
 	assert.DeepEqual(t, "f", opt.Tag("e"))
 	assert.DeepEqual(t, time.Second, opt.DialTimeout())
 	assert.DeepEqual(t, time.Second, opt.ReadTimeout())
 	assert.DeepEqual(t, time.Second, opt.WriteTimeout())
+	assert.DeepEqual(t, true, opt.SenseContextCancel())
 	assert.True(t, opt.IsSD())
 }
 

pkg/network/netpoll/connection.go

@@ -21,10 +21,11 @@ import (
 	"strings"
 	"syscall"
 
+	"github.com/cloudwego/netpoll"
+
 	errs "github.com/cloudwego/hertz/pkg/common/errors"
 	"github.com/cloudwego/hertz/pkg/common/hlog"
 	"github.com/cloudwego/hertz/pkg/network"
-	"github.com/cloudwego/netpoll"
 )
 
 type Conn struct {

pkg/network/netpoll/connection_test.go

@@ -23,8 +23,9 @@ import (
 	"testing"
 	"time"
 
-	"github.com/cloudwego/hertz/pkg/common/test/assert"
 	"github.com/cloudwego/netpoll"
+
+	"github.com/cloudwego/hertz/pkg/common/test/assert"
 )
 
 func TestReadBytes(t *testing.T) {

pkg/protocol/http1/client.go

@@ -398,7 +398,7 @@ func (c *HostClient) Do(ctx context.Context, req *protocol.Request, resp *protoc
 		default:
 		}
 
-		canIdempotentRetry, err = c.do(req, resp)
+		canIdempotentRetry, err = c.do(ctx, req, resp)
 		// If there is no custom retry and err is equal to nil, the loop simply exits.
 		if err == nil && isDefaultRetryFunc {
 			if connAttempts != 0 {
@@ -456,22 +456,98 @@ func (c *HostClient) PendingRequests() int {
 	return int(atomic.LoadInt32(&c.pendingRequests))
 }
 
-func (c *HostClient) do(req *protocol.Request, resp *protocol.Response) (bool, error) {
+func (c *HostClient) do(ctx context.Context, req *protocol.Request, resp *protocol.Response) (bool, error) {
 	nilResp := false
 	if resp == nil {
 		nilResp = true
 		resp = protocol.AcquireResponse()
 	}
+	defer func() {
+		if nilResp {
+			protocol.ReleaseResponse(resp)
+		}
+	}()
 
-	canIdempotentRetry, err := c.doNonNilReqResp(req, resp)
+	c.checkAndConfigureReqResp(req, resp)
 
-	if nilResp {
-		protocol.ReleaseResponse(resp)
+	// whole request timeout
+	o := req.Options()
+	deadline := time.Time{}
+	if v := o.RequestTimeout(); v > 0 {
+		deadline = o.StartTime().Add(v)
 	}
 
+	// dial starts
+	timeout := calcTimeout(deadline, minTimeout(c.DialTimeout, o.DialTimeout()))
+	if timeout < 0 {
+		return false, errTimeout
+	}
+	cc, inPool, err := c.acquireConn(timeout)
+	if err != nil {
+		return false, err
+	}
+	conn := cc.c
+	resp.ParseNetAddr(conn)
+	// force handshake using dial timeout
+	if c.IsTLS && timeout > 0 {
+		// NOTE: Handshake() here is optional as Write would tirigger handshake
+		// but for tls handshake, it writes and reads, and we need to set deadline for that.
+		tlsconn, ok := conn.(network.ConnTLSer)
+		if ok {
+			// currently netpoll doesn't support conn.SetDeadline nor tls, but crypto/tls.Conn does.
+			// in case netpoll supports tls in the future, may need to change this to
+			// call both conn.SetReadTimeout, and conn.SetWriteTimeout
+			err = conn.SetDeadline(time.Now().Add(timeout))
+			if err == nil {
+				err = tlsconn.Handshake()
+				// NOTE: no need conn.SetDeadline(time.Time{})?
+				// we always reset before Write and Read
+			}
+			if err != nil {
+				c.closeConn(cc)
+				return true, err
+			}
+		}
+	}
+
+	// write req
+	canIdempotentRetry, resetConnection, shouldFinish, err := c.writeReq(cc, inPool, req, resp, deadline)
+	if shouldFinish {
+		return canIdempotentRetry, err
+	}
+	if (c.ClientOptions.SenseContextCancel || o.SenseContextCancel()) && ctx.Done() != nil {
+		doChan := make(chan doResult)
+		go func() {
+			canRetry, dErr := c.readResp(cc, inPool, resetConnection, req, resp, deadline)
+			doChan <- doResult{
+				canIdempotentRetry: canRetry,
+				err:                dErr,
+			}
+		}()
+		select {
+		case <-ctx.Done():
+			// Close connection directly and wait for readResp releasing connection resource
+			cc.c.Close()
+			// wait for goroutine exited
+			// since the lifecycle of context has finished,
+			// there is no need to retry
+			<-doChan
+			return false, ctx.Err()
+		case res := <-doChan:
+			return res.canIdempotentRetry, res.err
+		}
+	}
+
+	canIdempotentRetry, err = c.readResp(cc, inPool, resetConnection, req, resp, deadline)
+
 	return canIdempotentRetry, err
 }
 
+type doResult struct {
+	canIdempotentRetry bool
+	err                error
+}
+
 func minTimeout(d0, d1 time.Duration) time.Duration {
 	if d0 <= 0 {
 		if d1 <= 0 {
@@ -516,7 +592,7 @@ func calcTimeout(deadline time.Time, timeout time.Duration) time.Duration {
 	return timeout
 }
 
-func (c *HostClient) doNonNilReqResp(req *protocol.Request, resp *protocol.Response) (bool, error) {
+func (c *HostClient) checkAndConfigureReqResp(req *protocol.Request, resp *protocol.Response) {
 	if req == nil {
 		panic("BUG: req cannot be nil")
 	}
@@ -536,46 +612,16 @@ func (c *HostClient) doNonNilReqResp(req *protocol.Request, resp *protocol.Respo
 	if c.DisablePathNormalizing {
 		req.URI().DisablePathNormalizing = true
 	}
+}
 
-	o := req.Options()
-	deadline := time.Time{}
-	if v := o.RequestTimeout(); v > 0 {
-		deadline = o.StartTime().Add(v)
-	}
-
-	// dial starts
-
-	timeout := calcTimeout(deadline, minTimeout(c.DialTimeout, o.DialTimeout()))
-	if timeout < 0 {
-		return false, errTimeout
-	}
-	cc, inPool, err := c.acquireConn(timeout)
-	if err != nil {
-		return false, err
-	}
+func (c *HostClient) writeReq(cc *clientConn, connInPool bool, req *protocol.Request, resp *protocol.Response, deadline time.Time) (
+	canIdempotentRetry bool,
+	resetConnection bool,
+	shouldFinish bool,
+	err error,
+) {
 	conn := cc.c
-	resp.ParseNetAddr(conn)
-
-	if c.IsTLS && timeout > 0 { // force handshake using dial timeout
-		// NOTE: Handshake() here is optional as Write would tirigger handshake
-		// but for tls handshake, it writes and reads, and we need to set deadline for that.
-		tlsconn, ok := conn.(network.ConnTLSer)
-		if ok {
-			// currently netpoll doesn't support conn.SetDeadline nor tls, but crypto/tls.Conn does.
-			// in case netpoll supports tls in the future, may need to change this to
-			// call both conn.SetReadTimeout, and conn.SetWriteTimeout
-			err := conn.SetDeadline(time.Now().Add(timeout))
-			if err == nil {
-				err = tlsconn.Handshake()
-				// NOTE: no need conn.SetDeadline(time.Time{})?
-				// we always reset before Write and Read
-			}
-			if err != nil {
-				c.closeConn(cc)
-				return true, err
-			}
-		}
-	}
+	o := req.Options()
 
 	usingProxy := false
 	if c.ProxyURI != nil && bytes.Equal(req.Scheme(), bytestr.StrHTTP) {
@@ -584,18 +630,16 @@ func (c *HostClient) doNonNilReqResp(req *protocol.Request, resp *protocol.Respo
 	}
 
 	// write starts
-
-	timeout = calcTimeout(deadline, minTimeout(c.WriteTimeout, o.WriteTimeout()))
+	timeout := calcTimeout(deadline, minTimeout(c.WriteTimeout, o.WriteTimeout()))
 	if timeout < 0 {
 		c.closeConn(cc)
-		return false, errTimeout
+		return false, resetConnection, true, errTimeout
 	}
 	if err = conn.SetWriteTimeout(timeout); err != nil {
 		c.closeConn(cc)
-		return true, err
+		return true, resetConnection, true, err
 	}
 
-	resetConnection := false
 	if c.MaxConnDuration > 0 && time.Since(cc.createdTime) > c.MaxConnDuration && !req.ConnectionClose() {
 		req.SetConnectionClose()
 		resetConnection = true
@@ -625,7 +669,7 @@ func (c *HostClient) doNonNilReqResp(req *protocol.Request, resp *protocol.Respo
 			err = errNorm.ToHertzError(err)
 		}
 		if !errors.Is(err, errs.ErrConnectionClosed) {
-			return true, err
+			return true, resetConnection, true, err
 		}
 
 		// introduced by https://github.com/cloudwego/hertz/pull/412
@@ -635,27 +679,34 @@ func (c *HostClient) doNonNilReqResp(req *protocol.Request, resp *protocol.Respo
 		// NOTE: can't use deadline since it likely already exceeded deadline when write
 		timeout = minTimeout(50*time.Millisecond, minTimeout(c.ReadTimeout, o.ReadTimeout()))
 		if conn.SetReadTimeout(timeout) != nil {
-			return true, err
+			return true, resetConnection, true, err
 		}
 		zr := c.acquireReader(conn)
 		defer zr.Release()
 		if respI.ReadHeaderAndLimitBody(resp, zr, c.MaxResponseBodySize) == nil {
 			if code := resp.StatusCode(); code >= 400 && code < 600 {
 				// strictly for 4xx only, but 5xx is also acceptable.
 				// both can be considered better response rather than write err
-				return false, nil
+				return false, resetConnection, true, nil
 			}
 		}
 
-		if inPool {
+		if connInPool {
 			err = errs.ErrBadPoolConn
 		}
-		return true, err
+		return true, resetConnection, true, err
 	}
+	return false, resetConnection, false, nil
+}
 
-	// read starts
+func (c *HostClient) readResp(cc *clientConn, connInPool, resetConnection bool, req *protocol.Request, resp *protocol.Response, deadline time.Time) (bool, error) {
+	var err error
+	customSkipBody := resp.SkipBody
+	conn := cc.c
+	o := req.Options()
 
-	timeout = calcTimeout(deadline, minTimeout(c.ReadTimeout, o.ReadTimeout()))
+	// read starts
+	timeout := calcTimeout(deadline, minTimeout(c.ReadTimeout, o.ReadTimeout()))
 	if timeout < 0 {
 		c.closeConn(cc)
 		return false, errTimeout
@@ -683,7 +734,7 @@ func (c *HostClient) doNonNilReqResp(req *protocol.Request, resp *protocol.Respo
 	if err != nil {
 		zr.Release() //nolint:errcheck
 		c.closeConn(cc)
-		if inPool && (err == io.EOF || err == syscall.ECONNRESET) {
+		if connInPool && (err == io.EOF || err == syscall.ECONNRESET) {
 			return true, errs.ErrBadPoolConn
 		}
 		// if this is not a pooled connection,
@@ -1454,4 +1505,8 @@ type ClientOptions struct {
 
 	// StateObserve execution interval
 	ObservationInterval time.Duration
+
+	// Whether to let the ctx passed into Do control the lifecycle of the request,
+	// so that when the ctx is canceled, the corresponding request will also be finished
+	SenseContextCancel bool
 }