Repository: grafana/dskit (go)

Issue:
ring/util.go WaitRingStability slow to respond
In Cortex the test TestStoreGateway_InitialSyncWithWaitRingStability/shuffle_sharding_strategy,_20_gateways,_RF_=_3,_SS_=_3_(bucket_index_enabled_=_false)
 takes 30 seconds. And most of it is spent in WaitRingStability, which is weird as everything runs locally in the same process.

Goal: improve the responsiveness of the function if possible

Pull request description:
ring/util.go: make WaitRingStability more responsive
**What this PR does**:

In some use cases, such as the Cortex store gateway, it is enough for the
user to know that the token distribution didn't change and it is not
important to know if some member changed between allowed states.

This patch implements the new WaitRingTokensStability function to
be able to wait for tokens stability and ignore allowed state changes.

The original motivation was that store gateway tests took a minute longer
then strictly necessary.

**Which issue(s) this PR fixes**:

Fixes #96 

**Checklist**
- [+] Tests updated
- [+] `CHANGELOG.md` updated - the order of entries should be `[CHANGE]`, `[FEATURE]`, `[ENHANCEMENT]`, `[BUGFIX]`


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index fd9ad94..58d50f2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -23,5 +23,6 @@
 * [ENHANCEMENT] Memberlist: prepare the data to send on the write before starting counting the elapsed time for `-memberlist.packet-write-timeout`, in order to reduce chances we hit the timeout when sending a packet to other node. #89
 * [ENHANCEMENT] flagext: for cases such as `DeprecatedFlag()` that need a logger, add RegisterFlagsWithLogger. #80
 * [ENHANCEMENT] Added option to BasicLifecycler to keep instance in the ring when stopping. #97
+* [ENHANCEMENT] Add WaitRingTokensStability function to ring, to be able to wait on ring stability excluding allowed state transitions. #95
 * [BUGFIX] spanlogger: Support multiple tenant IDs. #59
 * [BUGFIX] Memberlist: fixed corrupted packets when sending compound messages with more than 255 messages or messages bigger than 64KB. #85
