u-desk/internal/dbclient/cache.go

package dbclient

import (
	"crypto/sha256"
	"fmt"
	"sync"
	"time"
)

// QueryCache 查询缓存
type QueryCache struct {
	items      map[string]*CachedQuery
	size       int
	ttl        time.Duration
	mu         sync.RWMutex
	stopCh     chan struct{}
	wg         sync.WaitGroup

	// 智能缓存策略
	hitRate       float64              // 缓存命中率
	hitCount      int64                // 命中次数
	missCount     int64                // 未命中次数
	evictionCount int64                // 驱逐次数
	hotQueries    map[string]bool      // 热点查询标记
	cooldowns     map[string]time.Time // 冷却时间（避免频繁驱逐）

	// 内存限制
	maxMemoryBytes int64 // 缓存最大内存（字节），默认 100MB
	usedMemory     int64 // 当前估算内存使用量
}

// NewQueryCache 创建新的查询缓存
func NewQueryCache(size int, ttl time.Duration) *QueryCache {
	cache := &QueryCache{
		items:          make(map[string]*CachedQuery),
		size:           size,
		ttl:            ttl,
		stopCh:         make(chan struct{}),
		hitRate:        0.0,
		hitCount:       0,
		missCount:      0,
		evictionCount:  0,
		hotQueries:     make(map[string]bool),
		cooldowns:      make(map[string]time.Time),
		maxMemoryBytes: 100 * 1024 * 1024, // 默认 100MB
	}

	// 启动清理协程
	cache.StartCleanup()

	// 启动统计协程
	cache.StartStatsCollection()

	return cache
}

// Get 从缓存中获取查询结果
func (c *QueryCache) Get(params QueryParams) (*CachedQuery, error) {
	key := c.generateKey(params)

	c.mu.RLock()
	item, exists := c.items[key]
	if !exists {
		c.missCount++
		_, inCooldown := c.cooldowns[key]
		if inCooldown && time.Now().Before(c.cooldowns[key]) {
			c.mu.RUnlock()
			return nil, ErrCacheCooldown
		}
		c.mu.RUnlock()
		return nil, ErrCacheNotFound
	}

	// 检查是否过期
	if time.Now().After(item.ExpiryTime) {
		if c.isHotQuery(key) {
			c.mu.RUnlock()
			c.mu.Lock()
			item.ExpiryTime = time.Now().Add(c.ttl)
			c.hitCount++
			c.markAsHot(key)
			c.mu.Unlock()
			return item, nil
		}
		c.mu.RUnlock()
		c.mu.Lock()
		delete(c.items, key)
		c.evictionCount++
		c.missCount++
		c.mu.Unlock()
		return nil, ErrCacheExpired
	}

	// 命中
	c.hitCount++
	needsMark := !c.hotQueries[key]
	c.mu.RUnlock()

	if needsMark {
		c.mu.Lock()
		c.markAsHot(key)
		c.mu.Unlock()
	}

	return item, nil
}

// Set 将查询结果存入缓存
func (c *QueryCache) Set(params QueryParams, item *CachedQuery) {
	key := c.generateKey(params)

	// 估算条目内存大小
	itemSize := c.estimateSize(params, item)

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

	// 更新统计
	c.recordQueryAttempt(key)

	// 如果超过内存限制，执行驱逐直到有空间
	for c.usedMemory+itemSize > c.maxMemoryBytes && len(c.items) > 0 {
		c.smartEvict(key)
	}

	// 如果条目数已满，执行智能驱逐
	if len(c.items) >= c.size {
		c.smartEvict(key)
	}

	// 如果已有旧条目，先减去旧的大小
	if old, exists := c.items[key]; exists {
		c.usedMemory -= c.estimateItemSize(old)
	}

	c.items[key] = item
	c.usedMemory += itemSize

	// 标记为热点查询
	c.markAsHot(key)
}

