vendor: google.golang.org/grpc v1.58.3

Signed-off-by: Sebastiaan van Stijn <github@gone.nl>

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

@@ -28,26 +28,36 @@ import (
 	"google.golang.org/grpc/internal/channelz"
 	istatus "google.golang.org/grpc/internal/status"
 	"google.golang.org/grpc/internal/transport"
+	"google.golang.org/grpc/stats"
 	"google.golang.org/grpc/status"
 )
 
 // pickerWrapper is a wrapper of balancer.Picker. It blocks on certain pick
 // actions and unblock when there's a picker update.
 type pickerWrapper struct {
-	mu         sync.Mutex
-	done       bool
-	blockingCh chan struct{}
-	picker     balancer.Picker
+	mu            sync.Mutex
+	done          bool
+	idle          bool
+	blockingCh    chan struct{}
+	picker        balancer.Picker
+	statsHandlers []stats.Handler // to record blocking picker calls
 }
 
-func newPickerWrapper() *pickerWrapper {
-	return &pickerWrapper{blockingCh: make(chan struct{})}
+func newPickerWrapper(statsHandlers []stats.Handler) *pickerWrapper {
+	return &pickerWrapper{
+		blockingCh:    make(chan struct{}),
+		statsHandlers: statsHandlers,
+	}
 }
 
 // updatePicker is called by UpdateBalancerState. It unblocks all blocked pick.
 func (pw *pickerWrapper) updatePicker(p balancer.Picker) {
 	pw.mu.Lock()
-	if pw.done {
+	if pw.done || pw.idle {
+		// There is a small window where a picker update from the LB policy can
+		// race with the channel going to idle mode. If the picker is idle here,
+		// it is because the channel asked it to do so, and therefore it is sage
+		// to ignore the update from the LB policy.
 		pw.mu.Unlock()
 		return
 	}
@@ -63,10 +73,8 @@ func (pw *pickerWrapper) updatePicker(p balancer.Picker) {
 //   - wraps the done function in the passed in result to increment the calls
 //     failed or calls succeeded channelz counter before invoking the actual
 //     done function.
-func doneChannelzWrapper(acw *acBalancerWrapper, result *balancer.PickResult) {
-	acw.mu.Lock()
-	ac := acw.ac
-	acw.mu.Unlock()
+func doneChannelzWrapper(acbw *acBalancerWrapper, result *balancer.PickResult) {
+	ac := acbw.ac
 	ac.incrCallsStarted()
 	done := result.Done
 	result.Done = func(b balancer.DoneInfo) {
@@ -92,6 +100,7 @@ func (pw *pickerWrapper) pick(ctx context.Context, failfast bool, info balancer.
 	var ch chan struct{}
 
 	var lastPickErr error
+
 	for {
 		pw.mu.Lock()
 		if pw.done {
@@ -126,6 +135,20 @@ func (pw *pickerWrapper) pick(ctx context.Context, failfast bool, info balancer.
 			continue
 		}
 
+		// If the channel is set, it means that the pick call had to wait for a
+		// new picker at some point. Either it's the first iteration and this
+		// function received the first picker, or a picker errored with
+		// ErrNoSubConnAvailable or errored with failfast set to false, which
+		// will trigger a continue to the next iteration. In the first case this
+		// conditional will hit if this call had to block (the channel is set).
+		// In the second case, the only way it will get to this conditional is
+		// if there is a new picker.
+		if ch != nil {
+			for _, sh := range pw.statsHandlers {
+				sh.HandleRPC(ctx, &stats.PickerUpdated{})
+			}
+		}
+
 		ch = pw.blockingCh
 		p := pw.picker
 		pw.mu.Unlock()
@@ -152,14 +175,14 @@ func (pw *pickerWrapper) pick(ctx context.Context, failfast bool, info balancer.
 			return nil, balancer.PickResult{}, status.Error(codes.Unavailable, err.Error())
 		}
 
-		acw, ok := pickResult.SubConn.(*acBalancerWrapper)
+		acbw, ok := pickResult.SubConn.(*acBalancerWrapper)
 		if !ok {
 			logger.Errorf("subconn returned from pick is type %T, not *acBalancerWrapper", pickResult.SubConn)
 			continue
 		}
-		if t := acw.getAddrConn().getReadyTransport(); t != nil {
+		if t := acbw.ac.getReadyTransport(); t != nil {
 			if channelz.IsOn() {
-				doneChannelzWrapper(acw, &pickResult)
+				doneChannelzWrapper(acbw, &pickResult)
 				return t, pickResult, nil
 			}
 			return t, pickResult, nil
@@ -187,6 +210,25 @@ func (pw *pickerWrapper) close() {
 	close(pw.blockingCh)
 }
 
+func (pw *pickerWrapper) enterIdleMode() {
+	pw.mu.Lock()
+	defer pw.mu.Unlock()
+	if pw.done {
+		return
+	}
+	pw.idle = true
+}
+
+func (pw *pickerWrapper) exitIdleMode() {
+	pw.mu.Lock()
+	defer pw.mu.Unlock()
+	if pw.done {
+		return
+	}
+	pw.blockingCh = make(chan struct{})
+	pw.idle = false
+}
+
 // dropError is a wrapper error that indicates the LB policy wishes to drop the
 // RPC and not retry it.
 type dropError struct {