Fixed shuffle sharding consistency when zone-awareness is enabled and the shard size is increased or instances in a new zone are added (#3299)

* Fixed shuffle sharding consistency when zone-awareness is enabled and the shard size is increased or instances in a new zone are added

Signed-off-by: Marco Pracucci <[email protected]>

* Slightly more permissive deviance

Signed-off-by: Marco Pracucci <[email protected]>

* Cleaned up ring unit tests

Signed-off-by: Marco Pracucci <[email protected]>

* Use YoloBuf to reduce allocations in ShuffleShardSeed()

Signed-off-by: Marco Pracucci <[email protected]>

Author: pracucci

CHANGELOG.md

@@ -84,6 +84,7 @@
 * [BUGFIX] Experimental Alertmanager API: Do not allow empty Alertmanager configurations or bad template filenames to be submitted through the configuration API. #3185
 * [BUGFIX] Reduce failures to update heartbeat when using Consul. #3259
 * [BUGFIX] When using ruler sharding, moving all user rule groups from ruler to a different one and then back could end up with some user groups not being evaluated at all. #3235
+* [BUGFIX] Fixed shuffle sharding consistency when zone-awareness is enabled and the shard size is increased or instances in a new zone are added. #3299
 * [BUGFIX] Use a valid grpc header when logging IP addresses. #3307
 * [BUGFIX] Fixed the metric `cortex_prometheus_rule_group_duration_seconds` in the Ruler, it wouldn't report any values. #3310
 * [BUGFIX] Fixed gRPC connections leaking in rulers when rulers sharding is enabled and APIs called. #3314

pkg/ingester/active_series.go

@@ -5,12 +5,13 @@ import (
 	"math"
 	"sync"
 	"time"
-	"unsafe"
 
 	"github.com/cespare/xxhash"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/pkg/labels"
 	"go.uber.org/atomic"
+
+	"github.com/cortexproject/cortex/pkg/util"
 )
 
 const (
@@ -69,19 +70,15 @@ func fingerprint(series labels.Labels) uint64 {
 
 	sum.Reset()
 	for _, label := range series {
-		_, _ = sum.Write(yoloBuf(label.Name))
+		_, _ = sum.Write(util.YoloBuf(label.Name))
 		_, _ = sum.Write(sep)
-		_, _ = sum.Write(yoloBuf(label.Value))
+		_, _ = sum.Write(util.YoloBuf(label.Value))
 		_, _ = sum.Write(sep)
 	}
 
 	return sum.Sum64()
 }
 
-func yoloBuf(s string) []byte {
-	return *((*[]byte)(unsafe.Pointer(&s)))
-}
-
 // Purge removes expired entries from the cache. This function should be called
 // periodically to avoid memory leaks.
 func (c *ActiveSeries) Purge(keepUntil time.Time) {

pkg/ingester/active_series_test.go

@@ -211,9 +211,3 @@ func benchmarkPurge(b *testing.B, twice bool) {
 		}
 	}
 }
-
-func TestYoloBuf(t *testing.T) {
-	s := yoloBuf("hello world")
-
-	require.Equal(t, []byte("hello world"), s)
-}

pkg/querier/frontend/frontend_querier_queues.go

@@ -89,7 +89,7 @@ func (q *queues) getOrAddQueue(userID string, maxQueriers int) chan *request {
 	if uq == nil {
 		uq = &userQueue{
 			ch:    make(chan *request, q.maxUserQueueSize),
-			seed:  util.ShuffleShardSeed(userID),
+			seed:  util.ShuffleShardSeed(userID, ""),
 			index: -1,
 		}
 		q.userQueues[userID] = uq

pkg/ring/ring.go

@@ -471,9 +471,6 @@ func (r *Ring) ShuffleShard(identifier string, size int) ReadRing {
 		return cached
 	}
 
-	// Initialise the random generator used to select instances in the ring.
-	random := rand.New(rand.NewSource(util.ShuffleShardSeed(identifier)))
-
 	var result *Ring
 
 	// This deferred function will store newly computed ring into cache.
@@ -515,6 +512,12 @@ func (r *Ring) ShuffleShard(identifier string, size int) ReadRing {
 			tokens = r.ringTokens
 		}
 
+		// Initialise the random generator used to select instances in the ring.
+		// Since we consider each zone like an independent ring, we have to use dedicated
+		// pseudo-random generator for each zone, in order to guarantee the "consistency"
+		// property when the shard size changes or a new zone is added.
+		random := rand.New(rand.NewSource(util.ShuffleShardSeed(identifier, zone)))
+
 		// To select one more instance while guaranteeing the "consistency" property,
 		// we do pick a random value from the generator and resolve uniqueness collisions
 		// (if any) continuing walking the ring.

pkg/ring/ring_test.go

@@ -333,7 +333,7 @@ func TestRing_ShuffleShard(t *testing.T) {
 
 	for testName, testData := range tests {
 		t.Run(testName, func(t *testing.T) {
-			// Init the ring description.
+			// Init the ring.
 			ringDesc := &Desc{Ingesters: testData.ringInstances}
 			for id, instance := range ringDesc.Ingesters {
 				instance.Timestamp = time.Now().Unix()
@@ -512,8 +512,8 @@ func TestRing_ShuffleShard_Shuffling(t *testing.T) {
 	maxCombinations := int(math.Pow(float64(numTenants), 2)) - numTenants
 	for numMatching, probability := range theoreticalMatchings {
 		// We allow a max deviance of 10% compared to the theoretical probability,
-		// clamping it between 1% and 0.1% boundaries.
-		maxDeviance := math.Min(1, math.Max(0.1, probability*0.1))
+		// clamping it between 1% and 0.2% boundaries.
+		maxDeviance := math.Min(1, math.Max(0.2, probability*0.1))
 
 		actual := (float64(distribution[numMatching]) / float64(maxCombinations)) * 100
 		assert.InDelta(t, probability, actual, maxDeviance, "numMatching: %d", numMatching)
@@ -609,6 +609,156 @@ func TestRing_ShuffleShard_Consistency(t *testing.T) {
 	}
 }
 
+func TestRing_ShuffleShard_ConsistencyOnShardSizeChanged(t *testing.T) {
+	// Create 30 instances in 3 zones.
+	ringInstances := map[string]IngesterDesc{}
+	for i := 0; i < 30; i++ {
+		name, desc := generateRingInstance(i, i%3)
+		ringInstances[name] = desc
+	}
+
+	// Init the ring.
+	ringDesc := &Desc{Ingesters: ringInstances}
+	ring := Ring{
+		cfg: Config{
+			HeartbeatTimeout:     time.Hour,
+			ZoneAwarenessEnabled: true,
+		},
+		ringDesc:         ringDesc,
+		ringTokens:       ringDesc.getTokens(),
+		ringTokensByZone: ringDesc.getTokensByZone(),
+		ringZones:        getZones(ringDesc.getTokensByZone()),
+		strategy:         &DefaultReplicationStrategy{},
+	}
+
+	// Get the replication set with shard size = 3.
+	firstShard := ring.ShuffleShard("tenant-id", 3)
+	assert.Equal(t, 3, firstShard.IngesterCount())
+
+	firstSet, err := firstShard.GetAll(Read)
+	require.NoError(t, err)
+
+	// Increase shard size to 6.
+	secondShard := ring.ShuffleShard("tenant-id", 6)
+	assert.Equal(t, 6, secondShard.IngesterCount())
+
+	secondSet, err := secondShard.GetAll(Read)
+	require.NoError(t, err)
+
+	for _, firstInstance := range firstSet.Ingesters {
+		assert.True(t, secondSet.Includes(firstInstance.Addr), "new replication set is expected to include previous instance %s", firstInstance.Addr)
+	}
+
+	// Increase shard size to 9.
+	thirdShard := ring.ShuffleShard("tenant-id", 9)
+	assert.Equal(t, 9, thirdShard.IngesterCount())
+
+	thirdSet, err := thirdShard.GetAll(Read)
+	require.NoError(t, err)
+
+	for _, secondInstance := range secondSet.Ingesters {
+		assert.True(t, thirdSet.Includes(secondInstance.Addr), "new replication set is expected to include previous instance %s", secondInstance.Addr)
+	}
+
+	// Decrease shard size to 6.
+	fourthShard := ring.ShuffleShard("tenant-id", 6)
+	assert.Equal(t, 6, fourthShard.IngesterCount())
+
+	fourthSet, err := fourthShard.GetAll(Read)
+	require.NoError(t, err)
+
+	// We expect to have the same exact instances we had when the shard size was 6.
+	for _, secondInstance := range secondSet.Ingesters {
+		assert.True(t, fourthSet.Includes(secondInstance.Addr), "new replication set is expected to include previous instance %s", secondInstance.Addr)
+	}
+
+	// Decrease shard size to 3.
+	fifthShard := ring.ShuffleShard("tenant-id", 3)
+	assert.Equal(t, 3, fifthShard.IngesterCount())
+
+	fifthSet, err := fifthShard.GetAll(Read)
+	require.NoError(t, err)
+
+	// We expect to have the same exact instances we had when the shard size was 3.
+	for _, firstInstance := range firstSet.Ingesters {
+		assert.True(t, fifthSet.Includes(firstInstance.Addr), "new replication set is expected to include previous instance %s", firstInstance.Addr)
+	}
+}
+
+func TestRing_ShuffleShard_ConsistencyOnZonesChanged(t *testing.T) {
+	// Create 20 instances in 2 zones.
+	ringInstances := map[string]IngesterDesc{}
+	for i := 0; i < 20; i++ {
+		name, desc := generateRingInstance(i, i%2)
+		ringInstances[name] = desc
+	}
+
+	// Init the ring.
+	ringDesc := &Desc{Ingesters: ringInstances}
+	ring := Ring{
+		cfg: Config{
+			HeartbeatTimeout:     time.Hour,
+			ZoneAwarenessEnabled: true,
+		},
+		ringDesc:         ringDesc,
+		ringTokens:       ringDesc.getTokens(),
+		ringTokensByZone: ringDesc.getTokensByZone(),
+		ringZones:        getZones(ringDesc.getTokensByZone()),
+		strategy:         &DefaultReplicationStrategy{},
+	}
+
+	// Get the replication set with shard size = 2.
+	firstShard := ring.ShuffleShard("tenant-id", 2)
+	assert.Equal(t, 2, firstShard.IngesterCount())
+
+	firstSet, err := firstShard.GetAll(Read)
+	require.NoError(t, err)
+
+	// Increase shard size to 4.
+	secondShard := ring.ShuffleShard("tenant-id", 4)
+	assert.Equal(t, 4, secondShard.IngesterCount())
+
+	secondSet, err := secondShard.GetAll(Read)
+	require.NoError(t, err)
+
+	for _, firstInstance := range firstSet.Ingesters {
+		assert.True(t, secondSet.Includes(firstInstance.Addr), "new replication set is expected to include previous instance %s", firstInstance.Addr)
+	}
+
+	// Scale up cluster, adding 10 instances in 1 new zone.
+	for i := 20; i < 30; i++ {
+		name, desc := generateRingInstance(i, 2)
+		ringInstances[name] = desc
+	}
+
+	ring.ringDesc.Ingesters = ringInstances
+	ring.ringTokens = ringDesc.getTokens()
+	ring.ringTokensByZone = ringDesc.getTokensByZone()
+	ring.ringZones = getZones(ringDesc.getTokensByZone())
+
+	// Increase shard size to 6.
+	thirdShard := ring.ShuffleShard("tenant-id", 6)
+	assert.Equal(t, 6, thirdShard.IngesterCount())
+
+	thirdSet, err := thirdShard.GetAll(Read)
+	require.NoError(t, err)
+
+	for _, secondInstance := range secondSet.Ingesters {
+		assert.True(t, thirdSet.Includes(secondInstance.Addr), "new replication set is expected to include previous instance %s", secondInstance.Addr)
+	}
+
+	// Increase shard size to 9.
+	fourthShard := ring.ShuffleShard("tenant-id", 9)
+	assert.Equal(t, 9, fourthShard.IngesterCount())
+
+	fourthSet, err := fourthShard.GetAll(Read)
+	require.NoError(t, err)
+
+	for _, thirdInstance := range thirdSet.Ingesters {
+		assert.True(t, fourthSet.Includes(thirdInstance.Addr), "new replication set is expected to include previous instance %s", thirdInstance.Addr)
+	}
+}
+
 func BenchmarkRing_ShuffleShard(b *testing.B) {
 	for _, numInstances := range []int{50, 100, 1000} {
 		for _, numZones := range []int{1, 3} {

pkg/ruler/ruler_test.go

@@ -585,7 +585,7 @@ func TestSharding(t *testing.T) {
 
 // User shuffle shard token.
 func userToken(user string, skip int) uint32 {
-	r := rand.New(rand.NewSource(util.ShuffleShardSeed(user)))
+	r := rand.New(rand.NewSource(util.ShuffleShardSeed(user, "")))
 
 	for ; skip > 0; skip-- {
 		_ = r.Uint32()

pkg/util/shard.go

@@ -5,11 +5,19 @@ import (
 	"encoding/binary"
 )
 
+var (
+	seedSeparator = []byte{0}
+)
+
 // ShuffleShardSeed returns seed for random number generator, computed from provided identifier.
-func ShuffleShardSeed(identifier string) int64 {
+func ShuffleShardSeed(identifier, zone string) int64 {
 	// Use the identifier to compute an hash we'll use to seed the random.
 	hasher := md5.New()
-	hasher.Write([]byte(identifier)) // nolint:errcheck
+	hasher.Write(YoloBuf(identifier)) // nolint:errcheck
+	if zone != "" {
+		hasher.Write(seedSeparator) // nolint:errcheck
+		hasher.Write(YoloBuf(zone)) // nolint:errcheck
+	}
 	checksum := hasher.Sum(nil)
 
 	// Generate the seed based on the first 64 bits of the checksum.

pkg/util/yolo.go

@@ -0,0 +1,7 @@
+package util
+
+import "unsafe"
+
+func YoloBuf(s string) []byte {
+	return *((*[]byte)(unsafe.Pointer(&s)))
+}

pkg/util/yolo_test.go

@@ -0,0 +1,13 @@
+package util
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestYoloBuf(t *testing.T) {
+	s := YoloBuf("hello world")
+
+	require.Equal(t, []byte("hello world"), s)
+}