client: block RPCs early until the resolver has returned addresses

Author: dfawley

Diff for: clientconn.go

@@ -129,6 +129,7 @@ func DialContext(ctx context.Context, target string, opts ...DialOption) (conn *
 		dopts:          defaultDialOptions(),
 		blockingpicker: newPickerWrapper(),
 		czData:         new(channelzData),
+		resolvedAddrs:  make(chan struct{}),
 	}
 	cc.retryThrottler.Store((*retryThrottler)(nil))
 	cc.ctx, cc.cancel = context.WithCancel(context.Background())
@@ -399,6 +400,10 @@ type ClientConn struct {
 	balancerWrapper *ccBalancerWrapper
 	retryThrottler  atomic.Value
 
+	resolvedAddrsOnce sync.Once
+	resolvedAddrs     chan struct{}
+	hasResolvedAddrs  int32
+
 	channelzID int64 // channelz unique identification number
 	czData     *channelzData
 }
@@ -444,6 +449,26 @@ func (cc *ClientConn) scWatcher() {
 	}
 }
 
+// waitForResolvedAddrs blocks until the resolver has provided addresses or the
+// context expires.  Returns nil unless the context expires first; otherwise
+// returns a status error based on the context.
+func (cc *ClientConn) waitForResolvedAddrs(ctx context.Context) error {
+	// This is on the RPC path, so we use an atomic to avoid the need to do a
+	// more-expensive "select" below after the resolver has returned once.
+	if atomic.LoadInt32(&cc.hasResolvedAddrs) != 0 {
+		return nil
+	}
+	select {
+	case <-cc.resolvedAddrs:
+		atomic.StoreInt32(&cc.hasResolvedAddrs, 1)
+		return nil
+	case <-ctx.Done():
+		return status.FromContextError(ctx.Err()).Err()
+	case <-cc.ctx.Done():
+		return ErrClientConnClosing
+	}
+}
+
 func (cc *ClientConn) handleResolvedAddrs(addrs []resolver.Address, err error) {
 	cc.mu.Lock()
 	defer cc.mu.Unlock()
@@ -457,6 +482,7 @@ func (cc *ClientConn) handleResolvedAddrs(addrs []resolver.Address, err error) {
 	}
 
 	cc.curAddresses = addrs
+	cc.resolvedAddrsOnce.Do(func() { close(cc.resolvedAddrs) })
 
 	if cc.dopts.balancerBuilder == nil {
 		// Only look at balancer types and switch balancer if balancer dial

Diff for: stream.go

@@ -160,6 +160,11 @@ func newClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, meth
 		}()
 	}
 	c := defaultCallInfo()
+	// Provide an opportunity for the first RPC to see the first service config
+	// provided by the resolver.
+	if err := cc.waitForResolvedAddrs(ctx); err != nil {
+		return err
+	}
 	mc := cc.GetMethodConfig(method)
 	if mc.WaitForReady != nil {
 		c.failFast = !*mc.WaitForReady

Diff for: test/end2end_test.go

@@ -7103,3 +7103,55 @@ func (lis notifyingListener) Accept() (net.Conn, error) {
 	defer lis.connEstablished.Fire()
 	return lis.Listener.Accept()
 }
+
+func TestRPCWaitsForResolver(t *testing.T) {
+	te := testServiceConfigSetup(t, tcpClearRREnv)
+	te.startServer(&testServer{security: tcpClearRREnv.security})
+	defer te.tearDown()
+	r, rcleanup := manual.GenerateAndRegisterManualResolver()
+	defer rcleanup()
+
+	te.resolverScheme = r.Scheme()
+	te.nonBlockingDial = true
+	cc := te.clientConn()
+	tc := testpb.NewTestServiceClient(cc)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
+	// With no resolved addresses yet, this will timeout.
+	if _, err := tc.EmptyCall(ctx, &testpb.Empty{}); status.Code(err) != codes.DeadlineExceeded {
+		t.Fatalf("TestService/EmptyCall(_, _) = _, %v, want _, %s", err, codes.DeadlineExceeded)
+	}
+
+	ctx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	go func() {
+		time.Sleep(time.Second)
+		r.NewServiceConfig(`{
+		    "methodConfig": [
+		        {
+		            "name": [
+		                {
+		                    "service": "grpc.testing.TestService",
+		                    "method": "UnaryCall"
+		                }
+		            ],
+                    "maxRequestMessageBytes": 0
+		        }
+		    ]
+		}`)
+		r.NewAddress([]resolver.Address{{Addr: te.srvAddr}})
+	}()
+	// We wait a second before providing a service config and resolving
+	// addresses.  So this will wait for that and then honor the
+	// maxRequestMessageBytes it contains.
+	if _, err := tc.UnaryCall(ctx, &testpb.SimpleRequest{ResponseType: testpb.PayloadType_UNCOMPRESSABLE}); status.Code(err) != codes.ResourceExhausted {
+		t.Fatalf("TestService/UnaryCall(_, _) = _, %v, want _, nil", err)
+	}
+	if got := ctx.Err(); got != nil {
+		t.Fatalf("ctx.Err() = %v; want nil (deadline should be set short by service config)", got)
+	}
+	if _, err := tc.UnaryCall(ctx, &testpb.SimpleRequest{}); err != nil {
+		t.Fatalf("TestService/UnaryCall(_, _) = _, %v, want _, nil", err)
+	}
+}