stats: Add RPC event for blocking for a picker update

clientconn.go

@@ -349,7 +349,7 @@ func (cc *ClientConn) exitIdleMode() error {
 	cc.idlenessState = ccIdlenessStateExitingIdle
 	exitedIdle := false
 	if cc.blockingpicker == nil {
-		cc.blockingpicker = newPickerWrapper()
+		cc.blockingpicker = newPickerWrapper(cc.dopts.copts.StatsHandlers)
 	} else {
 		cc.blockingpicker.exitIdleMode()
 		exitedIdle = true

picker_wrapper.go

@@ -28,21 +28,26 @@ 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
-	idle       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.
@@ -125,6 +130,9 @@ func (pw *pickerWrapper) pick(ctx context.Context, failfast bool, info balancer.
 					return nil, balancer.PickResult{}, status.Error(codes.Canceled, errStr)
 				}
 			case <-ch:
+				for _, sh := range pw.statsHandlers {
+					sh.HandleRPC(ctx, &stats.PickerUpdated{})
+				}
 			}
 			continue
 		}

picker_wrapper_test.go

@@ -66,7 +66,7 @@ func (p *testingPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error
 }
 
 func (s) TestBlockingPickTimeout(t *testing.T) {
-	bp := newPickerWrapper()
+	bp := newPickerWrapper(nil)
 	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
 	defer cancel()
 	if _, _, err := bp.pick(ctx, true, balancer.PickInfo{}); status.Code(err) != codes.DeadlineExceeded {
@@ -75,7 +75,7 @@ func (s) TestBlockingPickTimeout(t *testing.T) {
 }
 
 func (s) TestBlockingPick(t *testing.T) {
-	bp := newPickerWrapper()
+	bp := newPickerWrapper(nil)
 	// All goroutines should block because picker is nil in bp.
 	var finishedCount uint64
 	for i := goroutineCount; i > 0; i-- {
@@ -94,7 +94,7 @@ func (s) TestBlockingPick(t *testing.T) {
 }
 
 func (s) TestBlockingPickNoSubAvailable(t *testing.T) {
-	bp := newPickerWrapper()
+	bp := newPickerWrapper(nil)
 	var finishedCount uint64
 	bp.updatePicker(&testingPicker{err: balancer.ErrNoSubConnAvailable, maxCalled: goroutineCount})
 	// All goroutines should block because picker returns no subConn available.
@@ -114,7 +114,7 @@ func (s) TestBlockingPickNoSubAvailable(t *testing.T) {
 }
 
 func (s) TestBlockingPickTransientWaitforready(t *testing.T) {
-	bp := newPickerWrapper()
+	bp := newPickerWrapper(nil)
 	bp.updatePicker(&testingPicker{err: balancer.ErrTransientFailure, maxCalled: goroutineCount})
 	var finishedCount uint64
 	// All goroutines should block because picker returns transientFailure and
@@ -135,7 +135,7 @@ func (s) TestBlockingPickTransientWaitforready(t *testing.T) {
 }
 
 func (s) TestBlockingPickSCNotReady(t *testing.T) {
-	bp := newPickerWrapper()
+	bp := newPickerWrapper(nil)
 	bp.updatePicker(&testingPicker{sc: testSCNotReady, maxCalled: goroutineCount})
 	var finishedCount uint64
 	// All goroutines should block because subConn is not ready.

stats/opencensus/trace.go

@@ -99,9 +99,11 @@ func populateSpan(ctx context.Context, rs stats.RPCStats, ti *traceInfo) {
 			trace.BoolAttribute("Client", rs.Client),
 			trace.BoolAttribute("FailFast", rs.FailFast),
 		)
+	case *stats.PickerUpdated:
+		span.Annotate(nil, "Delayed LB pick complete")
 	case *stats.InPayload:
 		// message id - "must be calculated as two different counters starting
-		// from 1 one for sent messages and one for received messages."
+		// from one for sent messages and one for received messages."
 		mi := atomic.AddUint32(&ti.countRecvMsg, 1)
 		span.AddMessageReceiveEvent(int64(mi), int64(rs.Length), int64(rs.CompressedLength))
 	case *stats.OutPayload:

stats/stats.go

@@ -59,6 +59,16 @@ func (s *Begin) IsClient() bool { return s.Client }
 
 func (s *Begin) isRPCStats() {}
 
+// PickerUpdated represents an event that the picker finished picking if the
+// ClientConn blocks on picker picking while performing a RPC.
+type PickerUpdated struct{}
+
+// IsClient indicates if the stats information is from client side. Only Client
+// Side interfaces with a Picker, thus always returns true.
+func (*PickerUpdated) IsClient() bool { return true }
+
+func (*PickerUpdated) isRPCStats() {}
+
 // InPayload contains the information for an incoming payload.
 type InPayload struct {
 	// Client is true if this InPayload is from client side.

test/retry_test.go

@@ -543,6 +543,8 @@ func (s) TestRetryStats(t *testing.T) {
 	handler.mu.Lock()
 	want := []stats.RPCStats{
 		&stats.Begin{},
+		&stats.PickerUpdated{},
+		&stats.PickerUpdated{},
 		&stats.OutHeader{FullMethod: "/grpc.testing.TestService/EmptyCall"},
 		&stats.OutPayload{WireLength: 5},
 		&stats.End{},
@@ -579,7 +581,7 @@ func (s) TestRetryStats(t *testing.T) {
 	// There is a race between receiving the payload (triggered by the
 	// application / gRPC library) and receiving the trailer (triggered at the
 	// transport layer).  Adjust the received stats accordingly if necessary.
-	const tIdx, pIdx = 13, 14
+	const tIdx, pIdx = 15, 16
 	_, okT := handler.s[tIdx].(*stats.InTrailer)
 	_, okP := handler.s[pIdx].(*stats.InPayload)
 	if okT && okP {
@@ -594,7 +596,11 @@ func (s) TestRetryStats(t *testing.T) {
 			t.Fatalf("at position %v: got %T; want %T", i, s, w)
 		}
 		wv, sv := reflect.ValueOf(w).Elem(), reflect.ValueOf(s).Elem()
-
+		if sv.NumField() == 0 {
+			// Resolver blocking and Picker blocking events have no fields,
+			// isClient() always returns true.
+			continue
+		}
 		// Validate that Client is always true
 		if sv.FieldByName("Client").Interface().(bool) != true {
 			t.Fatalf("at position %v: got Client=false; want true", i)
@@ -620,8 +626,8 @@ func (s) TestRetryStats(t *testing.T) {
 	}
 
 	// Validate timings between last Begin and preceding End.
-	end := handler.s[8].(*stats.End)
-	begin := handler.s[9].(*stats.Begin)
+	end := handler.s[10].(*stats.End)
+	begin := handler.s[11].(*stats.Begin)
 	diff := begin.BeginTime.Sub(end.EndTime)
 	if diff < 10*time.Millisecond || diff > 50*time.Millisecond {
 		t.Fatalf("pushback time before final attempt = %v; want ~10ms", diff)