From 534d64673fd38fc89a3862a3c164ba43f85f80bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E7=A5=A5=E8=B6=85?= <iwind.liu@gmail.com>
Date: Mon, 9 Oct 2023 12:48:30 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E5=86=85=E5=AD=98=E7=BC=93?=
 =?UTF-8?q?=E5=AD=98=E7=9B=B8=E5=85=B3=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/caches/memory_fragment_pool.go |  4 +---
 internal/caches/writer_memory.go        | 22 +++++++++++++---------
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/internal/caches/memory_fragment_pool.go b/internal/caches/memory_fragment_pool.go
index 7de47f3..b556e8e 100644
--- a/internal/caches/memory_fragment_pool.go
+++ b/internal/caches/memory_fragment_pool.go
@@ -126,9 +126,7 @@ func (this *MemoryFragmentPool) Get(expectSize int64) (resultBytes []byte, ok bo
 		return
 	}
 
-	if expectSize < minMemoryFragmentPoolItemSize {
-		return
-	}
+	// DO NOT check min segment size
 
 	this.mu.RLock()
 
diff --git a/internal/caches/writer_memory.go b/internal/caches/writer_memory.go
index 93e5795..a9bf54f 100644
--- a/internal/caches/writer_memory.go
+++ b/internal/caches/writer_memory.go
@@ -32,23 +32,20 @@ func NewMemoryWriter(memoryStorage *MemoryStorage, key string, expiredAt int64,
 		ModifiedAt: fasttime.Now().Unix(),
 		Status:     status,
 	}
-	if expectedBodySize > 0 {
+	if expectedBodySize > 0 && expectedBodySize <= maxMemoryFragmentPoolItemSize {
 		bodyBytes, ok := SharedFragmentMemoryPool.Get(expectedBodySize) // try to reuse memory
 		if ok {
 			valueItem.BodyValue = bodyBytes
+			valueItem.IsPrepared = true
 		} else {
-			if isDirty {
-				if expectedBodySize >= minMemoryFragmentPoolItemSize {
-					SharedFragmentMemoryPool.IncreaseNew()
-				}
+			if expectedBodySize <= (16 << 20) {
 				var allocSize = (expectedBodySize/16384 + 1) * 16384
 				valueItem.BodyValue = make([]byte, allocSize)[:expectedBodySize]
-			} else {
-				valueItem.BodyValue = make([]byte, expectedBodySize)
+				valueItem.IsPrepared = true
+
+				SharedFragmentMemoryPool.IncreaseNew()
 			}
 		}
-
-		valueItem.IsPrepared = true
 	}
 	var w = &MemoryWriter{
 		storage:          memoryStorage,
@@ -170,6 +167,13 @@ func (this *MemoryWriter) Discard() error {
 
 	this.storage.locker.Lock()
 	delete(this.storage.valuesMap, this.hash)
+
+	if this.item != nil &&
+		!this.item.isReferring &&
+		cap(this.item.BodyValue) >= minMemoryFragmentPoolItemSize {
+		SharedFragmentMemoryPool.Put(this.item.BodyValue)
+	}
+
 	this.storage.locker.Unlock()
 	return nil
 }