HATracker: Add a local cache warmup on start

Signed-off-by: SungJin1212 <tjdwls1201@gmail.com>

Author: SungJin1212

CHANGELOG.md

@@ -13,6 +13,7 @@
 * [FEATURE] Distributor: Add a per-tenant flag `-distributor.enable-type-and-unit-labels` that enables adding `__unit__` and `__type__` labels for remote write v2 and OTLP requests. This is a breaking change; the `-distributor.otlp.enable-type-and-unit-labels` flag is now deprecated, operates as a no-op, and has been consolidated into this new flag. #7077
 * [FEATURE] Querier: Add experimental projection pushdown support in Parquet Queryable. #7152
 * [FEATURE] Ingester: Add experimental active series queried metric. #7173
+* [ENHANCEMENT] HATracker: Add a local cache warmup on startup to prevent KV store operations. #7213
 * [ENHANCEMENT] Ingester: Add support for ingesting Native Histogram with Custom Buckets. #7191
 * [ENHANCEMENT] Ingester: Optimize labels out-of-order (ooo) check by allowing the iteration to terminate immediately upon finding the first unsorted label. #7186
 * [ENHANCEMENT] Distributor: Skip attaching `__unit__` and `__type__` labels when `-distributor.enable-type-and-unit-labels` is enabled, as these are appended from metadata. #7145

pkg/ha/ha_tracker.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"flag"
 	"fmt"
+	"maps"
 	"math/rand"
 	"slices"
 	"strings"
@@ -222,10 +223,84 @@ func NewHATracker(cfg HATrackerConfig, limits HATrackerLimits, trackerStatusConf
 		t.client = client
 	}
 
-	t.Service = services.NewBasicService(nil, t.loop, nil)
+	t.Service = services.NewBasicService(t.syncKVStoreToLocalMap, t.loop, nil)
 	return t, nil
 }
 
+// syncKVStoreToLocalMap warms up the local cache by fetching all active entries from the KV store.
+func (c *HATracker) syncKVStoreToLocalMap(ctx context.Context) error {
+	if !c.cfg.EnableHATracker {
+		return nil
+	}
+
+	start := time.Now()
+	level.Info(c.logger).Log("msg", "starting HA tracker cache warmup")
+
+	keys, err := c.client.List(ctx, "")
+	if err != nil {
+		level.Error(c.logger).Log("msg", "failed to list keys during HA tracker cache warmup", "err", err)
+		return err
+	}
+
+	if len(keys) == 0 {
+		level.Info(c.logger).Log("msg", "HA tracker cache warmup finished", "reason", "no keys found in KV store")
+		return nil
+	}
+
+	// create temporarily map
+	tempElected := make(map[string]ReplicaDesc, len(keys))
+	tempReplicaGroups := make(map[string]map[string]struct{})
+	successCount := 0
+
+	for _, key := range keys {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
+		val, err := c.client.Get(ctx, key)
+		if err != nil {
+			level.Warn(c.logger).Log("msg", "failed to fetch key during cache warmup", "key", key, "err", err)
+			continue
+		}
+
+		desc, ok := val.(*ReplicaDesc)
+		if !ok || desc == nil || desc.DeletedAt > 0 {
+			continue
+		}
+
+		user, cluster, keyHasSeparator := strings.Cut(key, "/")
+		if !keyHasSeparator {
+			continue
+		}
+
+		tempElected[key] = *desc
+		if tempReplicaGroups[user] == nil {
+			tempReplicaGroups[user] = make(map[string]struct{})
+		}
+		tempReplicaGroups[user][cluster] = struct{}{}
+		successCount++
+	}
+
+	c.electedLock.Lock()
+
+	// Update local map
+	maps.Copy(c.elected, tempElected)
+	for user, clusters := range tempReplicaGroups {
+		if c.replicaGroups[user] == nil {
+			c.replicaGroups[user] = make(map[string]struct{})
+		}
+		for cluster := range clusters {
+			c.replicaGroups[user][cluster] = struct{}{}
+		}
+	}
+	c.electedLock.Unlock()
+
+	c.updateUserReplicaGroupCount()
+
+	level.Info(c.logger).Log("msg", "HA tracker cache warmup completed", "duration", time.Since(start), "synced keys", successCount)
+	return nil
+}
+
 // Follows pattern used by ring for WatchKey.
 func (c *HATracker) loop(ctx context.Context) error {
 	if !c.cfg.EnableHATracker {

pkg/ha/ha_tracker_test.go

@@ -758,6 +758,108 @@ func TestCheckReplicaCleanup(t *testing.T) {
 	))
 }
 
