Use WaitRingTokensAndZonesStability instead of WaitRingStability

Signed-off-by: Justin Jung <[email protected]>

Author: justinjung04

pkg/ring/replication_set.go

@@ -145,10 +145,9 @@ func HasReplicationSetChangedWithoutState(before, after ReplicationSet) bool {
 	})
 }
 
-// HasReplicationSetChangedWithoutStateAndAddress returns false if two replications sets
-// are the same (with possibly different timestamps, instance states and address),
-// true if they differ in any other way (number of instances, tokens, or zones).
-func HasReplicationSetChangedWithoutStateAndAddress(before, after ReplicationSet) bool {
+// HasReplicationSetTokensOrZonesChanged returns true if two replications sets
+// are different in number of instances, tokens and zones, false otherwise.
+func HasReplicationSetTokensOrZonesChanged(before, after ReplicationSet) bool {
 	return hasReplicationSetChangedExcluding(before, after, func(i *InstanceDesc) {
 		i.Timestamp = 0
 		i.State = PENDING

pkg/ring/replication_set_test.go

@@ -341,7 +341,7 @@ func TestHasReplicationSetChangedWithoutStateAndAddress_IgnoresTimeStampAndState
 	// Only testing difference to underlying Equal function
 	for testName, testData := range replicationSetChangesTestCases {
 		t.Run(testName, func(t *testing.T) {
-			assert.Equal(t, testData.expectHasReplicationSetChangedWithoutStateAndAddress, HasReplicationSetChangedWithoutStateAndAddress(replicationSetChangesInitialState, testData.nextState), "HasReplicationSetChangedWithoutStateAndAddress wrong result")
+			assert.Equal(t, testData.expectHasReplicationSetChangedWithoutStateAndAddress, HasReplicationSetTokensOrZonesChanged(replicationSetChangesInitialState, testData.nextState), "HasReplicationSetChangedWithoutStateAndAddress wrong result")
 		})
 	}
 }

pkg/ring/util.go

@@ -99,13 +99,11 @@ func WaitRingStability(ctx context.Context, r *Ring, op Operation, minStability,
 	return waitStability(ctx, r, op, minStability, maxWaiting, HasReplicationSetChanged)
 }
 
-// WaitRingTokensStability waits for the Ring to be unchanged at
-// least for minStability time period, excluding transitioning between
-// allowed states (e.g. JOINING->ACTIVE if allowed by op).
-// This can be used to avoid wasting resources on moving data around
+// WaitRingTokensAndZonesStability waits for the Ring tokens and zones to be unchanged at least
+// for minStability time period. This can be used to avoid wasting resources on moving data around
 // due to multiple changes in the Ring.
-func WaitRingTokensStability(ctx context.Context, r *Ring, op Operation, minStability, maxWaiting time.Duration) error {
-	return waitStability(ctx, r, op, minStability, maxWaiting, HasReplicationSetChangedWithoutState)
+func WaitRingTokensAndZonesStability(ctx context.Context, r *Ring, op Operation, minStability, maxWaiting time.Duration) error {
+	return waitStability(ctx, r, op, minStability, maxWaiting, HasReplicationSetTokensOrZonesChanged)
 }
 
 func waitStability(ctx context.Context, r *Ring, op Operation, minStability, maxWaiting time.Duration, isChanged func(ReplicationSet, ReplicationSet) bool) error {

pkg/ring/util_test.go

@@ -138,7 +138,7 @@ func TestWaitRingStability_ShouldReturnAsSoonAsMinStabilityIsReachedOnNoChanges(
 	assert.Less(t, elapsedTime, 2*minStability)
 }
 
-func TestWaitRingTokensStability_ShouldReturnAsSoonAsMinStabilityIsReachedOnNoChanges(t *testing.T) {
+func TestWaitRingTokensAndZonesStability_ShouldReturnAsSoonAsMinStabilityIsReachedOnNoChanges(t *testing.T) {
 	t.Parallel()
 
 	const (
@@ -149,7 +149,7 @@ func TestWaitRingTokensStability_ShouldReturnAsSoonAsMinStabilityIsReachedOnNoCh
 	ring := createStartingRing()
 
 	startTime := time.Now()
-	require.NoError(t, WaitRingTokensStability(context.Background(), ring, Reporting, minStability, maxWaiting))
+	require.NoError(t, WaitRingTokensAndZonesStability(context.Background(), ring, Reporting, minStability, maxWaiting))
 	elapsedTime := time.Since(startTime)
 
 	assert.GreaterOrEqual(t, elapsedTime, minStability)
@@ -195,7 +195,7 @@ func TestWaitRingStability_ShouldReturnOnceMinStabilityOfInstancesHasBeenReached
 	assert.LessOrEqual(t, elapsedTime, minStability+addInstanceAfter+3*time.Second)
 }
 
-func TestWaitRingTokensStability_ShouldReturnOnceMinStabilityOfInstancesHasBeenReached(t *testing.T) {
+func TestWaitRingTokensAndZonesStability_ShouldReturnOnceMinStabilityOfInstancesHasBeenReached(t *testing.T) {
 	t.Parallel()
 
 	const (
@@ -210,7 +210,7 @@ func TestWaitRingTokensStability_ShouldReturnOnceMinStabilityOfInstancesHasBeenR
 	addInstanceAfterSomeTime(ring, addInstanceAfter)
 
 	startTime := time.Now()
-	require.NoError(t, WaitRingTokensStability(context.Background(), ring, Reporting, minStability, maxWaiting))
+	require.NoError(t, WaitRingTokensAndZonesStability(context.Background(), ring, Reporting, minStability, maxWaiting))
 	elapsedTime := time.Since(startTime)
 
 	assert.GreaterOrEqual(t, elapsedTime, minStability+addInstanceAfter)
@@ -264,7 +264,7 @@ func TestWaitRingStability_ShouldReturnErrorIfInstancesAddedAndMaxWaitingIsReach
 	assert.GreaterOrEqual(t, elapsedTime, maxWaiting)
 }
 
-func TestWaitRingTokensStability_ShouldReturnErrorIfInstancesAddedAndMaxWaitingIsReached(t *testing.T) {
+func TestWaitRingTokensAndZonesStability_ShouldReturnErrorIfInstancesAddedAndMaxWaitingIsReached(t *testing.T) {
 	t.Parallel()
 
 	const (
@@ -278,7 +278,7 @@ func TestWaitRingTokensStability_ShouldReturnErrorIfInstancesAddedAndMaxWaitingI
 	defer close(done)
 
 	startTime := time.Now()
-	require.Equal(t, context.DeadlineExceeded, WaitRingTokensStability(context.Background(), ring, Reporting, minStability, maxWaiting))
+	require.Equal(t, context.DeadlineExceeded, WaitRingTokensAndZonesStability(context.Background(), ring, Reporting, minStability, maxWaiting))
 	elapsedTime := time.Since(startTime)
 
 	assert.GreaterOrEqual(t, elapsedTime, maxWaiting)
@@ -335,7 +335,7 @@ func TestWaitRingStability_ShouldReturnErrorIfInstanceStateIsChangingAndMaxWaiti
 	assert.GreaterOrEqual(t, elapsedTime, maxWaiting)
 }
 
-func TestWaitRingTokensStability_ShouldReturnOnceMinStabilityOfInstancesHasBeenReachedWhileStateCanChange(t *testing.T) {
+func TestWaitRingTokensAndZonesStability_ShouldReturnOnceMinStabilityOfInstancesHasBeenReachedWhileStateCanChange(t *testing.T) {
 	t.Parallel()
 
 	const (
@@ -350,7 +350,7 @@ func TestWaitRingTokensStability_ShouldReturnOnceMinStabilityOfInstancesHasBeenR
 	defer close(done)
 
 	startTime := time.Now()
-	require.NoError(t, WaitRingTokensStability(context.Background(), ring, Reporting, minStability, maxWaiting))
+	require.NoError(t, WaitRingTokensAndZonesStability(context.Background(), ring, Reporting, minStability, maxWaiting))
 	elapsedTime := time.Since(startTime)
 
 	assert.GreaterOrEqual(t, elapsedTime, minStability)

pkg/storegateway/bucket_stores.go

@@ -66,9 +66,6 @@ type BucketStores struct {
 	storesMu sync.RWMutex
 	stores   map[string]*store.BucketStore
 
-	// List of users assigned to the BucketStores
-	userIDs []string
-
 	// Metrics.
 	syncTimes         prometheus.Histogram
 	syncLastSuccess   prometheus.Gauge
@@ -143,7 +140,7 @@ func NewBucketStores(cfg tsdb.BlocksStorageConfig, shardingStrategy ShardingStra
 func (u *BucketStores) InitialSync(ctx context.Context) error {
 	level.Info(u.logger).Log("msg", "synchronizing TSDB blocks for all users")
 
-	if err := u.syncUsersBlocksWithRetries(ctx, true, func(ctx context.Context, s *store.BucketStore) error {
+	if err := u.syncUsersBlocksWithRetries(ctx, func(ctx context.Context, s *store.BucketStore) error {
 		return s.InitialSync(ctx)
 	}); err != nil {
 		level.Warn(u.logger).Log("msg", "failed to synchronize TSDB blocks", "err", err)
@@ -155,13 +152,13 @@ func (u *BucketStores) InitialSync(ctx context.Context) error {
 }
 
 // SyncBlocks synchronizes the stores state with the Bucket store for every user.
-func (u *BucketStores) SyncBlocks(ctx context.Context, shouldUpdateUserList bool) error {
-	return u.syncUsersBlocksWithRetries(ctx, shouldUpdateUserList, func(ctx context.Context, s *store.BucketStore) error {
+func (u *BucketStores) SyncBlocks(ctx context.Context) error {
+	return u.syncUsersBlocksWithRetries(ctx, func(ctx context.Context, s *store.BucketStore) error {
 		return s.SyncBlocks(ctx)
 	})
 }
 
-func (u *BucketStores) syncUsersBlocksWithRetries(ctx context.Context, shouldUpdateUserList bool, f func(context.Context, *store.BucketStore) error) error {
+func (u *BucketStores) syncUsersBlocksWithRetries(ctx context.Context, f func(context.Context, *store.BucketStore) error) error {
 	retries := backoff.New(ctx, backoff.Config{
 		MinBackoff: 1 * time.Second,
 		MaxBackoff: 10 * time.Second,
@@ -170,7 +167,7 @@ func (u *BucketStores) syncUsersBlocksWithRetries(ctx context.Context, shouldUpd
 
 	var lastErr error
 	for retries.Ongoing() {
-		lastErr = u.syncUsersBlocks(ctx, shouldUpdateUserList, f)
+		lastErr = u.syncUsersBlocks(ctx, f)
 		if lastErr == nil {
 			return nil
 		}
@@ -185,7 +182,7 @@ func (u *BucketStores) syncUsersBlocksWithRetries(ctx context.Context, shouldUpd
 	return lastErr
 }
 
-func (u *BucketStores) syncUsersBlocks(ctx context.Context, shouldUpdateUserList bool, f func(context.Context, *store.BucketStore) error) (returnErr error) {
+func (u *BucketStores) syncUsersBlocks(ctx context.Context, f func(context.Context, *store.BucketStore) error) (returnErr error) {
 	defer func(start time.Time) {
 		u.syncTimes.Observe(time.Since(start).Seconds())
 		if returnErr == nil {
@@ -213,11 +210,7 @@ func (u *BucketStores) syncUsersBlocks(ctx context.Context, shouldUpdateUserList
 
 	includeUserIDs := make(map[string]struct{})
 
-	if shouldUpdateUserList {
-		u.userIDs = u.shardingStrategy.FilterUsers(ctx, scannedUserIDs)
-	}
-
-	for _, userID := range u.userIDs {
+	for _, userID := range u.shardingStrategy.FilterUsers(ctx, scannedUserIDs) {
 		includeUserIDs[userID] = struct{}{}
 	}
 

pkg/storegateway/bucket_stores_test.go

@@ -209,7 +209,7 @@ func TestBucketStores_SyncBlocks(t *testing.T) {
 
 	// Generate another block and sync blocks again.
 	generateStorageBlock(t, storageDir, userID, metricName, 100, 200, 15)
-	require.NoError(t, stores.SyncBlocks(ctx, false))
+	require.NoError(t, stores.SyncBlocks(ctx))
 
 	seriesSet, warnings, err = querySeries(stores, userID, metricName, 150, 180)
 	require.NoError(t, err)
@@ -257,27 +257,16 @@ func TestBucketStores_syncUsersBlocks(t *testing.T) {
 		shouldUpdateUserList bool
 	}{
 		"when sharding is disabled all users should be synced": {
-			shardingStrategy:     NewNoShardingStrategy(),
-			expectedStores:       3,
-			shouldUpdateUserList: true,
+			shardingStrategy: NewNoShardingStrategy(),
+			expectedStores:   3,
 		},
 		"when sharding is enabled only stores for filtered users should be created": {
 			shardingStrategy: func() ShardingStrategy {
 				s := &mockShardingStrategy{}
 				s.On("FilterUsers", mock.Anything, allUsers).Return([]string{"user-1", "user-2"})
 				return s
 			}(),
-			expectedStores:       2,
-			shouldUpdateUserList: true,
-		},
-		"should not sync any blocks if shouldUpdateUserList was never true": {
-			shardingStrategy: func() ShardingStrategy {
-				s := &mockShardingStrategy{}
-				s.On("FilterUsers", mock.Anything, allUsers).Return([]string{"user-1", "user-2"})
-				return s
-			}(),
-			expectedStores:       0,
-			shouldUpdateUserList: false,
+			expectedStores: 2,
 		},
 	}
 
@@ -294,7 +283,7 @@ func TestBucketStores_syncUsersBlocks(t *testing.T) {
 
 			// Sync user stores and count the number of times the callback is called.
 			var storesCount atomic.Int32
-			err = stores.syncUsersBlocks(context.Background(), testData.shouldUpdateUserList, func(ctx context.Context, bs *store.BucketStore) error {
+			err = stores.syncUsersBlocks(context.Background(), func(ctx context.Context, bs *store.BucketStore) error {
 				storesCount.Inc()
 				return nil
 			})
@@ -505,7 +494,7 @@ func TestBucketStores_deleteLocalFilesForExcludedTenants(t *testing.T) {
 
 	// Single user left in shard.
 	sharding.users = []string{user1}
-	require.NoError(t, stores.SyncBlocks(ctx, true))
+	require.NoError(t, stores.SyncBlocks(ctx))
 	require.Equal(t, []string{user1}, getUsersInDir(t, cfg.BucketStore.SyncDir))
 
 	require.NoError(t, testutil.GatherAndCompare(reg, strings.NewReader(`
@@ -522,7 +511,7 @@ func TestBucketStores_deleteLocalFilesForExcludedTenants(t *testing.T) {
 
 	// No users left in this shard.
 	sharding.users = nil
-	require.NoError(t, stores.SyncBlocks(ctx, true))
+	require.NoError(t, stores.SyncBlocks(ctx))
 	require.Equal(t, []string(nil), getUsersInDir(t, cfg.BucketStore.SyncDir))
 
 	require.NoError(t, testutil.GatherAndCompare(reg, strings.NewReader(`
@@ -536,7 +525,7 @@ func TestBucketStores_deleteLocalFilesForExcludedTenants(t *testing.T) {
 
 	// We can always get user back.
 	sharding.users = []string{user1}
-	require.NoError(t, stores.SyncBlocks(ctx, true))
+	require.NoError(t, stores.SyncBlocks(ctx))
 	require.Equal(t, []string{user1}, getUsersInDir(t, cfg.BucketStore.SyncDir))
 
 	require.NoError(t, testutil.GatherAndCompare(reg, strings.NewReader(`

pkg/storegateway/gateway.go

@@ -285,16 +285,15 @@ func (g *StoreGateway) running(ctx context.Context) error {
 	for {
 		select {
 		case <-syncTicker.C:
-			g.syncStores(ctx, syncReasonPeriodic, false)
+			g.syncStores(ctx, ringLastState, syncReasonPeriodic)
 		case <-ringTickerChan:
 			// We ignore the error because in case of error it will return an empty
 			// replication set which we use to compare with the previous state.
 			currRingState, _ := g.ring.GetAllHealthy(BlocksOwnerSync) // nolint:errcheck
 
 			// Ignore address when comparing to avoid block re-sync if tokens are persisted with tokens_file_path
-			if ring.HasReplicationSetChangedWithoutStateAndAddress(ringLastState, currRingState) {
-				ringLastState = currRingState
-				g.syncStores(ctx, syncReasonRingChange, true)
+			if ring.HasReplicationSetTokensOrZonesChanged(ringLastState, currRingState) {
+				ringLastState = g.syncStores(ctx, currRingState, syncReasonRingChange)
 			}
 		case <-ctx.Done():
 			return nil
@@ -311,19 +310,23 @@ func (g *StoreGateway) stopping(_ error) error {
 	return nil
 }
 
-func (g *StoreGateway) syncStores(ctx context.Context, reason string, shouldUpdateUserList bool) {
+func (g *StoreGateway) syncStores(ctx context.Context, currRingState ring.ReplicationSet, reason string) ring.ReplicationSet {
+	ringLastState := currRingState
 	if g.gatewayCfg.ShardingEnabled && g.gatewayCfg.ShardingRing.WaitStabilityMinDuration > 0 {
 		g.waitRingStability(ctx, reason)
+		ringLastState, _ = g.ring.GetAllHealthy(BlocksOwnerSync)
 	}
 
 	level.Info(g.logger).Log("msg", "synchronizing TSDB blocks for all users", "reason", reason)
 	g.bucketSync.WithLabelValues(reason).Inc()
 
-	if err := g.stores.SyncBlocks(ctx, shouldUpdateUserList); err != nil {
+	if err := g.stores.SyncBlocks(ctx); err != nil {
 		level.Warn(g.logger).Log("msg", "failed to synchronize TSDB blocks", "reason", reason, "err", err)
 	} else {
 		level.Info(g.logger).Log("msg", "successfully synchronized TSDB blocks for all users", "reason", reason)
 	}
+
+	return ringLastState
 }
 
 func (g *StoreGateway) Series(req *storepb.SeriesRequest, srv storegatewaypb.StoreGateway_SeriesServer) error {
@@ -368,7 +371,7 @@ func (g *StoreGateway) waitRingStability(ctx context.Context, reason string) {
 	maxWaiting := g.gatewayCfg.ShardingRing.WaitStabilityMaxDuration
 
 	level.Info(g.logger).Log("msg", "waiting until store-gateway ring topology is stable", "min_waiting", minWaiting.String(), "max_waiting", maxWaiting.String(), "reason", reason)
-	if err := ring.WaitRingStability(ctx, g.ring, BlocksOwnerSync, minWaiting, maxWaiting); err != nil {
+	if err := ring.WaitRingTokensAndZonesStability(ctx, g.ring, BlocksOwnerSync, minWaiting, maxWaiting); err != nil {
 		level.Warn(g.logger).Log("msg", "store-gateway ring topology is not stable after the max waiting time, proceeding anyway", "reason", reason)
 	} else {
 		level.Info(g.logger).Log("msg", "store-gateway ring topology is stable", "reason", reason)

pkg/storegateway/gateway_test.go

@@ -796,7 +796,8 @@ func TestStoreGateway_syncStoresShouldWaitRingStability(t *testing.T) {
 	}))
 
 	syncCtx, cancelCtx := context.WithTimeout(ctx, 3*time.Second)
-	g.syncStores(syncCtx, "test", true)
+	ringState, _ := g.ring.GetAllHealthy(BlocksOwnerSync)
+	g.syncStores(syncCtx, ringState, "test")
 	cancelCtx()
 
 	// No blocks should be synced
@@ -806,63 +807,6 @@ func TestStoreGateway_syncStoresShouldWaitRingStability(t *testing.T) {
 	assert.Equal(t, float64(0), metrics.GetSumOfGauges("cortex_bucket_store_blocks_loaded"))
 }
 
-func TestStoreGateway_PeriodicSyncShouldNotUpdateUserList(t *testing.T) {
-	registeredAt := time.Now()
-	bucketClient, storageDir := cortex_testutil.PrepareFilesystemBucket(t)
-
-	// This tests uses real TSDB blocks. 24h time range, 2h block range period
-	now := time.Now()
-	mockTSDB(t, path.Join(storageDir, "user-1"), 1, 12, now.Add(-24*time.Hour).Unix()*1000, now.Unix()*1000)
-
-	ctx := context.Background()
-	gatewayCfg := mockGatewayConfig()
-	gatewayCfg.ShardingEnabled = true
-	gatewayCfg.ShardingRing.RingCheckPeriod = time.Hour // Do not trigger the ring change sync in this test.
-
-	storageCfg := mockStorageConfig(t)
-	storageCfg.BucketStore.SyncInterval = 100 * time.Millisecond
-
-	reg := prometheus.NewPedanticRegistry()
-	ringStore, closer := consul.NewInMemoryClient(ring.GetCodec(), log.NewNopLogger(), nil)
-	t.Cleanup(func() { assert.NoError(t, closer.Close()) })
-
-	g, err := newStoreGateway(gatewayCfg, storageCfg, bucketClient, ringStore, defaultLimitsOverrides(t), mockLoggingLevel(), log.NewNopLogger(), reg)
-	require.NoError(t, err)
-
-	require.NoError(t, ringStore.CAS(ctx, RingKey, func(in interface{}) (interface{}, bool, error) {
-		ringDesc := ring.GetOrCreateRingDesc(in)
-		ringDesc.AddIngester("instance-1", "127.0.0.1", "", ring.Tokens{1, 2, 3}, ring.ACTIVE, registeredAt)
-		return ringDesc, true, nil
-	}))
-
-	require.NoError(t, services.StartAndAwaitRunning(ctx, g))
-	defer services.StopAndAwaitTerminated(ctx, g) //nolint:errcheck
-
-	// Assert on the initial state.
-	regs := util.NewUserRegistries()
-	regs.AddUserRegistry("test", reg)
-	metrics := regs.BuildMetricFamiliesPerUser()
-	assert.Equal(t, float64(12), metrics.GetSumOfGauges("cortex_bucket_store_blocks_loaded"))
-
-	// Add more blocks, and extra user
-	mockTSDB(t, path.Join(storageDir, "user-1"), 1, 3, now.Unix()*1000, now.Add(6*time.Hour).Unix()*1000)
-	mockTSDB(t, path.Join(storageDir, "user-2"), 1, 3, now.Unix()*1000, now.Add(6*time.Hour).Unix()*1000)
-
-	// Verify only user-1 blocks are synced
-	time.Sleep(time.Second)
-	metrics = regs.BuildMetricFamiliesPerUser()
-	assert.Equal(t, float64(12+3), metrics.GetSumOfGauges("cortex_bucket_store_blocks_loaded"))
-
-	// Force update user list
-	g.syncStores(ctx, "test", true)
-	g.waitRingStability(ctx, "test")
-
-	// Verify user-2 blocks are also synced
-	time.Sleep(time.Second)
-	metrics = regs.BuildMetricFamiliesPerUser()
-	assert.Equal(t, float64(12+3+3), metrics.GetSumOfGauges("cortex_bucket_store_blocks_loaded"))
-}
-
 func TestStoreGateway_RingLifecyclerShouldAutoForgetUnhealthyInstances(t *testing.T) {
 	const unhealthyInstanceID = "unhealthy-id"
 	const heartbeatTimeout = time.Minute