New implementation of roundrobin and pickfirst (#1506)

balancer_v1_wrapper.go

@@ -23,6 +23,7 @@ import (
 
 	"golang.org/x/net/context"
 	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/grpclog"
 	"google.golang.org/grpc/resolver"
@@ -45,6 +46,8 @@ func (bwb *balancerWrapperBuilder) Build(cc balancer.ClientConn, opts balancer.B
 		startCh:   make(chan struct{}),
 		conns:     make(map[resolver.Address]balancer.SubConn),
 		connSt:    make(map[balancer.SubConn]*scState),
+		csEvltr:   &connectivityStateEvaluator{},
+		state:     connectivity.Idle,
 	}
 	cc.UpdateBalancerState(connectivity.Idle, bw)
 	go bw.lbWatcher()
@@ -67,6 +70,10 @@ type balancerWrapper struct {
 
 	cc balancer.ClientConn
 
+	// To aggregate the connectivity state.
+	csEvltr *connectivityStateEvaluator
+	state   connectivity.State
+
 	mu     sync.Mutex
 	conns  map[resolver.Address]balancer.SubConn
 	connSt map[balancer.SubConn]*scState
@@ -134,7 +141,7 @@ func (bw *balancerWrapper) lbWatcher() {
 				newAddr := resolver.Address{
 					Addr:       a.Addr,
 					Type:       resolver.Backend, // All addresses from balancer are all backends.
-					ServerName: "",               // TODO(bar) support servername.
+					ServerName: "",
 					Metadata:   a.Metadata,
 				}
 				newAddrs = append(newAddrs, newAddr)
@@ -173,7 +180,7 @@ func (bw *balancerWrapper) lbWatcher() {
 				resAddrs[resolver.Address{
 					Addr:       a.Addr,
 					Type:       resolver.Backend, // All addresses from balancer are all backends.
-					ServerName: "",               // TODO(bar) support servername.
+					ServerName: "",
 					Metadata:   a.Metadata,
 				}] = a
 			}
@@ -187,7 +194,7 @@ func (bw *balancerWrapper) lbWatcher() {
 				if _, ok := resAddrs[a]; !ok {
 					del = append(del, c)
 					delete(bw.conns, a)
-					delete(bw.connSt, c)
+					// Keep the state of this sc in bw.connSt until its state becomes Shutdown.
 				}
 			}
 			bw.mu.Unlock()
@@ -230,12 +237,18 @@ func (bw *balancerWrapper) HandleSubConnStateChange(sc balancer.SubConn, s conne
 		scSt.down = bw.balancer.Up(scSt.addr)
 	} else if oldS == connectivity.Ready && s != connectivity.Ready {
 		if scSt.down != nil {
-			scSt.down(errConnClosing) // TODO(bar) what error to use?
+			scSt.down(errConnClosing)
 		}
 	}
-	// The connectivity state is ignored by clientConn now.
-	// TODO(bar) use the aggregated connectivity state.
-	bw.cc.UpdateBalancerState(connectivity.Ready, bw)
+	sa := bw.csEvltr.recordTransition(oldS, s)
+	if bw.state != sa {
+		bw.state = sa
+	}
+	bw.cc.UpdateBalancerState(bw.state, bw)
+	if s == connectivity.Shutdown {
+		// Remove state for this sc.
+		delete(bw.connSt, sc)
+	}
 	return
 }
 
@@ -276,27 +289,79 @@ func (bw *balancerWrapper) Pick(ctx context.Context, opts balancer.PickOptions)
 	if err != nil {
 		return nil, nil, err
 	}
-	var put func(balancer.DoneInfo)
+	var done func(balancer.DoneInfo)
 	if p != nil {
-		put = func(i balancer.DoneInfo) { p() }
+		done = func(i balancer.DoneInfo) { p() }
 	}
 	var sc balancer.SubConn
+	bw.mu.Lock()
+	defer bw.mu.Unlock()
 	if bw.pickfirst {
-		bw.mu.Lock()
 		// Get the first sc in conns.
 		for _, sc = range bw.conns {
 			break
 		}
-		bw.mu.Unlock()
 	} else {
-		bw.mu.Lock()
-		sc = bw.conns[resolver.Address{
+		var ok bool
+		sc, ok = bw.conns[resolver.Address{
 			Addr:       a.Addr,
 			Type:       resolver.Backend,
-			ServerName: "", // TODO(bar) support servername.
+			ServerName: "",
 			Metadata:   a.Metadata,
 		}]
-		bw.mu.Unlock()
+		if !ok && failfast {
+			return nil, nil, Errorf(codes.Unavailable, "there is no connection available")
+		}
+		if s, ok := bw.connSt[sc]; failfast && (!ok || s.s != connectivity.Ready) {
+			// If the returned sc is not ready and RPC is failfast,
+			// return error, and this RPC will fail.
+			return nil, nil, Errorf(codes.Unavailable, "there is no connection available")
+		}
 	}
-	return sc, put, nil
+
+	return sc, done, nil
+}
+
+// connectivityStateEvaluator gets updated by addrConns when their
+// states transition, based on which it evaluates the state of
+// ClientConn.
+type connectivityStateEvaluator struct {
+	mu                  sync.Mutex
+	numReady            uint64 // Number of addrConns in ready state.
+	numConnecting       uint64 // Number of addrConns in connecting state.
+	numTransientFailure uint64 // Number of addrConns in transientFailure.
+}
+
+// recordTransition records state change happening in every subConn and based on
+// that it evaluates what aggregated state should be.
+// It can only transition between Ready, Connecting and TransientFailure. Other states,
+// Idle and Shutdown are transitioned into by ClientConn; in the begining of the connection
+// before any subConn is created ClientConn is in idle state. In the end when ClientConn
+// closes it is in Shutdown state.
+// TODO Note that in later releases, a ClientConn with no activity will be put into an Idle state.
+func (cse *connectivityStateEvaluator) recordTransition(oldState, newState connectivity.State) connectivity.State {
+	cse.mu.Lock()
+	defer cse.mu.Unlock()
+
+	// Update counters.
+	for idx, state := range []connectivity.State{oldState, newState} {
+		updateVal := 2*uint64(idx) - 1 // -1 for oldState and +1 for new.
+		switch state {
+		case connectivity.Ready:
+			cse.numReady += updateVal
+		case connectivity.Connecting:
+			cse.numConnecting += updateVal
+		case connectivity.TransientFailure:
+			cse.numTransientFailure += updateVal
+		}
+	}
+
+	// Evaluate.
+	if cse.numReady > 0 {
+		return connectivity.Ready
+	}
+	if cse.numConnecting > 0 {
+		return connectivity.Connecting
+	}
+	return connectivity.TransientFailure
 }