Skip to content

Commit

Permalink
fix: call conn.release method in ext.ReleaseBodystream (#1252)
Browse files Browse the repository at this point in the history
  • Loading branch information
Duslia authored Jan 2, 2025
1 parent a859259 commit 5729635
Show file tree
Hide file tree
Showing 5 changed files with 23 additions and 7 deletions.
6 changes: 4 additions & 2 deletions pkg/common/test/mock/network.go
Original file line number Diff line number Diff line change
Expand Up @@ -325,7 +325,8 @@ func (m *Conn) AddCloseCallback(callback netpoll.CloseCallback) error {
}

type StreamConn struct {
Data []byte
HasReleased bool
Data []byte
}

func NewStreamConn() *StreamConn {
Expand Down Expand Up @@ -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 {
Expand Down
6 changes: 3 additions & 3 deletions pkg/common/test/mock/network_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
})
Expand Down
1 change: 0 additions & 1 deletion pkg/protocol/client/client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down
16 changes: 15 additions & 1 deletion pkg/protocol/http1/ext/stream.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions pkg/protocol/http1/req/request_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down

0 comments on commit 5729635

Please sign in to comment.