Additional tests and optimisations

Author: bajb

cache.go

@@ -14,7 +14,8 @@ type CacheMap struct {
 
 // A "thread" safe string to anything map
 type CacheMapShared struct {
-	shutdown     chan bool
+	shutdown     chan struct{}
+	ticker       *time.Ticker
 	cleanupCycle time.Duration
 	items        map[string]*Item
 	sync.RWMutex // Read Write mutex, guards access to internal map.
@@ -43,9 +44,20 @@ func (m CacheMap) Close() {
 	}
 }
 
-// Close stops the cleanup
-func (ms CacheMapShared) Close() {
-	ms.shutdown <- true
+// Close stops the cleanup background goroutine and ticker for this shard
+func (ms *CacheMapShared) Close() {
+	// Protect against double-close: close is safe only once; recover if already closed.
+	if ms.shutdown != nil {
+		select {
+		case <-ms.shutdown:
+			// already closed
+		default:
+			close(ms.shutdown)
+		}
+	}
+	if ms.ticker != nil {
+		ms.ticker.Stop()
+	}
 }
 
 // Returns shard under given key
@@ -130,14 +142,14 @@ func (m CacheMap) Remove(key string) {
 }
 
 // Removes an element from the map
-func (ms CacheMapShared) Remove(key string) {
+func (ms *CacheMapShared) Remove(key string) {
 	ms.Lock()
 	ms.remove(key)
 	ms.Unlock()
 }
 
 // Removes an element from the map
-func (ms CacheMapShared) remove(key string) {
+func (ms *CacheMapShared) remove(key string) {
 	if itm, ok := ms.items[key]; ok && itm.onDelete != nil {
 		itm.onDelete(itm)
 	}

cache_bench_test.go

@@ -0,0 +1,94 @@
+package ttlmap_test
+
+import (
+	"math/rand"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/packaged/ttlmap"
+)
+
+// helper to prepopulate cache with n keys
+func prepopulate(c ttlmap.CacheMap, n int) []string {
+	keys := make([]string, n)
+	for i := 0; i < n; i++ {
+		k := "k" + strconv.Itoa(i)
+		c.Set(k, i, nil)
+		keys[i] = k
+	}
+	return keys
+}
+
+func BenchmarkSet(b *testing.B) {
+	cache := ttlmap.New(ttlmap.WithCleanupDuration(time.Hour)) // disable cleanup overhead
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		cache.Set("key"+strconv.Itoa(i), i, nil)
+	}
+}
+
+func BenchmarkGetHit(b *testing.B) {
+	cache := ttlmap.New(ttlmap.WithCleanupDuration(time.Hour))
+	keys := prepopulate(cache, 1024)
+	// rotate through keys
+	b.ResetTimer()
+	idx := 0
+	for i := 0; i < b.N; i++ {
+		if idx == len(keys) {
+			idx = 0
+		}
+		cache.Get(keys[idx])
+		idx++
+	}
+}
+
+func BenchmarkGetMiss(b *testing.B) {
+	cache := ttlmap.New(ttlmap.WithCleanupDuration(time.Hour))
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		cache.Get("missing-" + strconv.Itoa(i))
+	}
+}
+
+func BenchmarkParallelSet(b *testing.B) {
+	cache := ttlmap.New(ttlmap.WithCleanupDuration(time.Hour))
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		// local counter to avoid contention on atomics
+		j := 0
+		for pb.Next() {
+			cache.Set("pset-"+strconv.Itoa(j), j, nil)
+			j++
+		}
+	})
+}
+
+func BenchmarkParallelGetHit(b *testing.B) {
+	cache := ttlmap.New(ttlmap.WithCleanupDuration(time.Hour))
+	keys := prepopulate(cache, 4096)
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		r := rand.New(rand.NewSource(time.Now().UnixNano()))
+		for pb.Next() {
+			cache.Get(keys[r.Intn(len(keys))])
+		}
+	})
+}
+
+func BenchmarkMixedRW(b *testing.B) {
+	cache := ttlmap.New(ttlmap.WithCleanupDuration(time.Hour))
+	keys := prepopulate(cache, 2048)
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		r := rand.New(rand.NewSource(time.Now().UnixNano()))
+		for pb.Next() {
+			if r.Intn(10) == 0 { // ~10% writes
+				k := keys[r.Intn(len(keys))]
+				cache.Set(k, r.Int(), nil)
+			} else {
+				cache.Get(keys[r.Intn(len(keys))])
+			}
+		}
+	})
+}

