diff --git a/pkg/queryfrontend/cache.go b/pkg/queryfrontend/cache.go new file mode 100644 index 0000000..eae145c --- /dev/null +++ b/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 +} diff --git a/pkg/queryfrontend/cache_test.go b/pkg/queryfrontend/cache_test.go new file mode 100644 index 0000000..f324e83 --- /dev/null +++ b/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() +}