优化内存写入速度

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()

internal/caches/writer_memory.go

@@ -16,6 +16,8 @@ type MemoryWriter struct {
 	bodySize   int64
 	status     int
 	isDirty    bool
+
+	expectedBodySize int64
 	maxSize          int64
 
 	hash    uint64
@@ -24,18 +26,24 @@ 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,
-		},
+		item:             valueItem,
 		status:           status,
 		isDirty:          isDirty,
+		expectedBodySize: expectedBodySize,
 		maxSize:          maxSize,
 		endFunc:          endFunc,
 	}
@@ -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))
+	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 在指定位置写入数据