Repository: etcd-io/etcd
File: client/v3/client.go (Go)

Code before:
// dialWithBalancer dials the client's current load balanced resolver group.  The scheme of the host
// of the provided endpoint determines the scheme used for all endpoints of the client connection.
func (c *Client) dialWithBalancer(dopts ...grpc.DialOption) (*grpc.ClientConn, error) {
	creds := c.credentialsForEndpoint(c.Endpoints()[0])
	opts := append(dopts, grpc.WithResolvers(c.resolver))
	return c.dial(creds, opts...)
}

// dial configures and dials any grpc balancer target.
func (c *Client) dial(creds grpccredentials.TransportCredentials, dopts ...grpc.DialOption) (*grpc.ClientConn, error) {
	opts := c.dialSetupOpts(creds, dopts...)

	if c.authTokenBundle != nil {
		opts = append(opts, grpc.WithPerRPCCredentials(c.authTokenBundle.PerRPCCredentials()))
	}

	opts = append(opts, c.cfg.DialOptions...)

	target := fmt.Sprintf("%s://%p/%s", resolver.Schema, c, authority(c.endpoints[0]))
	conn, err := grpc.NewClient(target, opts...)
	if err != nil {
		return nil, err
	}
	if dialTo := c.cfg.DialTimeout; dialTo > 0 {
		var cancel context.CancelFunc
		dctx, cancel := context.WithTimeout(c.ctx, dialTo)
		defer cancel()

		if err := waitForConnection(dctx, conn); err != nil {
			conn.Close()
			return nil, err
		}
	}
	return conn, nil
}

func waitForConnection(ctx context.Context, conn *grpc.ClientConn) error {
	cli := healthpb.NewHealthClient(conn)

	// Use WaitForReady to wait until the connection is ready. The health check
	// may return Unimplemented if the server does not expose the health endpoint,
	// or FailedPrecondition if the leader has not yet applied the configuration
	// change that enables it. In both cases, we can still treat the connection as
	// healthy enough to proceed.
	//
	// The downside is reduced error clarity: we may surface context deadline
	// exceeded instead of the underlying cause (for example, connection refused
	// because the endpoint is not ready, or an x509 error from TLS misconfiguration).
	// Users need to enable --debug mode to see the original error.
	//
	// Use withMax to disable retrying on Unimplemented, so that we can

Review comment (suggestion):
Maybe we could use `WithTimeoutCause` to return some custom "dial timeout" instead of generic "context timeout"

Code after:
// dialWithBalancer dials the client's current load balanced resolver group.  The scheme of the host
// of the provided endpoint determines the scheme used for all endpoints of the client connection.
func (c *Client) dialWithBalancer(dopts ...grpc.DialOption) (*grpc.ClientConn, error) {
	creds := c.credentialsForEndpoint(c.Endpoints()[0])
	opts := append(dopts, grpc.WithResolvers(c.resolver))
	return c.dial(creds, opts...)
}

// dial configures and dials any grpc balancer target.
func (c *Client) dial(creds grpccredentials.TransportCredentials, dopts ...grpc.DialOption) (*grpc.ClientConn, error) {
	opts := c.dialSetupOpts(creds, dopts...)

	if c.authTokenBundle != nil {
		opts = append(opts, grpc.WithPerRPCCredentials(c.authTokenBundle.PerRPCCredentials()))
	}

	opts = append(opts, c.cfg.DialOptions...)

	target := fmt.Sprintf("%s://%p/%s", resolver.Schema, c, authority(c.endpoints[0]))
	conn, err := grpc.NewClient(target, opts...)
	if err != nil {
		return nil, err
	}
	if dialTimeout := c.cfg.DialTimeout; dialTimeout > 0 {
		dctx, cancel := context.WithTimeout(c.ctx, dialTimeout)
		defer cancel()

		if err := waitForConnection(dctx, conn); err != nil {
			conn.Close()
			return nil, err
		}
	}
	return conn, nil
}

func waitForConnection(ctx context.Context, conn *grpc.ClientConn) error {
	cli := healthpb.NewHealthClient(conn)

	// Use WaitForReady to wait until the connection is ready. The health check
	// may return Unimplemented if the server does not expose the health endpoint,
	// or FailedPrecondition if the leader has not yet applied the configuration
	// change that enables it. In both cases, we can still treat the connection as
	// healthy enough to proceed.
	//
	// Use withMax to disable retrying on Unimplemented, so that we can
	// return the original error immediately.
	_, err := cli.Check(ctx, &healthpb.HealthCheckRequest{}, grpc.WaitForReady(true), withMax(0))
	if err == nil {
		return nil
	}
	if cerr := ctx.Err(); cerr != nil {