Distributor might start serving writes with an empty HA tracke…

CHANGELOG.md

@@ -28,6 +28,7 @@
 * [BUGFIX] Fix issue where all incoming HTTP requests have duplicate trace spans. #6920
 * [BUGFIX] Querier: do not retry requests to store-gateway when a query gets canceled. #6934
 * [BUGFIX] Querier: return 499 status code instead of 500 when a request to remote read endpoint gets canceled. #6934
+* [BUGFIX] Distributor: distributor might start serving writes with an empty HA tracker state. #5796
 * [BUGFIX] Querier: fix issue where `-querier.max-fetched-series-per-query` is not applied to `/series` endpoint if the series are loaded from ingesters. #7055
 
 ### Mixin

pkg/distributor/ha_tracker.go

@@ -202,10 +202,55 @@ func newHATracker(cfg HATrackerConfig, limits haTrackerLimits, reg prometheus.Re
 		t.client = client
 	}
 
-	t.Service = services.NewBasicService(nil, t.loop, nil)
+	t.Service = services.NewBasicService(t.syncHAStateOnStart, t.loop, nil)
 	return t, nil
 }
 
+// Will update the internal cache to be in sync with the KV Store
+func (h *haTracker) syncHAStateOnStart(ctx context.Context) error {
+	if !h.cfg.EnableHATracker {
+		return nil
+	}
+
+	// haTracker holds HATrackerConfig with the prefix.
+	keys, err := h.client.List(ctx, "")
+	if err != nil {
+		level.Warn(h.logger).Log("msg", "syncHAStateOnStart: failed to list replica keys", "err", err)
+		return err
+	}
+
+	if len(keys) == 0 {
+		level.Warn(h.logger).Log("msg", "syncHAStateOnStart: no keys for HA tracker prefix", "err", err)
+		return nil
+	}
+
+	// TODO what is the benefit to allowing an a concurrent/parallel process here
+	// What is the amount fo keys we are getting back? Is it inefficient to serially allow this behavior?
+	for i := 0; i < len(keys); i++ {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
+		val, err := h.client.Get(ctx, keys[i])
+		if err != nil {
+			level.Warn(h.logger).Log("msg", "syncHAStateOnStart: failed to get replica value", "key", keys[i], "err", err)
+			return err
+		}
+
+		desc, ok := val.(*ReplicaDesc)
+		if !ok {
+			level.Error(h.logger).Log("msg", "syncHAStateOnStart: got invalid replica descriptor", "key", keys[i])
+			continue
+		}
+
+		if notProcessed := h.processKVStoreEntry(keys[i], desc); !notProcessed {
+			level.Warn(h.logger).Log("msg", "syncHAStateOnStart: failed to processed replica value to cache", "key", keys[i], "err", err)
+		}
+	}
+
+	return nil
+}
+
 // Follows pattern used by ring for WatchKey.
 func (h *haTracker) loop(ctx context.Context) error {
 	if !h.cfg.EnableHATracker {
@@ -226,39 +271,11 @@ func (h *haTracker) loop(ctx context.Context) error {
 	// The KVStore config we gave when creating h should have contained a prefix,
 	// which would have given us a prefixed KVStore client. So, we can pass empty string here.
 	h.client.WatchPrefix(ctx, "", func(key string, value interface{}) bool {
-		replica := value.(*ReplicaDesc)
-		segments := strings.SplitN(key, "/", 2)
-
-		// Valid key would look like cluster/replica, and a key without a / such as `ring` would be invalid.
-		if len(segments) != 2 {
-			return true
-		}
-
-		user := segments[0]
-		cluster := segments[1]
-
-		if replica.DeletedAt > 0 {
-			h.electedReplicaChanges.DeleteLabelValues(user, cluster)
-			h.electedReplicaTimestamp.DeleteLabelValues(user, cluster)
-
-			h.electedLock.Lock()
-			defer h.electedLock.Unlock()
-			userClusters := h.clusters[user]
-			if userClusters != nil {
-				delete(userClusters, cluster)
-				if len(userClusters) == 0 {
-					delete(h.clusters, user)
-				}
-			}
-			return true
+		replica, ok := value.(*ReplicaDesc)
+		if !ok {
+			return false
 		}
-
-		// Store the received information into our cache
-		h.electedLock.Lock()
-		h.updateCache(user, cluster, replica)
-		h.electedLock.Unlock()
-		h.electedReplicaPropagationTime.Observe(time.Since(timestamp.Time(replica.ReceivedAt)).Seconds())
-		return true
+		return h.processKVStoreEntry(key, replica)
 	})
 
 	wg.Wait()
@@ -274,6 +291,41 @@ const (
 	deletionTimeout = 30 * time.Minute
 )
 
+func (h *haTracker) processKVStoreEntry(key string, replica *ReplicaDesc) bool {
+	segments := strings.SplitN(key, "/", 2)
+
+	// Valid key would look like cluster/replica, and a key without a / such as `ring` would be invalid.
+	if len(segments) != 2 {
+		return true
+	}
+
+	user := segments[0]
+	cluster := segments[1]
+
+	if replica.DeletedAt > 0 {
+		h.electedReplicaChanges.DeleteLabelValues(user, cluster)
+		h.electedReplicaTimestamp.DeleteLabelValues(user, cluster)
+
+		h.electedLock.Lock()
+		defer h.electedLock.Unlock()
+		userClusters := h.clusters[user]
+		if userClusters != nil {
+			delete(userClusters, cluster)
+			if len(userClusters) == 0 {
+				delete(h.clusters, user)
+			}
+		}
+		return true
+	}
+
+	// Store the received information into our cache
+	h.electedLock.Lock()
+	h.updateCache(user, cluster, replica)
+	h.electedLock.Unlock()
+	h.electedReplicaPropagationTime.Observe(time.Since(timestamp.Time(replica.ReceivedAt)).Seconds())
+
+	return true
+}
 func (h *haTracker) updateKVLoop(ctx context.Context) {
 	cleanupTick := time.NewTicker(util.DurationWithJitter(cleanupCyclePeriod, cleanupCycleJitterVariance))
 	defer cleanupTick.Stop()

pkg/distributor/ha_tracker_test.go

@@ -154,7 +154,6 @@ func TestWatchPrefixAssignment(t *testing.T) {
 	require.NoError(t, err)
 	require.NoError(t, services.StartAndAwaitRunning(context.Background(), c))
 	defer services.StopAndAwaitTerminated(context.Background(), c) //nolint:errcheck
-
 	// Write the first time.
 	now := time.Now()
 
@@ -792,6 +791,58 @@ func TestCheckReplicaCleanup(t *testing.T) {
 	))
 }
 
+// TestStartingSync tests that values are synced to KVStore on startup
+func TestStartingSync(t *testing.T) {
+	t.Run("check that if a second replica is set up during preLoad that it does not write to it since its not the leader", func(t *testing.T) {
+		cluster := "c1"
+		replica := "r1"
+		var c *haTracker
+		var err error
+		var now time.Time
+
+		codec := GetReplicaDescCodec()
+		kvStore, closer := consul.NewInMemoryClient(codec, log.NewNopLogger(), nil)
+		t.Cleanup(func() { assert.NoError(t, closer.Close()) })
+
+		mock := kv.PrefixClient(kvStore, "prefix")
+		c, err = newHATracker(HATrackerConfig{
+			EnableHATracker:        true,
+			KVStore:                kv.Config{Mock: mock},
+			UpdateTimeout:          15 * time.Second,
+			UpdateTimeoutJitterMax: 0,
+			FailoverTimeout:        time.Millisecond * 2,
+		}, trackerLimits{maxClusters: 100}, nil, log.NewNopLogger())
+		require.NoError(t, err)
+		require.NoError(t, services.StartAndAwaitRunning(context.Background(), c))
+
+		now = time.Now()
+		err = c.checkReplica(context.Background(), "user", cluster, replica, now)
+		assert.NoError(t, err)
+
+		// Check to see if the value in the trackers cache is correct.
+		checkReplicaTimestamp(t, time.Second, c, "user", cluster, replica, now)
+		services.StopAndAwaitTerminated(context.Background(), c) //nolint:errcheck
+
+		replicaTwo := "r2"
+		c, err = newHATracker(HATrackerConfig{
+			EnableHATracker:        true,
+			KVStore:                kv.Config{Mock: mock},
+			UpdateTimeout:          15 * time.Second,
+			UpdateTimeoutJitterMax: 0,
+			FailoverTimeout:        time.Millisecond * 2,
+		}, trackerLimits{maxClusters: 100}, nil, log.NewNopLogger())
+
+		require.NoError(t, err)
+		require.NoError(t, services.StartAndAwaitRunning(context.Background(), c))
+		defer services.StopAndAwaitTerminated(context.Background(), c) //nolint:errcheck
+
+		now = time.Now()
+		err = c.checkReplica(context.Background(), "user", cluster, replicaTwo, now)
+		assert.Error(t, err)
+		assert.ErrorIs(t, err, replicasDidNotMatchError{replica: "r2", elected: "r1"})
+	})
+}
+
 func checkUserClusters(t *testing.T, duration time.Duration, c *haTracker, user string, expectedClusters int) {
 	t.Helper()
 	test.Poll(t, duration, nil, func() interface{} {