Skip to content

Commit

Permalink
fix: call conn.release method in ext.ReleaseBodystream
Browse files Browse the repository at this point in the history
  • Loading branch information
Duslia committed Dec 31, 2024
1 parent 4cc53fe commit e64b1b1
Show file tree
Hide file tree
Showing 4 changed files with 23 additions and 6 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
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 e64b1b1

Please sign in to comment.