Added support for Haystack tracing

vendor/google.golang.org/grpc/balancer_v1_wrapper.go

@@ -19,16 +19,14 @@
 package grpc
 
 import (
+	"context"
 	"strings"
 	"sync"
 
-	"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"
-	"google.golang.org/grpc/status"
 )
 
 type balancerWrapperBuilder struct {
@@ -55,7 +53,7 @@ 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{},
+		csEvltr:    &balancer.ConnectivityStateEvaluator{},
 		state:      connectivity.Idle,
 	}
 	cc.UpdateBalancerState(connectivity.Idle, bw)
@@ -80,10 +78,6 @@ type balancerWrapper struct {
 	cc         balancer.ClientConn
 	targetAddr string // Target without the scheme.
 
-	// To aggregate the connectivity state.
-	csEvltr *connectivityStateEvaluator
-	state   connectivity.State
-
 	mu     sync.Mutex
 	conns  map[resolver.Address]balancer.SubConn
 	connSt map[balancer.SubConn]*scState
@@ -92,6 +86,10 @@ type balancerWrapper struct {
 	// - NewSubConn is created, cc wants to notify balancer of state changes;
 	// - Build hasn't return, cc doesn't have access to balancer.
 	startCh chan struct{}
+
+	// To aggregate the connectivity state.
+	csEvltr *balancer.ConnectivityStateEvaluator
+	state   connectivity.State
 }
 
 // lbWatcher watches the Notify channel of the balancer and manages
@@ -248,7 +246,7 @@ func (bw *balancerWrapper) HandleSubConnStateChange(sc balancer.SubConn, s conne
 			scSt.down(errConnClosing)
 		}
 	}
-	sa := bw.csEvltr.recordTransition(oldS, s)
+	sa := bw.csEvltr.RecordTransition(oldS, s)
 	if bw.state != sa {
 		bw.state = sa
 	}
@@ -283,9 +281,8 @@ func (bw *balancerWrapper) Close() {
 }
 
 // The picker is the balancerWrapper itself.
-// Pick should never return ErrNoSubConnAvailable.
 // It either blocks or returns error, consistent with v1 balancer Get().
-func (bw *balancerWrapper) Pick(ctx context.Context, opts balancer.PickOptions) (balancer.SubConn, func(balancer.DoneInfo), error) {
+func (bw *balancerWrapper) Pick(ctx context.Context, opts balancer.PickOptions) (sc balancer.SubConn, done func(balancer.DoneInfo), err error) {
 	failfast := true // Default failfast is true.
 	if ss, ok := rpcInfoFromContext(ctx); ok {
 		failfast = ss.failfast
@@ -294,79 +291,51 @@ func (bw *balancerWrapper) Pick(ctx context.Context, opts balancer.PickOptions)
 	if err != nil {
 		return nil, nil, err
 	}
-	var done func(balancer.DoneInfo)
 	if p != nil {
-		done = func(i balancer.DoneInfo) { p() }
+		done = func(balancer.DoneInfo) { p() }
+		defer func() {
+			if err != nil {
+				p()
+			}
+		}()
 	}
-	var sc balancer.SubConn
+
 	bw.mu.Lock()
 	defer bw.mu.Unlock()
 	if bw.pickfirst {
 		// Get the first sc in conns.
-		for _, sc = range bw.conns {
-			break
-		}
-	} else {
-		var ok bool
-		sc, ok = bw.conns[resolver.Address{
-			Addr:       a.Addr,
-			Type:       resolver.Backend,
-			ServerName: "",
-			Metadata:   a.Metadata,
-		}]
-		if !ok && failfast {
-			return nil, nil, status.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, status.Errorf(codes.Unavailable, "there is no connection available")
+		for _, sc := range bw.conns {
+			return sc, done, nil
 		}
+		return nil, nil, balancer.ErrNoSubConnAvailable
+	}
+	sc, ok1 := bw.conns[resolver.Address{
+		Addr:       a.Addr,
+		Type:       resolver.Backend,
+		ServerName: "",
+		Metadata:   a.Metadata,
+	}]
+	s, ok2 := bw.connSt[sc]
+	if !ok1 || !ok2 {
+		// This can only happen due to a race where Get() returned an address
+		// that was subsequently removed by Notify.  In this case we should
+		// retry always.
+		return nil, nil, balancer.ErrNoSubConnAvailable
+	}
+	switch s.s {
+	case connectivity.Ready, connectivity.Idle:
+		return sc, done, nil
+	case connectivity.Shutdown, connectivity.TransientFailure:
+		// If the returned sc has been shut down or is in transient failure,
+		// return error, and this RPC will fail or wait for another picker (if
+		// non-failfast).
+		return nil, nil, balancer.ErrTransientFailure
+	default:
+		// For other states (connecting or unknown), the v1 balancer would
+		// traditionally wait until ready and then issue the RPC.  Returning
+		// ErrNoSubConnAvailable will be a slight improvement in that it will
+		// allow the balancer to choose another address in case others are
+		// connected.
+		return nil, nil, balancer.ErrNoSubConnAvailable
 	}
-
-	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 beginning 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
 }