diff --git a/cache/memory/memory.go b/cache/memory/memory.go index e53f1ba..8951b01 100644 --- a/cache/memory/memory.go +++ b/cache/memory/memory.go @@ -46,7 +46,6 @@ type Cache struct { memoryUsage int64 maxMemorySize int64 maxCacheSize int64 - sync.RWMutex } func New(globalTTL time.Duration) *Cache { @@ -167,26 +166,21 @@ func (c *Cache) Set(key string, value []byte, ttl time.Duration) { entry.MemorySize = int64(len(key) + len(entry.Value) + approxEntryOverhead) } - // Check if this is a new entry or an update. The exists-check, counter - // update and store must be atomic together: without the lock, two - // concurrent Sets of the same brand-new key both observe it as absent and - // both increment entryCount/memoryUsage, leaving the counters inflated - // for a single stored entry. - c.Lock() - oldEntry, exists := c.entries.Load(key) - if exists { - // Update memory usage: subtract old entry size, add new entry size - oldCacheEntry := oldEntry.(CacheEntry) - atomic.AddInt64(&c.memoryUsage, -oldCacheEntry.MemorySize) + // Swap is a single atomic operation on the sync.Map: it stores the new + // entry and reports the previous value (if any) in one step, so the + // entryCount/memoryUsage deltas below always match what actually + // happened to the map, even when Set races with Delete/evict on the same + // key. This keeps Set lock-free instead of serializing every write. + prev, loaded := c.entries.Swap(key, entry) + if loaded { + // Existing entry replaced: subtract its memory size. + prevEntry := prev.(CacheEntry) + atomic.AddInt64(&c.memoryUsage, -prevEntry.MemorySize) } else { - // New entry + // New entry. atomic.AddInt64(&c.entryCount, 1) } - - // Add new entry's memory size to total atomic.AddInt64(&c.memoryUsage, entry.MemorySize) - c.entries.Store(key, entry) - c.Unlock() } func (c *Cache) Get(key string) ([]byte, bool) { @@ -282,9 +276,14 @@ func (c *Cache) CleanExpiredEntries() { c.entries.Range(func(key, value any) bool { entry := value.(CacheEntry) if entry.ExpiresAt.Before(now) { - if _, exists := c.entries.LoadAndDelete(key); exists { + // Subtract the MemorySize of the value LoadAndDelete actually + // removed, not the Range snapshot: a concurrent Set of the same + // key between the Range read and the delete would otherwise + // account for the wrong entry's size. + if removed, exists := c.entries.LoadAndDelete(key); exists { + removedEntry := removed.(CacheEntry) atomic.AddInt64(&c.entryCount, -1) - atomic.AddInt64(&c.memoryUsage, -entry.MemorySize) + atomic.AddInt64(&c.memoryUsage, -removedEntry.MemorySize) } } return true @@ -342,7 +341,16 @@ func (c *Cache) evictToFreeMemory(bytesToFree int64) { // previous code capped the candidate set at maxCacheSize/5 and stopped at // the end of that subset, so a large write could leave the cache over its // memory limit. - entries := make([]keyMemorySize, 0, int(c.maxCacheSize)) + // + // Size the slice from the live entry count, not the configured + // maxCacheSize: with a large configured limit (e.g. millions of entries) + // and a near-empty cache, pre-allocating by maxCacheSize would allocate + // far more memory than the cache actually holds. + liveEntries := atomic.LoadInt64(&c.entryCount) + if liveEntries < 0 { + liveEntries = 0 + } + entries := make([]keyMemorySize, 0, liveEntries) c.entries.Range(func(k, v any) bool { key := k.(string) entry := v.(CacheEntry) diff --git a/cache/memory/memory_additional_test.go b/cache/memory/memory_additional_test.go index 5ede35f..872f4f0 100644 --- a/cache/memory/memory_additional_test.go +++ b/cache/memory/memory_additional_test.go @@ -1,6 +1,7 @@ package libpack_cache_memory import ( + "fmt" "io" "sync" "testing" @@ -148,6 +149,95 @@ func TestMemoryCacheConcurrentSetSameNewKey_CounterNotDoubled(t *testing.T) { } } +func TestMemoryCacheConcurrentSetDeleteEvict_CountersMatchActualState(t *testing.T) { + // Regression for M4 (evictToFreeMemory pre-allocation sized from live + // entryCount, not configured maxCacheSize) and S2 (lock-free Set via + // sync.Map.Swap, so Set no longer double-accounts when it races with + // Delete/evictOldest/evictToFreeMemory/CleanExpiredEntries on the same + // key). + // + // The specific key each goroutine touches, and which operation wins a + // given race, is inherently nondeterministic. What must always hold, + // deterministically, once every goroutine has joined, is the invariant + // that entryCount/memoryUsage are derived state: they must exactly match + // the cache's real final contents. Recomputing "expected" from the + // actual final map (the known final key set) and diffing against the + // atomic counters catches any lost or double-counted update. + const ( + numKeys = 50 + numWorkers = 16 + numOpsEach = 300 + ) + + // Small maxCacheSize so Set's own eviction path (evictOldest) also fires + // under the concurrent load, in addition to the evictors driven directly + // below. + cache := NewWithSize(DefaultTestExpiration, DefaultMaxMemorySize, 40) + defer cache.Shutdown() + + keyFor := func(i int) string { return fmt.Sprintf("hammer-key-%d", i) } + + // Pre-seed so evictors have entries to act on from the start. + for i := range numKeys { + cache.Set(keyFor(i), []byte("seed-value"), DefaultTestExpiration) + } + + var wg sync.WaitGroup + for w := range numWorkers { + wg.Add(1) + go func(worker int) { + defer wg.Done() + for op := range numOpsEach { + k := keyFor((worker + op) % numKeys) + switch op % 3 { + case 0: + cache.Set(k, []byte("value"), DefaultTestExpiration) + case 1: + cache.Delete(k) + case 2: + cache.evictOldest(2) + } + } + }(w) + } + + wg.Add(2) + go func() { + defer wg.Done() + for range 100 { + cache.evictToFreeMemory(128) + } + }() + go func() { + defer wg.Done() + for range 100 { + cache.CleanExpiredEntries() + } + }() + + wg.Wait() + + // Stop the background cleanup goroutine before recomputing expected + // state, so it cannot race with the Range below (nothing has expired + // yet - DefaultTestExpiration is 5s - but this removes any doubt). + cache.Shutdown() + + var wantCount, wantMemory int64 + cache.entries.Range(func(_, v any) bool { + entry := v.(CacheEntry) + wantCount++ + wantMemory += entry.MemorySize + return true + }) + + if got := cache.CountQueries(); got != wantCount { + t.Fatalf("entryCount = %d, want %d (actual live entries after join)", got, wantCount) + } + if got := cache.GetMemoryUsage(); got != wantMemory { + t.Fatalf("memoryUsage = %d, want %d (sum of live entries' MemorySize after join)", got, wantMemory) + } +} + func TestMemoryCacheClose_IdempotentAndStillUsable(t *testing.T) { // Compile-time guard: the memory backend must satisfy io.Closer so // cache.Shutdown() can stop its cleanup goroutine (it otherwise leaks