From b468eb0c1e509df063774372a5dfce216bc98132 Mon Sep 17 00:00:00 2001 From: laurentketterle-hub Date: Thu, 6 Aug 2026 17:48:24 +0200 Subject: [PATCH] Fix cache key collision in Query Frontend with epoch-aware shard keys During tenant shuffle-sharding rebalance, the Query Frontend's cache key generation did not account for the active shard configuration, causing stale cached fragments from old shard assignments to be merged with new fragments. This resulted in partial/incomplete query results returned with HTTP 200 OK. Changes: - Add ShardConfig and CacheKey types with epoch-aware hashing (types.go) - Implement QueryRangeCache middleware that incorporates shard epoch in cache keys, bypasses cache during transitions, and supports proactive stale-epoch invalidation (cache.go) - Add comprehensive unit tests verifying key rotation, transition bypass, stale invalidation, and hash determinism (cache_test.go) - Update main.go with a demo showcasing the fix The cache key now includes a hash of the tenant's shard configuration, ensuring that when shards change (scale-out/scale-in/rebalance), cache keys naturally rotate and old fragments are never mixed with new ones. Closes #1 Signed-off-by: laurentketterle-hub --- go.mod | 3 + main.go | 71 +++++++- pkg/queryfrontend/queryrange/cache.go | 124 +++++++++++++ pkg/queryfrontend/queryrange/cache_test.go | 198 +++++++++++++++++++++ pkg/queryfrontend/queryrange/types.go | 108 +++++++++++ 5 files changed, 502 insertions(+), 2 deletions(-) create mode 100644 go.mod create mode 100644 pkg/queryfrontend/queryrange/cache.go create mode 100644 pkg/queryfrontend/queryrange/cache_test.go create mode 100644 pkg/queryfrontend/queryrange/types.go diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..d427693 --- /dev/null +++ b/go.mod @@ -0,0 +1,3 @@ +module github.com/laurentketterle-hub/Grafana-Mimir + +go 1.26.5 diff --git a/main.go b/main.go index 49f4dee..f48cd5e 100644 --- a/main.go +++ b/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") } diff --git a/pkg/queryfrontend/queryrange/cache.go b/pkg/queryfrontend/queryrange/cache.go new file mode 100644 index 0000000..dd52991 --- /dev/null +++ b/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) +} diff --git a/pkg/queryfrontend/queryrange/cache_test.go b/pkg/queryfrontend/queryrange/cache_test.go new file mode 100644 index 0000000..452a4b7 --- /dev/null +++ b/pkg/queryfrontend/queryrange/cache_test.go @@ -0,0 +1,198 @@ +package queryrange + +import ( + "context" + "testing" + "time" +) + +func TestCacheKeyChangesWithShardConfig(t *testing.T) { + shardState := NewShardState() + cache := 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 + + // Set initial shard configuration: 3 shards + shardState.SetConfig(ShardConfig{ + TenantID: tenantID, + Shards: []ShardID{"shard-1", "shard-2", "shard-3"}, + Version: 1, + }) + shardState.CompleteTransition(tenantID) + + key1, err := cache.BuildCacheKey(tenantID, query, start, end, step) + if err != nil { + t.Fatalf("failed to build cache key: %v", err) + } + + // Cache a result under the old shard config + ctx := context.Background() + cache.Set(ctx, key1, []byte("result-v1")) + + // Verify the cached result can be retrieved + result, found := cache.Get(ctx, key1) + if !found { + t.Error("expected cache hit under same shard config") + } + if string(result) != "result-v1" { + t.Errorf("expected 'result-v1', got %q", string(result)) + } + + // Now change shard configuration: scale out to 5 shards + shardState.SetConfig(ShardConfig{ + TenantID: tenantID, + Shards: []ShardID{"shard-1", "shard-2", "shard-3", "shard-4", "shard-5"}, + Version: 2, + }) + + key2, err := cache.BuildCacheKey(tenantID, query, start, end, step) + if err != nil { + t.Fatalf("failed to build cache key after rebalance: %v", err) + } + + // Cache keys must differ when shard config changes + if key1.String() == key2.String() { + t.Errorf("cache keys must differ after shard rebalance! old: %s new: %s", key1.String(), key2.String()) + } + + // During transition, cache should be bypassed for data completeness + _, found = cache.Get(ctx, key1) + if found { + t.Error("during shard transition, cache should be bypassed to guarantee data completeness") + } + + // After transition completes, stale entries from old epoch remain but + // new queries use the new epoch key. InvalidateStaleEpoch can proactively + // clean up old entries. Without invalidation, old epoch entries can still + // be retrieved with their original keys (no collision with new keys). + shardState.CompleteTransition(tenantID) + + // New queries use the new key with current epoch + cache.Set(ctx, key2, []byte("result-v2")) + result2, found := cache.Get(ctx, key2) + if !found { + t.Error("expected cache hit with new shard config") + } + if string(result2) != "result-v2" { + t.Errorf("expected 'result-v2', got %q", string(result2)) + } + + // Proactive invalidation of stale epoch entries + cache.InvalidateStaleEpoch(tenantID) + _, found = cache.Get(ctx, key1) + if found { + t.Error("after InvalidateStaleEpoch, old cache key should not return results") + } +} + +func TestCacheKeyIncludesShardEpoch(t *testing.T) { + shardState := NewShardState() + tenantID := "tenant-2" + query := "sum(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 := 30 * time.Second + + shardState.SetConfig(ShardConfig{TenantID: tenantID, Shards: []ShardID{"a", "b"}, Version: 1}) + shardState.CompleteTransition(tenantID) + cache := NewQueryRangeCache(shardState) + key1, _ := cache.BuildCacheKey(tenantID, query, start, end, step) + + shardState.SetConfig(ShardConfig{TenantID: tenantID, Shards: []ShardID{"a", "b"}, Version: 3}) + shardState.CompleteTransition(tenantID) + key2, _ := cache.BuildCacheKey(tenantID, query, start, end, step) + + if key1.String() == key2.String() { + t.Error("cache keys must differ when version/epoch changes, even with same shard list") + } +} + +func TestDifferentTenantsHaveDifferentKeys(t *testing.T) { + shardState := NewShardState() + cache := NewQueryRangeCache(shardState) + query := "up" + 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 := time.Minute + + shardState.SetConfig(ShardConfig{TenantID: "tenant-A", Shards: []ShardID{"s1", "s2"}, Version: 1}) + shardState.CompleteTransition("tenant-A") + shardState.SetConfig(ShardConfig{TenantID: "tenant-B", Shards: []ShardID{"s1", "s2"}, Version: 1}) + shardState.CompleteTransition("tenant-B") + + keyA, _ := cache.BuildCacheKey("tenant-A", query, start, end, step) + keyB, _ := cache.BuildCacheKey("tenant-B", query, start, end, step) + + if keyA.String() == keyB.String() { + t.Error("different tenants must have different cache keys") + } +} + +func TestInvalidateStaleEpoch(t *testing.T) { + shardState := NewShardState() + cache := NewQueryRangeCache(shardState) + tenantID := "tenant-3" + ctx := context.Background() + + shardState.SetConfig(ShardConfig{TenantID: tenantID, Shards: []ShardID{"s1", "s2"}, Version: 1}) + shardState.CompleteTransition(tenantID) + key1, _ := cache.BuildCacheKey(tenantID, "query1", time.Now(), time.Now().Add(time.Hour), time.Minute) + cache.Set(ctx, key1, []byte("data1")) + + shardState.SetConfig(ShardConfig{TenantID: tenantID, Shards: []ShardID{"s1", "s2", "s3"}, Version: 2}) + shardState.CompleteTransition(tenantID) + + if cache.Size() != 1 { + t.Errorf("expected 1 entry before invalidation, got %d", cache.Size()) + } + cache.InvalidateStaleEpoch(tenantID) + if cache.Size() != 0 { + t.Errorf("expected 0 entries after stale epoch invalidation, got %d", cache.Size()) + } +} + +func TestCacheBypassDuringTransition(t *testing.T) { + shardState := NewShardState() + cache := NewQueryRangeCache(shardState) + tenantID := "tenant-4" + ctx := context.Background() + + shardState.SetConfig(ShardConfig{TenantID: tenantID, Shards: []ShardID{"s1"}, Version: 1}) + shardState.CompleteTransition(tenantID) + key, _ := cache.BuildCacheKey(tenantID, "query", time.Now(), time.Now().Add(time.Hour), time.Minute) + cache.Set(ctx, key, []byte("cached-data")) + + shardState.SetConfig(ShardConfig{TenantID: tenantID, Shards: []ShardID{"s1", "s2"}, Version: 2}) + // NOT calling CompleteTransition - tenant is transitioning + _, found := cache.Get(ctx, key) + if found { + t.Error("cache must be bypassed during shard transition to prevent partial results") + } + + shardState.CompleteTransition(tenantID) + newKey, _ := cache.BuildCacheKey(tenantID, "query", time.Now(), time.Now().Add(time.Hour), time.Minute) + cache.Set(ctx, newKey, []byte("new-data")) + result, found := cache.Get(ctx, newKey) + if !found { + t.Error("cache should work after transition completes") + } + if string(result) != "new-data" { + t.Errorf("expected 'new-data', got %q", string(result)) + } +} + +func TestShardConfigHashIsDeterministic(t *testing.T) { + cfg1 := ShardConfig{TenantID: "t1", Shards: []ShardID{"shard-b", "shard-a", "shard-c"}, Version: 1} + cfg2 := ShardConfig{TenantID: "t1", Shards: []ShardID{"shard-c", "shard-b", "shard-a"}, Version: 1} + if cfg1.Hash() != cfg2.Hash() { + t.Error("hash must be order-independent: same shards in different order must produce same hash") + } + cfg3 := ShardConfig{TenantID: "t1", Shards: []ShardID{"shard-a", "shard-b", "shard-c"}, Version: 2} + if cfg1.Hash() == cfg3.Hash() { + t.Error("hash must change when version changes") + } +} diff --git a/pkg/queryfrontend/queryrange/types.go b/pkg/queryfrontend/queryrange/types.go new file mode 100644 index 0000000..1135724 --- /dev/null +++ b/pkg/queryfrontend/queryrange/types.go @@ -0,0 +1,108 @@ +package queryrange + +import ( + "crypto/sha256" + "fmt" + "sort" + "strings" + "sync" + "time" +) + +// ShardID represents a single shard identifier in the tenant's shard configuration. +type ShardID string + +// ShardConfig represents the active shard assignment for a tenant. +// When this configuration changes (e.g., during scale-out/scale-in), +// the epoch changes and cache keys are naturally rotated. +type ShardConfig struct { + TenantID string + Shards []ShardID + Version int64 // monotonically increasing version/epoch +} + +// Hash returns a deterministic hash of the shard configuration. +// Changes when shards are added, removed, or reordered. +func (sc ShardConfig) Hash() string { + shardIDs := make([]string, len(sc.Shards)) + for i, s := range sc.Shards { + shardIDs[i] = string(s) + } + sort.Strings(shardIDs) + input := fmt.Sprintf("%s:%s:%d", sc.TenantID, strings.Join(shardIDs, ","), sc.Version) + h := sha256.Sum256([]byte(input)) + return fmt.Sprintf("%x", h[:8]) // first 8 bytes as hex +} + +// CacheKey represents a fully-qualified cache key for a query fragment. +// Incorporates tenant, query, time range, step, and shard configuration hash +// to prevent collisions during shard rebalancing. +type CacheKey struct { + TenantID string + Query string + Start time.Time + End time.Time + Step time.Duration + ShardEpoch string // hash of the shard configuration at time of caching +} + +// String returns the canonical cache key string. +func (ck CacheKey) String() string { + return fmt.Sprintf("%s:%s:%d:%d:%d:%s", + ck.TenantID, + ck.Query, + ck.Start.Unix(), + ck.End.Unix(), + ck.Step.Milliseconds(), + ck.ShardEpoch, + ) +} + +// ShardState tracks the current shard configuration and epoch for a tenant. +type ShardState struct { + mu sync.RWMutex + configs map[string]ShardConfig // tenantID -> config + transitioning map[string]bool // tenantID -> is transitioning +} + +// NewShardState creates a new shard state tracker. +func NewShardState() *ShardState { + return &ShardState{ + configs: make(map[string]ShardConfig), + transitioning: make(map[string]bool), + } +} + +// SetConfig updates the shard configuration for a tenant and marks it as transitioning. +func (ss *ShardState) SetConfig(config ShardConfig) { + ss.mu.Lock() + defer ss.mu.Unlock() + + old, exists := ss.configs[config.TenantID] + if !exists || old.Hash() != config.Hash() { + ss.transitioning[config.TenantID] = true + } + ss.configs[config.TenantID] = config +} + +// GetConfig returns the current shard configuration for a tenant. +func (ss *ShardState) GetConfig(tenantID string) (ShardConfig, bool) { + ss.mu.RLock() + defer ss.mu.RUnlock() + cfg, ok := ss.configs[tenantID] + return cfg, ok +} + +// IsTransitioning returns true if the tenant is currently undergoing a shard rebalance. +func (ss *ShardState) IsTransitioning(tenantID string) bool { + ss.mu.RLock() + defer ss.mu.RUnlock() + return ss.transitioning[tenantID] +} + +// CompleteTransition marks a tenant's shard transition as complete. +func (ss *ShardState) CompleteTransition(tenantID string) { + ss.mu.Lock() + defer ss.mu.Unlock() + ss.transitioning[tenantID] = false +}