From 5ea2a5a2339161cf4f5c7642398312c03d884914 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Leszczy=C5=84ski?= <2000michal@wp.pl>
Date: Wed, 13 Mar 2024 15:45:55 +0100
Subject: [PATCH] feat(repair): use RingDescriber per table basis

This is required for tablet tables which have per table ring description.
---
 pkg/service/repair/generator.go | 15 ++++----
 pkg/service/repair/plan.go      | 62 ++++++++++++++++-----------------
 2 files changed, 37 insertions(+), 40 deletions(-)

diff --git a/pkg/service/repair/generator.go b/pkg/service/repair/generator.go
index 8c5cff0d85..651a94365e 100644
--- a/pkg/service/repair/generator.go
+++ b/pkg/service/repair/generator.go
@@ -107,21 +107,18 @@ func (g *generator) Run(ctx context.Context) error {
 	g.logger.Info(ctx, "Start generator")
 	var genErr error
 
+	ringDescriber := scyllaclient.NewRingDescriber(ctx, g.client)
 	for _, ksp := range g.plan.Keyspaces {
-		if !g.shouldGenerate() {
-			break
-		}
-
-		ring, err := g.client.DescribeVnodeRing(ctx, ksp.Keyspace)
-		if err != nil {
-			return errors.Wrap(err, "describe ring")
-		}
-
 		for _, tp := range ksp.Tables {
 			if !g.shouldGenerate() {
 				break
 			}
 
+			ring, err := ringDescriber.DescribeRing(ctx, ksp.Keyspace, tp.Table)
+			if err != nil {
+				return errors.Wrap(err, "describe ring")
+			}
+
 			tg := g.newTableGenerator(ksp.Keyspace, tp, ring)
 			// All errors are logged, so in order to reduce clutter,
 			// return only the first one.
diff --git a/pkg/service/repair/plan.go b/pkg/service/repair/plan.go
index 36685e4596..4d749d622c 100644
--- a/pkg/service/repair/plan.go
+++ b/pkg/service/repair/plan.go
@@ -63,42 +63,40 @@ func newPlan(ctx context.Context, target Target, client *scyllaclient.Client) (*
 		maxP     int
 	)
 
+	ringDescriber := scyllaclient.NewRingDescriber(ctx, client)
 	for _, u := range target.Units {
-		ring, err := client.DescribeVnodeRing(ctx, u.Keyspace)
-		if err != nil {
-			return nil, errors.Wrapf(err, "keyspace %s: get ring description", u.Keyspace)
-		}
-		// Allow repairing single node cluster for better UX.
-		if len(status) > 1 && !ShouldRepairRing(ring, target.DC, target.Host) {
-			continue
-		}
-
-		// Update max parallel
-		maxP = max(maxP, MaxRingParallel(ring, target.DC))
-
-		// Update ranges and hosts
-		rangesCnt := 0
-		replicaSetCnt := 0
-		for _, rep := range ring.ReplicaTokens {
-			filtered := filterReplicaSet(rep.ReplicaSet, ring.HostDC, target)
-			if len(filtered) == 0 {
+		var tables []tablePlan
+		for _, t := range u.Tables {
+			ring, err := ringDescriber.DescribeRing(ctx, u.Keyspace, t)
+			if err != nil {
+				return nil, errors.Wrapf(err, "keyspace %s.%s: get ring description", u.Keyspace, t)
+			}
+			// Allow repairing single node cluster for better UX.
+			if len(status) > 1 && !ShouldRepairRing(ring, target.DC, target.Host) {
 				continue
 			}
 
-			replicaSetCnt++
-			allHosts.Add(filtered...)
+			// Update max parallel
+			maxP = max(maxP, MaxRingParallel(ring, target.DC))
+
+			// Update ranges and hosts
+			rangesCnt := 0
+			replicaSetCnt := 0
+			for _, rep := range ring.ReplicaTokens {
+				filtered := filterReplicaSet(rep.ReplicaSet, ring.HostDC, target)
+				if len(filtered) == 0 {
+					continue
+				}
+
+				replicaSetCnt++
+				allHosts.Add(filtered...)
 
-			for _, h := range filtered {
-				for _, t := range u.Tables {
+				for _, h := range filtered {
 					ranges[newHostKsTable(h, u.Keyspace, t)] += len(rep.Ranges)
 				}
+				rangesCnt += len(rep.Ranges)
 			}
-			rangesCnt += len(rep.Ranges)
-		}
 
-		// Update table plan
-		var tables []tablePlan
-		for _, t := range u.Tables {
 			tables = append(tables, tablePlan{
 				Table:         t,
 				ReplicaSetCnt: replicaSetCnt,
@@ -106,10 +104,12 @@ func newPlan(ctx context.Context, target Target, client *scyllaclient.Client) (*
 			})
 		}
 
-		ks = append(ks, keyspacePlan{
-			Keyspace: u.Keyspace,
-			Tables:   tables,
-		})
+		if len(tables) > 0 {
+			ks = append(ks, keyspacePlan{
+				Keyspace: u.Keyspace,
+				Tables:   tables,
+			})
+		}
 	}
 
 	if len(ks) == 0 {