Spaces:
Running
Running
Download internal/cache/cache.go from AmaniQuery/amaniquery-agent: direct link, hf CLI and curl.
- Browser
- Download file 16.4 kB
-
https://huggingface.co/spaces/AmaniQuery/amaniquery-agent/resolve/main/internal/cache/cache.go
- Command line
-
hf download hf://spaces/AmaniQuery/amaniquery-agent/internal/cache/cache.go
-
curl -L -o cache.go https://huggingface.co/spaces/AmaniQuery/amaniquery-agent/resolve/main/internal/cache/cache.go
16.4 kB
| // Package cache provides multi-tier caching implementation with local LRU and Redis. | |
| // Production-grade caching with L1 (in-memory LRU), L2 (Redis), and stampede protection. | |
| package cache | |
| import ( | |
| "context" | |
| "encoding/json" | |
| "fmt" | |
| "strings" | |
| "sync" | |
| "time" | |
| "github.com/redis/go-redis/v9" | |
| "golang.org/x/sync/singleflight" | |
| ) | |
| const ( | |
| DefaultRegion = "nbo" | |
| DefaultTier = "L1" | |
| KeySeparator = ":" | |
| ) | |
| var ( | |
| ErrCacheMiss = &CacheMissError{} | |
| ErrCircuitOpen = fmt.Errorf("circuit breaker is open") | |
| ) | |
| type CacheMissError struct{} | |
| func (e *CacheMissError) Error() string { return "cache miss" } | |
| type CacheMetrics struct { | |
| mu sync.RWMutex | |
| LocalHits int64 | |
| LocalMisses int64 | |
| RedisHits int64 | |
| RedisMisses int64 | |
| StampedeHits int64 | |
| } | |
| func (m *CacheMetrics) Snapshot() CacheMetrics { | |
| m.mu.RLock() | |
| defer m.mu.RUnlock() | |
| return CacheMetrics{ | |
| LocalHits: m.LocalHits, | |
| LocalMisses: m.LocalMisses, | |
| RedisHits: m.RedisHits, | |
| RedisMisses: m.RedisMisses, | |
| StampedeHits: m.StampedeHits, | |
| } | |
| } | |
| type Config struct { | |
| RedisURL string `yaml:"redisUrl"` | |
| RedisAddr string `yaml:"redisAddr"` | |
| RedisPassword string `yaml:"redisPassword"` | |
| RedisDB int `yaml:"redisDb"` | |
| LocalSize int `yaml:"localSize"` | |
| TTL time.Duration `yaml:"ttl"` | |
| MaxRetries int `yaml:"maxRetries"` | |
| PoolSize int `yaml:"poolSize"` | |
| MaxEntrySize int `yaml:"maxEntrySize"` | |
| KeyPrefix string `yaml:"keyPrefix"` | |
| Region string `yaml:"region"` | |
| // PrefixIndexEnabled maintains a per-namespace Redis SET index of cache keys | |
| // so DeleteByPrefix can invalidate in O(matched) instead of O(keyspace). | |
| // See docs/architecture/cdn-invalidation-prefix-index.md. | |
| PrefixIndexEnabled bool `yaml:"prefixIndexEnabled"` | |
| } | |
| type CachedResponse struct { | |
| Status int `json:"status"` | |
| Headers map[string]string `json:"headers"` | |
| Body json.RawMessage `json:"body"` | |
| Compressed bool `json:"compressed"` | |
| Timestamp time.Time `json:"timestamp"` | |
| } | |
| func MarshalResponse(status int, headers map[string]string, body []byte) []byte { | |
| entry := CachedResponse{ | |
| Status: status, | |
| Headers: headers, | |
| Body: json.RawMessage(body), | |
| Timestamp: time.Now(), | |
| } | |
| data, _ := json.Marshal(entry) | |
| return data | |
| } | |
| func UnmarshalResponse(data []byte) (*CachedResponse, error) { | |
| var entry CachedResponse | |
| if err := json.Unmarshal(data, &entry); err != nil { | |
| return nil, err | |
| } | |
| return &entry, nil | |
| } | |
| type MultiTierCache struct { | |
| local *LRUCache | |
| redis *redis.Client | |
| ttl time.Duration | |
| keyPrefix string | |
| maxEntrySize int | |
| region string | |
| prefixIndex bool | |
| sf singleflight.Group | |
| metrics *CacheMetrics | |
| healthy bool | |
| mu sync.RWMutex | |
| } | |
| func New(cfg Config) (*MultiTierCache, error) { | |
| region := cfg.Region | |
| if region == "" { | |
| region = DefaultRegion | |
| } | |
| keyPrefix := cfg.KeyPrefix | |
| if keyPrefix == "" { | |
| keyPrefix = fmt.Sprintf("%s%s%s%s", DefaultTier, KeySeparator, region, KeySeparator) | |
| } | |
| ttl := cfg.TTL | |
| if ttl == 0 { | |
| ttl = 5 * time.Minute | |
| } | |
| localSize := cfg.LocalSize | |
| if localSize == 0 { | |
| localSize = 10000 | |
| } | |
| poolSize := cfg.PoolSize | |
| if poolSize == 0 { | |
| poolSize = 20 | |
| } | |
| c := &MultiTierCache{ | |
| local: NewLRUCache(localSize), | |
| ttl: ttl, | |
| keyPrefix: keyPrefix, | |
| maxEntrySize: cfg.MaxEntrySize, | |
| region: region, | |
| prefixIndex: cfg.PrefixIndexEnabled, | |
| metrics: &CacheMetrics{}, | |
| healthy: false, | |
| } | |
| var redisClient *redis.Client | |
| var redisErr error | |
| switch { | |
| case cfg.RedisURL != "": | |
| var opt *redis.Options | |
| opt, redisErr = redis.ParseURL(cfg.RedisURL) | |
| if redisErr == nil { | |
| opt.MaxRetries = cfg.MaxRetries | |
| opt.PoolSize = poolSize | |
| redisClient = redis.NewClient(opt) | |
| } | |
| case cfg.RedisAddr != "": | |
| redisClient = redis.NewClient(&redis.Options{ | |
| Addr: cfg.RedisAddr, | |
| Password: cfg.RedisPassword, | |
| DB: cfg.RedisDB, | |
| MaxRetries: cfg.MaxRetries, | |
| PoolSize: poolSize, | |
| MinIdleConns: 5, | |
| DialTimeout: 5 * time.Second, | |
| ReadTimeout: 3 * time.Second, | |
| WriteTimeout: 3 * time.Second, | |
| }) | |
| } | |
| if redisClient != nil { | |
| ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) | |
| defer cancel() | |
| if err := redisClient.Ping(ctx).Err(); err != nil { | |
| redisClient.Close() | |
| redisClient = nil | |
| } else { | |
| c.healthy = true | |
| } | |
| } | |
| c.redis = redisClient | |
| return c, nil | |
| } | |
| func (c *MultiTierCache) Healthy() bool { | |
| c.mu.RLock() | |
| defer c.mu.RUnlock() | |
| return c.healthy | |
| } | |
| // Keys are stored verbatim (callers build them via NormalizeKey). The previous | |
| // implementation hashed keys a SECOND time here, making the stored key prefix- | |
| // opaque — SCAN-pattern invalidation (DeleteByPrefix, admin invalidate-pattern) | |
| // could never match and silently no-op'd. | |
| func (c *MultiTierCache) Get(ctx context.Context, key string) ([]byte, error) { | |
| if value, found := c.local.Get(key); found { | |
| c.metrics.mu.Lock() | |
| c.metrics.LocalHits++ | |
| c.metrics.mu.Unlock() | |
| EmitCacheHitsLocal(c.region) | |
| return value, nil | |
| } | |
| c.metrics.mu.Lock() | |
| c.metrics.LocalMisses++ | |
| c.metrics.mu.Unlock() | |
| EmitCacheMissesLocal(c.region) | |
| if c.redis != nil { | |
| if !c.Healthy() { | |
| return nil, ErrCacheMiss | |
| } | |
| value, err := c.redis.Get(ctx, key).Bytes() | |
| if err == nil { | |
| c.metrics.mu.Lock() | |
| c.metrics.RedisHits++ | |
| c.metrics.mu.Unlock() | |
| EmitCacheHitsRedis(c.region) | |
| // Mirror into L1 with the REMAINING Redis TTL, not the full default: | |
| // short-TTL entries (e.g. 5s error-caches) must not be re-promoted to | |
| // a 5-minute L1 lifetime, or L1 would serve data past Redis expiry. | |
| promoteTTL := c.ttl | |
| if ttl, terr := c.redis.TTL(ctx, key).Result(); terr == nil && ttl > 0 && ttl < promoteTTL { | |
| promoteTTL = ttl | |
| } | |
| c.local.Set(key, value, promoteTTL) | |
| return value, nil | |
| } | |
| c.metrics.mu.Lock() | |
| c.metrics.RedisMisses++ | |
| c.metrics.mu.Unlock() | |
| EmitCacheMissesRedis(c.region) | |
| } | |
| return nil, ErrCacheMiss | |
| } | |
| func (c *MultiTierCache) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error { | |
| if c.maxEntrySize > 0 && len(value) > c.maxEntrySize { | |
| return nil | |
| } | |
| if ttl == 0 { | |
| ttl = c.ttl | |
| } | |
| c.local.Set(key, value, ttl) | |
| if c.redis != nil && c.Healthy() { | |
| // Index maintenance is atomic with the data write (same MULTI/EXEC) so | |
| // the index can never miss a just-written key. | |
| if c.prefixIndex { | |
| if ns := c.namespaceOf(key); ns != "" { | |
| pipe := c.redis.TxPipeline() | |
| pipe.Set(ctx, key, value, ttl) | |
| pipe.SAdd(ctx, c.indexKey(ns), key) | |
| pipe.Expire(ctx, c.indexKey(ns), ttl) | |
| _, _ = pipe.Exec(ctx) | |
| return nil | |
| } | |
| } | |
| _ = c.redis.Set(ctx, key, value, ttl).Err() | |
| } | |
| return nil | |
| } | |
| func (c *MultiTierCache) Delete(ctx context.Context, key string) error { | |
| c.local.Delete(key) | |
| if c.redis != nil && c.Healthy() { | |
| if c.prefixIndex { | |
| if ns := c.namespaceOf(key); ns != "" { | |
| c.redis.SRem(ctx, c.indexKey(ns), key) | |
| } | |
| } | |
| return c.redis.Del(ctx, key).Err() | |
| } | |
| return nil | |
| } | |
| func (c *MultiTierCache) GetOrSet(ctx context.Context, key string, fn func() ([]byte, error), ttl time.Duration) ([]byte, error) { | |
| if cached, err := c.Get(ctx, key); err == nil { | |
| return cached, nil | |
| } | |
| val, err, _ := c.sf.Do(key, func() (interface{}, error) { | |
| data, fetchErr := fn() | |
| if fetchErr != nil { | |
| return nil, fetchErr | |
| } | |
| c.metrics.mu.Lock() | |
| c.metrics.StampedeHits++ | |
| c.metrics.mu.Unlock() | |
| EmitStampedeHit(c.region) | |
| _ = c.Set(ctx, key, data, ttl) | |
| return data, nil | |
| }) | |
| if err != nil { | |
| return nil, err | |
| } | |
| return val.([]byte), nil | |
| } | |
| // DeleteByPrefix removes every cached entry whose key starts with prefix — from | |
| // both the local L1 tier and Redis. prefix may be a full key prefix | |
| // ("L1:nbo:rag:") or a bare tail ("rag:") resolved against this cache's | |
| // keyPrefix. An empty prefix purges the whole cache keyspace. | |
| // | |
| // When the prefix index is enabled the deletion runs in O(matched) via the | |
| // per-namespace SET index (single atomic Lua call); SCAN remains the fallback | |
| // for patterns with no index entry (e.g. keys written before the flag). | |
| func (c *MultiTierCache) DeleteByPrefix(ctx context.Context, prefix string) error { | |
| pattern := c.prefixPattern(prefix) // glob for Redis SCAN | |
| localPrefix := strings.TrimSuffix(pattern, "*") // literal prefix for L1 | |
| // Always clear the local L1 tier, regardless of Redis health. | |
| c.local.DeletePrefix(localPrefix) | |
| if c.redis == nil || !c.Healthy() { | |
| return nil | |
| } | |
| if c.prefixIndex { | |
| handled, err := c.deleteByPrefixIndexed(ctx, prefix, localPrefix) | |
| if err != nil { | |
| return err | |
| } | |
| if handled { | |
| return nil | |
| } | |
| } | |
| // SCAN fallback (also the default path when the index is disabled). | |
| var cursor uint64 | |
| for { | |
| keys, next, err := c.redis.Scan(ctx, cursor, pattern, 100).Result() | |
| if err != nil { | |
| return err | |
| } | |
| if len(keys) > 0 { | |
| if err := c.redis.Del(ctx, keys...).Err(); err != nil { | |
| return err | |
| } | |
| } | |
| cursor = next | |
| if cursor == 0 { | |
| break | |
| } | |
| } | |
| return nil | |
| } | |
| // indexKey returns the Redis SET key indexing the cache keys of a namespace. | |
| func (c *MultiTierCache) indexKey(namespace string) string { | |
| return "idx:cache:" + c.region + ":" + namespace | |
| } | |
| // namespaceOf extracts the logical namespace (the segment right after the key | |
| // prefix) from a readable key, e.g. "L1:nbo:rag:..." -> "rag". | |
| func (c *MultiTierCache) namespaceOf(key string) string { | |
| rest := strings.TrimPrefix(key, c.keyPrefix) | |
| if i := strings.Index(rest, ":"); i > 0 { | |
| return rest[:i] | |
| } | |
| return rest | |
| } | |
| // namespaceOfPrefix maps an invalidation prefix to its index namespace. Bare | |
| // tails ("rag:") and full-key prefixes ("L1:nbo:rag:") both resolve to "rag"; | |
| // empty/keyspace-wide prefixes have no namespace (whole-cache SCAN purge). | |
| func (c *MultiTierCache) namespaceOfPrefix(prefix string) string { | |
| if prefix == "" { | |
| return "" | |
| } | |
| if strings.HasPrefix(prefix, c.keyPrefix) { | |
| rest := strings.TrimPrefix(prefix, c.keyPrefix) | |
| if i := strings.Index(rest, ":"); i > 0 { | |
| return rest[:i] | |
| } | |
| return "" | |
| } | |
| if i := strings.Index(prefix, ":"); i > 0 { | |
| return prefix[:i] | |
| } | |
| return prefix | |
| } | |
| // prefixIndexScript atomically deletes cache keys matching a literal prefix and | |
| // removes them from the namespace index in a single round trip. DEL/SREM are | |
| // chunked because Lua's unpack() overflows past ~8000 arguments — a hot legal | |
| // namespace can hold far more keys than that. | |
| var prefixIndexScript = redis.NewScript(` | |
| local members = redis.call('SMEMBERS', KEYS[1]) | |
| local batch = {} | |
| local removed = 0 | |
| local function flush() | |
| if #batch > 0 then | |
| redis.call('DEL', unpack(batch)) | |
| redis.call('SREM', KEYS[1], unpack(batch)) | |
| removed = removed + #batch | |
| batch = {} | |
| end | |
| end | |
| for _, k in ipairs(members) do | |
| if string.find(k, ARGV[1], 1, true) == 1 then | |
| table.insert(batch, k) | |
| if #batch >= 500 then | |
| flush() | |
| end | |
| end | |
| end | |
| flush() | |
| return removed | |
| `) | |
| // deleteByPrefixIndexed removes matching keys through the per-namespace SET | |
| // index in one atomic Lua call. Returns handled=true when the index was | |
| // authoritative (nothing left for SCAN to find); false asks the caller to fall | |
| // back to SCAN (unknown namespace with no index entry). | |
| func (c *MultiTierCache) deleteByPrefixIndexed(ctx context.Context, prefix, localPrefix string) (bool, error) { | |
| ns := c.namespaceOfPrefix(prefix) | |
| if ns == "" { | |
| return false, nil | |
| } | |
| idxKey := c.indexKey(ns) | |
| // Match keys by their RESOLVED literal prefix so bare tails and full-key | |
| // prefixes behave identically to the SCAN path. | |
| n, err := prefixIndexScript.Run(ctx, c.redis, []string{idxKey}, localPrefix).Int() | |
| if err != nil { | |
| // Script failure must not leave stale cache: fall back to SCAN so the | |
| // invalidation still happens (slower, but correct). | |
| return false, nil | |
| } | |
| exists, err := c.redis.Exists(ctx, idxKey).Result() | |
| if err != nil { | |
| return false, nil | |
| } | |
| // n>0: matched & removed. exists>0: index is live with no matches | |
| // (authoritative no-op). Both absent: namespace never indexed — fall back. | |
| return n > 0 || exists > 0, nil | |
| } | |
| // prefixPattern resolves a caller-supplied prefix into a SCAN glob. | |
| func (c *MultiTierCache) prefixPattern(prefix string) string { | |
| if prefix == "" { | |
| return c.keyPrefix + "*" | |
| } | |
| if strings.HasPrefix(prefix, c.keyPrefix) { | |
| return prefix + "*" | |
| } | |
| return c.keyPrefix + prefix + "*" | |
| } | |
| // NormalizeKey builds a READABLE, prefix-preserving key from parts. Keys are | |
| // stored verbatim. The previous implementation hashed the whole key, which | |
| // destroyed the logical prefix — pattern invalidation (DeleteByPrefix, admin | |
| // invalidate-pattern) could never match and silently no-op'd. Readable keys | |
| // make prefix/pattern invalidation work end-to-end (queries, users, models are | |
| // all addressable via glob) and are debuggable in Redis. Redis supports keys up | |
| // to 512MB, and typical cache keys here are well under 1KB. | |
| func (c *MultiTierCache) NormalizeKey(parts ...string) string { | |
| var sb strings.Builder | |
| sb.WriteString(c.keyPrefix) | |
| for i, p := range parts { | |
| if i > 0 { | |
| sb.WriteString(KeySeparator) | |
| } | |
| sb.WriteString(strings.TrimSpace(p)) | |
| } | |
| return sb.String() | |
| } | |
| func (c *MultiTierCache) GetJSON(ctx context.Context, key string, v interface{}) error { | |
| data, err := c.Get(ctx, key) | |
| if err != nil { | |
| return err | |
| } | |
| return json.Unmarshal(data, v) | |
| } | |
| func (c *MultiTierCache) SetJSON(ctx context.Context, key string, v interface{}, ttl time.Duration) error { | |
| data, err := json.Marshal(v) | |
| if err != nil { | |
| return err | |
| } | |
| return c.Set(ctx, key, data, ttl) | |
| } | |
| func (c *MultiTierCache) GetMetrics() CacheMetrics { | |
| return c.metrics.Snapshot() | |
| } | |
| func (c *MultiTierCache) Close() error { | |
| if c.redis != nil { | |
| return c.redis.Close() | |
| } | |
| return nil | |
| } | |
| // LRUCache implements a simple LRU cache with per-entry TTL. Entries are | |
| // evicted when they expire, so L1 can never serve data past its TTL (previously | |
| // L1 held entries indefinitely until capacity eviction — stale responses could | |
| // be served long after their Redis expiry). | |
| type LRUCache struct { | |
| mu sync.RWMutex | |
| capacity int | |
| items map[string]*lruItem | |
| head *lruItem | |
| tail *lruItem | |
| } | |
| type lruItem struct { | |
| key string | |
| value []byte | |
| expiresAt time.Time | |
| prev *lruItem | |
| next *lruItem | |
| } | |
| func NewLRUCache(capacity int) *LRUCache { | |
| return &LRUCache{ | |
| capacity: capacity, | |
| items: make(map[string]*lruItem), | |
| } | |
| } | |
| func (c *LRUCache) Get(key string) ([]byte, bool) { | |
| c.mu.Lock() | |
| defer c.mu.Unlock() | |
| item, found := c.items[key] | |
| if !found { | |
| return nil, false | |
| } | |
| if !item.expiresAt.IsZero() && time.Now().After(item.expiresAt) { | |
| c.removeItem(item) | |
| delete(c.items, key) | |
| return nil, false | |
| } | |
| c.moveToFront(item) | |
| return item.value, true | |
| } | |
| func (c *LRUCache) Set(key string, value []byte, ttl time.Duration) { | |
| c.mu.Lock() | |
| defer c.mu.Unlock() | |
| var expiresAt time.Time | |
| if ttl > 0 { | |
| expiresAt = time.Now().Add(ttl) | |
| } | |
| if item, found := c.items[key]; found { | |
| item.value = value | |
| item.expiresAt = expiresAt | |
| c.moveToFront(item) | |
| return | |
| } | |
| item := &lruItem{key: key, value: value, expiresAt: expiresAt} | |
| c.items[key] = item | |
| c.addToFront(item) | |
| if len(c.items) > c.capacity { | |
| c.evictLRU() | |
| } | |
| } | |
| func (c *LRUCache) Delete(key string) { | |
| c.mu.Lock() | |
| defer c.mu.Unlock() | |
| if item, found := c.items[key]; found { | |
| c.removeItem(item) | |
| delete(c.items, key) | |
| } | |
| } | |
| // DeletePrefix removes all entries whose key starts with prefix. | |
| func (c *LRUCache) DeletePrefix(prefix string) { | |
| c.mu.Lock() | |
| defer c.mu.Unlock() | |
| for key, item := range c.items { | |
| if strings.HasPrefix(key, prefix) { | |
| c.removeItem(item) | |
| delete(c.items, key) | |
| } | |
| } | |
| } | |
| func (c *LRUCache) moveToFront(item *lruItem) { | |
| if item == c.head { | |
| return | |
| } | |
| c.removeItem(item) | |
| c.addToFront(item) | |
| } | |
| func (c *LRUCache) addToFront(item *lruItem) { | |
| item.prev = nil | |
| item.next = c.head | |
| if c.head != nil { | |
| c.head.prev = item | |
| } | |
| c.head = item | |
| if c.tail == nil { | |
| c.tail = item | |
| } | |
| } | |
| func (c *LRUCache) removeItem(item *lruItem) { | |
| if item.prev != nil { | |
| item.prev.next = item.next | |
| } else { | |
| c.head = item.next | |
| } | |
| if item.next != nil { | |
| item.next.prev = item.prev | |
| } else { | |
| c.tail = item.prev | |
| } | |
| } | |
| func (c *LRUCache) evictLRU() { | |
| if c.tail == nil { | |
| return | |
| } | |
| delete(c.items, c.tail.key) | |
| c.removeItem(c.tail) | |
| } | |