Add topoproto.IsTabletWithinKeyRanges(...) helper

go/vt/topo/topoproto/tablet.go

@@ -30,6 +30,7 @@ import (
 
 	"vitess.io/vitess/go/netutil"
 	"vitess.io/vitess/go/sets"
+	"vitess.io/vitess/go/vt/key"
 	"vitess.io/vitess/go/vt/vterrors"
 
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
@@ -295,3 +296,24 @@ func IsServingType(tabletType topodatapb.TabletType) bool {
 		return false
 	}
 }
+
+// IsTabletWithinKeyRanges returns true if the provided tablet is within a slice of *topodatapb.KeyRange.
+func IsTabletWithinKeyRanges(tablet *topodatapb.Tablet, keyRanges []*topodatapb.KeyRange) bool {
+	tabletKeyRange := tablet.GetKeyRange()
+	for _, keyRange := range keyRanges {
+		if key.KeyRangeContainsKeyRange(keyRange, tabletKeyRange) {
+			return true
+		}
+	}
+	return false
+}
+
+// IsTabletWithinKeyRangesString returns true if the provided tablet is within a provided string of
+// one or more shard ranges. An error is returned if the shard range spec string cannot be parsed.
+func IsTabletWithinKeyRangesString(tablet *topodatapb.Tablet, keyRangesString string) (bool, error) {
+	keyRanges, err := key.ParseShardingSpec(keyRangesString)
+	if err != nil {
+		return false, err
+	}
+	return IsTabletWithinKeyRanges(tablet, keyRanges), nil
+}

go/vt/topo/topoproto/tablet_test.go

@@ -21,7 +21,9 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
+	"vitess.io/vitess/go/vt/key"
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
 )
 
@@ -129,3 +131,29 @@ func TestIsTabletsInList(t *testing.T) {
 		})
 	}
 }
+
+func TestIsTabletWithinKeyRangesString(t *testing.T) {
+	tablet := &topodatapb.Tablet{
+		Keyspace: "ks",
+		KeyRange: key.NewKeyRange([]byte{0x00}, []byte{0x40}),
+	}
+
+	{
+		// match
+		found, err := IsTabletWithinKeyRangesString(tablet, "-40")
+		require.NoError(t, err)
+		require.True(t, found)
+	}
+	{
+		// no match
+		found, err := IsTabletWithinKeyRangesString(tablet, "80-")
+		require.NoError(t, err)
+		require.False(t, found)
+	}
+	{
+		// bad shard spec
+		found, err := IsTabletWithinKeyRangesString(tablet, "!!!won'tparse!!!")
+		require.Error(t, err)
+		require.False(t, found)
+	}
+}

go/vt/vtorc/config/config.go

@@ -17,12 +17,17 @@
 package config
 
 import (
+	"fmt"
+	"strings"
 	"time"
 
 	"github.com/spf13/pflag"
 
 	"vitess.io/vitess/go/viperutil"
+	"vitess.io/vitess/go/vt/key"
+	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/servenv"
+	"vitess.io/vitess/go/vt/topo/topoproto"
 )
 
 var configurationLoaded = make(chan bool)
@@ -412,3 +417,40 @@
 func WaitForConfigurationToBeLoaded() {
 	<-configurationLoaded
 }
+
+// ParseClustersToWatch parses a []string of keyspace/shards to watch
+// into a map of keyspaces with one or more key ranges to watch.
+func ParseClustersToWatch(clustersToWatch []string) (map[string][]*topodatapb.KeyRange, error) {
+	keyRangesByKS := make(map[string][]*topodatapb.KeyRange)
+	for _, ks := range clustersToWatch {
+		if strings.Contains(ks, "/") && !strings.HasSuffix(ks, "/") {
+			// Validate keyspace/shard parses.
+			k, s, err := topoproto.ParseKeyspaceShard(ks)
+			if err != nil {
+				log.Errorf("Could not parse keyspace/shard %q: %+v", ks, err)
+				continue
+			}
+			if !key.IsValidKeyRange(s) {
+				return keyRangesByKS, fmt.Errorf("invalid key range %q while parsing clusters to watch", s)
+			}
+			// Parse the shard name into key range value.
+			keyRanges, err := key.ParseShardingSpec(s)
+			if err != nil {
+				return keyRangesByKS, fmt.Errorf("could not parse shard name %q: %+v", s, err)
+			}
+			keyRangesByKS[k] = append(shardsToWatch[k], keyRanges...)
+		} else {
+			// Remove trailing slash if exists.
+			ks = strings.TrimSuffix(ks, "/")
+			// We store the entire range of key range if nothing is specified.
+			keyRangesByKS[ks] = []*topodatapb.KeyRange{key.NewCompleteKeyRange()}
+		}
+	}
+	return keyRangesByKS, nil
+}
+
+// ParseClustersToWatchString parses a string of keyspace/shards to watch
+// into a map of keyspaces with one or more key ranges to watch.
+func ParseClustersToWatchString(clustersToWatch string) (map[string][]*topodatapb.KeyRange, error) {
+	return ParseClustersToWatch(strings.Split(clustersToWatch, ","))
+}

