add buffer pool for proxy forwarder (#16942)

This commit is contained in:
Harshavardhana 2023-04-06 15:54:12 -07:00 committed by GitHub
parent 4c204707fd
commit 47b7469a60
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -24,6 +24,7 @@ import (
"net/http/httputil" "net/http/httputil"
"net/url" "net/url"
"strings" "strings"
"sync"
"time" "time"
) )
@ -50,6 +51,33 @@ func NewForwarder(f *Forwarder) *Forwarder {
return f return f
} }
type bufPool struct {
sz int
pool sync.Pool
}
func (b *bufPool) Put(buf []byte) {
if cap(buf) < b.sz || cap(buf) > b.sz*2 {
// Buffer too small or will likely leak memory after being expanded.
// Drop it.
return
}
b.pool.Put(&buf)
}
func (b *bufPool) Get() []byte {
bufp := b.pool.Get().(*[]byte)
return (*bufp)[:b.sz]
}
func newBufPool(sz int) httputil.BufferPool {
return &bufPool{sz: sz, pool: sync.Pool{
New: func() interface{} {
buf := make([]byte, sz)
return &buf
},
}}
}
// ServeHTTP forwards HTTP traffic using the configured transport // ServeHTTP forwards HTTP traffic using the configured transport
func (f *Forwarder) ServeHTTP(w http.ResponseWriter, inReq *http.Request) { func (f *Forwarder) ServeHTTP(w http.ResponseWriter, inReq *http.Request) {
outReq := new(http.Request) outReq := new(http.Request)
@ -59,6 +87,7 @@ func (f *Forwarder) ServeHTTP(w http.ResponseWriter, inReq *http.Request) {
Director: func(req *http.Request) { Director: func(req *http.Request) {
f.modifyRequest(req, inReq.URL) f.modifyRequest(req, inReq.URL)
}, },
BufferPool: newBufPool(128 << 10),
Transport: f.RoundTripper, Transport: f.RoundTripper,
FlushInterval: defaultFlushInterval, FlushInterval: defaultFlushInterval,
ErrorHandler: f.customErrHandler, ErrorHandler: f.customErrHandler,