From 8227592451b70e96127af526e80868ef1c33902c Mon Sep 17 00:00:00 2001
From: DarthSim <darthsim@gmail.com>
Date: Mon, 12 Jun 2023 22:00:55 +0300
Subject: [PATCH] Optimize bufpool

---
 bufpool/bufpool.go | 24 ++++++++++++++++++++++--
 imagedata/read.go  |  4 +---
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/bufpool/bufpool.go b/bufpool/bufpool.go
index 6eb97c92..2cd0f054 100644
--- a/bufpool/bufpool.go
+++ b/bufpool/bufpool.go
@@ -106,7 +106,7 @@ func (p *Pool) calibrateAndClean() {
 	for i := 0; i < steps; i++ {
 		callsSum += p.tmpCalls[i]
 
-		if callsSum > defSum || defStep < 0 {
+		if defStep < 0 && callsSum > defSum {
 			defStep = i
 		}
 
@@ -219,7 +219,7 @@ func (p *Pool) Get(size int, grow bool) *bytes.Buffer {
 
 	growSize := p.defaultSize
 	if grow {
-		growSize = imath.Max(size, growSize)
+		growSize = imath.Max(p.normalizeCap(size), growSize)
 	}
 
 	// Grow the buffer only if we know the requested size and it is smaller than
@@ -261,6 +261,26 @@ func (p *Pool) Put(buf *bytes.Buffer) {
 	p.insert(buf)
 }
 
+// GrowBuffer growth capacity of the buffer to the normalized provided value
+func (p *Pool) GrowBuffer(buf *bytes.Buffer, cap int) {
+	cap = p.normalizeCap(cap)
+	if buf.Cap() < cap {
+		buf.Grow(cap - buf.Len())
+	}
+}
+
+func (p *Pool) normalizeCap(cap int) int {
+	// Don't normalize cap if it's larger than maxSize
+	// since we'll throw this buf out anyway
+	maxSize := int(atomic.LoadUint64(&p.maxSize))
+	if maxSize > 0 && cap > maxSize {
+		return cap
+	}
+
+	ind := index(cap)
+	return imath.Max(cap, minSize<<ind)
+}
+
 func saveEntry(e *entry) {
 	e.buf = nil
 	e.next = nil
diff --git a/imagedata/read.go b/imagedata/read.go
index 103c1aae..531528e3 100644
--- a/imagedata/read.go
+++ b/imagedata/read.go
@@ -52,9 +52,7 @@ func readAndCheckImage(r io.Reader, contentLength int, secopts security.Options)
 		return nil, wrapError(err)
 	}
 
-	if contentLength > buf.Cap() {
-		buf.Grow(contentLength - buf.Len())
-	}
+	downloadBufPool.GrowBuffer(buf, contentLength)
 
 	if err = br.Flush(); err != nil {
 		buf.Reset()