Agent-System/internal/ctxmgr/cache.go

package ctxmgr

import (
	"crypto/sha256"
	"encoding/hex"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// 缓存层级常量定义
const (
	LayerC1Search    = "C1" // 检索缓存 (kb.Search / SearchChatMemory), 默认 TTL 10m
	LayerC2Rewrite   = "C2" // 改写缓存 (RewriteResult), 默认 TTL 30m
	LayerC3Embedding = "C3" // 文本向量化缓存, 默认 TTL 60m
	LayerC4Context   = "C4" // 上下文组装产物, 默认 TTL 5m (黑板/任务状态更新主动失效)
	LayerC5Summary   = "C5" // 会话滚动摘要, 默认 TTL 5m (摘要生成主动失效)

	DefaultMaxEntriesPerLayer = 1000
)

type cacheEntry struct {
	value     any
	expiresAt time.Time
	costMs    int64
	hits      int64
}

// CacheStats 缓存运行监控指标
type CacheStats struct {
	Hits    int64   `json:"hits"`
	Misses  int64   `json:"misses"`
	SavedMs int64   `json:"saved_ms"`
	HitRate float64 `json:"hit_rate"`
}

// MultiLevelCache 统一 LRU+TTL 内存多级缓存系统
type MultiLevelCache struct {
	mu           sync.RWMutex
	layers       map[string]map[string]*cacheEntry
	defaultTTLs  map[string]time.Duration
	maxEntries   int
	enabled      bool
	totalHits    atomic.Int64
	totalMisses  atomic.Int64
	totalSavedMs atomic.Int64
}

var globalCache *MultiLevelCache
var globalCacheOnce sync.Once

// GetGlobalCache 获取全局多级缓存单例
func GetGlobalCache() *MultiLevelCache {
	globalCacheOnce.Do(func() {
		globalCache = NewMultiLevelCache()
	})
	return globalCache
}

// NewMultiLevelCache 创建多级缓存实例
func NewMultiLevelCache() *MultiLevelCache {
	enabled := true
	if env := strings.ToLower(strings.TrimSpace(os.Getenv("CACHE_ENABLED"))); env == "false" || env == "0" || env == "off" {
		enabled = false
	}

	return &MultiLevelCache{
		layers: map[string]map[string]*cacheEntry{
			LayerC1Search:    make(map[string]*cacheEntry),
			LayerC2Rewrite:   make(map[string]*cacheEntry),
			LayerC3Embedding: make(map[string]*cacheEntry),
			LayerC4Context:   make(map[string]*cacheEntry),
			LayerC5Summary:   make(map[string]*cacheEntry),
		},
		defaultTTLs: map[string]time.Duration{
			LayerC1Search:    10 * time.Minute,
			LayerC2Rewrite:   30 * time.Minute,
			LayerC3Embedding: 60 * time.Minute,
			LayerC4Context:   5 * time.Minute,
			LayerC5Summary:   5 * time.Minute,
		},
		maxEntries: DefaultMaxEntriesPerLayer,
		enabled:    enabled,
	}
}

// SetEnabled 动态控制缓存开关（用于测试或运行时降级）
func (c *MultiLevelCache) SetEnabled(enabled bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.enabled = enabled
}

// IsEnabled 检查缓存是否启用
func (c *MultiLevelCache) IsEnabled() bool {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.enabled
}

// HashKey 计算多字段 SHA256 字符串指纹
func HashKey(parts ...string) string {
	h := sha256.New()
	for _, p := range parts {
		h.Write([]byte(p))
		h.Write([]byte("|"))
	}
	return hex.EncodeToString(h.Sum(nil))
}

// Get 读取指定层级的缓存
func (c *MultiLevelCache) Get(layer string, key string) (any, bool) {
	if c == nil || !c.IsEnabled() || key == "" {
		return nil, false
	}

	c.mu.RLock()
	layerMap, ok := c.layers[layer]
	if !ok {
		c.mu.RUnlock()
		c.totalMisses.Add(1)
		return nil, false
	}

	entry, found := layerMap[key]
	if !found {
		c.mu.RUnlock()
		c.totalMisses.Add(1)
		return nil, false
	}

	if time.Now().After(entry.expiresAt) {
		c.mu.RUnlock()
		// 惰性驱逐
		c.Invalidate(layer, key)
		c.totalMisses.Add(1)
		return nil, false
	}

	c.mu.RUnlock()

	entry.hits++
	c.totalHits.Add(1)
	if entry.costMs > 0 {
		c.totalSavedMs.Add(entry.costMs)
	}
	return entry.value, true
}

// Set 写入指定层级的缓存
func (c *MultiLevelCache) Set(layer string, key string, val any, costMs int64) {
	if c == nil || !c.IsEnabled() || key == "" || val == nil {
		return
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	layerMap, ok := c.layers[layer]
	if !ok {
		layerMap = make(map[string]*cacheEntry)
		c.layers[layer] = layerMap
	}

	// 容量保护：超过单层最大容量时清除超限旧条目
	if len(layerMap) >= c.maxEntries {
		now := time.Now()
		for k, v := range layerMap {
			if now.After(v.expiresAt) {
				delete(layerMap, k)
			}
		}
		if len(layerMap) >= c.maxEntries {
			// 若无过期则随机驱逐10%
			n := 0
			for k := range layerMap {
				delete(layerMap, k)
				n++
				if n >= c.maxEntries/10 {
					break
				}
			}
		}
	}

	ttl := c.defaultTTLs[layer]
	if ttl <= 0 {
		ttl = 10 * time.Minute
	}

	layerMap[key] = &cacheEntry{
		value:     val,
		expiresAt: time.Now().Add(ttl),
		costMs:    costMs,
		hits:      0,
	}
}

// Invalidate 失效指定层级的某一个 Key
func (c *MultiLevelCache) Invalidate(layer string, key string) {
	if c == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if layerMap, ok := c.layers[layer]; ok {
		delete(layerMap, key)
	}
}

// InvalidateLayer 整体清空某个层级（如黑板更新时清空 C4）
func (c *MultiLevelCache) InvalidateLayer(layer string) {
	if c == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if _, ok := c.layers[layer]; ok {
		c.layers[layer] = make(map[string]*cacheEntry)
	}
}

// Clear 清空全部层级与统计
func (c *MultiLevelCache) Clear() {
	if c == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	for l := range c.layers {
		c.layers[l] = make(map[string]*cacheEntry)
	}
	c.totalHits.Store(0)
	c.totalMisses.Store(0)
	c.totalSavedMs.Store(0)
}

// Stats 获取缓存性能与命中统计指标
func (c *MultiLevelCache) Stats() CacheStats {
	if c == nil {
		return CacheStats{}
	}
	hits := c.totalHits.Load()
	misses := c.totalMisses.Load()
	saved := c.totalSavedMs.Load()
	rate := 0.0
	if total := hits + misses; total > 0 {
		rate = float64(hits) / float64(total)
	}
	return CacheStats{
		Hits:    hits,
		Misses:  misses,
		SavedMs: saved,
		HitRate: rate,
	}
}