From def45d7dfce52a7075b6e4cfe2bd946d14c817c0 Mon Sep 17 00:00:00 2001
From: Michael Demmer <mdemmer@slack-corp.com>
Date: Thu, 23 Jan 2025 08:58:38 -0800
Subject: [PATCH] add num_backup_conns option to force discovery of
 non-zone-local vtgates

---
 go/vt/vtgateproxy/discovery.go   | 28 ++++++++++++++++++++--------
 go/vt/vtgateproxy/vtgateproxy.go |  2 ++
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/go/vt/vtgateproxy/discovery.go b/go/vt/vtgateproxy/discovery.go
index 34bafa5cbd7..78ea8d179b4 100644
--- a/go/vt/vtgateproxy/discovery.go
+++ b/go/vt/vtgateproxy/discovery.go
@@ -84,6 +84,7 @@ type JSONGateResolverBuilder struct {
 	affinityField  string
 	affinityValue  string
 	numConnections int
+	numBackupConns int
 
 	mu        sync.RWMutex
 	targets   map[string][]targetHost
@@ -115,6 +116,7 @@ func RegisterJSONGateResolver(
 	affinityField string,
 	affinityValue string,
 	numConnections int,
+	numBackupConns int,
 ) (*JSONGateResolverBuilder, error) {
 	jsonDiscovery := &JSONGateResolverBuilder{
 		targets:        map[string][]targetHost{},
@@ -125,6 +127,7 @@ func RegisterJSONGateResolver(
 		affinityField:  affinityField,
 		affinityValue:  affinityValue,
 		numConnections: numConnections,
+		numBackupConns: numBackupConns,
 		sorter:         newShuffleSorter(),
 	}
 
@@ -265,7 +268,7 @@ func (b *JSONGateResolverBuilder) parse() (bool, error) {
 		return false, fmt.Errorf("error parsing JSON discovery file %s: %v", b.jsonPath, err)
 	}
 
-	var targets = map[string][]targetHost{}
+	var allTargets = map[string][]targetHost{}
 	for _, host := range hosts {
 		hostname, hasHostname := host["host"]
 		address, hasAddress := host[b.addressField]
@@ -312,7 +315,7 @@ func (b *JSONGateResolverBuilder) parse() (bool, error) {
 		}
 
 		target := targetHost{hostname.(string), fmt.Sprintf("%s:%s", address, port), poolType.(string), affinity.(string), affinity == b.affinityValue}
-		targets[target.PoolType] = append(targets[target.PoolType], target)
+		allTargets[target.PoolType] = append(allTargets[target.PoolType], target)
 	}
 
 	// If a pool disappears, the metric will not record this unless all counts
@@ -322,16 +325,25 @@ func (b *JSONGateResolverBuilder) parse() (bool, error) {
 	// targets and only resetting pools which disappear.
 	targetCount.ResetAll()
 
-	for poolType := range targets {
-		b.sorter.shuffleSort(targets[poolType])
-		if len(targets[poolType]) > *numConnections {
-			targets[poolType] = targets[poolType][:b.numConnections]
+	var selected = map[string][]targetHost{}
+
+	for poolType := range allTargets {
+		b.sorter.shuffleSort(allTargets[poolType])
+
+		// try to pick numConnections from the front of the list (local zone) and numBackupConnections
+		// from the tail (remote zone). if that's not possible, just take the whole set
+		if len(allTargets[poolType]) >= b.numConnections+b.numBackupConns {
+			remoteOffset := len(allTargets[poolType]) - b.numBackupConns
+			selected[poolType] = append(allTargets[poolType][:b.numConnections], allTargets[poolType][remoteOffset:]...)
+		} else {
+			selected[poolType] = allTargets[poolType]
 		}
-		targetCount.Set(poolType, int64(len(targets[poolType])))
+
+		targetCount.Set(poolType, int64(len(selected[poolType])))
 	}
 
 	b.mu.Lock()
-	b.targets = targets
+	b.targets = selected
 	b.mu.Unlock()
 
 	return true, nil
diff --git a/go/vt/vtgateproxy/vtgateproxy.go b/go/vt/vtgateproxy/vtgateproxy.go
index c1d67b459b9..c7ff8f2c78e 100644
--- a/go/vt/vtgateproxy/vtgateproxy.go
+++ b/go/vt/vtgateproxy/vtgateproxy.go
@@ -53,6 +53,7 @@ const (
 var (
 	vtgateHostsFile = flag.String("vtgate_hosts_file", "", "json file describing the host list to use for vtgate:// resolution")
 	numConnections  = flag.Int("num_connections", 4, "number of outbound GPRC connections to maintain")
+	numBackupConns  = flag.Int("num_backup_conns", 1, "number of backup remote-zone GPRC connections to maintain")
 	poolTypeField   = flag.String("pool_type_field", "", "Field name used to specify the target vtgate type and filter the hosts")
 	affinityField   = flag.String("affinity_field", "", "Attribute (JSON file) used to specify the routing affinity , e.g. 'az_id'")
 	affinityValue   = flag.String("affinity_value", "", "Value to match for routing affinity , e.g. 'use-az1'")
@@ -233,6 +234,7 @@ func Init() {
 		*affinityField,
 		*affinityValue,
 		*numConnections,
+		*numBackupConns,
 	)
 
 	if err != nil {