Fix bug where adding/removing instance with RF > len(zones) will introduce more than one member change in Get replication set (#5429)

* Fix bug where adding/removing instance with RF > len(zones) will introduce more than one member change in Get replication set

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

* Add tests

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

* Add changelog

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

* Add more tests

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

* Update variable name

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

* Add extended replica set test for ring.Get

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

* Fix batch benchmark tests

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

---------

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

Author: justinjung04

CHANGELOG.md

@@ -43,6 +43,7 @@
 * [BUGFIX] Ring: Allow RF greater than number of zones to select more than one instance per zone #5411
 * [BUGFIX] Distributor: Fix potential data corruption in cases of timeout between distributors and ingesters. #5422
 * [BUGFIX] Store Gateway: Fix bug in store gateway ring comparison logic. #5426
+* [BUGFIX] Ring: Fix bug in consistency of Get func in a scaling zone-aware ring. #5429
 
 ## 1.15.1 2023-04-26
 

pkg/compactor/shuffle_sharding_grouper_test.go

@@ -751,7 +751,7 @@ func (r *RingMock) Collect(ch chan<- prometheus.Metric) {}
 
 func (r *RingMock) Describe(ch chan<- *prometheus.Desc) {}
 
-func (r *RingMock) Get(key uint32, op ring.Operation, bufDescs []ring.InstanceDesc, bufHosts, bufZones []string) (ring.ReplicationSet, error) {
+func (r *RingMock) Get(key uint32, op ring.Operation, bufDescs []ring.InstanceDesc, bufHosts []string, bufZones map[string]int) (ring.ReplicationSet, error) {
 	args := r.Called(key, op, bufDescs, bufHosts, bufZones)
 	return args.Get(0).(ring.ReplicationSet), args.Error(1)
 }

pkg/ring/batch.go

@@ -78,10 +78,10 @@ func DoBatch(ctx context.Context, op Operation, r ReadRing, keys []uint32, callb
 	var (
 		bufDescs [GetBufferSize]InstanceDesc
 		bufHosts [GetBufferSize]string
-		bufZones [GetBufferSize]string
+		bufZones = make(map[string]int, GetZoneSize)
 	)
 	for i, key := range keys {
-		replicationSet, err := r.Get(key, op, bufDescs[:0], bufHosts[:0], bufZones[:0])
+		replicationSet, err := r.Get(key, op, bufDescs[:0], bufHosts[:0], bufZones)
 		if err != nil {
 			cleanup()
 			return err

pkg/ring/ring.go

@@ -31,6 +31,10 @@ const (
 	// GetBufferSize is the suggested size of buffers passed to Ring.Get(). It's based on
 	// a typical replication factor 3, plus extra room for a JOINING + LEAVING instance.
 	GetBufferSize = 5
+
+	// GetZoneSize is the suggested size of zone map passed to Ring.Get(). It's based on
+	// a typical replication factor 3.
+	GetZoneSize = 3
 )
 
 // ReadRing represents the read interface to the ring.
@@ -39,7 +43,7 @@ type ReadRing interface {
 	// Get returns n (or more) instances which form the replicas for the given key.
 	// bufDescs, bufHosts and bufZones are slices to be overwritten for the return value
 	// to avoid memory allocation; can be nil, or created with ring.MakeBuffersForGet().
-	Get(key uint32, op Operation, bufDescs []InstanceDesc, bufHosts, bufZones []string) (ReplicationSet, error)
+	Get(key uint32, op Operation, bufDescs []InstanceDesc, bufHosts []string, bufZones map[string]int) (ReplicationSet, error)
 
 	// GetAllHealthy returns all healthy instances in the ring, for the given operation.
 	// This function doesn't check if the quorum is honored, so doesn't fail if the number
@@ -340,24 +344,30 @@ func (r *Ring) updateRingState(ringDesc *Desc) {
 }
 
 // Get returns n (or more) instances which form the replicas for the given key.
-func (r *Ring) Get(key uint32, op Operation, bufDescs []InstanceDesc, bufHosts, bufZones []string) (ReplicationSet, error) {
+// This implementation guarantees:
+// - Stability: given the same ring, two invocations returns the same set for same operation.
+// - Consistency: adding/removing 1 instance from the ring returns set with no more than 1 difference for same operation.
+func (r *Ring) Get(key uint32, op Operation, bufDescs []InstanceDesc, bufHosts []string, bufZones map[string]int) (ReplicationSet, error) {
 	r.mtx.RLock()
 	defer r.mtx.RUnlock()
 	if r.ringDesc == nil || len(r.ringTokens) == 0 {
 		return ReplicationSet{}, ErrEmptyRing
 	}
 
 	var (
-		replicationFactor = r.cfg.ReplicationFactor
-		instances         = bufDescs[:0]
-		start             = searchToken(r.ringTokens, key)
-		iterations        = 0
+		replicationFactor      = r.cfg.ReplicationFactor
+		instances              = bufDescs[:0]
+		start                  = searchToken(r.ringTokens, key)
+		iterations             = 0
+		maxInstancePerZone     = replicationFactor / len(r.ringZones)
+		zonesWithExtraInstance = replicationFactor % len(r.ringZones)
 
 		// We use a slice instead of a map because it's faster to search within a
 		// slice than lookup a map for a very low number of items.
-		distinctHosts = bufHosts[:0]
-		distinctZones = bufZones[:0]
+		distinctHosts       = bufHosts[:0]
+		numOfInstanceByZone = resetZoneMap(bufZones)
 	)
+
 	for i := start; len(distinctHosts) < replicationFactor && iterations < len(r.ringTokens); i++ {
 		iterations++
 		// Wrap i around in the ring.
@@ -370,14 +380,20 @@ func (r *Ring) Get(key uint32, op Operation, bufDescs []InstanceDesc, bufHosts,
 			return ReplicationSet{}, ErrInconsistentTokensInfo
 		}
 
-		// We want n *distinct* instances && distinct zones.
+		// We want n *distinct* instances.
 		if util.StringsContain(distinctHosts, info.InstanceID) {
 			continue
 		}
 
 		// Ignore if the instances don't have a zone set.
 		if r.cfg.ZoneAwarenessEnabled && info.Zone != "" {
-			if util.StringsContain(distinctZones, info.Zone) {
+			maxNumOfInstance := maxInstancePerZone
+			// If we still have room for zones with extra instance, increase the instance threshold by 1
+			if zonesWithExtraInstance > 0 {
+				maxNumOfInstance++
+			}
+
+			if numOfInstanceByZone[info.Zone] >= maxNumOfInstance {
 				continue
 			}
 		}
@@ -392,11 +408,14 @@ func (r *Ring) Get(key uint32, op Operation, bufDescs []InstanceDesc, bufHosts,
 		} else if r.cfg.ZoneAwarenessEnabled && info.Zone != "" {
 			// We should only add the zone if we are not going to extend,
 			// as we want to extend the instance in the same AZ.
-			distinctZones = append(distinctZones, info.Zone)
-
-			if len(distinctZones) == len(r.ringZones) {
-				// reset the zones to repeatedly get hosts from distinct zones
-				distinctZones = distinctZones[:0]
+			if numOfInstance, ok := numOfInstanceByZone[info.Zone]; !ok {
+				numOfInstanceByZone[info.Zone] = 1
+			} else if numOfInstance < maxInstancePerZone {
+				numOfInstanceByZone[info.Zone]++
+			} else {
+				// This zone will have an extra instance
+				numOfInstanceByZone[info.Zone]++
+				zonesWithExtraInstance--
 			}
 		}