// smartEvict 智能驱逐策略
func (c *QueryCache) smartEvict(newKey string) {
	if len(c.items) == 0 {
		return
	}
	// LRU + LFU 混合策略
	var evictKey string
	var worstScore float64 = -1

	for key, item := range c.items {
		if key == newKey {
			continue
		}

		score := c.calculateEvictionScore(key, item)
		if score > worstScore {
			worstScore = score
			evictKey = key
		}
	}

	if evictKey != "" {
		if evicted, exists := c.items[evictKey]; exists {
			c.usedMemory -= c.estimateItemSize(evicted)
		}
		c.cooldowns[evictKey] = time.Now().Add(1 * time.Minute)
		delete(c.items, evictKey)
		c.evictionCount++
	}
}

// calculateEvictionScore 计算驱逐分数（越低越适合保留）
func (c *QueryCache) calculateEvictionScore(key string, item *CachedQuery) float64 {
	now := time.Now()

	// 基础分数
	score := 1.0

	// 热点查询加分（优先保留）
	if c.isHotQuery(key) {
		score -= 0.5
	}

	// 接近过期的加分（优先驱逐即将过期的）
	if item.ExpiryTime.Sub(now) < c.ttl/2 {
		score += 0.3
	}

	// 最近使用的加分（优先保留最近使用的）
	if !item.LastUsed.IsZero() {
		recency := now.Sub(item.LastUsed)
		if recency < 5*time.Minute {
			score -= 0.2
		}
	}

	return score
}

// isHotQuery 检查是否为热点查询
func (c *QueryCache) isHotQuery(key string) bool {
	return c.hotQueries[key]
}

// markAsHot 标记为热点查询
func (c *QueryCache) markAsHot(key string) {
	c.hotQueries[key] = true
}

// cleanupHotMarkers 清理热点标记
func (c *QueryCache) cleanupHotMarkers() {
	now := time.Now()
	for key := range c.hotQueries {
		// 清理超过10分钟未使用的热点标记
		if item, exists := c.items[key]; exists {
			if now.Sub(item.LastUsed) > 10*time.Minute {
				delete(c.hotQueries, key)
			}
		} else {
			delete(c.hotQueries, key)
		}
	}
}

// recordQueryAttempt 记录查询尝试
func (c *QueryCache) recordQueryAttempt(key string) {
	// 更新命中率
	c.updateHitRate()

	// 更新最后使用时间
	if item, exists := c.items[key]; exists {
		item.LastUsed = time.Now()
	}
}

// updateHitRate 更新命中率
func (c *QueryCache) updateHitRate() {
	total := c.hitCount + c.missCount
	if total > 0 {
		c.hitRate = float64(c.hitCount) / float64(total)
	}
}

// Delete 从缓存中删除指定查询
func (c *QueryCache) Delete(params QueryParams) {
	key := c.generateKey(params)

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

	if item, exists := c.items[key]; exists {
		c.usedMemory -= c.estimateItemSize(item)
		delete(c.items, key)
	}
}

// Clear 清空整个缓存
func (c *QueryCache) Clear() {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.items = make(map[string]*CachedQuery)
	c.usedMemory = 0
}

// Size 获取缓存大小
func (c *QueryCache) Size() int {
	c.mu.RLock()
	defer c.mu.RUnlock()

	return len(c.items)
}

// CleanupExpired 清理过期的缓存条目
func (c *QueryCache) CleanupExpired() {
	c.mu.Lock()
	defer c.mu.Unlock()

	now := time.Now()
	for key, item := range c.items {
		if now.After(item.ExpiryTime) {
			c.usedMemory -= c.estimateItemSize(item)
			delete(c.items, key)
		}
	}
}

// Keys 获取缓存中所有的键
func (c *QueryCache) Keys() []string {
	c.mu.RLock()
	defer c.mu.RUnlock()

	keys := make([]string, 0, len(c.items))
	for key := range c.items {
		keys = append(keys, key)
	}
	return keys
}

// Stats 获取缓存统计信息
func (c *QueryCache) Stats() CacheStats {
	c.mu.RLock()
	defer c.mu.RUnlock()

	now := time.Now()
	expired := 0
	active := 0

	for _, item := range c.items {
		if now.After(item.ExpiryTime) {
			expired++
		} else {
			active++
		}
	}

	return CacheStats{
		TotalItems:    len(c.items),
		ActiveItems:   active,
		ExpiredItems:  expired,
		Size:          c.size,
		TTL:           c.ttl,
		HitRate:       c.hitRate,
		HitCount:      c.hitCount,
		MissCount:     c.missCount,
		EvictionCount: c.evictionCount,
		HotQueries:    len(c.hotQueries),
	}
}

