From 825b1d3d0a11bccc41fdc933c4e64eb21ba1b439 Mon Sep 17 00:00:00 2001 From: darkweak Date: Mon, 6 Oct 2025 19:52:44 +0200 Subject: [PATCH 1/2] fix(memory-bloat): try to solve some issues --- pkg/api/souin.go | 3 +- pkg/middleware/middleware.go | 73 +++++++++------------------- pkg/middleware/writer.go | 94 +++++++++++++++++++++++++++--------- 3 files changed, 93 insertions(+), 77 deletions(-) diff --git a/pkg/api/souin.go b/pkg/api/souin.go index 131723edb..796810776 100644 --- a/pkg/api/souin.go +++ b/pkg/api/souin.go @@ -186,7 +186,7 @@ func EvictMapping(current types.Storer) { } } - if updated { + if updated && len(mapping.GetMapping()) > 0 { v, e := proto.Marshal(mapping) if e != nil { fmt.Println("Impossible to re-encode the mapping", core.MappingKeyPrefix+k) @@ -199,7 +199,6 @@ func EvictMapping(current types.Storer) { current.Delete(core.MappingKeyPrefix + k) } } - time.Sleep(time.Minute) } func (s *SouinAPI) purgeMapping() { diff --git a/pkg/middleware/middleware.go b/pkg/middleware/middleware.go index 9b4848039..aedaa913e 100644 --- a/pkg/middleware/middleware.go +++ b/pkg/middleware/middleware.go @@ -58,6 +58,8 @@ func registerMappingKeysEviction(logger core.Logger, storers []types.Storer) { logger.Debugf("run mapping eviction for storer %s", current.Name()) api.EvictMapping(current) + + time.Sleep(time.Minute) } }(storer) } @@ -342,8 +344,7 @@ func (s *SouinBaseHandler) Store( if res.Header.Get("Content-Length") == "" { res.Header.Set("Content-Length", fmt.Sprint(bLen)) } - respBodyMaxSize := int(s.Configuration.GetDefaultCache().GetMaxBodyBytes()) - if respBodyMaxSize > 0 && bLen > respBodyMaxSize { + if customWriter.maxSizeReached { customWriter.Header().Set("Cache-Status", status+"; detail=UPSTREAM-RESPONSE-TOO-LARGE; key="+rfc.GetCacheKeyFromCtx(rq.Context())) return nil @@ -516,9 +517,7 @@ func (s *SouinBaseHandler) Upstream( } err := s.Store(customWriter, rq, requestCc, cachedKey, uri) - defer customWriter.handleBuffer(func(b *bytes.Buffer) { - b.Reset() - }) + defer customWriter.resetBuffer() return singleflightValue{ body: customWriter.Buf.Bytes(), @@ -581,9 +580,7 @@ func (s *SouinBaseHandler) Revalidate(validator *core.Revalidator, next handlerF statusCode := customWriter.GetStatusCode() if err == nil { if validator.IfUnmodifiedSincePresent && statusCode != http.StatusNotModified { - customWriter.handleBuffer(func(b *bytes.Buffer) { - b.Reset() - }) + customWriter.resetBuffer() customWriter.Rw.WriteHeader(http.StatusPreconditionFailed) return nil, errors.New("") @@ -591,9 +588,7 @@ func (s *SouinBaseHandler) Revalidate(validator *core.Revalidator, next handlerF if validator.IfModifiedSincePresent { if lastModified, err := time.Parse(time.RFC1123, customWriter.Header().Get("Last-Modified")); err == nil && validator.IfModifiedSince.Sub(lastModified) > 0 { - customWriter.handleBuffer(func(b *bytes.Buffer) { - b.Reset() - }) + customWriter.resetBuffer() customWriter.Rw.WriteHeader(http.StatusNotModified) return nil, errors.New("") @@ -615,9 +610,8 @@ func (s *SouinBaseHandler) Revalidate(validator *core.Revalidator, next handlerF ), ) - defer customWriter.handleBuffer(func(b *bytes.Buffer) { - b.Reset() - }) + defer customWriter.resetBuffer() + return singleflightValue{ body: customWriter.Buf.Bytes(), headers: customWriter.Header().Clone(), @@ -789,7 +783,7 @@ func (s *SouinBaseHandler) ServeHTTP(rw http.ResponseWriter, rq *http.Request, n bufPool.Reset() defer s.bufPool.Put(bufPool) - customWriter := NewCustomWriter(req, rw, bufPool) + customWriter := NewCustomWriter(req, rw, bufPool, int(s.Configuration.GetDefaultCache().GetMaxBodyBytes())) customWriter.Headers.Add("Range", req.Header.Get("Range")) // req.Header.Del("Range") @@ -838,18 +832,14 @@ func (s *SouinBaseHandler) ServeHTTP(rw http.ResponseWriter, rq *http.Request, n } if validator.NotModified { customWriter.WriteHeader(http.StatusNotModified) - customWriter.handleBuffer(func(b *bytes.Buffer) { - b.Reset() - }) + customWriter.resetBuffer() _, _ = customWriter.Send() return nil } customWriter.WriteHeader(response.StatusCode) - customWriter.handleBuffer(func(b *bytes.Buffer) { - _, _ = io.Copy(b, response.Body) - }) + _, _ = customWriter.copyToBuffer(response.Body) _, _ = customWriter.Send() return nil @@ -875,9 +865,7 @@ func (s *SouinBaseHandler) ServeHTTP(rw http.ResponseWriter, rq *http.Request, n } customWriter.WriteHeader(response.StatusCode) s.Configuration.GetLogger().Debugf("Serve from cache %+v", req) - customWriter.handleBuffer(func(b *bytes.Buffer) { - _, _ = io.Copy(b, response.Body) - }) + _, _ = customWriter.copyToBuffer(response.Body) _, err := customWriter.Send() prometheus.Increment(prometheus.CachedResponseCounter) @@ -897,11 +885,9 @@ func (s *SouinBaseHandler) ServeHTTP(rw http.ResponseWriter, rq *http.Request, n } customWriter.WriteHeader(response.StatusCode) rfc.HitStaleCache(&response.Header) - customWriter.handleBuffer(func(b *bytes.Buffer) { - _, _ = io.Copy(b, response.Body) - }) + _, _ = customWriter.copyToBuffer(response.Body) _, err := customWriter.Send() - customWriter = NewCustomWriter(req, rw, bufPool) + customWriter = NewCustomWriter(req, rw, bufPool, int(s.Configuration.GetDefaultCache().GetMaxBodyBytes())) go func(v *core.Revalidator, goCw *CustomWriter, goRq *http.Request, goNext func(http.ResponseWriter, *http.Request) error, goCc *cacheobject.RequestCacheDirectives, goCk string, goUri string) { _ = s.Revalidate(v, goNext, goCw, goRq, goCc, goCk, goUri) }(validator, customWriter, req, next, requestCc, cachedKey, uri) @@ -923,18 +909,13 @@ func (s *SouinBaseHandler) ServeHTTP(rw http.ResponseWriter, rq *http.Request, n response.Header.Set("Cache-Status", response.Header.Get("Cache-Status")+code) maps.Copy(customWriter.Header(), response.Header) customWriter.WriteHeader(response.StatusCode) - customWriter.handleBuffer(func(b *bytes.Buffer) { - b.Reset() - _, _ = io.Copy(b, response.Body) - }) + _, _ = customWriter.resetAndCopyToBuffer(response.Body) _, err := customWriter.Send() return err } rw.WriteHeader(http.StatusGatewayTimeout) - customWriter.handleBuffer(func(b *bytes.Buffer) { - b.Reset() - }) + customWriter.resetBuffer() _, err := customWriter.Send() return err @@ -945,9 +926,7 @@ func (s *SouinBaseHandler) ServeHTTP(rw http.ResponseWriter, rq *http.Request, n rfc.SetCacheStatusHeader(response, storerName) customWriter.WriteHeader(response.StatusCode) maps.Copy(customWriter.Header(), response.Header) - customWriter.handleBuffer(func(b *bytes.Buffer) { - _, _ = io.Copy(b, response.Body) - }) + _, _ = customWriter.copyToBuffer(response.Body) _, _ = customWriter.Send() return err @@ -956,9 +935,7 @@ func (s *SouinBaseHandler) ServeHTTP(rw http.ResponseWriter, rq *http.Request, n if statusCode != http.StatusNotModified && validator.Matched { customWriter.WriteHeader(http.StatusNotModified) - customWriter.handleBuffer(func(b *bytes.Buffer) { - b.Reset() - }) + customWriter.resetBuffer() _, _ = customWriter.Send() return err @@ -973,9 +950,7 @@ func (s *SouinBaseHandler) ServeHTTP(rw http.ResponseWriter, rq *http.Request, n customWriter.WriteHeader(response.StatusCode) rfc.HitStaleCache(&response.Header) maps.Copy(customWriter.Header(), response.Header) - customWriter.handleBuffer(func(b *bytes.Buffer) { - _, _ = io.Copy(b, response.Body) - }) + _, _ = customWriter.copyToBuffer(response.Body) _, err := customWriter.Send() return err @@ -989,9 +964,7 @@ func (s *SouinBaseHandler) ServeHTTP(rw http.ResponseWriter, rq *http.Request, n customWriter.WriteHeader(response.StatusCode) rfc.HitStaleCache(&response.Header) maps.Copy(customWriter.Header(), response.Header) - customWriter.handleBuffer(func(b *bytes.Buffer) { - _, _ = io.Copy(b, response.Body) - }) + _, _ = customWriter.copyToBuffer(response.Body) _, err := customWriter.Send() return err @@ -1015,10 +988,8 @@ func (s *SouinBaseHandler) ServeHTTP(rw http.ResponseWriter, rq *http.Request, n response.Header.Set("Cache-Status", response.Header.Get("Cache-Status")+code) maps.Copy(customWriter.Header(), response.Header) customWriter.WriteHeader(response.StatusCode) - customWriter.handleBuffer(func(b *bytes.Buffer) { - b.Reset() - _, _ = io.Copy(b, response.Body) - }) + _, _ = customWriter.resetAndCopyToBuffer(response.Body) + _, err := customWriter.Send() return err diff --git a/pkg/middleware/writer.go b/pkg/middleware/writer.go index c4e836861..c15f35f9e 100644 --- a/pkg/middleware/writer.go +++ b/pkg/middleware/writer.go @@ -3,12 +3,14 @@ package middleware import ( "bytes" "fmt" + "io" "net/http" "strconv" "strings" "sync" "github.com/darkweak/go-esi/esi" + "github.com/darkweak/souin/context" "github.com/darkweak/souin/pkg/rfc" ) @@ -19,32 +21,53 @@ type SouinWriterInterface interface { var _ SouinWriterInterface = (*CustomWriter)(nil) -func NewCustomWriter(rq *http.Request, rw http.ResponseWriter, b *bytes.Buffer) *CustomWriter { +func NewCustomWriter(rq *http.Request, rw http.ResponseWriter, b *bytes.Buffer, maxSize int) *CustomWriter { return &CustomWriter{ - statusCode: 200, - Buf: b, - Req: rq, - Rw: rw, - Headers: http.Header{}, - mutex: sync.Mutex{}, + statusCode: 200, + Buf: b, + Req: rq, + Rw: rw, + Headers: http.Header{}, + mutex: sync.Mutex{}, + maxSize: maxSize, + maxSizeReached: false, } } // CustomWriter handles the response and provide the way to cache the value type CustomWriter struct { - Buf *bytes.Buffer - Rw http.ResponseWriter - Req *http.Request - Headers http.Header - headersSent bool - mutex sync.Mutex - statusCode int + Buf *bytes.Buffer + Rw http.ResponseWriter + Req *http.Request + Headers http.Header + headersSent bool + mutex sync.Mutex + statusCode int + maxSize int + maxSizeReached bool } -func (r *CustomWriter) handleBuffer(callback func(*bytes.Buffer)) { +func (r *CustomWriter) resetBuffer() { r.mutex.Lock() - callback(r.Buf) - r.mutex.Unlock() + defer r.mutex.Unlock() + + r.Buf.Reset() +} + +func (r *CustomWriter) copyToBuffer(src io.Reader) (int64, error) { + r.mutex.Lock() + defer r.mutex.Unlock() + + return io.Copy(r.Buf, src) +} + +func (r *CustomWriter) resetAndCopyToBuffer(src io.Reader) (int64, error) { + r.mutex.Lock() + defer r.mutex.Unlock() + + r.Buf.Reset() + + return io.Copy(r.Buf, src) } // Header will write the response headers @@ -79,10 +102,34 @@ func (r *CustomWriter) WriteHeader(code int) { // Write will write the response body func (r *CustomWriter) Write(b []byte) (int, error) { - r.handleBuffer(func(actual *bytes.Buffer) { - actual.Grow(len(b)) - _, _ = actual.Write(b) - }) + r.mutex.Lock() + defer r.mutex.Unlock() + + if r.maxSizeReached { + return r.Rw.Write(b) + } + + if r.maxSize > 0 && (r.Buf.Len()+len(b)) > r.maxSize { + r.maxSizeReached = true + + if !r.headersSent && r.Req.Context().Err() == nil { + r.Rw.Header().Set( + "Cache-Status", + fmt.Sprintf( + "%s; fwd=uri-miss; detail=UPSTREAM-RESPONSE-TOO-LARGE; key=%s", + r.Req.Context().Value(context.CacheName), + rfc.GetCacheKeyFromCtx(r.Req.Context()), + ), + ) + } + + _, _ = r.Rw.Write(r.Buf.Bytes()) + + r.Buf.Reset() + } + + r.Buf.Grow(len(b)) + _, _ = r.Buf.Write(b) return len(b), nil } @@ -142,9 +189,8 @@ func parseRange(rangeHeaders []string, contentRange string) ([]rangeValue, range // Send delays the response to handle Cache-Status func (r *CustomWriter) Send() (int, error) { - defer r.handleBuffer(func(b *bytes.Buffer) { - b.Reset() - }) + defer r.resetBuffer() + storedLength := r.Header().Get(rfc.StoredLengthHeader) if storedLength != "" { r.Header().Set("Content-Length", storedLength) From 8f2ce4c8a7bfc95e4297d0f1ec011c839c98c4b0 Mon Sep 17 00:00:00 2001 From: darkweak Date: Sat, 11 Oct 2025 19:03:20 +0200 Subject: [PATCH 2/2] fix(memory-bloat): close response body and use a string Builder for Range requests --- pkg/middleware/middleware.go | 9 +++++++++ pkg/middleware/writer.go | 9 +++++---- plugins/caddy/admin.go | 2 +- 3 files changed, 15 insertions(+), 5 deletions(-) diff --git a/pkg/middleware/middleware.go b/pkg/middleware/middleware.go index aedaa913e..e72b6178c 100644 --- a/pkg/middleware/middleware.go +++ b/pkg/middleware/middleware.go @@ -817,6 +817,15 @@ func (s *SouinBaseHandler) ServeHTTP(rw http.ResponseWriter, rq *http.Request, n backfillIds++ } + defer func() { + if fresh != nil { + _ = fresh.Body.Close() + } + if stale != nil { + _ = stale.Body.Close() + } + }() + headerName, _ := s.SurrogateKeyStorer.GetSurrogateControl(customWriter.Header()) if fresh != nil && (!modeContext.Strict || rfc.ValidateCacheControl(fresh, requestCc)) { go func() { diff --git a/pkg/middleware/writer.go b/pkg/middleware/writer.go index c15f35f9e..5d86939ff 100644 --- a/pkg/middleware/writer.go +++ b/pkg/middleware/writer.go @@ -202,7 +202,7 @@ func (r *CustomWriter) Send() (int, error) { if r.Headers.Get("Range") != "" { - var bufStr string + bufStr := new(strings.Builder) mimeType := r.Header().Get("Content-Type") r.WriteHeader(http.StatusPartialContent) @@ -244,16 +244,17 @@ func (r *CustomWriter) Send() (int, error) { content = content[:header.to-header.from] } - bufStr += fmt.Sprintf(` + bufStr.WriteString(fmt.Sprintf(` %s Content-Type: %s Content-Range: bytes %d-%d/%d %s -`, separator, mimeType, header.from, header.to, r.Buf.Len(), content) +`, separator, mimeType, header.from, header.to, r.Buf.Len(), content)) } - result = []byte(bufStr + separator + "--") + bufStr.WriteString(separator + "--") + result = []byte(bufStr.String()) } } diff --git a/plugins/caddy/admin.go b/plugins/caddy/admin.go index 70c7c5253..fda530d0a 100644 --- a/plugins/caddy/admin.go +++ b/plugins/caddy/admin.go @@ -78,7 +78,7 @@ func (a *adminAPI) Provision(ctx caddy.Context) error { return nil } -// Routes returns the admin routes. +// Routes return the admin routes. func (a *adminAPI) Routes() []caddy.AdminRoute { return []caddy.AdminRoute{ {