+func TestHATracker_CacheWarmupOnStart(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	reg := prometheus.NewPedanticRegistry()
+
+	codec := GetReplicaDescCodec()
+	kvStore, closer := consul.NewInMemoryClient(codec, log.NewNopLogger(), nil)
+	t.Cleanup(func() { assert.NoError(t, closer.Close()) })
+
+	mockKV := kv.PrefixClient(kvStore, "prefix")
+
+	// CAS valid entry
+	user1 := "user1"
+	clusterUser1 := "clusterUser1"
+	key1 := fmt.Sprintf("%s/%s", user1, clusterUser1)
+	desc1 := &ReplicaDesc{
+		Replica:    "replica-0",
+		ReceivedAt: timestamp.FromTime(time.Now()),
+	}
+
+	err := mockKV.CAS(ctx, key1, func(_ any) (any, bool, error) {
+		return desc1, true, nil
+	})
+	require.NoError(t, err)
+
+	user2 := "user2"
+	clusterUser2 := "clusterUser2"
+	key2 := fmt.Sprintf("%s/%s", user2, clusterUser2)
+	desc2 := &ReplicaDesc{
+		Replica:    "replica-0",
+		ReceivedAt: timestamp.FromTime(time.Now()),
+	}
+	err = mockKV.CAS(ctx, key2, func(_ any) (any, bool, error) {
+		return desc2, true, nil
+	})
+	require.NoError(t, err)
+
+	// CAS deleted entry
+	clusterDeleted := "clusterDeleted"
+	keyDeleted := fmt.Sprintf("%s/%s", user1, clusterDeleted)
+	descDeleted := &ReplicaDesc{
+		Replica:    "replica-old",
+		ReceivedAt: timestamp.FromTime(time.Now()),
+		DeletedAt:  timestamp.FromTime(time.Now()), // Marked as deleted
+	}
+	err = mockKV.CAS(ctx, keyDeleted, func(_ any) (any, bool, error) {
+		return descDeleted, true, nil
+	})
+	require.NoError(t, err)
+
+	cfg := HATrackerConfig{
+		EnableHATracker:        true,
+		KVStore:                kv.Config{Mock: mockKV}, // Use the seeded KV
+		UpdateTimeout:          time.Second,
+		UpdateTimeoutJitterMax: 0,
+		FailoverTimeout:        time.Second,
+	}
+
+	tracker, err := NewHATracker(cfg, trackerLimits{maxReplicaGroups: 100}, haTrackerStatusConfig, prometheus.WrapRegistererWithPrefix("cortex_", reg), "test-ha-tracker", log.NewNopLogger())
+	require.NoError(t, err)
+
+	// Start ha tracker
+	require.NoError(t, services.StartAndAwaitRunning(ctx, tracker))
+	defer services.StopAndAwaitTerminated(ctx, tracker) // nolint:errcheck
+
+	tracker.electedLock.Lock()
+	// Check local cache updated
+	desc1Cached, ok := tracker.elected[key1]
+	require.True(t, ok)
+	require.Equal(t, desc1.Replica, desc1Cached.Replica)
+
+	_, ok = tracker.elected[keyDeleted]
+	require.False(t, ok)
+
+	desc2Cached, ok := tracker.elected[key2]
+	require.True(t, ok)
+	require.Equal(t, desc2.Replica, desc2Cached.Replica)
+
+	// user1 should have 1 group (clusterUser1), ignoring clusterDeleted
+	require.NotNil(t, tracker.replicaGroups[user1])
+	require.Equal(t, 1, len(tracker.replicaGroups[user1]))
+	_, hasClusterUser1 := tracker.replicaGroups[user1][clusterUser1]
+	require.True(t, hasClusterUser1)
+
+	// user2 should have 1 group (clusterUser2), ignoring clusterDeleted
+	require.NotNil(t, tracker.replicaGroups[user2])
+	require.Equal(t, 1, len(tracker.replicaGroups[user2]))
+	_, hasClusterUser2 := tracker.replicaGroups[user2][clusterUser2]
+	require.True(t, hasClusterUser2)
+
+	tracker.electedLock.Unlock()
+
+	// Check metric updated
+	require.NoError(t, testutil.GatherAndCompare(reg, strings.NewReader(`
+		# HELP cortex_ha_tracker_user_replica_group_count Number of HA replica groups tracked for each user.
+		# TYPE cortex_ha_tracker_user_replica_group_count gauge
+		cortex_ha_tracker_user_replica_group_count{user="user1"} 1
+		cortex_ha_tracker_user_replica_group_count{user="user2"} 1
+	`), "cortex_ha_tracker_user_replica_group_count",
+	))
+}
+
 func checkUserReplicaGroups(t *testing.T, duration time.Duration, c *HATracker, user string, expectedReplicaGroups int) {
 	t.Helper()
 	test.Poll(t, duration, nil, func() any {