优化内存缓存限制

internal/caches/storage_memory.go

@@ -30,6 +30,8 @@ type MemoryItem struct {
 	Status      int
 	IsDone      bool
 	ModifiedAt  int64
+
+	TotalSize int64
 }
 
 func (this *MemoryItem) IsExpired() bool {
@@ -50,7 +52,7 @@ type MemoryStorage struct {
 
 	purgeTicker *utils.Ticker
 
-	totalSize     int64
+	usedSize      int64
 	writingKeyMap map[string]zero.Zero // key => bool
 
 	ignoreKeys *setutils.FixedSet
@@ -84,13 +86,6 @@ func NewMemoryStorage(policy *serverconfigs.HTTPCachePolicy, parentStorage Stora
 func (this *MemoryStorage) Init() error {
 	_ = this.list.Init()
 
-	this.list.OnAdd(func(item *Item) {
-		atomic.AddInt64(&this.totalSize, item.TotalSize())
-	})
-	this.list.OnRemove(func(item *Item) {
-		atomic.AddInt64(&this.totalSize, -item.TotalSize())
-	})
-
 	this.initPurgeTicker()
 
 	// 启动定时Flush memory to disk任务
@@ -205,12 +200,12 @@ func (this *MemoryStorage) openWriter(key string, expiresAt int64, status int, h
 	}
 
 	// 检查是否超出最大值
-	capacityBytes := this.memoryCapacityBytes()
+	var capacityBytes = this.memoryCapacityBytes()
 	if bodySize < 0 {
 		bodySize = 0
 	}
-	if capacityBytes > 0 && capacityBytes <= this.totalSize+bodySize {
-		return nil, NewCapacityError("write memory cache failed: over memory size: " + strconv.FormatInt(capacityBytes, 10) + ", current size: " + strconv.FormatInt(this.totalSize, 10) + " bytes")
+	if capacityBytes > 0 && capacityBytes <= atomic.LoadInt64(&this.usedSize)+bodySize {
+		return nil, NewCapacityError("write memory cache failed: over memory size: " + strconv.FormatInt(capacityBytes, 10) + ", current size: " + strconv.FormatInt(this.usedSize, 10) + " bytes")
 	}
 
 	// 先删除
@@ -220,10 +215,17 @@ func (this *MemoryStorage) openWriter(key string, expiresAt int64, status int, h
 	}
 
 	isWriting = true
-	return NewMemoryWriter(this, key, expiresAt, status, isDirty, maxSize, func() {
+	return NewMemoryWriter(this, key, expiresAt, status, isDirty, maxSize, func(valueItem *MemoryItem) {
 		this.locker.Lock()
 		delete(this.writingKeyMap, key)
 		this.locker.Unlock()
+
+		if valueItem != nil && valueItem.IsDone {
+			valueItem.TotalSize = int64(len(valueItem.HeaderValue) + len(valueItem.BodyValue) + len(key) + 256 /** meta size **/)
+
+			atomic.AddInt64(&this.usedSize, valueItem.TotalSize)
+			runtime.SetFinalizer(valueItem, this.valueItemFinalizer)
+		}
 	}), nil
 }
 
@@ -252,7 +254,7 @@ func (this *MemoryStorage) CleanAll() error {
 	this.locker.Lock()
 	this.valuesMap = map[uint64]*MemoryItem{}
 	_ = this.list.Reset()
-	atomic.StoreInt64(&this.totalSize, 0)
+	atomic.StoreInt64(&this.usedSize, 0)
 	this.locker.Unlock()
 	return nil
 }
@@ -363,6 +365,11 @@ func (this *MemoryStorage) CanUpdatePolicy(newPolicy *serverconfigs.HTTPCachePol
 
 // AddToList 将缓存添加到列表
 func (this *MemoryStorage) AddToList(item *Item) {
+	// skip added item
+	if item.MetaSize > 0 {
+		return
+	}
+
 	item.MetaSize = int64(len(item.Key)) + 128 /** 128是我们评估的数据结构的长度 **/
 	var hash = types.String(this.hash(item.Key))
 
@@ -380,7 +387,7 @@ func (this *MemoryStorage) TotalDiskSize() int64 {
 
 // TotalMemorySize 内存尺寸
 func (this *MemoryStorage) TotalMemorySize() int64 {
-	return atomic.LoadInt64(&this.totalSize)
+	return atomic.LoadInt64(&this.usedSize)
 }
 
 // IgnoreKey 忽略某个Key，即不缓存某个Key
@@ -563,17 +570,29 @@ func (this *MemoryStorage) memoryCapacityBytes() int64 {
 	if this.policy == nil {
 		return 0
 	}
-	c1 := int64(0)
-	if this.policy.Capacity != nil {
-		c1 = this.policy.Capacity.Bytes()
-	}
+
 	if SharedManager.MaxMemoryCapacity != nil {
-		c2 := SharedManager.MaxMemoryCapacity.Bytes()
-		if c2 > 0 {
-			return c2
+		var capacityBytes = SharedManager.MaxMemoryCapacity.Bytes()
+		if capacityBytes > 0 {
+			return capacityBytes
 		}
 	}
-	return c1
+
+	var capacity = this.policy.Capacity // copy
+	if capacity != nil {
+		var capacityBytes = capacity.Bytes()
+		if capacityBytes > 0 {
+			return capacityBytes
+		}
+	}
+
+	// half of the system memory
+	var memoryGB = utils.SystemMemoryGB()
+	if memoryGB < 1 {
+		memoryGB = 1
+	}
+
+	return int64(memoryGB) << 30 / 2
 }
 
 func (this *MemoryStorage) deleteWithoutLocker(key string) error {
@@ -604,3 +623,7 @@ func (this *MemoryStorage) initPurgeTicker() {
 		}
 	})
 }
+
+func (this *MemoryStorage) valueItemFinalizer(valueItem *MemoryItem) {
+	atomic.AddInt64(&this.usedSize, -valueItem.TotalSize)
+}

internal/caches/writer_memory.go

@@ -20,12 +20,12 @@ type MemoryWriter struct {
 
 	hash    uint64
 	item    *MemoryItem
-	endFunc func()
+	endFunc func(valueItem *MemoryItem)
 	once    sync.Once
 }
 
-func NewMemoryWriter(memoryStorage *MemoryStorage, key string, expiredAt int64, status int, isDirty bool, maxSize int64, endFunc func()) *MemoryWriter {
-	w := &MemoryWriter{
+func NewMemoryWriter(memoryStorage *MemoryStorage, key string, expiredAt int64, status int, isDirty bool, maxSize int64, endFunc func(valueItem *MemoryItem)) *MemoryWriter {
+	var w = &MemoryWriter{
 		storage:   memoryStorage,
 		key:       key,
 		expiredAt: expiredAt,
@@ -39,6 +39,7 @@ func NewMemoryWriter(memoryStorage *MemoryStorage, key string, expiredAt int64,
 		maxSize: maxSize,
 		endFunc: endFunc,
 	}
+
 	w.hash = w.calculateHash(key)
 
 	return w
@@ -87,7 +88,7 @@ func (this *MemoryWriter) BodySize() int64 {
 func (this *MemoryWriter) Close() error {
 	// 需要在Locker之外
 	defer this.once.Do(func() {
-		this.endFunc()
+		this.endFunc(this.item)
 	})
 
 	if this.item == nil {
@@ -102,7 +103,6 @@ func (this *MemoryWriter) Close() error {
 			select {
 			case this.storage.dirtyChan <- this.key:
 			default:
-
 			}
 		}
 	}
@@ -115,7 +115,7 @@ func (this *MemoryWriter) Close() error {
 func (this *MemoryWriter) Discard() error {
 	// 需要在Locker之外
 	defer this.once.Do(func() {
-		this.endFunc()
+		this.endFunc(nil)
 	})
 
 	this.storage.locker.Lock()

internal/utils/fs/status.go

@@ -4,6 +4,7 @@ package fsutils
 
 import (
 	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
+	"github.com/iwind/TeaGo/Tea"
 	"sync/atomic"
 	"time"
 )
@@ -67,6 +68,10 @@ func DiskIsFast() bool {
 }
 
 func DiskIsExtremelyFast() bool {
+	// 在开发环境下返回false，以便于测试
+	if Tea.IsTesting() {
+		return false
+	}
 	return DiskSpeed == SpeedExtremelyFast
 }