waf-platform/EdgeNode/internal/utils/ttlcache/piece.go

package ttlcache

import (
	"github.com/TeaOSLab/EdgeNode/internal/utils/expires"
	"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
	"github.com/TeaOSLab/EdgeNode/internal/utils/zero"
	"sync"
)

type Piece[T any] struct {
	m           map[uint64]*Item[T]
	expiresList *expires.List
	maxItems    int
	lastGCTime  int64

	locker sync.RWMutex
}

func NewPiece[T any](maxItems int) *Piece[T] {
	return &Piece[T]{
		m:           map[uint64]*Item[T]{},
		expiresList: expires.NewSingletonList(),
		maxItems:    maxItems,
	}
}

func (this *Piece[T]) Add(key uint64, item *Item[T]) (ok bool) {
	this.locker.RLock()
	if this.maxItems > 0 && len(this.m) >= this.maxItems {
		this.locker.RUnlock()
		return
	}
	this.locker.RUnlock()

	this.locker.Lock()
	oldItem, exists := this.m[key]
	if exists && oldItem.expiresAt == item.expiresAt {
		this.locker.Unlock()
		return true
	}
	this.m[key] = item
	this.locker.Unlock()

	this.expiresList.Add(key, item.expiresAt)

	return true
}

func (this *Piece[T]) IncreaseInt64(key uint64, delta T, expiresAt int64, extend bool) (result T) {
	this.locker.Lock()
	item, ok := this.m[key]
	if ok && item.expiresAt > fasttime.Now().Unix() {
		int64Value, isInt64 := any(item.Value).(int64)
		if isInt64 {
			result = any(int64Value + any(delta).(int64)).(T)
		}
		item.Value = result
		if extend {
			item.expiresAt = expiresAt
		}
		this.expiresList.Add(key, expiresAt)
	} else {
		if len(this.m) < this.maxItems {
			result = delta
			this.m[key] = &Item[T]{
				Value:     delta,
				expiresAt: expiresAt,
			}
			this.expiresList.Add(key, expiresAt)
		}
	}
	this.locker.Unlock()

	return
}

func (this *Piece[T]) Delete(key uint64) {
	this.expiresList.Remove(key)

	this.locker.Lock()
	delete(this.m, key)
	this.locker.Unlock()
}

func (this *Piece[T]) Read(key uint64) (item *Item[T]) {
	this.locker.RLock()
	item = this.m[key]
	if item != nil && item.expiresAt < fasttime.Now().Unix() {
		item = nil
	}
	this.locker.RUnlock()

	return
}

func (this *Piece[T]) Count() (count int) {
	this.locker.RLock()
	count = len(this.m)
	this.locker.RUnlock()
	return
}

// SampleExpireRate 采样检查过期率（优化：避免全量遍历）
func (this *Piece[T]) SampleExpireRate(sampleSize int) int {
	this.locker.RLock()
	defer this.locker.RUnlock()

	if len(this.m) == 0 {
		return 0
	}

	var currentTime = fasttime.Now().Unix()
	var expiredCount = 0
	var checkedCount = 0

	// 采样检查
	for _, item := range this.m {
		if checkedCount >= sampleSize {
			break
		}
		checkedCount++
		if item.expiresAt <= currentTime {
			expiredCount++
		}
	}

	return expiredCount
}

func (this *Piece[T]) GC() int {
	var currentTime = fasttime.Now().Unix()

	// 优化：如果上次 GC 时间很近，跳过（避免频繁 GC）
	if this.lastGCTime > 0 && currentTime-this.lastGCTime < 1 {
		return 0
	}

	if this.lastGCTime == 0 {
		this.lastGCTime = currentTime - 3600
	}

	var minTime = this.lastGCTime
	var maxTime = currentTime
	if minTime > maxTime {
		// 过去的时间比现在大，则从这一秒重新开始
		minTime = maxTime
	}

	// 优化：批量收集过期项，减少锁持有时间
	var expiredKeys []uint64
	for i := minTime; i <= maxTime; i++ {
		var itemMap = this.expiresList.GC(i)
		if len(itemMap) > 0 {
			for key := range itemMap {
				expiredKeys = append(expiredKeys, key)
			}
		}
	}

	// 批量删除
	var expiredCount = len(expiredKeys)
	if expiredCount > 0 {
		// 重新构建 itemMap
		itemMap := make(expires.ItemMap)
		for _, key := range expiredKeys {
			itemMap[key] = zero.New()
		}
		this.gcItemMap(itemMap)
	}

	this.lastGCTime = currentTime
	return expiredCount
}

func (this *Piece[T]) Clean() {
	this.locker.Lock()
	this.m = map[uint64]*Item[T]{}
	this.locker.Unlock()

	this.expiresList.Clean()
}

func (this *Piece[T]) Destroy() {
	this.locker.Lock()
	this.m = nil
	this.expiresList.Clean()
	this.locker.Unlock()

	this.expiresList.Clean()
}

func (this *Piece[T]) gcItemMap(itemMap expires.ItemMap) {
	this.locker.Lock()
	for key := range itemMap {
		delete(this.m, key)
	}
	this.locker.Unlock()
}