Merge pull request #882 from menghanl/blocking_balancer

Move balancer initialization into a goroutine

balancer.go

@@ -206,6 +206,11 @@ func (rr *roundRobin) watchAddrUpdates() error {
 }
 
 func (rr *roundRobin) Start(target string) error {
+	rr.mu.Lock()
+	defer rr.mu.Unlock()
+	if rr.done {
+		return ErrClientConnClosing
+	}
 	if rr.r == nil {
 		// If there is no name resolver installed, it is not needed to
 		// do name resolution. In this case, target is added into rr.addrs

clientconn.go

@@ -271,16 +271,17 @@ func DialContext(ctx context.Context, target string, opts ...DialOption) (conn *
 		cc.dopts.bs = DefaultBackoffConfig
 	}
 
-	var (
+	var ok bool
-		ok    bool
+	waitC := make(chan error, 1)
-		addrs []Address
+	go func() {
-	)
+		var addrs []Address
 		if cc.dopts.balancer == nil {
 			// Connect to target directly if balancer is nil.
 			addrs = append(addrs, Address{Addr: target})
 		} else {
 			if err := cc.dopts.balancer.Start(target); err != nil {
-			return nil, err
+				waitC <- err
+				return
 			}
 			ch := cc.dopts.balancer.Notify()
 			if ch == nil {
@@ -289,12 +290,11 @@ func DialContext(ctx context.Context, target string, opts ...DialOption) (conn *
 			} else {
 				addrs, ok = <-ch
 				if !ok || len(addrs) == 0 {
-				return nil, errNoAddr
+					waitC <- errNoAddr
+					return
 				}
 			}
 		}
-	waitC := make(chan error, 1)
-	go func() {
 		for _, a := range addrs {
 			if err := cc.resetAddrConn(a, false, nil); err != nil {
 				waitC <- err

clientconn_test.go

@@ -75,7 +75,7 @@ func TestTLSServerNameOverwrite(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Failed to create credentials %v", err)
 	}
-	conn, err := Dial("Non-Existent.Server:80", WithTransportCredentials(creds), WithTimeout(time.Millisecond))
+	conn, err := Dial("Non-Existent.Server:80", WithTransportCredentials(creds))
 	if err != nil {
 		t.Fatalf("Dial(_, _) = _, %v, want _, <nil>", err)
 	}
@@ -93,6 +93,43 @@ func TestDialContextCancel(t *testing.T) {
 	}
 }
 
+// blockingBalancer mimics the behavior of balancers whose initialization takes a long time.
+// In this test, reading from blockingBalancer.Notify() blocks forever.
+type blockingBalancer struct {
+	ch chan []Address
+}
+
+func newBlockingBalancer() Balancer {
+	return &blockingBalancer{ch: make(chan []Address)}
+}
+func (b *blockingBalancer) Start(target string) error {
+	return nil
+}
+func (b *blockingBalancer) Up(addr Address) func(error) {
+	return nil
+}
+func (b *blockingBalancer) Get(ctx context.Context, opts BalancerGetOptions) (addr Address, put func(), err error) {
+	return Address{}, nil, nil
+}
+func (b *blockingBalancer) Notify() <-chan []Address {
+	return b.ch
+}
+func (b *blockingBalancer) Close() error {
+	close(b.ch)
+	return nil
+}
+
+func TestDialWithBlockingBalancer(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	dialDone := make(chan struct{})
+	go func() {
+		DialContext(ctx, "Non-Existent.Server:80", WithBlock(), WithInsecure(), WithBalancer(newBlockingBalancer()))
+		close(dialDone)
+	}()
+	cancel()
+	<-dialDone
+}
+
 func TestCredentialsMisuse(t *testing.T) {
 	tlsCreds, err := credentials.NewClientTLSFromFile(tlsDir+"ca.pem", "x.test.youtube.com")
 	if err != nil {