From f54ffa14aca15d9115587283ec57e4259aa67cec Mon Sep 17 00:00:00 2001 From: Rennbon <343688972@qq.com> Date: Thu, 3 Mar 2022 20:46:32 +0800 Subject: [PATCH] feature: Keep the memory usage of the service at a stable level (#1216) * feature: add request and response body size limit, it prevents the large body from slowly stretching the memory of the entire service * fix: http.go fmt * refact: optimize code naming * Update http.go Co-authored-by: Erik Dubbelboer * Update http.go Co-authored-by: Erik Dubbelboer * Update http.go Co-authored-by: Erik Dubbelboer * Update http.go Co-authored-by: Erik Dubbelboer * Update http.go Co-authored-by: Erik Dubbelboer Co-authored-by: Erik Dubbelboer --- http.go | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/http.go b/http.go index d61eb622be..47431cdcd8 100644 --- a/http.go +++ b/http.go @@ -17,6 +17,18 @@ import ( "github.com/valyala/bytebufferpool" ) +var ( + requestBodyPoolSizeLimit = -1 + responseBodyPoolSizeLimit = -1 +) + +// SetBodySizePoolLimit set the max body size for bodies to be returned to the pool. +// If the body size is larger it will be released instead of put back into the pool for reuse. +func SetBodySizePoolLimit(reqBodyLimit, respBodyLimit int) { + requestBodyPoolSizeLimit = reqBodyLimit + responseBodyPoolSizeLimit = respBodyLimit +} + // Request represents HTTP request. // // It is forbidden copying Request instances. Create new instances @@ -957,6 +969,9 @@ func readMultipartForm(r io.Reader, boundary string, size, maxInMemoryFileSize i // Reset clears request contents. func (req *Request) Reset() { + if requestBodyPoolSizeLimit >= 0 && req.body != nil { + req.ReleaseBody(requestBodyPoolSizeLimit) + } req.Header.Reset() req.resetSkipHeader() req.timeout = 0 @@ -986,6 +1001,9 @@ func (req *Request) RemoveMultipartFormFiles() { // Reset clears response contents. func (resp *Response) Reset() { + if responseBodyPoolSizeLimit >= 0 && resp.body != nil { + resp.ReleaseBody(responseBodyPoolSizeLimit) + } resp.Header.Reset() resp.resetSkipHeader() resp.SkipBody = false