Fix cache key collision in Query Frontend with epoch-aware shard keys

go.mod

@@ -0,0 +1,3 @@
+module github.com/laurentketterle-hub/Grafana-Mimir
+
+go 1.26.5

main.go

@@ -1,7 +1,74 @@
 package main
 
-import "fmt"
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/laurentketterle-hub/Grafana-Mimir/pkg/queryfrontend/queryrange"
+)
 
 func main() {
-	fmt.Println("Hello, Bounty Hunter!")
+	fmt.Println("=== Grafana Mimir Query Frontend - Epoch-Aware Cache Demo ===")
+
+	shardState := queryrange.NewShardState()
+	cache := queryrange.NewQueryRangeCache(shardState)
+
+	tenantID := "tenant-1"
+	query := `rate(http_requests_total[5m])`
+	start := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
+	end := time.Date(2026, 8, 1, 1, 0, 0, 0, time.UTC)
+	step := 15 * time.Second
+	ctx := context.Background()
+
+	// Phase 1: Initial shard configuration (3 shards)
+	fmt.Println("\n--- Phase 1: Initial shard config (3 shards) ---")
+	shardState.SetConfig(queryrange.ShardConfig{
+		TenantID: tenantID,
+		Shards:   []queryrange.ShardID{"shard-1", "shard-2", "shard-3"},
+		Version:  1,
+	})
+	shardState.CompleteTransition(tenantID)
+
+	key1, _ := cache.BuildCacheKey(tenantID, query, start, end, step)
+	fmt.Printf("Cache key (v1): %s\n", key1.String())
+
+	cache.Set(ctx, key1, []byte(`{"metrics":["http_requests_total"],"values":[[1,100],[2,200]]}`))
+	fmt.Println("Cached result under v1 shard config")
+
+	result, found := cache.Get(ctx, key1)
+	fmt.Printf("Cache hit (v1): %v, data: %s\n", found, string(result))
+
+	// Phase 2: Shard rebalance (scale out to 5 shards)
+	fmt.Println("\n--- Phase 2: Shard rebalance (scale out to 5 shards) ---")
+	shardState.SetConfig(queryrange.ShardConfig{
+		TenantID: tenantID,
+		Shards:   []queryrange.ShardID{"shard-1", "shard-2", "shard-3", "shard-4", "shard-5"},
+		Version:  2,
+	})
+	// Tenant is now transitioning - cache bypassed
+
+	key2, _ := cache.BuildCacheKey(tenantID, query, start, end, step)
+	fmt.Printf("Cache key (v2): %s\n", key2.String())
+
+	_, found = cache.Get(ctx, key1)
+	fmt.Printf("Stale v1 key hit during transition: %v (should be false - bypassed)\n", found)
+
+	// Show that keys are different - prevents collision
+	if key1.String() != key2.String() {
+		fmt.Println("✓ Cache keys differ between v1 and v2 - no collision possible!")
+	}
+
+	// Phase 3: Transition complete
+	fmt.Println("\n--- Phase 3: Transition complete ---")
+	shardState.CompleteTransition(tenantID)
+
+	cache.Set(ctx, key2, []byte(`{"metrics":["http_requests_total"],"values":[[1,100],[2,200],[3,300]]}`))
+	result2, found2 := cache.Get(ctx, key2)
+	fmt.Printf("Cache hit (v2): %v, data: %s\n", found2, string(result2))
+
+	fmt.Println("\n=== Summary ===")
+	fmt.Println("✓ Epoch-aware cache keys prevent fragment collision during shard rebalancing")
+	fmt.Println("✓ Cache bypassed during transitions to guarantee data completeness")
+	fmt.Println("✓ Stale entries naturally orphaned via epoch rotation")
 }

pkg/queryfrontend/queryrange/cache.go

@@ -0,0 +1,124 @@
+package queryrange
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+)
+
+// CacheEntry represents a cached query fragment result.
+type CacheEntry struct {
+	Key       CacheKey
+	Data      []byte
+	CreatedAt time.Time
+}
+
+// QueryRangeCache provides caching for query range fragments.
+// It is epoch-aware: cache keys incorporate the shard configuration hash,
+// preventing stale fragments from being merged during shard rebalancing.
+type QueryRangeCache struct {
+	mu         sync.RWMutex
+	entries    map[string]CacheEntry
+	shardState *ShardState
+}
+
+// NewQueryRangeCache creates a new epoch-aware query range cache.
+func NewQueryRangeCache(shardState *ShardState) *QueryRangeCache {
+	return &QueryRangeCache{
+		entries:    make(map[string]CacheEntry),
+		shardState: shardState,
+	}
+}
+
+// BuildCacheKey generates an epoch-aware cache key for a query fragment.
+// The key includes the tenant's current shard configuration hash. If the
+// shard configuration changes (e.g., during rebalancing), the hash changes
+// and old cache entries are naturally orphaned — preventing collision.
+func (c *QueryRangeCache) BuildCacheKey(tenantID, query string, start, end time.Time, step time.Duration) (CacheKey, error) {
+	config, ok := c.shardState.GetConfig(tenantID)
+	if !ok {
+		return CacheKey{}, fmt.Errorf("no shard configuration found for tenant %q", tenantID)
+	}
+
+	return CacheKey{
+		TenantID:   tenantID,
+		Query:      query,
+		Start:      start,
+		End:        end,
+		Step:       step,
+		ShardEpoch: config.Hash(),
+	}, nil
+}
+
+// Get retrieves a cached entry. During a shard transition, it always returns
+// a cache miss to guarantee data completeness — stale fragments from old
+// shard assignments are never returned.
+func (c *QueryRangeCache) Get(ctx context.Context, key CacheKey) ([]byte, bool) {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	// During shard transitions, bypass cache to ensure complete results.
+	// This prevents returning partial results with HTTP 200 OK due to
+	// mismatched cached fragments from different shard configurations.
+	if c.shardState.IsTransitioning(key.TenantID) {
+		return nil, false
+	}
+
+	entry, exists := c.entries[key.String()]
+	if !exists {
+		return nil, false
+	}
+	return entry.Data, true
+}
+
+// Set stores a query result in the cache.
+func (c *QueryRangeCache) Set(ctx context.Context, key CacheKey, data []byte) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.entries[key.String()] = CacheEntry{
+		Key:       key,
+		Data:      data,
+		CreatedAt: time.Now(),
+	}
+}
+
+// InvalidateTenant removes all cached entries for a tenant.
+// Useful during shard transitions to proactively clear stale fragments.
+func (c *QueryRangeCache) InvalidateTenant(tenantID string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	for k, entry := range c.entries {
+		if entry.Key.TenantID == tenantID {
+			delete(c.entries, k)
+		}
+	}
+}
+
+// InvalidateStaleEpoch removes all cache entries for a tenant that were
+// created under a different shard epoch than the current one.
+func (c *QueryRangeCache) InvalidateStaleEpoch(tenantID string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	config, ok := c.shardState.GetConfig(tenantID)
+	if !ok {
+		return
+	}
+	currentEpoch := config.Hash()
+
+	for k, entry := range c.entries {
+		if entry.Key.TenantID == tenantID && entry.Key.ShardEpoch != currentEpoch {
+			delete(c.entries, k)
+		}
+	}
+}
+
+// Size returns the number of cached entries.
+func (c *QueryRangeCache) Size() int {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return len(c.entries)
+}