diff --git a/internal/nodes/http_writer.go b/internal/nodes/http_writer.go index ea0fcaa..d94881c 100644 --- a/internal/nodes/http_writer.go +++ b/internal/nodes/http_writer.go @@ -22,13 +22,22 @@ import ( "net/http" "path/filepath" "strings" + "sync/atomic" ) +// 限制WebP能够同时使用的Buffer内存使用量 +const webpMaxBufferSize int64 = 1_000_000_000 + +var webpTotalBufferSize int64 = 0 +var webpBufferPool = utils.NewBufferPool(1024) + // HTTPWriter 响应Writer type HTTPWriter struct { req *HTTPRequest writer http.ResponseWriter + size int64 + webpIsEncoding bool webpBuffer *bytes.Buffer webpIsWriting bool @@ -83,6 +92,7 @@ func (this *HTTPWriter) SetCompression(config *serverconfigs.HTTPCompressionConf // Prepare 准备输出 // 缓存不调用此函数 func (this *HTTPWriter) Prepare(size int64, status int) { + this.size = size this.statusCode = status if status == http.StatusOK { @@ -240,8 +250,18 @@ func (this *HTTPWriter) SetOk() { // Close 关闭 func (this *HTTPWriter) Close() { + if this.webpIsEncoding { + defer func() { + atomic.AddInt64(&webpTotalBufferSize, -this.size*32) + webpBufferPool.Put(this.webpBuffer) + }() + } + // webp writer if this.isOk && this.webpIsEncoding { + var bufferLen = int64(this.webpBuffer.Len()) + atomic.AddInt64(&webpTotalBufferSize, bufferLen*8) + imageData, _, err := image.Decode(this.webpBuffer) if err != nil { _, _ = io.Copy(this.writer, this.webpBuffer) @@ -257,6 +277,7 @@ func (this *HTTPWriter) Close() { f = 100 } this.webpIsWriting = true + err = webp.Encode(this, imageData, &webp.Options{ Lossless: false, Quality: f, @@ -274,6 +295,9 @@ func (this *HTTPWriter) Close() { this.cacheWriter = nil } } + + atomic.AddInt64(&webpTotalBufferSize, -bufferLen*8) + this.webpBuffer.Reset() } // compression writer @@ -343,12 +367,15 @@ func (this *HTTPWriter) prepareWebP(size int64) { this.req.web.WebP.IsOn && this.req.web.WebP.MatchResponse(this.Header().Get("Content-Type"), size, filepath.Ext(this.req.requestPath()), this.req.Format) && this.req.web.WebP.MatchAccept(this.req.requestHeader("Accept")) && - len(this.writer.Header().Get("Content-Encoding")) == 0 { + len(this.writer.Header().Get("Content-Encoding")) == 0 && + atomic.LoadInt64(&webpTotalBufferSize) < webpMaxBufferSize { this.webpIsEncoding = true - this.webpBuffer = &bytes.Buffer{} + this.webpBuffer = webpBufferPool.Get() this.Header().Del("Content-Length") this.Header().Set("Content-Type", "image/webp") + + atomic.AddInt64(&webpTotalBufferSize, size*32) } } diff --git a/internal/utils/buffer_pool.go b/internal/utils/buffer_pool.go new file mode 100644 index 0000000..f08614c --- /dev/null +++ b/internal/utils/buffer_pool.go @@ -0,0 +1,48 @@ +// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved. + +package utils + +import "bytes" + +// BufferPool pool for get byte slice +type BufferPool struct { + c chan *bytes.Buffer +} + +// NewBufferPool 创建新对象 +func NewBufferPool(maxSize int) *BufferPool { + if maxSize <= 0 { + maxSize = 1024 + } + pool := &BufferPool{ + c: make(chan *bytes.Buffer, maxSize), + } + return pool +} + +// Get 获取一个新的Buffer +func (this *BufferPool) Get() (b *bytes.Buffer) { + select { + case b = <-this.c: + b.Reset() + default: + b = &bytes.Buffer{} + } + return +} + +// Put 放回一个使用过的byte slice +func (this *BufferPool) Put(b *bytes.Buffer) { + b.Reset() + + select { + case this.c <- b: + default: + // 已达最大容量,则抛弃 + } +} + +// Size 当前的数量 +func (this *BufferPool) Size() int { + return len(this.c) +}