go/vt/vtorc/logic/tablet_discovery.go

@@ -19,7 +19,6 @@
 import (
 	"context"
 	"errors"
-	"fmt"
 	"slices"
 	"strings"
 	"sync"
@@ -32,7 +31,6 @@
 
 	"vitess.io/vitess/go/stats"
 	"vitess.io/vitess/go/vt/external/golib/sqlutils"
-	"vitess.io/vitess/go/vt/key"
 	"vitess.io/vitess/go/vt/log"
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
 	"vitess.io/vitess/go/vt/topo"
@@ -103,35 +101,14 @@
 
 // initializeShardsToWatch parses the --clusters_to_watch flag-value
 // into a map of keyspace/shards.
-func initializeShardsToWatch() error {
-	shardsToWatch = make(map[string][]*topodatapb.KeyRange)
+func initializeShardsToWatch() (err error) {
 	if len(clustersToWatch) == 0 {
 		return nil
 	}
 
-	for _, ks := range clustersToWatch {
-		if strings.Contains(ks, "/") && !strings.HasSuffix(ks, "/") {
-			// Validate keyspace/shard parses.
-			k, s, err := topoproto.ParseKeyspaceShard(ks)
-			if err != nil {
-				log.Errorf("Could not parse keyspace/shard %q: %+v", ks, err)
-				continue
-			}
-			if !key.IsValidKeyRange(s) {
-				return fmt.Errorf("invalid key range %q while parsing clusters to watch", s)
-			}
-			// Parse the shard name into key range value.
-			keyRanges, err := key.ParseShardingSpec(s)
-			if err != nil {
-				return fmt.Errorf("could not parse shard name %q: %+v", s, err)
-			}
-			shardsToWatch[k] = append(shardsToWatch[k], keyRanges...)
-		} else {
-			// Remove trailing slash if exists.
-			ks = strings.TrimSuffix(ks, "/")
-			// We store the entire range of key range if nothing is specified.
-			shardsToWatch[ks] = []*topodatapb.KeyRange{key.NewCompleteKeyRange()}
-		}
+	shardsToWatch, err = config.ParseClustersToWatch(clustersToWatch)
+	if err != nil {
+		return err
 	}
 
 	if len(shardsToWatch) == 0 {
@@ -140,27 +117,20 @@
 	return nil
 }
 
-// shouldWatchTablet checks if the given tablet is part of the watch list.
-func shouldWatchTablet(tablet *topodatapb.Tablet) bool {
+// IsIncluded checks if the given tablet is part of the watch list.
+func (s shardsToWatchMap) IsIncluded(tablet *topodatapb.Tablet) bool {
 	// If we are watching all keyspaces, then we want to watch this tablet too.
-	if len(shardsToWatch) == 0 {
+	if len(s) == 0 {
 		return true
 	}
-	shardRanges, ok := shardsToWatch[tablet.GetKeyspace()]
+	shardRanges, ok := s[tablet.GetKeyspace()]
 	// If we don't have the keyspace in our map, then this tablet
 	// doesn't need to be watched.
 	if !ok {
 		return false
 	}
-	// Get the tablet's key range, and check if
-	// it is part of the shard ranges we are watching.
-	kr := tablet.GetKeyRange()
-	for _, shardRange := range shardRanges {
-		if key.KeyRangeContainsKeyRange(shardRange, kr) {
-			return true
-		}
-	}
-	return false
+	// Check if tablet is part of the shard ranges we are watching.
+	return topoproto.IsTabletWithinKeyRanges(tablet, shardRanges)
 }
 
 // OpenTabletDiscovery opens the vitess topo if enables and returns a ticker
@@ -249,7 +219,7 @@
 		matchedTablets := make([]*topo.TabletInfo, 0, len(tablets))
 		func() {
 			for _, t := range tablets {
 				if shouldWatchTablet(t.Tablet) {
 					matchedTablets = append(matchedTablets, t)
 				}
 			}