stats: Add RPC event for blocking for a picker update

clientconn.go

@@ -42,6 +42,7 @@ import (
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/resolver"
 	"google.golang.org/grpc/serviceconfig"
+	"google.golang.org/grpc/stats"
 	"google.golang.org/grpc/status"
 
 	_ "google.golang.org/grpc/balancer/roundrobin"           // To register roundrobin.
@@ -349,7 +350,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
@@ -743,6 +744,9 @@ func (cc *ClientConn) waitForResolvedAddrs(ctx context.Context) error {
 	}
 	select {
 	case <-cc.firstResolveEvent.Done():
+		for _, sh := range cc.dopts.copts.StatsHandlers {
+			sh.HandleRPC(ctx, &stats.ResolverResolved{})
+		}
 		return nil
 	case <-ctx.Done():
 		return status.FromContextError(ctx.Err()).Err()

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
 }
 
-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,10 @@ 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.PickerPicked{})
+				}
+
 			}
 			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,13 @@ func populateSpan(ctx context.Context, rs stats.RPCStats, ti *traceInfo) {
 			trace.BoolAttribute("Client", rs.Client),
 			trace.BoolAttribute("FailFast", rs.FailFast),
 		)
+	case *stats.ResolverResolved:
+		span.Annotate(nil, "Delayed name resolution complete")
+	case *stats.PickerPicked:
+		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,26 @@ func (s *Begin) IsClient() bool { return s.Client }
 
 func (s *Begin) isRPCStats() {}
 
+// ResolverResolved represents an event that the resolved finished resolving if
+// the ClientConn blocks on resolver resolution while performing a RPC.
+type ResolverResolved struct{}
+
+// IsClient indicates if the stats information is from client side. Only Client
+// Side interfaces with a resolver, thus always returns true.
+func (rr *ResolverResolved) IsClient() bool { return true }
+
+func (rr *ResolverResolved) isRPCStats() {}
+
+// PickerPicked represents an event that the picker finished picking if the ClientConn
+// blocks on picker picking while performing a RPC.
+type PickerPicked struct{}
+
+// IsClient indicates if the stats information is from client side. Only Client
+// Side interfaces with a Picker, thus always returns true.
+func (pp *PickerPicked) IsClient() bool { return true }
+
+func (pp *PickerPicked) isRPCStats() {}
+
 // InPayload contains the information for an incoming payload.
 type InPayload struct {
 	// Client is true if this InPayload is from client side.