dialOption: export WithContextDialer() (#2629)

fixes #2627

balancer/grpclb/grpclb_remote_balancer.go

@@ -279,10 +279,7 @@ func (lb *lbBalancer) dialRemoteLB(remoteLBName string) {
 		dopts = append(dopts, grpc.WithInsecure())
 	}
 	if lb.opt.Dialer != nil {
-		// WithDialer takes a different type of function, so we instead use a
-		// special DialOption here.
-		wcd := internal.WithContextDialer.(func(func(context.Context, string) (net.Conn, error)) grpc.DialOption)
-		dopts = append(dopts, wcd(lb.opt.Dialer))
+		dopts = append(dopts, grpc.WithContextDialer(lb.opt.Dialer))
 	}
 	// Explicitly set pickfirst as the balancer.
 	dopts = append(dopts, grpc.WithBalancerName(grpc.PickFirstBalancerName))

balancer/grpclb/grpclb_test.go

@@ -117,9 +117,9 @@ func (c *serverNameCheckCreds) OverrideServerName(s string) error {
 
 // fakeNameDialer replaces fakeName with localhost when dialing.
 // This will test that custom dialer is passed from Dial to grpclb.
-func fakeNameDialer(addr string, timeout time.Duration) (net.Conn, error) {
+func fakeNameDialer(ctx context.Context, addr string) (net.Conn, error) {
 	addr = strings.Replace(addr, fakeName, "localhost", 1)
-	return net.DialTimeout("tcp", addr, timeout)
+	return (&net.Dialer{}).DialContext(ctx, "tcp", addr)
 }
 
 // merge merges the new client stats into current stats.
@@ -382,7 +382,7 @@ func TestGRPCLB(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 	cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
-		grpc.WithTransportCredentials(&creds), grpc.WithDialer(fakeNameDialer))
+		grpc.WithTransportCredentials(&creds), grpc.WithContextDialer(fakeNameDialer))
 	if err != nil {
 		t.Fatalf("Failed to dial to the backend %v", err)
 	}
@@ -433,7 +433,7 @@ func TestGRPCLBWeighted(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 	cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
-		grpc.WithTransportCredentials(&creds), grpc.WithDialer(fakeNameDialer))
+		grpc.WithTransportCredentials(&creds), grpc.WithContextDialer(fakeNameDialer))
 	if err != nil {
 		t.Fatalf("Failed to dial to the backend %v", err)
 	}
@@ -498,7 +498,7 @@ func TestDropRequest(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 	cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
-		grpc.WithTransportCredentials(&creds), grpc.WithDialer(fakeNameDialer))
+		grpc.WithTransportCredentials(&creds), grpc.WithContextDialer(fakeNameDialer))
 	if err != nil {
 		t.Fatalf("Failed to dial to the backend %v", err)
 	}
@@ -585,7 +585,7 @@ func TestBalancerDisconnects(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 	cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
-		grpc.WithTransportCredentials(&creds), grpc.WithDialer(fakeNameDialer))
+		grpc.WithTransportCredentials(&creds), grpc.WithContextDialer(fakeNameDialer))
 	if err != nil {
 		t.Fatalf("Failed to dial to the backend %v", err)
 	}
@@ -666,7 +666,7 @@ func TestFallback(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 	cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
-		grpc.WithTransportCredentials(&creds), grpc.WithDialer(fakeNameDialer))
+		grpc.WithTransportCredentials(&creds), grpc.WithContextDialer(fakeNameDialer))
 	if err != nil {
 		t.Fatalf("Failed to dial to the backend %v", err)
 	}
@@ -760,7 +760,7 @@ func runAndGetStats(t *testing.T, drop bool, runRPCs func(*grpc.ClientConn)) *rp
 	cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
 		grpc.WithTransportCredentials(&creds),
 		grpc.WithPerRPCCredentials(failPreRPCCred{}),
-		grpc.WithDialer(fakeNameDialer))
+		grpc.WithContextDialer(fakeNameDialer))
 	if err != nil {
 		t.Fatalf("Failed to dial to the backend %v", err)
 	}

benchmark/benchmain/main.go

