From d5456a246b9531c9902ba1d1f6a0bb9d31681c62 Mon Sep 17 00:00:00 2001
From: Tim Vaillancourt <tim@timvaillancourt.com>
Date: Tue, 3 Dec 2024 02:21:07 +0100
Subject: [PATCH] `vtorc`: fetch shard names only every
 `--topo-information-refresh-duration`

Signed-off-by: Tim Vaillancourt <tim@timvaillancourt.com>
---
 go/vt/vtorc/logic/keyspace_shard_discovery.go | 44 ++++++++++++++++++-
 .../logic/keyspace_shard_discovery_test.go    | 17 +++++++
 go/vt/vtorc/logic/tablet_discovery.go         | 10 +----
 3 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/go/vt/vtorc/logic/keyspace_shard_discovery.go b/go/vt/vtorc/logic/keyspace_shard_discovery.go
index b1e93fe2a01..78d468c397f 100644
--- a/go/vt/vtorc/logic/keyspace_shard_discovery.go
+++ b/go/vt/vtorc/logic/keyspace_shard_discovery.go
@@ -22,12 +22,43 @@ import (
 	"strings"
 	"sync"
 
+	"vitess.io/vitess/go/stats"
 	"vitess.io/vitess/go/vt/log"
-
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/vtorc/inst"
 )
 
+var (
+	// keyspaceShardNames stores the current names of shards by keyspace.
+	keyspaceShardNames         = make(map[string][]string)
+	keyspaceShardNamesMu       sync.Mutex
+	statsKeyspaceShardsWatched = stats.NewGaugesFuncWithMultiLabels("KeyspaceShardsWatched",
+		"The keyspace/shards watched by VTOrc",
+		[]string{"Keyspace", "Shard"},
+		getKeyspaceShardsStats,
+	)
+)
+
+// getKeyspaceShardsStats returns the current keyspace/shards watched in stats format.
+func getKeyspaceShardsStats() map[string]int64 {
+	keyspaceShardNamesMu.Lock()
+	defer keyspaceShardNamesMu.Unlock()
+	keyspaceShards := make(map[string]int64)
+	for ks, shards := range keyspaceShardNames {
+		for _, shard := range shards {
+			keyspaceShards[ks+"."+shard] = 1
+		}
+	}
+	return keyspaceShards
+}
+
+// GetKeyspaceShardNames returns the names of the shards in a given keyspace.
+func GetKeyspaceShardNames(keyspaceName string) []string {
+	keyspaceShardNamesMu.Lock()
+	defer keyspaceShardNamesMu.Unlock()
+	return keyspaceShardNames[keyspaceName]
+}
+
 // RefreshAllKeyspacesAndShards reloads the keyspace and shard information for the keyspaces that vtorc is concerned with.
 func RefreshAllKeyspacesAndShards() {
 	var keyspaces []string
@@ -134,13 +165,22 @@ func refreshAllShards(ctx context.Context, keyspaceName string) error {
 		log.Error(err)
 		return err
 	}
-	for _, shardInfo := range shardInfos {
+
+	shardNames := make([]string, 0, len(shardInfos))
+	for shardName, shardInfo := range shardInfos {
 		err = inst.SaveShard(shardInfo)
 		if err != nil {
 			log.Error(err)
 			return err
 		}
+		shardNames = append(shardNames, shardName)
 	}
+	sort.Strings(shardNames)
+
+	keyspaceShardNamesMu.Lock()
+	defer keyspaceShardNamesMu.Unlock()
+	keyspaceShardNames[keyspaceName] = shardNames
+
 	return nil
 }
 
diff --git a/go/vt/vtorc/logic/keyspace_shard_discovery_test.go b/go/vt/vtorc/logic/keyspace_shard_discovery_test.go
index 097865db84a..5ecee3980a9 100644
--- a/go/vt/vtorc/logic/keyspace_shard_discovery_test.go
+++ b/go/vt/vtorc/logic/keyspace_shard_discovery_test.go
@@ -52,6 +52,12 @@ var (
 )
 
 func TestRefreshAllKeyspaces(t *testing.T) {
+	// reset keyspaceShardNames
+	keyspaceShardNames = make(map[string][]string)
+	defer func() {
+		keyspaceShardNames = make(map[string][]string)
+	}()
+
 	// Store the old flags and restore on test completion
 	oldTs := ts
 	oldClustersToWatch := clustersToWatch
@@ -119,6 +125,17 @@ func TestRefreshAllKeyspaces(t *testing.T) {
 	verifyKeyspaceInfo(t, "ks4", keyspaceDurabilityTest, "")
 	verifyPrimaryAlias(t, "ks4", "80-", "zone_ks4-0000000101", "")
 
+	// Confirm caching of shard names
+	require.Equal(t, map[string][]string{
+		"ks1": {"-80", "80-"},
+		"ks2": {"-80", "80-"},
+		"ks3": {"-80", "80-"},
+		"ks4": {"-80", "80-"},
+	}, keyspaceShardNames)
+	for _, ksName := range keyspaceNames {
+		require.Equal(t, []string{"-80", "80-"}, GetKeyspaceShardNames(ksName))
+	}
+	require.Len(t, GetKeyspaceShardNames("does-not-exist"), 0)
 }
 
 func TestRefreshKeyspace(t *testing.T) {
diff --git a/go/vt/vtorc/logic/tablet_discovery.go b/go/vt/vtorc/logic/tablet_discovery.go
index 1ee9e2cfefa..785e0bba506 100644
--- a/go/vt/vtorc/logic/tablet_discovery.go
+++ b/go/vt/vtorc/logic/tablet_discovery.go
@@ -118,15 +118,7 @@ func refreshTabletsUsing(loader func(tabletAlias string), forceRefresh bool) {
 				input := strings.Split(ks, "/")
 				keyspaceShards = append(keyspaceShards, &topo.KeyspaceShard{Keyspace: input[0], Shard: input[1]})
 			} else {
-				// Assume this is a keyspace and find all shards in keyspace
-				ctx, cancel := context.WithTimeout(context.Background(), topo.RemoteOperationTimeout)
-				defer cancel()
-				shards, err := ts.GetShardNames(ctx, ks)
-				if err != nil {
-					// Log the errr and continue
-					log.Errorf("Error fetching shards for keyspace: %v", ks)
-					continue
-				}
+				shards := GetKeyspaceShardNames(ks)
 				if len(shards) == 0 {
 					log.Errorf("Topo has no shards for ks: %v", ks)
 					continue