Repository: cloudwego/hertz
Pull request #1252 (2024-12-24 07:08:09)

fix: call conn.release method in ext.ReleaseBodystream
#### What type of PR is this?
<!--
Add one of the following kinds:

build: Changes that affect the build system or external dependencies (example scopes: gulp, broccoli, npm)
ci: Changes to our CI configuration files and scripts (example scopes: Travis, Circle, BrowserStack, SauceLabs)
docs: Documentation only changes
feat: A new feature
optimize: A new optimization
fix: A bug fix
perf: A code change that improves performance
refactor: A code change that neither fixes a bug nor adds a feature
style: Changes that do not affect the meaning of the code (white space, formatting, missing semi-colons, etc)
test: Adding missing tests or correcting existing tests
chore: Changes to the build process or auxiliary tools and libraries such as documentation generation
-->

#### Check the PR title.
<!--
The description of the title will be attached in Release Notes, 
so please describe it from user-oriented, what this PR does / why we need it.
Please check your PR title with the below requirements:
-->
- [ ] This PR title match the format: \<type\>(optional scope): \<description\>
- [ ] The description of this PR title is user-oriented and clear enough for others to understand.
- [ ] Attach the PR updating the user documentation if the current PR requires user awareness at the usage level. [User docs repo](https://github.com/cloudwego/cloudwego.github.io)


#### (Optional) Translate the PR title into Chinese.
在 ext.ReleaseBodystream 中调用 conn.Release() 方法

#### (Optional) More detailed description for this PR(en: English/zh: Chinese).
<!--
Provide more detailed info for review(e.g., it's recommended to provide perf data if this is a perf type PR).
-->
en: In a streaming scenario, if there is a large amount of remaining data on the connection, all of it will be read into memory when skipRest is performed, which may lead to an OOM error.
zh(optional): 在流式场景下，如果连接上剩余数据很多时，在 skipRest 时会全部读到内存中，可能导致 OOM。

#### (Optional) Which issue(s) this PR fixes:
<!--
Automatically closes linked issue when PR is merged.
Eg: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
-->

#### (Optional) The PR that updates user documentation:
<!--
If the current PR requires user awareness at the usage level, please submit a PR to update user docs. [User docs repo](https://github.com/cloudwego/cloudwego.github.io)
-->


Patch:
diff --git a/pkg/protocol/http1/ext/stream.go b/pkg/protocol/http1/ext/stream.go
index ae81b56..c771cb8 100644
--- a/pkg/protocol/http1/ext/stream.go
+++ b/pkg/protocol/http1/ext/stream.go
@@ -272,6 +272,12 @@ func (rs *bodyStream) skipRest() error {
 			if err != nil {
 				return err
 			}
+			// After Skip, the buffer needs to be released to prevent OOM if there are too much data on conn.
+			err = rs.reader.Release()
+			if err != nil {
+				return err
+			}
+
 		}
 	}
 	// max value of pSize is 8193, it's safe.
@@ -300,7 +306,15 @@ func (rs *bodyStream) skipRest() error {
 		if skip > needSkipLen {
 			skip = needSkipLen
 		}
-		rs.reader.Skip(skip)
+		err := rs.reader.Skip(skip)
+		if err != nil {
+			return err
+		}
+		// After Skip, the buffer needs to be released to prevent OOM if there are too much data on conn.
+		err = rs.reader.Release()
+		if err != nil {
+			return err
+		}
 		needSkipLen -= skip
 		if needSkipLen == 0 {
 			return nil


Test patch:
diff --git a/pkg/common/test/mock/network.go b/pkg/common/test/mock/network.go
index fe671ae..b3ae54b 100644
--- a/pkg/common/test/mock/network.go
+++ b/pkg/common/test/mock/network.go
@@ -325,7 +325,8 @@ func (m *Conn) AddCloseCallback(callback netpoll.CloseCallback) error {
 }
 
 type StreamConn struct {
-	Data []byte
+	HasReleased bool
+	Data        []byte
 }
 
 func NewStreamConn() *StreamConn {
@@ -354,7 +355,8 @@ func (m *StreamConn) Skip(n int) error {
 }
 
 func (m *StreamConn) Release() error {
-	panic("implement me")
+	m.HasReleased = true
+	return nil
 }
 
 func (m *StreamConn) Len() int {
diff --git a/pkg/common/test/mock/network_test.go b/pkg/common/test/mock/network_test.go
index 4c9c4cf..34b0769 100644
--- a/pkg/common/test/mock/network_test.go
+++ b/pkg/common/test/mock/network_test.go
@@ -174,13 +174,13 @@ func TestStreamConn(t *testing.T) {
 		assert.DeepEqual(t, cap(conn.Data), conn.Len())
 		err = conn.Skip(conn.Len() + 1)
 		assert.DeepEqual(t, "not enough data", err.Error())
+		err = conn.Release()
+		assert.DeepEqual(t, nil, err)
+		assert.DeepEqual(t, true, conn.HasReleased)
 	})
 
 	t.Run("TestNotImplement", func(t *testing.T) {
 		conn := NewStreamConn()
-		assert.Panic(t, func() {
-			conn.Release()
-		})
 		assert.Panic(t, func() {
 			conn.ReadByte()
 		})
diff --git a/pkg/protocol/client/client_test.go b/pkg/protocol/client/client_test.go
index 49e7e7d..2babab4 100644
--- a/pkg/protocol/client/client_test.go
+++ b/pkg/protocol/client/client_test.go
@@ -33,7 +33,6 @@ type MockDoer struct {
 }
 
 func (m *MockDoer) Do(ctx context.Context, req *protocol.Request, resp *protocol.Response) error {
-
 	// this is the real logic in (c *HostClient) doNonNilReqResp method
 	if len(req.Header.Host()) == 0 {
 		req.Header.SetHostBytes(req.URI().Host())
diff --git a/pkg/protocol/http1/req/request_test.go b/pkg/protocol/http1/req/request_test.go
index 0411187..e9451e7 100644
--- a/pkg/protocol/http1/req/request_test.go
+++ b/pkg/protocol/http1/req/request_test.go
@@ -1425,6 +1425,7 @@ func TestStreamNotEnoughData(t *testing.T) {
 	err = ext.ReleaseBodyStream(req.BodyStream())
 	assert.Nil(t, err)
 	assert.DeepEqual(t, 0, len(conn.Data))
+	assert.DeepEqual(t, true, conn.HasReleased)
 }
 
 func TestRequestBodyStreamWithTrailer(t *testing.T) {