// generateKey 生成缓存键
func (c *QueryCache) generateKey(params QueryParams) string {
	key := fmt.Sprintf("%s|%s|%d|%d|%s|%s|%s|%v",
		params.SQL, params.Database, params.Limit, params.Offset,
		params.Table, params.Where, params.SortBy, params.IsReadOnly)
	h := sha256.Sum256([]byte(key))
	return fmt.Sprintf("%x", h)
}

// evictOldest 删除最老的缓存条目
func (c *QueryCache) evictOldest() {
	var oldestKey string
	var oldestTime time.Time

	for key, item := range c.items {
		if oldestKey == "" || item.CreatedAt.Before(oldestTime) {
			oldestKey = key
			oldestTime = item.CreatedAt
		}
	}

	if oldestKey != "" {
		delete(c.items, oldestKey)
	}
}

// StartCleanup 启动清理协程
func (c *QueryCache) StartCleanup() {
	c.wg.Add(1)
	go func() {
		defer c.wg.Done()

		ticker := time.NewTicker(c.ttl / 2) // 每 TTL/2 时间检查一次
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
				c.CleanupExpired()
				c.cleanupCooldowns() // 清理冷却时间
			case <-c.stopCh:
				return
			}
		}
	}()
}

// StartStatsCollection 启动统计收集协程
func (c *QueryCache) StartStatsCollection() {
	c.wg.Add(1)
	go func() {
		defer c.wg.Done()

		ticker := time.NewTicker(1 * time.Minute) // 每分钟收集一次统计
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
				c.updateHitRate()
				c.cleanupHotMarkers()
			case <-c.stopCh:
				return
			}
		}
	}()
}

// cleanupCooldowns 清理冷却时间
func (c *QueryCache) cleanupCooldowns() {
	now := time.Now()
	for key, cooldown := range c.cooldowns {
		if now.After(cooldown) {
			delete(c.cooldowns, key)
		}
	}
}

// Stop 停止缓存清理
func (c *QueryCache) Stop() {
	close(c.stopCh)
	c.wg.Wait()
}

// CacheStats 缓存统计信息
type CacheStats struct {
	TotalItems    int
	ActiveItems   int
	ExpiredItems  int
	Size          int
	TTL           time.Duration
	HitRate       float64
	HitCount      int64
	MissCount     int64
	EvictionCount int64
	HotQueries    int
}

// 缓存错误定义
var (
	ErrCacheNotFound = &CacheError{Message: "缓存未找到"}
	ErrCacheExpired  = &CacheError{Message: "缓存已过期"}
	ErrCacheCooldown = &CacheError{Message: "查询在冷却中"}
)

// CacheError 缓存错误
type CacheError struct {
	Message string
}

func (e *CacheError) Error() string {
	return e.Message
}

// estimateSize 估算缓存条目的内存大小（字节）
func (c *QueryCache) estimateSize(params QueryParams, item *CachedQuery) int64 {
	size := int64(len(params.SQL) + len(params.Database) + len(params.Table) +
		len(params.Where) + len(params.SortBy))
	if item != nil && item.Result != nil {
		size += c.estimateItemSize(item)
	}
	return size
}

// estimateItemSize 估算 CachedQuery 的内存大小
func (c *QueryCache) estimateItemSize(item *CachedQuery) int64 {
	if item == nil || item.Result == nil {
		return 128 // 基础结构体大小
	}
	size := int64(128) // CachedQuery 结构体基础大小
	for _, row := range item.Result.Data {
		for _, v := range row {
			switch val := v.(type) {
			case string:
				size += int64(len(val))
			case []byte:
				size += int64(len(val))
			case nil:
				// 无额外开销
			default:
				size += 64 // 其他类型的估算值
			}
		}
	}
	size += int64(len(item.Result.Columns)) * 64 // 列名估算
	return size
}