Skip to content

Commit

Permalink
feat: bytebuferpoll use mcache replace
Browse files Browse the repository at this point in the history
  • Loading branch information
mac authored and mac committed Dec 17, 2023
1 parent 82f9bfa commit 53a7209
Show file tree
Hide file tree
Showing 5 changed files with 85 additions and 8 deletions.
2 changes: 1 addition & 1 deletion pkg/app/context.go
Original file line number Diff line number Diff line change
Expand Up @@ -582,7 +582,7 @@ func (ctx *RequestContext) FileAttachment(filepath, filename string) {

// SetBodyString sets response body to the given value.
func (ctx *RequestContext) SetBodyString(body string) {
ctx.Response.SetBodyString(body)
ctx.Response.SetBodyString(body, string(ctx.URI().RequestURI()))
}

// SetContentTypeBytes sets response Content-Type.
Expand Down
51 changes: 51 additions & 0 deletions pkg/app/server/hertz_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import (
"io/ioutil"
"net"
"net/http"
"runtime"
"strings"
"sync"
"sync/atomic"
Expand Down Expand Up @@ -248,6 +249,56 @@ func TestServer_Run(t *testing.T) {
_ = hertz.Shutdown(ctx)
}

func TestBigBodyBug(t *testing.T) {
runtime.GOMAXPROCS(3)
hertz := New(WithHostPorts("127.0.0.1:8888"))
hertz.GET("/test1", func(c context.Context, ctx *app.RequestContext) {
body := make([]byte, 1024*1024*9)
ctx.SetBodyString(string(body))
})
hertz.GET("/test2", func(c context.Context, ctx *app.RequestContext) {
body := make([]byte, 1024)
ctx.SetBodyString(string(body))
})
hertz.GET("/test3", func(c context.Context, ctx *app.RequestContext) {
body := make([]byte, 1024*2)
ctx.SetBodyString(string(body))
})
go hertz.Run()

go func() {
for i := 0; i < 2; i++ {
go func() {
for true {
http.Get("http://127.0.0.1:8888/test1")
}
}()
}
}()

go func() {
for i := 0; i < 5; i++ {
go func() {
for true {
http.Get("http://127.0.0.1:8888/test2")
}
}()
}
}()

go func() {
for i := 0; i < 5; i++ {
go func() {
for true {
http.Get("http://127.0.0.1:8888/test3")
}
}()
}
}()

<-make(chan struct{})
}

func TestNotAbsolutePath(t *testing.T) {
engine := New(WithHostPorts("127.0.0.1:9990"))
engine.POST("/", func(c context.Context, ctx *app.RequestContext) {
Expand Down
14 changes: 13 additions & 1 deletion pkg/common/bytebufferpool/pool.go
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@ import (
"sort"
"sync"
"sync/atomic"

"github.com/bytedance/gopkg/lang/mcache"
)

const (
Expand All @@ -54,7 +56,7 @@ const (
minSize = 1 << minBitSize
maxSize = 1 << (minBitSize + steps - 1)

calibrateCallsThreshold = 42000
calibrateCallsThreshold = 10 // Just for test
maxPercentile = 0.95
)

Expand Down Expand Up @@ -96,6 +98,16 @@ func (p *Pool) Get() *ByteBuffer {
}
}

func (p *Pool) GetWithSize(size int) *ByteBuffer {
return &ByteBuffer{
B: mcache.Malloc(size)[:0],
}
}

func (p *Pool) PutWithByte(b []byte) {
mcache.Free(b)
}

// Put returns byte buffer to the pool.
//
// ByteBuffer.B mustn't be touched after returning it to the pool.
Expand Down
2 changes: 1 addition & 1 deletion pkg/common/config/option.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ const (
defaultAddr = ":8888"
defaultNetwork = "tcp"
defaultBasePath = "/"
defaultMaxRequestBodySize = 4 * 1024 * 1024
defaultMaxRequestBodySize = 0
defaultWaitExitTimeout = time.Second * 5
defaultReadBufferSize = 4 * 1024
)
Expand Down
24 changes: 19 additions & 5 deletions pkg/protocol/response.go
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@
package protocol

import (
"fmt"
"io"
"net"
"sync"
Expand Down Expand Up @@ -138,9 +139,11 @@ func (resp *Response) SetConnectionClose() {
}

// SetBodyString sets response body.
func (resp *Response) SetBodyString(body string) {
resp.CloseBodyStream() //nolint:errcheck
resp.BodyBuffer().SetString(body) //nolint:errcheck
func (resp *Response) SetBodyString(body, url string) {
resp.CloseBodyStream() //nolint:errcheck
// resp.BodyBuffer().SetString(body) //nolint:errcheck
// resp.BodyBuffer(url).SetString(body) //nolint:errcheck
resp.BodyBufferWithSize(len(body), url).SetString(body) //nolint:errcheck
}

func (resp *Response) ConstructBodyStream(body *bytebufferpool.ByteBuffer, bodyStream io.Reader) {
Expand Down Expand Up @@ -299,7 +302,8 @@ func (resp *Response) ResetBody() {
resp.body.Reset()
return
}
responseBodyPool.Put(resp.body)
// responseBodyPool.Put(resp.body)
responseBodyPool.PutWithByte(resp.body.B)
resp.body = nil
}
}
Expand Down Expand Up @@ -385,11 +389,21 @@ func (resp *Response) CloseBodyStream() error {
return err
}

func (resp *Response) BodyBuffer() *bytebufferpool.ByteBuffer {
func (resp *Response) BodyBufferWithSize(size int, uri string) *bytebufferpool.ByteBuffer {
if resp.body == nil {
resp.body = responseBodyPool.GetWithSize(size)
}
resp.bodyRaw = nil
fmt.Printf("flipped url=%s resp Body len=%d cap=%d\n", uri, resp.body.Len(), cap(resp.body.B))
return resp.body
}

func (resp *Response) BodyBuffer(url ...string) *bytebufferpool.ByteBuffer {
if resp.body == nil {
resp.body = responseBodyPool.Get()
}
resp.bodyRaw = nil
fmt.Printf("flipped url=%s resp Body len=%d cap=%d\n", url, resp.body.Len(), cap(resp.body.B))
return resp.body
}

Expand Down

0 comments on commit 53a7209

Please sign in to comment.