缓存框架

2026-10-03 15:06:36 +08:00 · 2026-06-11 15:10:11 +08:00
parent e8e0326879
commit 0221d4de14
20 changed files with 1699 additions and 11 deletions
@@ -0,0 +1,430 @@
+// Package diskcache implements a platform-level disk-backed cache with size limit, TTL, and LRU eviction.
+package diskcache
+
+import (
+	"container/list"
+	"context"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"sort"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/Rain-kl/Wavelet/internal/db"
+	"github.com/Rain-kl/Wavelet/internal/model"
+	"github.com/peterbourgon/diskv/v3"
+)
+
+// ErrCacheMiss represents a cache miss.
+var ErrCacheMiss = errors.New("cache miss")
+
+// Constants for disk cache configuration and sizing
+const (
+	defaultCacheDir        = "uploads/diskcache"
+	headerSize             = 8 // 8 bytes metadata prefix for expiration UnixNano timestamp
+	defaultMaxSizeMB       = 100
+	defaultTTLMinutes      = 60
+	defaultCleanupInterval = 10
+)
+
+var (
+	globalCache     *DiskCache
+	globalCacheOnce sync.Once
+)
+
+// Status represents the runtime cache statistics.
+type Status struct {
+	TotalSize  int64  `json:"total_size"`
+	KeysCount  int    `json:"keys_count"`
+	MaxSizeMB  int64  `json:"max_size_mb"`
+	TTLMinutes int64  `json:"ttl_minutes"`
+	LRUEnabled bool   `json:"lru_enabled"`
+	BasePath   string `json:"base_path"`
+}
+
+// DiskCache implements the disk-backed cache with size limits, TTL, and LRU eviction.
+type DiskCache struct {
+	mu         sync.RWMutex
+	d          *diskv.Diskv
+	basePath   string
+	maxSize    int64 // in bytes
+	defaultTTL time.Duration
+	lruEnabled bool
+
+	// LRU and Size tracking
+	currentSize int64
+	items       map[string]*list.Element
+	evictList   *list.List
+}
+
+type cacheItem struct {
+	key       string
+	size      int64
+	expiredAt time.Time
+}
+
+// GetGlobalCache returns the global singleton DiskCache instance.
+func GetGlobalCache() *DiskCache {
+	globalCacheOnce.Do(func() {
+		globalCache = New(defaultCacheDir)
+		// Load initial configs from database
+		globalCache.ReloadConfig(context.Background())
+		// Start background routine to clean expired items every 10 minutes
+		go globalCache.startCleanupWorker(defaultCleanupInterval * time.Minute)
+	})
+	return globalCache
+}
+
+// New creates a new DiskCache instance.
+func New(basePath string) *DiskCache {
+	d := diskv.New(diskv.Options{
+		BasePath:     basePath,
+		Transform:    func(_ string) []string { return []string{} }, // flat structure for easy walk
+		CacheSizeMax: 1024 * 1024,                                   // 1MB in-memory cache size for diskv itself
+	})
+
+	c := &DiskCache{
+		d:          d,
+		basePath:   basePath,
+		maxSize:    defaultMaxSizeMB * 1024 * 1024,  // 100MB default
+		defaultTTL: defaultTTLMinutes * time.Minute, // 60 minutes default
+		lruEnabled: true,
+		items:      make(map[string]*list.Element),
+		evictList:  list.New(),
+	}
+
+	// Scan directory on startup to rebuild LRU and size tracking
+	_ = c.loadTracker()
+	return c
+}
+
+// Set stores a key-value pair in the cache with a specific TTL.
+// If ttl <= 0, the default cache TTL is used.
+func (c *DiskCache) Set(key string, value []byte, ttl time.Duration) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if ttl <= 0 {
+		ttl = c.defaultTTL
+	}
+
+	var expiredAt time.Time
+	if ttl > 0 {
+		expiredAt = time.Now().Add(ttl)
+	}
+
+	// Prepare data layout: 8 bytes expiration timestamp + raw payload
+	buf := make([]byte, headerSize+len(value))
+	var expNano int64
+	if !expiredAt.IsZero() {
+		expNano = expiredAt.UnixNano()
+	}
+	binary.BigEndian.PutUint64(buf[0:headerSize], uint64(expNano))
+	copy(buf[headerSize:], value)
+
+	// Write to diskv
+	if err := c.d.Write(key, buf); err != nil {
+		return fmt.Errorf("failed to write key to disk: %w", err)
+	}
+
+	// Get file size on disk (approximate)
+	size := int64(len(buf))
+
+	// Update memory tracker
+	if elem, ok := c.items[key]; ok {
+		item := elem.Value.(*cacheItem)
+		c.currentSize += size - item.size
+		item.size = size
+		item.expiredAt = expiredAt
+		c.evictList.MoveToFront(elem)
+	} else {
+		item := &cacheItem{
+			key:       key,
+			size:      size,
+			expiredAt: expiredAt,
+		}
+		elem := c.evictList.PushFront(item)
+		c.items[key] = elem
+		c.currentSize += size
+	}
+
+	// Evict items if size limit exceeded and LRU is enabled
+	c.evict()
+
+	return nil
+}
+
+// Get retrieves a key's value from the cache.
+func (c *DiskCache) Get(key string) ([]byte, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	elem, ok := c.items[key]
+	if !ok {
+		return nil, ErrCacheMiss
+	}
+
+	item := elem.Value.(*cacheItem)
+
+	// Check expiration
+	if !item.expiredAt.IsZero() && time.Now().After(item.expiredAt) {
+		// Lazily delete expired item
+		_ = c.deleteUnlocked(key)
+		return nil, ErrCacheMiss
+	}
+
+	// Read from diskv
+	data, err := c.d.Read(key)
+	if err != nil {
+		// Key exists in memory but not on disk, sync state
+		_ = c.deleteUnlocked(key)
+		return nil, ErrCacheMiss
+	}
+
+	if len(data) < headerSize {
+		_ = c.deleteUnlocked(key)
+		return nil, ErrCacheMiss
+	}
+
+	// Update LRU access order
+	c.evictList.MoveToFront(elem)
+
+	// Slice off the metadata header
+	return data[headerSize:], nil
+}
+
+// Delete removes a key-value pair from the cache.
+func (c *DiskCache) Delete(key string) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.deleteUnlocked(key)
+}
+
+func (c *DiskCache) deleteUnlocked(key string) error {
+	if elem, ok := c.items[key]; ok {
+		item := elem.Value.(*cacheItem)
+		c.currentSize -= item.size
+		c.evictList.Remove(elem)
+		delete(c.items, key)
+	}
+	return c.d.Erase(key)
+}
+
+// Clear flushes all cached elements.
+func (c *DiskCache) Clear() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.currentSize = 0
+	c.items = make(map[string]*list.Element)
+	c.evictList.Init()
+
+	return c.d.EraseAll()
+}
+
+// Status returns the cache status.
+func (c *DiskCache) Status() Status {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	return Status{
+		TotalSize:  c.currentSize,
+		KeysCount:  len(c.items),
+		MaxSizeMB:  c.maxSize / (1024 * 1024),
+		TTLMinutes: int64(c.defaultTTL.Minutes()),
+		LRUEnabled: c.lruEnabled,
+		BasePath:   c.basePath,
+	}
+}
+
+// ReloadConfig reloads policies from database configs dynamically.
+func (c *DiskCache) ReloadConfig(ctx context.Context) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	// Ensure DB is initialized before querying
+	if db.DB(ctx) == nil {
+		return
+	}
+
+	// 1. Max Size
+	var scMaxSize model.SystemConfig
+	maxSizeMB := int64(defaultMaxSizeMB)
+	if err := scMaxSize.GetByKey(ctx, model.ConfigKeyDiskCacheMaxSizeMB); err == nil && scMaxSize.Value != "" {
+		if val, err := strconv.ParseInt(scMaxSize.Value, 10, 64); err == nil && val > 0 {
+			maxSizeMB = val
+		}
+	}
+	c.maxSize = maxSizeMB * 1024 * 1024
+
+	// 2. Default TTL
+	var scTTL model.SystemConfig
+	ttlMinutes := int64(defaultTTLMinutes)
+	if err := scTTL.GetByKey(ctx, model.ConfigKeyDiskCacheTTLMinutes); err == nil && scTTL.Value != "" {
+		if val, err := strconv.ParseInt(scTTL.Value, 10, 64); err == nil && val >= 0 {
+			ttlMinutes = val
+		}
+	}
+	c.defaultTTL = time.Duration(ttlMinutes) * time.Minute
+
+	// 3. LRU Enabled
+	var scLRU model.SystemConfig
+	lruEnabled := true
+	if err := scLRU.GetByKey(ctx, model.ConfigKeyDiskCacheLRUEnabled); err == nil && scLRU.Value != "" {
+		if val, err := strconv.ParseBool(scLRU.Value); err == nil {
+			lruEnabled = val
+		}
+	}
+	c.lruEnabled = lruEnabled
+
+	// Apply eviction immediately under new configs
+	c.evict()
+}
+
+// evict evicts oldest items if current size exceeds maxSize and LRU is enabled.
+func (c *DiskCache) evict() {
+	if !c.lruEnabled {
+		return
+	}
+
+	for c.currentSize > c.maxSize && c.evictList.Len() > 0 {
+		elem := c.evictList.Back()
+		if elem == nil {
+			break
+		}
+		item := elem.Value.(*cacheItem)
+		key := item.key
+
+		// Remove from memory
+		c.currentSize -= item.size
+		c.evictList.Remove(elem)
+		delete(c.items, key)
+
+		// Delete from disk
+		_ = c.d.Erase(key)
+	}
+}
+
+// loadTracker walks the directory to rebuild the LRU and size tracking structures.
+func (c *DiskCache) loadTracker() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.currentSize = 0
+	c.items = make(map[string]*list.Element)
+	c.evictList = list.New()
+
+	files, err := os.ReadDir(c.basePath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return nil
+		}
+		return err
+	}
+
+	type tempItem struct {
+		key       string
+		size      int64
+		expiredAt time.Time
+		mtime     time.Time
+	}
+	var loadedItems []tempItem
+
+	for _, file := range files {
+		if file.IsDir() {
+			continue
+		}
+		name := file.Name()
+		// Skip temporary files
+		if strings.HasPrefix(name, ".") || strings.Contains(name, "temp") {
+			continue
+		}
+
+		info, err := file.Info()
+		if err != nil {
+			continue
+		}
+
+		filePath := filepath.Join(c.basePath, name)
+		// #nosec G304
+		f, err := os.Open(filePath)
+		if err != nil {
+			continue
+		}
+
+		var expiredAt time.Time
+		var expNano int64
+		err = binary.Read(f, binary.BigEndian, &expNano)
+		_ = f.Close()
+		if err != nil {
+			// Corrupt metadata header: delete file
+			_ = os.Remove(filePath)
+			continue
+		}
+
+		if expNano > 0 {
+			expiredAt = time.Unix(0, expNano)
+			// Expired: delete file
+			if time.Now().After(expiredAt) {
+				_ = os.Remove(filePath)
+				continue
+			}
+		}
+
+		loadedItems = append(loadedItems, tempItem{
+			key:       name,
+			size:      info.Size(),
+			expiredAt: expiredAt,
+			mtime:     info.ModTime(),
+		})
+	}
+
+	// Sort loaded items by modification time ascending (oldest first)
+	sort.Slice(loadedItems, func(i, j int) bool {
+		return loadedItems[i].mtime.Before(loadedItems[j].mtime)
+	})
+
+	// Populate LRU (PushFront so that newest items are at the front, oldest at the back)
+	for _, item := range loadedItems {
+		entry := &cacheItem{
+			key:       item.key,
+			size:      item.size,
+			expiredAt: item.expiredAt,
+		}
+		element := c.evictList.PushFront(entry)
+		c.items[item.key] = element
+		c.currentSize += item.size
+	}
+
+	return nil
+}
+
+// startCleanupWorker periodically cleans up expired cache items.
+func (c *DiskCache) startCleanupWorker(interval time.Duration) {
+	ticker := time.NewTicker(interval)
+	for range ticker.C {
+		c.cleanExpired()
+	}
+}
+
+// cleanExpired scans memory for expired items and removes them.
+func (c *DiskCache) cleanExpired() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	now := time.Now()
+	for key, elem := range c.items {
+		item := elem.Value.(*cacheItem)
+		if !item.expiredAt.IsZero() && now.After(item.expiredAt) {
+			c.currentSize -= item.size
+			c.evictList.Remove(elem)
+			delete(c.items, key)
+			_ = c.d.Erase(key)
+		}
+	}
+}
@@ -0,0 +1,190 @@
+// Copyright 2026 Arctel.net
+// SPDX-License-Identifier: Apache-2.0
+
+package diskcache
+
+import (
+	"bytes"
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/Rain-kl/Wavelet/internal/db"
+	"github.com/Rain-kl/Wavelet/internal/model"
+	"github.com/Rain-kl/Wavelet/internal/testhelper"
+)
+
+func TestDiskCacheBasic(t *testing.T) {
+	testDir := "uploads/test_diskcache_basic"
+	defer func() { _ = os.RemoveAll(testDir) }()
+	_ = os.RemoveAll(testDir)
+
+	c := New(testDir)
+	defer func() { _ = c.Clear() }()
+
+	key := "key1"
+	val := []byte("value1")
+
+	// Get non-existent
+	_, err := c.Get(key)
+	if err != ErrCacheMiss {
+		t.Fatalf("expected ErrCacheMiss, got %v", err)
+	}
+
+	// Set & Get
+	err = c.Set(key, val, 10*time.Second)
+	if err != nil {
+		t.Fatalf("failed to set cache: %v", err)
+	}
+
+	got, err := c.Get(key)
+	if err != nil {
+		t.Fatalf("failed to get cache: %v", err)
+	}
+
+	if !bytes.Equal(got, val) {
+		t.Errorf("expected %s, got %s", val, got)
+	}
+
+	// Delete
+	err = c.Delete(key)
+	if err != nil {
+		t.Fatalf("failed to delete: %v", err)
+	}
+
+	_, err = c.Get(key)
+	if err != ErrCacheMiss {
+		t.Errorf("expected ErrCacheMiss after delete, got %v", err)
+	}
+}
+
+func TestDiskCacheTTL(t *testing.T) {
+	testDir := "uploads/test_diskcache_ttl"
+	defer func() { _ = os.RemoveAll(testDir) }()
+	_ = os.RemoveAll(testDir)
+
+	c := New(testDir)
+	defer func() { _ = c.Clear() }()
+
+	key := "ttlkey"
+	val := []byte("ttlval")
+
+	// Set with 200ms TTL
+	err := c.Set(key, val, 200*time.Millisecond)
+	if err != nil {
+		t.Fatalf("failed to set: %v", err)
+	}
+
+	// Immediate Get should succeed
+	got, err := c.Get(key)
+	if err != nil {
+		t.Fatalf("failed to get: %v", err)
+	}
+	if !bytes.Equal(got, val) {
+		t.Errorf("expected %s, got %s", val, got)
+	}
+
+	// Sleep 250ms to expire
+	time.Sleep(250 * time.Millisecond)
+
+	// Get should fail with cache miss
+	_, err = c.Get(key)
+	if err != ErrCacheMiss {
+		t.Errorf("expected ErrCacheMiss after TTL expiration, got %v", err)
+	}
+}
+
+func TestDiskCacheLRUEviction(t *testing.T) {
+	testDir := "uploads/test_diskcache_lru"
+	defer func() { _ = os.RemoveAll(testDir) }()
+	_ = os.RemoveAll(testDir)
+
+	c := New(testDir)
+	defer func() { _ = c.Clear() }()
+
+	// Force a very small max size of 20 bytes for testing (8 bytes header + payload)
+	// So 2 items of 2 bytes payload = 2 * (8 + 2) = 20 bytes max.
+	c.maxSize = 20
+	c.lruEnabled = true
+
+	// Write item 1: 8 + 2 = 10 bytes
+	err := c.Set("k1", []byte("v1"), 0)
+	if err != nil {
+		t.Fatalf("failed to set k1: %v", err)
+	}
+
+	// Write item 2: 8 + 2 = 10 bytes
+	err = c.Set("k2", []byte("v2"), 0)
+	if err != nil {
+		t.Fatalf("failed to set k2: %v", err)
+	}
+
+	// Both should exist
+	if _, err := c.Get("k1"); err != nil {
+		t.Errorf("k1 should exist: %v", err)
+	}
+	if _, err := c.Get("k2"); err != nil {
+		t.Errorf("k2 should exist: %v", err)
+	}
+
+	// Write item 3: 8 + 2 = 10 bytes -> total size would be 30, exceeding 20.
+	// This should evict the oldest item. Since k1 was accessed, but then k2 was accessed,
+	// wait, let's access k1 again to make it the most recently used, so k2 becomes oldest!
+	_, _ = c.Get("k1") // k1 is now MRU, k2 is LRU
+
+	err = c.Set("k3", []byte("v3"), 0)
+	if err != nil {
+		t.Fatalf("failed to set k3: %v", err)
+	}
+
+	// k2 should be evicted, k1 and k3 should exist
+	_, err = c.Get("k2")
+	if err != ErrCacheMiss {
+		t.Errorf("expected k2 to be evicted, got error %v", err)
+	}
+
+	if _, err := c.Get("k1"); err != nil {
+		t.Errorf("k1 should still exist: %v", err)
+	}
+
+	if _, err := c.Get("k3"); err != nil {
+		t.Errorf("k3 should exist: %v", err)
+	}
+}
+
+func TestDiskCacheReloadConfig(t *testing.T) {
+	dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
+	defer cleanup()
+
+	testDir := "uploads/test_diskcache_reload"
+	defer func() { _ = os.RemoveAll(testDir) }()
+	_ = os.RemoveAll(testDir)
+
+	c := New(testDir)
+	defer func() { _ = c.Clear() }()
+
+	// Update DB config values
+	dbConn.Model(&model.SystemConfig{}).Where("key = ?", model.ConfigKeyDiskCacheMaxSizeMB).Update("value", "250")
+	dbConn.Model(&model.SystemConfig{}).Where("key = ?", model.ConfigKeyDiskCacheTTLMinutes).Update("value", "120")
+	dbConn.Model(&model.SystemConfig{}).Where("key = ?", model.ConfigKeyDiskCacheLRUEnabled).Update("value", "false")
+
+	// Invalidate Redis config cache to force DB reload
+	if db.Redis != nil {
+		db.Redis.Del(context.Background(), db.PrefixedKey(model.SystemConfigRedisHashKey))
+	}
+
+	// Reload config
+	c.ReloadConfig(context.Background())
+
+	status := c.Status()
+	if status.MaxSizeMB != 250 {
+		t.Errorf("expected MaxSizeMB to be 250, got %d", status.MaxSizeMB)
+	}
+	if status.TTLMinutes != 120 {
+		t.Errorf("expected TTLMinutes to be 120, got %d", status.TTLMinutes)
+	}
+	if status.LRUEnabled != false {
+		t.Errorf("expected LRUEnabled to be false, got %t", status.LRUEnabled)
+	}
+}