diff --git a/balancer/balancer.go b/balancer/balancer.go index 67518de9a1cf..c9bb59320a43 100644 --- a/balancer/balancer.go +++ b/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() diff --git a/picker_wrapper.go b/picker_wrapper.go index a2575c9637b0..f9625496c403 100644 --- a/picker_wrapper.go +++ b/picker_wrapper.go @@ -165,6 +165,11 @@ func (bp *pickerWrapper) pick(ctx context.Context, failfast bool, opts balancer. } return t, done, nil } + if done != nil { + // Calling done with nil error, no bytes sent and no bytes received. + // DoneInfo with default value works. + done(balancer.DoneInfo{}) + } 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 diff --git a/test/balancer_test.go b/test/balancer_test.go index 5215a8461654..c8c422afa206 100644 --- a/test/balancer_test.go +++ b/test/balancer_test.go @@ -103,10 +103,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 +184,24 @@ 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. + finished := make(chan struct{}) + go func() { + for i := 0; i < 20; i++ { + tc.UnaryCall(ctx, &testpb.SimpleRequest{}) + } + close(finished) + }() + te.srv.Stop() + <-finished + if len(b.pickOptions) != len(b.doneInfo) { + t.Fatalf("Got %d picks, %d doneInfo, want equal amount", len(b.pickOptions), len(b.doneInfo)) + } }