add firstready balancer (#349)

* add firstready balancer

* fail if an invalid balancer is picked

* maintain the currentConn even if another one becomes ready

* remove unnecessary globals

* normalize the names

go/vt/vtgateproxy/firstready_balancer.go

@@ -0,0 +1,98 @@
+package vtgateproxy
+
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+// The firstready balancer implements the GRPC load balancer abstraction by
+// routing all queries to the first available target in the list returned from
+// discovery.
+//
+// Similar to the builtin "round_robin" balancer, the base functionality takes care
+// of establishing subconns to all targets in the list and keeping them in the "ready"
+// state, so all we have to do is pick the first available one from the set.
+//
+// This is in contrast to the `pick_first` balancer which only establishes the subconn
+// to a single target at a time and is therefore subject to undesirable behaviors if,
+// for example, the first host in the set is unreachable for some time, but not declared
+// down.
+// https://github.com/grpc/grpc-go/blob/master/pickfirst.go
+
+import (
+	"sync"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	"vitess.io/vitess/go/vt/log"
+)
+
+// newBuilder creates a new first_ready balancer builder.
+func newBuilder() balancer.Builder {
+	return base.NewBalancerBuilder("first_ready", &frPickerBuilder{}, base.Config{HealthCheck: true})
+}
+
+func init() {
+	balancer.Register(newBuilder())
+}
+
+// frPickerBuilder implements both the Builder and the Picker interfaces.
+//
+// Once a conn is chosen and is in the ready state, it will remain as the
+// active subconn even if other connections become available.
+type frPickerBuilder struct {
+	mu          sync.Mutex
+	currentConn balancer.SubConn
+}
+
+func (f *frPickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
+	log.V(100).Infof("first_ready: Build called with info: %v", info)
+
+	if len(info.ReadySCs) == 0 {
+		return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
+	}
+
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	// If we've already chosen a subconn, and it is still in the ready list, then
+	// no need to change state
+	if f.currentConn != nil {
+		log.V(100).Infof("first_ready: currentConn is active, checking if still ready")
+		for sc := range info.ReadySCs {
+			if f.currentConn == sc {
+				log.V(100).Infof("first_ready: currentConn still active - not changing")
+				return f
+			}
+		}
+	}
+
+	// Otherwise either we don't have an active conn or the conn we were using is
+	// no longer active, so pick an arbitrary new one out of the map.
+	log.V(100).Infof("first_ready: currentConn is not active, picking a new one")
+	for sc := range info.ReadySCs {
+		f.currentConn = sc
+		break
+	}
+
+	return f
+}
+
+// Pick simply returns the currently chosen conn
+func (f *frPickerBuilder) Pick(balancer.PickInfo) (balancer.PickResult, error) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	return balancer.PickResult{SubConn: f.currentConn}, nil
+}

go/vt/vtgateproxy/vtgateproxy.go

@@ -21,6 +21,7 @@ package vtgateproxy
 import (
 	"context"
 	"flag"
+	"fmt"
 	"io"
 	"net/url"
 	"strings"
@@ -55,6 +56,7 @@ var (
 	affinityValue   = flag.String("affinity_value", "", "Value to match for routing affinity , e.g. 'use-az1'")
 	addressField    = flag.String("address_field", "address", "field name in the json file containing the address")
 	portField       = flag.String("port_field", "port", "field name in the json file containing the port")
+	balancerType    = flag.String("balancer", "round_robin", "load balancing algorithm to use")
 
 	timings = stats.NewTimings("Timings", "proxy timings by operation", "operation")
 
@@ -201,9 +203,19 @@ func (proxy *VTGateProxy) StreamExecute(ctx context.Context, session *vtgateconn
 }
 
 func Init() {
-	log.V(100).Infof("Registering GRPC dial options")
+	log.Infof("registering GRPC dial options: balancer type %s", *balancerType)
+
+	switch *balancerType {
+	case "round_robin":
+	case "first_ready":
+	case "pick_first":
+		break
+	default:
+		log.Fatalf("invalid balancer type %s", *balancerType)
+	}
+
 	grpcclient.RegisterGRPCDialOptions(func(opts []grpc.DialOption) ([]grpc.DialOption, error) {
-		return append(opts, grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`)), nil
+		return append(opts, grpc.WithDefaultServiceConfig(fmt.Sprintf(`{"loadBalancingConfig": [{"%s":{}}]}`, *balancerType))), nil
 	})
 
 	RegisterJSONGateResolver(