From 8a0318b4f35504756992638aa992737bb20b9524 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: Sun, 1 Oct 2023 15:06:58 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E5=86=85=E5=AD=98=E5=86=99?=
 =?UTF-8?q?=E5=85=A5=E9=80=9F=E5=BA=A6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/caches/storage_memory.go |  5 ++-
 internal/caches/writer_memory.go  | 59 +++++++++++++++++++++----------
 2 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/internal/caches/storage_memory.go b/internal/caches/storage_memory.go
index d3e8528..4d3c1e3 100644
--- a/internal/caches/storage_memory.go
+++ b/internal/caches/storage_memory.go
@@ -32,6 +32,9 @@ type MemoryItem struct {
 	ModifiedAt  int64
 
 	TotalSize int64
+
+	IsPrepared  bool
+	WriteOffset int64
 }
 
 func (this *MemoryItem) IsExpired() bool {
@@ -215,7 +218,7 @@ func (this *MemoryStorage) openWriter(key string, expiresAt int64, status int, h
 	}
 
 	isWriting = true
-	return NewMemoryWriter(this, key, expiresAt, status, isDirty, maxSize, func(valueItem *MemoryItem) {
+	return NewMemoryWriter(this, key, expiresAt, status, isDirty, bodySize, maxSize, func(valueItem *MemoryItem) {
 		this.locker.Lock()
 		delete(this.writingKeyMap, key)
 		this.locker.Unlock()
diff --git a/internal/caches/writer_memory.go b/internal/caches/writer_memory.go
index 3ad813d..21d2d8d 100644
--- a/internal/caches/writer_memory.go
+++ b/internal/caches/writer_memory.go
@@ -16,7 +16,9 @@ type MemoryWriter struct {
 	bodySize   int64
 	status     int
 	isDirty    bool
-	maxSize    int64
+
+	expectedBodySize int64
+	maxSize          int64
 
 	hash    uint64
 	item    *MemoryItem
@@ -24,20 +26,26 @@ type MemoryWriter struct {
 	once    sync.Once
 }
 
-func NewMemoryWriter(memoryStorage *MemoryStorage, key string, expiredAt int64, status int, isDirty bool, maxSize int64, endFunc func(valueItem *MemoryItem)) *MemoryWriter {
+func NewMemoryWriter(memoryStorage *MemoryStorage, key string, expiredAt int64, status int, isDirty bool, expectedBodySize int64, maxSize int64, endFunc func(valueItem *MemoryItem)) *MemoryWriter {
+	var valueItem = &MemoryItem{
+		ExpiresAt:  expiredAt,
+		ModifiedAt: time.Now().Unix(),
+		Status:     status,
+	}
+	if expectedBodySize > 0 {
+		valueItem.BodyValue = make([]byte, expectedBodySize)
+		valueItem.IsPrepared = true
+	}
 	var w = &MemoryWriter{
-		storage:   memoryStorage,
-		key:       key,
-		expiredAt: expiredAt,
-		item: &MemoryItem{
-			ExpiresAt:  expiredAt,
-			ModifiedAt: time.Now().Unix(),
-			Status:     status,
-		},
-		status:  status,
-		isDirty: isDirty,
-		maxSize: maxSize,
-		endFunc: endFunc,
+		storage:          memoryStorage,
+		key:              key,
+		expiredAt:        expiredAt,
+		item:             valueItem,
+		status:           status,
+		isDirty:          isDirty,
+		expectedBodySize: expectedBodySize,
+		maxSize:          maxSize,
+		endFunc:          endFunc,
 	}
 
 	w.hash = w.calculateHash(key)
@@ -54,17 +62,32 @@ func (this *MemoryWriter) WriteHeader(data []byte) (n int, err error) {
 
 // Write 写入数据
 func (this *MemoryWriter) Write(data []byte) (n int, err error) {
-	this.bodySize += int64(len(data))
-	this.item.BodyValue = append(this.item.BodyValue, data...)
+	var l = len(data)
+	if l == 0 {
+		return
+	}
+
+	if this.item.IsPrepared {
+		if this.item.WriteOffset+int64(l) > this.expectedBodySize {
+			err = ErrWritingUnavailable
+			return
+		}
+		copy(this.item.BodyValue[this.item.WriteOffset:], data)
+		this.item.WriteOffset += int64(l)
+	} else {
+		this.item.BodyValue = append(this.item.BodyValue, data...)
+	}
+
+	this.bodySize += int64(l)
 
 	// 检查尺寸
 	if this.maxSize > 0 && this.bodySize > this.maxSize {
 		err = ErrEntityTooLarge
 		this.storage.IgnoreKey(this.key, this.maxSize)
-		return len(data), err
+		return l, err
 	}
 
-	return len(data), nil
+	return l, nil
 }
 
 // WriteAt 在指定位置写入数据