Using subring on getShardedRules when shuffle sharding (#4466)

* Using subring on getShardedRules when shuffle sharding

Signed-off-by: Alan Protasio <[email protected]>

Author: alanprot

CHANGELOG.md

@@ -39,6 +39,7 @@
 * [ENHANCEMENT] Updated Prometheus to include changes from prometheus/prometheus#9083. Now whenever `/labels` API calls include matchers, blocks store is queried for `LabelNames` with matchers instead of `Series` calls which was inefficient. #4380
 * [ENHANCEMENT] Exemplars are now emitted for all gRPC calls and many operations tracked by histograms. #4462
 * [ENHANCEMENT] New options `-server.http-listen-network` and `-server.grpc-listen-network` allow binding as 'tcp4' or 'tcp6'. #4462
+* [ENHANCEMENT] Rulers: Using shuffle sharding subring on GetRules API. #4466
 * [BUGFIX] Fixes a panic in the query-tee when comparing result. #4465
 * [BUGFIX] Frontend: Fixes @ modifier functions (start/end) when splitting queries by time. #4464
 * [BUGFIX] Compactor: compactor will no longer try to compact blocks that are already marked for deletion. Previously compactor would consider blocks marked for deletion within `-compactor.deletion-delay / 2` period as eligible for compaction. #4328

pkg/ruler/ruler.go

@@ -651,7 +651,7 @@ func (r *Ruler) GetRules(ctx context.Context) ([]*GroupStateDesc, error) {
 	}
 
 	if r.cfg.EnableSharding {
-		return r.getShardedRules(ctx)
+		return r.getShardedRules(ctx, userID)
 	}
 
 	return r.getLocalRules(userID)
@@ -744,8 +744,14 @@ func (r *Ruler) getLocalRules(userID string) ([]*GroupStateDesc, error) {
 	return groupDescs, nil
 }
 
-func (r *Ruler) getShardedRules(ctx context.Context) ([]*GroupStateDesc, error) {
-	rulers, err := r.ring.GetReplicationSetForOperation(RingOp)
+func (r *Ruler) getShardedRules(ctx context.Context, userID string) ([]*GroupStateDesc, error) {
+	ring := ring.ReadRing(r.ring)
+
+	if shardSize := r.limits.RulerTenantShardSize(userID); shardSize > 0 && r.cfg.ShardingStrategy == util.ShardingStrategyShuffle {
+		ring = r.ring.ShuffleShard(userID, shardSize)
+	}
+
+	rulers, err := ring.GetReplicationSetForOperation(RingOp)
 	if err != nil {
 		return nil, err
 	}

pkg/ruler/ruler_test.go

@@ -427,11 +427,14 @@ func TestGetRules(t *testing.T) {
 					require.NoError(t, err)
 					require.Equal(t, len(allRulesByUser[u]), len(rules))
 					if tc.sharding {
-						mockPoolLClient := r.clientsPool.(*mockRulerClientsPool)
-
-						// Right now we are calling all rules even with shuffle sharding
-						require.Equal(t, int32(len(rulerAddrMap)), mockPoolLClient.numberOfCalls.Load())
-						mockPoolLClient.numberOfCalls.Store(0)
+						mockPoolClient := r.clientsPool.(*mockRulerClientsPool)
+
+						if tc.shardingStrategy == util.ShardingStrategyShuffle {
+							require.Equal(t, int32(tc.shuffleShardSize), mockPoolClient.numberOfCalls.Load())
+						} else {
+							require.Equal(t, int32(len(rulerAddrMap)), mockPoolClient.numberOfCalls.Load())
+						}
+						mockPoolClient.numberOfCalls.Store(0)
 					}
 				})
 			}