balancer: make sure non-nil done returned by Pick is called

balancer/balancer.go

@@ -215,8 +215,10 @@ type Picker interface {
 	//
 	// If a SubConn is returned:
 	// - If it is READY, gRPC will send the RPC on it;
-	// - If it is not ready, or becomes not ready after it's returned, gRPC will block
-	//   until UpdateBalancerState() is called and will call pick on the new picker.
+	// - If it is not ready, or becomes not ready after it's returned, gRPC will
+	//   block until UpdateBalancerState() is called and will call pick on the
+	//   new picker. The done function returned from Pick(), if not nil, will be
+	//   called with nil error, no bytes sent and no bytes received.
 	//
 	// If the returned error is not nil:
 	// - If the error is ErrNoSubConnAvailable, gRPC will block until UpdateBalancerState()

picker_wrapper.go

@@ -165,6 +165,14 @@ func (bp *pickerWrapper) pick(ctx context.Context, failfast bool, opts balancer.
 			}
 			return t, done, nil
 		}
+		if done != nil {
+			done(balancer.DoneInfo{
+				Err:           nil,
+				Trailer:       nil,
+				BytesSent:     false,
+				BytesReceived: false,
+			})
+		}
 		grpclog.Infof("blockingPicker: the picked transport is not ready, loop back to repick")
 		// If ok == false, ac.state is not READY.
 		// A valid picker always returns READY subConn. This means the state of ac

test/balancer_test.go

@@ -21,6 +21,7 @@ package test
 import (
 	"context"
 	"reflect"
+	"sync"
 	"testing"
 	"time"
 
@@ -103,10 +104,10 @@ type picker struct {
 }
 
 func (p *picker) Pick(ctx context.Context, opts balancer.PickOptions) (balancer.SubConn, func(balancer.DoneInfo), error) {
-	p.bal.pickOptions = append(p.bal.pickOptions, opts)
 	if p.err != nil {
 		return nil, nil, p.err
 	}
+	p.bal.pickOptions = append(p.bal.pickOptions, opts)
 	return p.sc, func(d balancer.DoneInfo) { p.bal.doneInfo = append(p.bal.doneInfo, d) }, nil
 }
 
@@ -184,4 +185,25 @@ func testPickAndDone(t *testing.T, e env) {
 	if len(b.doneInfo) < 2 || !reflect.DeepEqual(b.doneInfo[1].Trailer, testTrailerMetadata) {
 		t.Fatalf("b.doneInfo = %v; want b.doneInfo[1].Trailer = %v", b.doneInfo, testTrailerMetadata)
 	}
+	if len(b.pickOptions) != len(b.doneInfo) {
+		t.Fatalf("Got %d picks, but %d doneInfo, want equal amount", len(b.pickOptions), len(b.doneInfo))
+	}
+	// To test done() is always called, even if it's returned with a non-Ready
+	// SubConn.
+	//
+	// Stop server and at the same time send RPCs. There are chances that picker
+	// is not updated in time, causing a non-Ready SubConn to be returned.
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		for i := 0; i < 20; i++ {
+			tc.UnaryCall(ctx, &testpb.SimpleRequest{})
+		}
+		wg.Done()
+	}()
+	te.srv.Stop()
+	wg.Wait()
+	if len(b.pickOptions) != len(b.doneInfo) {
+		t.Fatalf("Got %d picks, %d doneInfo, want equal amount", len(b.pickOptions), len(b.doneInfo))
+	}
 }