fix: topology-aware cache keys with shard epoch to prevent collision during rebalance

pkg/queryfrontend/cache.go

@@ -0,0 +1,71 @@
+package queryfrontend
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"fmt"
+	"sync"
+)
+
+// Request represents a PromQL query request.
+type Request struct {
+	Query    string
+	Start    int64
+	End      int64
+	Step     int64
+	TenantID string
+	Hash     string
+}
+
+// ShardConfig represents the active shard assignment for a tenant.
+type ShardConfig struct {
+	TenantID    string
+	ShardIDs    []int
+	Epoch       int64
+	RoutingHash string
+}
+
+// ShardStateManager manages shard configurations atomically.
+type ShardStateManager struct {
+	mu      sync.RWMutex
+	configs map[string]*ShardConfig
+}
+
+func NewShardStateManager() *ShardStateManager {
+	return &ShardStateManager{configs: make(map[string]*ShardConfig)}
+}
+
+func (m *ShardStateManager) GetShardConfig(tenantID string) *ShardConfig {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return m.configs[tenantID]
+}
+
+func (m *ShardStateManager) UpdateShardConfig(cfg *ShardConfig) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.configs[cfg.TenantID] = cfg
+}
+
+// GenerateKey creates a topology-aware cache key incorporating shard epoch.
+// Format: tenantID:epoch:routingHash:requestHash
+func GenerateKey(req Request, shardCfg *ShardConfig) string {
+	if shardCfg == nil {
+		return fmt.Sprintf("%s:0:default:%s", req.TenantID, hashRequest(req))
+	}
+	return fmt.Sprintf("%s:%d:%s:%s", req.TenantID, shardCfg.Epoch, shardCfg.RoutingHash, hashRequest(req))
+}
+
+func hashRequest(req Request) string {
+	h := sha256.New()
+	fmt.Fprintf(h, "%s:%d:%d:%d:%s", req.Query, req.Start, req.End, req.Step, req.Hash)
+	return hex.EncodeToString(h.Sum(nil))[:16]
+}
+
+// ValidateConsistency checks if shard config changed mid-flight.
+func ValidateConsistency(startEpoch int64, currentCfg *ShardConfig) bool {
+	if currentCfg == nil {
+		return startEpoch == 0
+	}
+	return startEpoch == currentCfg.Epoch
+}

pkg/queryfrontend/cache_test.go

@@ -0,0 +1,45 @@
+package queryfrontend
+
+import (
+	"sync"
+	"testing"
+)
+
+func TestCacheKeysDifferByShardConfig(t *testing.T) {
+	req := Request{Query: "up", Start: 1000, End: 2000, Step: 15, TenantID: "tenant-A", Hash: "abc"}
+	cfgOld := &ShardConfig{TenantID: "tenant-A", Epoch: 1, RoutingHash: "old"}
+	cfgNew := &ShardConfig{TenantID: "tenant-A", Epoch: 2, RoutingHash: "new"}
+	if GenerateKey(req, cfgOld) == GenerateKey(req, cfgNew) {
+		t.Error("keys should differ when shard config changes")
+	}
+}
+
+func TestCacheKeysSameForSameConfig(t *testing.T) {
+	req := Request{Query: "up", Start: 1000, End: 2000, Step: 15, TenantID: "t", Hash: "x"}
+	cfg := &ShardConfig{TenantID: "t", Epoch: 1, RoutingHash: "h"}
+	if GenerateKey(req, cfg) != GenerateKey(req, cfg) {
+		t.Error("keys should match for same config")
+	}
+}
+
+func TestValidateConsistency(t *testing.T) {
+	cfg := &ShardConfig{Epoch: 5}
+	if !ValidateConsistency(5, cfg) {
+		t.Error("should be consistent when epochs match")
+	}
+	if ValidateConsistency(4, cfg) {
+		t.Error("should be inconsistent when epochs differ")
+	}
+}
+
+func TestShardStateManagerConcurrent(t *testing.T) {
+	mgr := NewShardStateManager()
+	mgr.UpdateShardConfig(&ShardConfig{TenantID: "x", Epoch: 1})
+	var wg sync.WaitGroup
+	for i := 0; i < 50; i++ {
+		wg.Add(2)
+		go func() { defer wg.Done(); _ = mgr.GetShardConfig("x") }()
+		go func() { defer wg.Done(); mgr.UpdateShardConfig(&ShardConfig{TenantID: "x", Epoch: 2}) }()
+	}
+	wg.Wait()
+}