cleanup.go

@@ -26,15 +26,20 @@ func (ms *CacheMapShared) Cleanup() {
 }
 
 func (ms *CacheMapShared) initCleanup(dur time.Duration) {
-	ticker := time.Tick(dur)
-	go (func() {
+	// Initialize shutdown channel once
+	if ms.shutdown == nil {
+		ms.shutdown = make(chan struct{})
+	}
+	// Use NewTicker so we can Stop it later to avoid ticker leaks
+	ms.ticker = time.NewTicker(dur)
+	go func() {
 		for {
 			select {
 			case <-ms.shutdown:
 				return
-			case <-ticker:
+			case <-ms.ticker.C:
 				ms.Cleanup()
 			}
 		}
-	})()
+	}()
 }

ttlmap_test.go

@@ -0,0 +1,105 @@
+package ttlmap_test
+
+import (
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/packaged/ttlmap"
+)
+
+func TestHasAndExpiry(t *testing.T) {
+	cache := ttlmap.New(ttlmap.WithDefaultTTL(50*time.Millisecond), ttlmap.WithCleanupDuration(10*time.Millisecond))
+	if cache.Has("missing") {
+		t.Fatalf("expected Has to be false for missing key")
+	}
+	cache.Set("a", 1, nil)
+	if !cache.Has("a") {
+		t.Fatalf("expected Has to be true right after Set")
+	}
+	time.Sleep(70 * time.Millisecond)
+	if cache.Has("a") {
+		t.Fatalf("expected Has to be false after TTL expiry")
+	}
+}
+
+func TestGetExpiryAndGetItemCopy(t *testing.T) {
+	ttl := 80 * time.Millisecond
+	cache := ttlmap.New(ttlmap.WithDefaultTTL(ttl), ttlmap.WithCleanupDuration(10*time.Millisecond))
+	cache.Set("k", "v", &ttl)
+
+	exp1 := cache.GetExpiry("k")
+	if exp1 == nil {
+		t.Fatalf("expected expiry for existing key")
+	}
+	// Missing key should return nil
+	if exp := cache.GetExpiry("missing"); exp != nil {
+		t.Fatalf("expected nil expiry for missing key")
+	}
+
+	// Get a copy and Touch it; the underlying stored item should not change expiry
+	itm, ok := cache.GetItem("k")
+	if !ok || itm == nil {
+		t.Fatalf("expected GetItem to succeed")
+	}
+	// Touch the returned copy
+	itm.Touch()
+	exp2 := cache.GetExpiry("k")
+	if exp2 == nil || !exp2.Equal(*exp1) {
+		t.Fatalf("expected stored item expiry to remain unchanged after touching copy")
+	}
+}
+
+func TestFlush(t *testing.T) {
+	cache := ttlmap.New(ttlmap.WithDefaultTTL(500 * time.Millisecond))
+	cache.Set("x", 1, nil)
+	cache.Set("y", 2, nil)
+	if !cache.Has("x") || !cache.Has("y") {
+		t.Fatalf("expected keys to exist before Flush")
+	}
+	cache.Flush()
+	if cache.Has("x") || cache.Has("y") {
+		t.Fatalf("expected keys to be removed after Flush")
+	}
+}
+
+func TestCloseIdempotent(t *testing.T) {
+	cache := ttlmap.New(ttlmap.WithCleanupDuration(5 * time.Millisecond))
+	// Should not panic on repeated Close
+	cache.Close()
+	cache.Close()
+}
+
+func TestBackgroundUpdateSingleFlight(t *testing.T) {
+	cache := ttlmap.New(ttlmap.WithDefaultTTL(20 * time.Millisecond))
+	key := "sf"
+	cache.Set(key, 1, nil)
+
+	var calls int
+	var mu sync.Mutex
+	start := make(chan struct{})
+	var wg sync.WaitGroup
+	// Launch many contenders; only one updater should run at a time for the key
+	n := 50
+	wg.Add(n)
+	for i := 0; i < n; i++ {
+		go func() {
+			defer wg.Done()
+			<-start
+			cache.BackgroundUpdate(key, func() (interface{}, error) {
+				mu.Lock()
+				calls++
+				mu.Unlock()
+				time.Sleep(5 * time.Millisecond)
+				return 2, nil
+			})
+		}()
+	}
+	close(start)
+	wg.Wait()
+
+	// Only one should have executed; the rest return immediately because the key is locked
+	if calls != 1 {
+		t.Fatalf("expected exactly one updater call, got %d", calls)
+	}
+}