hunjing/mcp/web-search/internal/cache/cache.go

// Package cache 提供搜索缓存（两级：内存 LRU + 可选 Redis）。
package cache

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	lru "github.com/hashicorp/golang-lru/v2"
	"github.com/redis/go-redis/v9"

	"github.com/chenxin/hunjing/mcp/web-search/internal/config"
	"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)

// entry 缓存条目。
type entry struct {
	result    *models.SearchResult
	expiresAt time.Time
}

// Manager 两级缓存管理器。
// L1: 内存 LRU（本地毫秒级，单进程共享）
// L2: Redis（跨实例共享，可按开关关闭）
type Manager struct {
	enabled bool
	ttl     time.Duration
	lru     *lru.Cache[string, *entry]
	rdb     *redis.Client
	redisDB int
	mu      sync.RWMutex
	l1Hits  int64
	l2Hits  int64
	misses  int64
}

// New 创建缓存管理器。若 Redis 启用则初始化 L2 连接。
func New(cfg config.CacheConfig) (*Manager, error) {
	if !cfg.Enabled {
		return &Manager{enabled: false}, nil
	}

	c, err := lru.New[string, *entry](cfg.MaxItems)
	if err != nil {
		return nil, fmt.Errorf("cache: create lru: %w", err)
	}

	m := &Manager{
		enabled: true,
		ttl:     cfg.TTL,
		lru:     c,
	}

	// 初始化 Redis 连接（仅 enabled 时）
	if cfg.Redis.Enabled {
		m.rdb = redis.NewClient(&redis.Options{
			Addr:     cfg.Redis.Addr,
			Password: cfg.Redis.Password,
			DB:       cfg.Redis.DB,
		})
		m.redisDB = cfg.Redis.DB
		// 探测 Redis 连通性
		if err := m.rdb.Ping(context.Background()).Err(); err != nil {
			return nil, fmt.Errorf("cache: redis ping failed: %w", err)
		}
	}

	return m, nil
}

// cacheKey 生成缓存键。格式：bowen:cache:{provider}:{hash[:8]}:{max_results}
func cacheKey(req *models.SearchRequest, provider models.ProviderKind) string {
	h := sha256.Sum256([]byte(fmt.Sprintf("%s:%s:%d", provider, req.Query, req.MaxResults)))
	return fmt.Sprintf("bowen:cache:%x:%d", h[:8], req.MaxResults)
}

// Get 获取缓存结果。
// 查询顺序：L1 (LRU) → L2 (Redis) → nil
// 若 L2 命中，自动升温到 L1。
func (m *Manager) Get(req *models.SearchRequest, provider models.ProviderKind) *models.SearchResult {
	if !m.enabled {
		return nil
	}
	key := cacheKey(req, provider)

	// L1: 内存 LRU
	if result := m.getL1(key); result != nil {
		m.mu.Lock()
		m.l1Hits++
		m.mu.Unlock()
		return result
	}

	// L2: Redis（如启用）
	if m.rdb != nil {
		if result := m.getL2(key); result != nil {
			// 升温到 L1
			m.setL1(key, result)
			m.mu.Lock()
			m.l2Hits++
			m.mu.Unlock()
			return result
		}
	}

	m.mu.Lock()
	m.misses++
	m.mu.Unlock()
	return nil
}

func (m *Manager) getL1(key string) *models.SearchResult {
	m.mu.RLock()
	e, ok := m.lru.Get(key)
	m.mu.RUnlock()
	if !ok {
		return nil
	}
	if time.Now().After(e.expiresAt) {
		m.mu.Lock()
		m.lru.Remove(key)
		m.mu.Unlock()
		return nil
	}
	result := *e.result
	result.Cached = true
	return &result
}

func (m *Manager) getL2(key string) *models.SearchResult {
	data, err := m.rdb.Get(context.Background(), key).Bytes()
	if err != nil {
		return nil
	}
	var result models.SearchResult
	if err := json.Unmarshal(data, &result); err != nil {
		return nil
	}
	result.Cached = true
	return &result
}

// Set 写入缓存（L1 + L2，若 Redis 启用则同步写 L2）。
func (m *Manager) Set(req *models.SearchRequest, provider models.ProviderKind, result *models.SearchResult) {
	if !m.enabled {
		return
	}
	key := cacheKey(req, provider)

	// 写 L1
	m.setL1(key, result)

	// 写 L2（异步 fire-and-forget，不阻塞）
	if m.rdb != nil {
		go func() {
			data, err := json.Marshal(result)
			if err != nil {
				return
			}
			m.rdb.Set(context.Background(), key, data, m.ttl)
		}()
	}
}

func (m *Manager) setL1(key string, result *models.SearchResult) {
	m.mu.Lock()
	defer m.mu.Unlock()
	r := *result
	m.lru.Add(key, &entry{
		result:    &r,
		expiresAt: time.Now().Add(m.ttl),
	})
}

// Stats 缓存统计（L1 命中、L2 命中、L1+L2 总未命中）。
func (m *Manager) Stats() (l1Hits, l2Hits, misses int64) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.l1Hits, m.l2Hits, m.misses
}

// HitRate 综合缓存命中率（L1+L2 命中 / 总请求）。
func (m *Manager) HitRate() float64 {
	m.mu.RLock()
	defer m.mu.RUnlock()
	total := m.l1Hits + m.l2Hits + m.misses
	if total == 0 {
		return 0
	}
	return float64(m.l1Hits+m.l2Hits) / float64(total)
}

// SetTTL 运行时更新 TTL（仅影响后续写入的条目）。
func (m *Manager) SetTTL(d time.Duration) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.ttl = d
}

// Clear 清空缓存。
func (m *Manager) Clear() {
	if !m.enabled {
		return
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	m.lru.Purge()
	m.l1Hits = 0
	m.l2Hits = 0
	m.misses = 0
}

// Close 关闭 Redis 连接。
func (m *Manager) Close() error {
	if m.rdb != nil {
		return m.rdb.Close()
	}
	return nil
}