maintain the currentConn even if another one becomes ready

go/vt/vtgateproxy/firstready_balancer.go

@@ -31,17 +31,20 @@ limitations under the License.
 // https://github.com/grpc/grpc-go/blob/master/pickfirst.go
 
 import (
+ "sync"
+
  "google.golang.org/grpc/balancer"
  "google.golang.org/grpc/balancer/base"
  "google.golang.org/grpc/grpclog"
+ "vitess.io/vitess/go/vt/log"
 )
 
 // Name is the name of first_ready balancer.
 const Name = "first_ready"
 
 var logger = grpclog.Component("firstready")
 
-// newBuilder creates a new roundrobin balancer builder.
+// newBuilder creates a new first_ready balancer builder.
 func newBuilder() balancer.Builder {
  return base.NewBalancerBuilder(Name, &frPickerBuilder{}, base.Config{HealthCheck: true})
 }
@@ -50,31 +53,52 @@ func init() {
  balancer.Register(newBuilder())
 }
 
-type frPickerBuilder struct{}
+// 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("firstreadyPicker: Build called with info: %v", info)
 
-func (*frPickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
- logger.Infof("firstreadyPicker: Build called with info: %v", info)
  if len(info.ReadySCs) == 0 {
  return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
  }
 
- var subConn balancer.SubConn
+ 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("firstreadyPicker: currentConn is active, checking if still ready")
+ for sc := range info.ReadySCs {
+ if f.currentConn == sc {
+ log.V(100).Infof("firstreadyPicker: 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("firstreadyPicker: currentConn is not active, picking a new one")
  for sc := range info.ReadySCs {
- subConn = sc
+ f.currentConn = sc
  break
  }
- return &frPicker{
- subConn: subConn,
- }
-}
 
-type frPicker struct {
- // subConn is the first ready subconn when this picker was
- // created. The slice is immutable. Each Get() will do a round robin
- // selection from it and return the selected SubConn.
- subConn balancer.SubConn
+ return f
 }
 
-func (p *frPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
- return balancer.PickResult{SubConn: p.subConn}, nil
+// 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
 }