diff --git a/ring/replication_set.go b/ring/replication_set.go
index 4748589..461429d 100644
--- a/ring/replication_set.go
+++ b/ring/replication_set.go
@@ -130,6 +130,24 @@ func (r ReplicationSet) GetAddressesWithout(exclude string) []string {
 // HasReplicationSetChanged returns true if two replications sets are the same (with possibly different timestamps),
 // false if they differ in any way (number of instances, instance states, tokens, zones, ...).
 func HasReplicationSetChanged(before, after ReplicationSet) bool {
+	return hasReplicationSetChangedExcluding(before, after, func(i *InstanceDesc) {
+		i.Timestamp = 0
+	})
+}
+
+// HasReplicationSetChangedWithoutState returns true if two replications sets
+// are the same (with possibly different timestamps and instance states),
+// false if they differ in any other way (number of instances, tokens, zones, ...).
+func HasReplicationSetChangedWithoutState(before, after ReplicationSet) bool {
+	return hasReplicationSetChangedExcluding(before, after, func(i *InstanceDesc) {
+		i.Timestamp = 0
+		i.State = PENDING
+	})
+}
+
+// Do comparison of replicasets, but apply a function first
+// to be able to exclude (reset) some values
+func hasReplicationSetChangedExcluding(before, after ReplicationSet, exclude func(*InstanceDesc)) bool {
 	beforeInstances := before.Instances
 	afterInstances := after.Instances
 
@@ -144,9 +162,8 @@ func HasReplicationSetChanged(before, after ReplicationSet) bool {
 		b := beforeInstances[i]
 		a := afterInstances[i]
 
-		// Exclude the heartbeat timestamp from the comparison.
-		b.Timestamp = 0
-		a.Timestamp = 0
+		exclude(&a)
+		exclude(&b)
 
 		if !b.Equal(a) {
 			return true
diff --git a/ring/util.go b/ring/util.go
index a836aa2..b39f2f2 100644
--- a/ring/util.go
+++ b/ring/util.go
@@ -96,6 +96,19 @@ func WaitInstanceState(ctx context.Context, r ReadRing, instanceID string, state
 // WaitRingStability monitors the ring topology for the provided operation and waits until it
 // keeps stable for at least minStability.
 func WaitRingStability(ctx context.Context, r *Ring, op Operation, minStability, maxWaiting time.Duration) error {
+	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
+// 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 waitStability(ctx context.Context, r *Ring, op Operation, minStability, maxWaiting time.Duration, isChanged func(ReplicationSet, ReplicationSet) bool) error {
 	// Configure the max waiting time as a context deadline.
 	ctx, cancel := context.WithTimeout(ctx, maxWaiting)
 	defer cancel()
@@ -117,7 +130,7 @@ func WaitRingStability(ctx context.Context, r *Ring, op Operation, minStability,
 			// replication set which we use to compare with the previous state.
 			currRingState, _ := r.GetAllHealthy(op) // nolint:errcheck
 
-			if HasReplicationSetChanged(ringLastState, currRingState) {
+			if isChanged(ringLastState, currRingState) {
 				ringLastState = currRingState
 				ringLastStateTs = time.Now()
 			} else if time.Since(ringLastStateTs) >= minStability {


Test patch:
diff --git a/ring/replication_set_test.go b/ring/replication_set_test.go
index 74ef17b..42ecc0f 100644
--- a/ring/replication_set_test.go
+++ b/ring/replication_set_test.go
@@ -241,3 +241,71 @@ func TestReplicationSet_Do(t *testing.T) {
 		})
 	}
 }
+
+var (
+	replicationSetChangesInitialState = ReplicationSet{
+		Instances: []InstanceDesc{
+			{Addr: "127.0.0.1"},
+			{Addr: "127.0.0.2"},
+			{Addr: "127.0.0.3"},
+		},
+	}
+	replicationSetChangesTestCases = map[string]struct {
+		nextState                                  ReplicationSet
+		expectHasReplicationSetChanged             bool
+		expectHasReplicationSetChangedWithoutState bool
+	}{
+		"timestamp changed": {
+			ReplicationSet{
+				Instances: []InstanceDesc{
+					{Addr: "127.0.0.1", Timestamp: time.Hour.Microseconds()},
+					{Addr: "127.0.0.2"},
+					{Addr: "127.0.0.3"},
+				},
+			},
+			false,
+			false,
+		},
+		"state changed": {
+			ReplicationSet{
+				Instances: []InstanceDesc{
+					{Addr: "127.0.0.1", State: PENDING},
+					{Addr: "127.0.0.2"},
+					{Addr: "127.0.0.3"},
+				},
+			},
+			true,
+			false,
+		},
+		"more instances": {
+			ReplicationSet{
+				Instances: []InstanceDesc{
+					{Addr: "127.0.0.1"},
+					{Addr: "127.0.0.2"},
+					{Addr: "127.0.0.3"},
+					{Addr: "127.0.0.4"},
+				},
+			},
+			true,
+			true,
+		},
+	}
+)
+
+func TestHasReplicationSetChanged_IgnoresTimeStamp(t *testing.T) {
+	// Only testing difference to underlying Equal function
+	for testName, testData := range replicationSetChangesTestCases {
+		t.Run(testName, func(t *testing.T) {
+			assert.Equal(t, testData.expectHasReplicationSetChanged, HasReplicationSetChanged(replicationSetChangesInitialState, testData.nextState), "HasReplicationSetChanged wrong result")
+		})
+	}
+}
+
+func TestHasReplicationSetChangedWithoutState_IgnoresTimeStampAndState(t *testing.T) {
+	// Only testing difference to underlying Equal function
+	for testName, testData := range replicationSetChangesTestCases {
+		t.Run(testName, func(t *testing.T) {
+			assert.Equal(t, testData.expectHasReplicationSetChangedWithoutState, HasReplicationSetChangedWithoutState(replicationSetChangesInitialState, testData.nextState), "HasReplicationSetChangedWithoutState wrong result")
+		})
+	}
+}
diff --git a/ring/util_test.go b/ring/util_test.go
index 9a4a69c..8a4563e 100644
--- a/ring/util_test.go
+++ b/ring/util_test.go
@@ -79,7 +79,7 @@ func TestGenerateTokens(t *testing.T) {
 	}
 }
 
-func TestGenerateTokensIgnoresOldTokens(t *testing.T) {
+func TestGenerateTokens_IgnoresOldTokens(t *testing.T) {
 	first := GenerateTokens(1000000, nil)
 	second := GenerateTokens(1000000, first)
 
@@ -96,14 +96,7 @@ func TestGenerateTokensIgnoresOldTokens(t *testing.T) {
 	}
 }
 
-func TestWaitRingStabilityShouldReturnAsSoonAsMinStabilityIsReachedOnNoChanges(t *testing.T) {
-	t.Parallel()
-
-	const (
-		minStability = 2 * time.Second
-		maxWaiting   = 10 * time.Second
-	)
-
+func createStartingRing() *Ring {
 	// Init the ring.
 	ringDesc := &Desc{Ingesters: map[string]InstanceDesc{
 		"instance-1": {Addr: "127.0.0.1", State: ACTIVE, Timestamp: time.Now().Unix()},
@@ -123,49 +116,53 @@ func TestWaitRingStabilityShouldReturnAsSoonAsMinStabilityIsReachedOnNoChanges(t
 		strategy:            NewDefaultReplicationStrategy(),
 	}
 
+	return ring
+}
+
+func TestWaitRingStability_ShouldReturnAsSoonAsMinStabilityIsReachedOnNoChanges(t *testing.T) {
+	t.Parallel()
+
+	const (
+		minStability = 2 * time.Second
+		maxWaiting   = 10 * time.Second
+	)
+
+	ring := createStartingRing()
+
 	startTime := time.Now()
 	require.NoError(t, WaitRingStability(context.Background(), ring, Reporting, minStability, maxWaiting))
 	elapsedTime := time.Since(startTime)
 
-	assert.InDelta(t, minStability, elapsedTime, float64(2*time.Second))
+	assert.GreaterOrEqual(t, elapsedTime, minStability)
+	assert.Less(t, elapsedTime, 2*minStability)
 }
 
-func TestWaitRingStabilityShouldReturnOnceMinStabilityHasBeenReached(t *testing.T) {
+func TestWaitRingTokensStability_ShouldReturnAsSoonAsMinStabilityIsReachedOnNoChanges(t *testing.T) {
 	t.Parallel()
 
 	const (
-		minStability     = 3 * time.Second
-		addInstanceAfter = 2 * time.Second
-		maxWaiting       = 15 * time.Second
+		minStability = 2 * time.Second
+		maxWaiting   = 10 * time.Second
 	)
 
-	// Init the ring.
-	ringDesc := &Desc{Ingesters: map[string]InstanceDesc{
-		"instance-1": {Addr: "instance-1", State: ACTIVE, Timestamp: time.Now().Unix()},
-		"instance-2": {Addr: "instance-2", State: PENDING, Timestamp: time.Now().Unix()},
-		"instance-3": {Addr: "instance-3", State: JOINING, Timestamp: time.Now().Unix()},
-		"instance-4": {Addr: "instance-4", State: LEAVING, Timestamp: time.Now().Unix()},
-		"instance-5": {Addr: "instance-5", State: ACTIVE, Timestamp: time.Now().Unix()},
-	}}
+	ring := createStartingRing()
 
-	ring := &Ring{
-		cfg:                 Config{HeartbeatTimeout: time.Minute},
-		ringDesc:            ringDesc,
-		ringTokens:          ringDesc.GetTokens(),
-		ringTokensByZone:    ringDesc.getTokensByZone(),
-		ringInstanceByToken: ringDesc.getTokensInfo(),
-		ringZones:           getZones(ringDesc.getTokensByZone()),
-		strategy:            NewDefaultReplicationStrategy(),
-	}
+	startTime := time.Now()
+	require.NoError(t, WaitRingTokensStability(context.Background(), ring, Reporting, minStability, maxWaiting))
+	elapsedTime := time.Since(startTime)
 
-	// Add 1 new instance after some time.
+	assert.GreaterOrEqual(t, elapsedTime, minStability)
+	assert.Less(t, elapsedTime, 2*minStability)
+}
+
+func addInstanceAfterSomeTime(ring *Ring, addInstanceAfter time.Duration) {
 	go func() {
 		time.Sleep(addInstanceAfter)
 
 		ring.mtx.Lock()
 		defer ring.mtx.Unlock()
-
-		instanceID := fmt.Sprintf("instance-%d", len(ringDesc.Ingesters)+1)
+		ringDesc := ring.ringDesc
+		instanceID := fmt.Sprintf("127.0.0.%d", len(ringDesc.Ingesters)+1)
 		ringDesc.Ingesters[instanceID] = InstanceDesc{Addr: instanceID, State: ACTIVE, Timestamp: time.Now().Unix()}
 		ring.ringDesc = ringDesc
 		ring.ringTokens = ringDesc.GetTokens()
@@ -173,45 +170,56 @@ func TestWaitRingStabilityShouldReturnOnceMinStabilityHasBeenReached(t *testing.
 		ring.ringInstanceByToken = ringDesc.getTokensInfo()
 		ring.ringZones = getZones(ringDesc.getTokensByZone())
 	}()
+}
+
+func TestWaitRingStability_ShouldReturnOnceMinStabilityOfInstancesHasBeenReached(t *testing.T) {
+	t.Parallel()
+
+	const (
+		minStability     = 3 * time.Second
+		addInstanceAfter = 2 * time.Second
+		maxWaiting       = 15 * time.Second
+	)
+
+	ring := createStartingRing()
+
+	// Add 1 new instance after some time.
+	addInstanceAfterSomeTime(ring, addInstanceAfter)
 
 	startTime := time.Now()
 	require.NoError(t, WaitRingStability(context.Background(), ring, Reporting, minStability, maxWaiting))
 	elapsedTime := time.Since(startTime)
 
-	assert.GreaterOrEqual(t, elapsedTime.Milliseconds(), (minStability + addInstanceAfter).Milliseconds())
-	assert.LessOrEqual(t, elapsedTime.Milliseconds(), (minStability + addInstanceAfter + 3*time.Second).Milliseconds())
+	assert.GreaterOrEqual(t, elapsedTime, minStability+addInstanceAfter)
+	assert.LessOrEqual(t, elapsedTime, minStability+addInstanceAfter+3*time.Second)
 }
 
-func TestWaitRingStabilityShouldReturnErrorIfMaxWaitingIsReached(t *testing.T) {
+func TestWaitRingTokensStability_ShouldReturnOnceMinStabilityOfInstancesHasBeenReached(t *testing.T) {
 	t.Parallel()
 
 	const (
-		minStability = 2 * time.Second
-		maxWaiting   = 7 * time.Second
+		minStability     = 3 * time.Second
+		addInstanceAfter = 2 * time.Second
+		maxWaiting       = 15 * time.Second
 	)
 
-	// Init the ring.
-	ringDesc := &Desc{Ingesters: map[string]InstanceDesc{
-		"instance-1": {Addr: "instance-1", State: ACTIVE, Timestamp: time.Now().Unix()},
-		"instance-2": {Addr: "instance-2", State: PENDING, Timestamp: time.Now().Unix()},
-		"instance-3": {Addr: "instance-3", State: JOINING, Timestamp: time.Now().Unix()},
-		"instance-4": {Addr: "instance-4", State: LEAVING, Timestamp: time.Now().Unix()},
-		"instance-5": {Addr: "instance-5", State: ACTIVE, Timestamp: time.Now().Unix()},
-	}}
+	ring := createStartingRing()
 
-	ring := &Ring{
-		cfg:                 Config{HeartbeatTimeout: time.Minute},
-		ringDesc:            ringDesc,
-		ringTokens:          ringDesc.GetTokens(),
-		ringTokensByZone:    ringDesc.getTokensByZone(),
-		ringInstanceByToken: ringDesc.getTokensInfo(),
-		ringZones:           getZones(ringDesc.getTokensByZone()),
-		strategy:            NewDefaultReplicationStrategy(),
-	}
+	// Add 1 new instance after some time.
+	addInstanceAfterSomeTime(ring, addInstanceAfter)
 
+	startTime := time.Now()
+	require.NoError(t, WaitRingTokensStability(context.Background(), ring, Reporting, minStability, maxWaiting))
+	elapsedTime := time.Since(startTime)
+
+	assert.GreaterOrEqual(t, elapsedTime, minStability+addInstanceAfter)
+	assert.LessOrEqual(t, elapsedTime, minStability+addInstanceAfter+3*time.Second)
+}
+
+func addInstancesPeriodically(ring *Ring) chan struct{} {
 	// Keep changing the ring.
 	done := make(chan struct{})
-	defer close(done)
+
 	go func() {
 		for {
 			select {
@@ -219,8 +227,8 @@ func TestWaitRingStabilityShouldReturnErrorIfMaxWaitingIsReached(t *testing.T) {
 				return
 			case <-time.After(time.Second):
 				ring.mtx.Lock()
-
-				instanceID := fmt.Sprintf("instance-%d", len(ringDesc.Ingesters)+1)
+				ringDesc := ring.ringDesc
+				instanceID := fmt.Sprintf("127.0.0.%d", len(ringDesc.Ingesters)+1)
 				ringDesc.Ingesters[instanceID] = InstanceDesc{Addr: instanceID, State: ACTIVE, Timestamp: time.Now().Unix()}
 				ring.ringDesc = ringDesc
 				ring.ringTokens = ringDesc.GetTokens()
@@ -232,15 +240,123 @@ func TestWaitRingStabilityShouldReturnErrorIfMaxWaitingIsReached(t *testing.T) {
 			}
 		}
 	}()
+	return done
+}
+
+func TestWaitRingStability_ShouldReturnErrorIfInstancesAddedAndMaxWaitingIsReached(t *testing.T) {
+	t.Parallel()
+
+	const (
+		minStability = 2 * time.Second
+		maxWaiting   = 7 * time.Second
+	)
+
+	ring := createStartingRing()
+
+	done := addInstancesPeriodically(ring)
+	defer close(done)
 
 	startTime := time.Now()
 	require.Equal(t, context.DeadlineExceeded, WaitRingStability(context.Background(), ring, Reporting, minStability, maxWaiting))
 	elapsedTime := time.Since(startTime)
 
-	assert.InDelta(t, maxWaiting, elapsedTime, float64(2*time.Second))
+	assert.GreaterOrEqual(t, elapsedTime, maxWaiting)
+}
+
+func TestWaitRingTokensStability_ShouldReturnErrorIfInstancesAddedAndMaxWaitingIsReached(t *testing.T) {
+	t.Parallel()
+
+	const (
+		minStability = 2 * time.Second
+		maxWaiting   = 7 * time.Second
+	)
+
+	ring := createStartingRing()
+
+	done := addInstancesPeriodically(ring)
+	defer close(done)
+
+	startTime := time.Now()
+	require.Equal(t, context.DeadlineExceeded, WaitRingTokensStability(context.Background(), ring, Reporting, minStability, maxWaiting))
+	elapsedTime := time.Since(startTime)
+
+	assert.GreaterOrEqual(t, elapsedTime, maxWaiting)
+}
+
+// Keep changing the ring in a way to avoid repeating the same set of states for at least 2 sec
+func changeStatePeriodically(ring *Ring) chan struct{} {
+	done := make(chan struct{})
+	go func() {
+		instanceToMutate := "instance-1"
+		states := []InstanceState{PENDING, JOINING, ACTIVE, LEAVING}
+		stateIdx := 0
+
+		for states[stateIdx] != ring.ringDesc.Ingesters[instanceToMutate].State {
+			stateIdx++
+		}
+
+		for {
+			select {
+			case <-done:
+				return
+			case <-time.After(time.Second):
+				stateIdx++
+				ring.mtx.Lock()
+				ringDesc := ring.ringDesc
+				desc := InstanceDesc{Addr: "127.0.0.1", State: states[stateIdx%len(states)], Timestamp: time.Now().Unix()}
+				ringDesc.Ingesters[instanceToMutate] = desc
+				ring.mtx.Unlock()
+			}
+		}
+	}()
+
+	return done
+}
+
+func TestWaitRingStability_ShouldReturnErrorIfInstanceStateIsChangingAndMaxWaitingIsReached(t *testing.T) {
+	t.Parallel()
+
+	const (
+		minStability = 2 * time.Second
+		maxWaiting   = 7 * time.Second
+	)
+
+	ring := createStartingRing()
+
+	// Keep changing the ring.
+	done := changeStatePeriodically(ring)
+	defer close(done)
+
+	startTime := time.Now()
+	require.Equal(t, context.DeadlineExceeded, WaitRingStability(context.Background(), ring, Reporting, minStability, maxWaiting))
+	elapsedTime := time.Since(startTime)
+
+	assert.GreaterOrEqual(t, elapsedTime, maxWaiting)
+}
+
+func TestWaitRingTokensStability_ShouldReturnOnceMinStabilityOfInstancesHasBeenReachedWhileStateCanChange(t *testing.T) {
+	t.Parallel()
+
+	const (
+		minStability = 2 * time.Second
+		maxWaiting   = 7 * time.Second
+	)
+
+	ring := createStartingRing()
+
+	// Keep changing the ring.
+	done := changeStatePeriodically(ring)
+	defer close(done)
+
+	startTime := time.Now()
+	require.NoError(t, WaitRingTokensStability(context.Background(), ring, Reporting, minStability, maxWaiting))
+	elapsedTime := time.Since(startTime)
+
+	assert.GreaterOrEqual(t, elapsedTime, minStability)
+	assert.Less(t, elapsedTime, 2*minStability)
 }
 
-func TestWaitInstanceStateTimeout(t *testing.T) {
+func TestWaitInstanceState_Timeout(t *testing.T) {
 	t.Parallel()
 
 	const (
@@ -260,7 +376,7 @@ func TestWaitInstanceStateTimeout(t *testing.T) {
 	ring.AssertCalled(t, "GetInstanceState", instanceID)
 }
 
-func TestWaitInstanceStateTimeoutOnError(t *testing.T) {
+func TestWaitInstanceState_TimeoutOnError(t *testing.T) {
 	t.Parallel()
 
 	const (
@@ -280,7 +396,7 @@ func TestWaitInstanceStateTimeoutOnError(t *testing.T) {
 	ring.AssertCalled(t, "GetInstanceState", instanceID)
 }
 
-func TestWaitInstanceStateExitsAfterActualStateEqualsState(t *testing.T) {
+func TestWaitInstanceState_ExitsAfterActualStateEqualsState(t *testing.T) {
 	t.Parallel()
 
 	const (