From d076024883ab8707d2fbb802e02db4c79df2941f Mon Sep 17 00:00:00 2001 From: slasher Date: Thu, 4 Sep 2025 14:31:45 +0800 Subject: [PATCH] fix(rpc): timeout with response in rpc client @formatter:off Signed-off-by: slasher --- blobstore/common/rpc/client.go | 67 +++++++++++++++++------------ blobstore/common/rpc/client_test.go | 23 +++++++--- go.mod | 1 - go.sum | 2 - 4 files changed, 56 insertions(+), 37 deletions(-) diff --git a/blobstore/common/rpc/client.go b/blobstore/common/rpc/client.go index e684243aa..da04cc465 100644 --- a/blobstore/common/rpc/client.go +++ b/blobstore/common/rpc/client.go @@ -22,6 +22,7 @@ import ( "net/http" urllib "net/url" "strings" + "sync" "time" "github.com/cubefs/cubefs/blobstore/common/crc32block" @@ -282,10 +283,6 @@ func (c *client) Close() { func (c *client) doWithCtx(ctx context.Context, req *http.Request) (resp *http.Response, err error) { req = req.WithContext(ctx) - if c.bandwidthBPMs > 0 && req.Body != nil { - t := req.ContentLength/c.bandwidthBPMs + c.bodyBaseTimeoutMs - req.Body = &timeoutReadCloser{timeoutMs: t, body: req.Body} - } resp, err = c.client.Do(req) if err != nil { span := trace.SpanFromContextSafe(ctx) @@ -293,8 +290,9 @@ func (c *client) doWithCtx(ctx context.Context, req *http.Request) (resp *http.R return } if c.bandwidthBPMs > 0 { - t := resp.ContentLength/c.bandwidthBPMs + c.bodyBaseTimeoutMs - resp.Body = &timeoutReadCloser{timeoutMs: t, body: resp.Body} + timeout := time.Millisecond * time.Duration(resp.ContentLength/c.bandwidthBPMs+c.bodyBaseTimeoutMs) + timer := time.NewTimer(time.Hour) + resp.Body = &timeoutReadCloser{body: resp.Body, timer: timer, timeout: timeout} } return } @@ -352,36 +350,49 @@ func ParseResponseErr(resp *http.Response) (err error) { return NewError(resp.StatusCode, resp.Status, fmt.Errorf("%d response", resp.StatusCode)) } +// http request body will not closed by net/http, it should close by caller. +// but response body must be closed after Client.Do. type timeoutReadCloser struct { body io.ReadCloser - timeoutMs int64 + timer *time.Timer + timeout time.Duration + closeOnce sync.Once } func (tr *timeoutReadCloser) Close() (err error) { - return tr.body.Close() + tr.closeOnce.Do(func() { + err = tr.body.Close() + tr.timer.Stop() + }) + return } -func (tr *timeoutReadCloser) Read(p []byte) (n int, err error) { - readOk := make(chan struct{}) - if tr.timeoutMs > 0 { - startTime := time.Now().UnixNano() / 1e6 - after := time.After(time.Millisecond * time.Duration(tr.timeoutMs)) - go func() { - n, err = tr.body.Read(p) - close(readOk) - }() - select { - case <-readOk: - // really cost time - tr.timeoutMs = tr.timeoutMs - (time.Now().UnixNano()/1e6 - startTime) - return - case <-after: - tr.body.Close() - return 0, ErrBodyReadTimeout - } +func (tr *timeoutReadCloser) Read(p []byte) (int, error) { + if tr.timeout <= 0 { + return 0, ErrBodyReadTimeout + } + tr.timer.Reset(tr.timeout) + + start := time.Now() + + var n int + var err error + readOk := make(chan struct{}) + go func() { + n, err = tr.body.Read(p) + close(readOk) + }() + + select { + case <-readOk: + tr.timeout -= time.Since(start) + return n, err + case <-tr.timer.C: + tr.timeout = 0 + tr.Close() // trigger to break Read + <-readOk + return 0, ErrBodyReadTimeout } - tr.body.Close() - return 0, ErrBodyReadTimeout } func serverCrcEncodeCheck(ctx context.Context, request *http.Request, resp *http.Response) (err error) { diff --git a/blobstore/common/rpc/client_test.go b/blobstore/common/rpc/client_test.go index 1fd06d45f..4d7b34156 100644 --- a/blobstore/common/rpc/client_test.go +++ b/blobstore/common/rpc/client_test.go @@ -67,13 +67,18 @@ func (s *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } } -type testTimeoutReader struct{} +type testTimeoutReader struct{ closeTimes int } func (t *testTimeoutReader) Read(p []byte) (n int, err error) { time.Sleep(time.Millisecond * 20) return 0, nil } +func (t *testTimeoutReader) Close() error { + t.closeTimes++ + return nil +} + func doAfterCrc(w http.ResponseWriter, req *http.Request) { dec := crc32block.NewBodyDecoder(req.Body) defer dec.Close() @@ -254,13 +259,19 @@ func TestClient_PostWithReadResponseTimeout(t *testing.T) { func TestTimeoutReadCloser_Read(t *testing.T) { // test for read data for input buffer timeout - readCloser := timeoutReadCloser{ - timeoutMs: 1, - body: io.NopCloser(&testTimeoutReader{}), - } + body := &testTimeoutReader{} + timer := time.NewTimer(time.Millisecond) + responseBody := timeoutReadCloser{body: body, timer: timer} res := make([]byte, 30) - _, err := readCloser.Read(res) + _, err := responseBody.Read(res) require.Equal(t, ErrBodyReadTimeout, err) + _, err = responseBody.Read(res) + require.Equal(t, ErrBodyReadTimeout, err) + require.NoError(t, responseBody.Close()) + require.NoError(t, responseBody.Close()) + require.NoError(t, responseBody.Close()) + require.Equal(t, 1, body.closeTimes) + require.False(t, responseBody.timer.Stop()) } func TestClient_Put(t *testing.T) { diff --git a/go.mod b/go.mod index 44fbe13d0..9651fae05 100644 --- a/go.mod +++ b/go.mod @@ -44,7 +44,6 @@ require ( github.com/spf13/cobra v1.2.1 github.com/stretchr/testify v1.8.4 github.com/tecbot/gorocksdb v0.0.0-20191217155057-f0fad39f321c - github.com/xtaci/smux v1.5.16 github.com/zeebo/xxh3 v1.0.2 go.etcd.io/etcd/raft/v3 v3.5.8 go.uber.org/automaxprocs v1.5.1 diff --git a/go.sum b/go.sum index db52d2577..93fb55cdd 100644 --- a/go.sum +++ b/go.sum @@ -575,8 +575,6 @@ github.com/xdg-go/scram v1.1.1/go.mod h1:RaEWvsqvNKKvBPvcKeFjrG2cJqOkHTiyTpzz23n github.com/xdg-go/stringprep v1.0.3/go.mod h1:W3f5j4i+9rC0kuIEJL0ky1VpHXQU3ocBgklLGvcBnW8= github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2/go.mod h1:UETIi67q53MR2AWcXfiuqkDkRtnGDLqkBTpCHuJHxtU= github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q= -github.com/xtaci/smux v1.5.16 h1:FBPYOkW8ZTjLKUM4LI4xnnuuDC8CQ/dB04HD519WoEk= -github.com/xtaci/smux v1.5.16/go.mod h1:OMlQbT5vcgl2gb49mFkYo6SMf+zP3rcjcwQz7ZU7IGY= github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=