@@ -194,20 +194,19 @@ func makeClient(benchFeatures stats.Features) (testpb.BenchmarkServiceClient, fu
 	if *useBufconn {
 		bcLis := bufconn.Listen(256 * 1024)
 		lis = bcLis
-		opts = append(opts, grpc.WithDialer(func(string, time.Duration) (net.Conn, error) {
-			return nw.TimeoutDialer(
-				func(string, string, time.Duration) (net.Conn, error) {
-					return bcLis.Dial()
-				})("", "", 0)
+		opts = append(opts, grpc.WithContextDialer(func(ctx context.Context, address string) (net.Conn, error) {
+			return nw.ContextDialer(func(context.Context, string, string) (net.Conn, error) {
+				return bcLis.Dial()
+			})(ctx, "", "")
 		}))
 	} else {
 		var err error
 		lis, err = net.Listen("tcp", "localhost:0")
 		if err != nil {
 			grpclog.Fatalf("Failed to listen: %v", err)
 		}
-		opts = append(opts, grpc.WithDialer(func(_ string, timeout time.Duration) (net.Conn, error) {
-			return nw.TimeoutDialer(net.DialTimeout)("tcp", lis.Addr().String(), timeout)
+		opts = append(opts, grpc.WithContextDialer(func(ctx context.Context, address string) (net.Conn, error) {
+			return nw.ContextDialer((&net.Dialer{}).DialContext)(ctx, "tcp", lis.Addr().String())
 		}))
 	}
 	lis = nw.Listener(lis)

benchmark/benchmark.go

@@ -320,8 +320,8 @@ func runUnary(b *testing.B, benchFeatures stats.Features) {
 	defer stopper()
 	conn := NewClientConn(
 		target, grpc.WithInsecure(),
-		grpc.WithDialer(func(address string, timeout time.Duration) (net.Conn, error) {
-			return nw.TimeoutDialer(net.DialTimeout)("tcp", address, timeout)
+		grpc.WithContextDialer(func(ctx context.Context, address string) (net.Conn, error) {
+			return nw.ContextDialer((&net.Dialer{}).DialContext)(ctx, "tcp", address)
 		}),
 	)
 	tc := testpb.NewBenchmarkServiceClient(conn)
@@ -374,8 +374,8 @@ func runStream(b *testing.B, benchFeatures stats.Features) {
 	defer stopper()
 	conn := NewClientConn(
 		target, grpc.WithInsecure(),
-		grpc.WithDialer(func(address string, timeout time.Duration) (net.Conn, error) {
-			return nw.TimeoutDialer(net.DialTimeout)("tcp", address, timeout)
+		grpc.WithContextDialer(func(ctx context.Context, address string) (net.Conn, error) {
+			return nw.ContextDialer((&net.Dialer{}).DialContext)(ctx, "tcp", address)
 		}),
 	)
 	tc := testpb.NewBenchmarkServiceClient(conn)

dialoptions.go

@@ -329,14 +329,17 @@ func WithTimeout(d time.Duration) DialOption {
 	})
 }
 
-func withContextDialer(f func(context.Context, string) (net.Conn, error)) DialOption {
+// WithContextDialer returns a DialOption that sets a dialer to create
+// connections. If FailOnNonTempDialError() is set to true, and an error is
+// returned by f, gRPC checks the error's Temporary() method to decide if it
+// should try to reconnect to the network address.
+func WithContextDialer(f func(context.Context, string) (net.Conn, error)) DialOption {
 	return newFuncDialOption(func(o *dialOptions) {
 		o.copts.Dialer = f
 	})
 }
 
 func init() {
-	internal.WithContextDialer = withContextDialer
 	internal.WithResolverBuilder = withResolverBuilder
 	internal.WithHealthCheckFunc = withHealthCheckFunc
 }
@@ -345,8 +348,10 @@ func init() {
 // network addresses. If FailOnNonTempDialError() is set to true, and an error
 // is returned by f, gRPC checks the error's Temporary() method to decide if it
 // should try to reconnect to the network address.
+//
+// Deprecated: use WithContextDialer instead
 func WithDialer(f func(string, time.Duration) (net.Conn, error)) DialOption {
-	return withContextDialer(
+	return WithContextDialer(
 		func(ctx context.Context, addr string) (net.Conn, error) {
 			if deadline, ok := ctx.Deadline(); ok {
 				return f(addr, time.Until(deadline))

internal/internal.go

@@ -26,8 +26,6 @@ import (
 )
 
 var (
-	// WithContextDialer is exported by dialoptions.go
-	WithContextDialer interface{} // func(context.Context, string) (net.Conn, error) grpc.DialOption
 	// WithResolverBuilder is exported by dialoptions.go
 	WithResolverBuilder interface{} // func (resolver.Builder) grpc.DialOption
 	// WithHealthCheckFunc is not exported by dialoptions.go

test/gracefulstop_test.go

@@ -76,15 +76,15 @@ func (d *delayListener) allowClientRead() {
 	d.cc.allowRead()
 }
 
-func (d *delayListener) Dial(to time.Duration) (net.Conn, error) {
+func (d *delayListener) Dial(ctx context.Context) (net.Conn, error) {
 	if d.dialed {
 		// Only hand out one connection (net.Dial can return more even after the
 		// listener is closed).  This is not thread-safe, but Dial should never be
 		// called concurrently in this environment.
 		return nil, fmt.Errorf("no more conns")
 	}
 	d.dialed = true
-	c, err := net.DialTimeout("tcp", d.Listener.Addr().String(), to)
+	c, err := (&net.Dialer{}).DialContext(ctx, "tcp", d.Listener.Addr().String())
 	if err != nil {
 		return nil, err
 	}
@@ -137,7 +137,7 @@ func (s) TestGracefulStop(t *testing.T) {
 		closeCalled:  make(chan struct{}),
 		allowCloseCh: make(chan struct{}),
 	}
-	d := func(_ string, to time.Duration) (net.Conn, error) { return dlis.Dial(to) }
+	d := func(ctx context.Context, _ string) (net.Conn, error) { return dlis.Dial(ctx) }
 	serverGotReq := make(chan struct{})
 
 	ss := &stubServer{
@@ -180,7 +180,7 @@ func (s) TestGracefulStop(t *testing.T) {
 	// even though GracefulStop has closed the listener.
 	ctx, dialCancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer dialCancel()
-	cc, err := grpc.DialContext(ctx, "", grpc.WithInsecure(), grpc.WithBlock(), grpc.WithDialer(d))
+	cc, err := grpc.DialContext(ctx, "", grpc.WithInsecure(), grpc.WithBlock(), grpc.WithContextDialer(d))
 	if err != nil {
 		dlis.allowClientRead()
 		t.Fatalf("grpc.Dial(%q) = %v", lis.Addr().String(), err)