vitess.io/vitess@v0.16.2/go/vt/vtctl/grpcvtctldclient/client.go (about)

     1  /*
     2  Copyright 2021 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  // Package grpcvtctldclient contains the gRPC version of the vtctld client
    18  // protocol.
    19  package grpcvtctldclient
    20  
    21  import (
    22  	"context"
    23  	"errors"
    24  	"fmt"
    25  
    26  	"google.golang.org/grpc"
    27  	"google.golang.org/grpc/connectivity"
    28  
    29  	"vitess.io/vitess/go/vt/grpcclient"
    30  	"vitess.io/vitess/go/vt/vtctl/grpcclientcommon"
    31  	"vitess.io/vitess/go/vt/vtctl/vtctldclient"
    32  
    33  	vtctlservicepb "vitess.io/vitess/go/vt/proto/vtctlservice"
    34  )
    35  
    36  var (
    37  	ErrConnectionShutdown = errors.New("gRPCVtctldClient in a SHUTDOWN state")
    38  	ErrConnectionTimeout  = errors.New("gRPC connection wait time exceeded")
    39  )
    40  
    41  const connClosedMsg = "grpc: the client connection is closed"
    42  
    43  type gRPCVtctldClient struct {
    44  	cc *grpc.ClientConn
    45  	c  vtctlservicepb.VtctldClient
    46  }
    47  
    48  //go:generate -command grpcvtctldclient go run ../vtctldclient/codegen
    49  //go:generate grpcvtctldclient --out client_gen.go
    50  
    51  func gRPCVtctldClientFactory(addr string) (vtctldclient.VtctldClient, error) {
    52  	opt, err := grpcclientcommon.SecureDialOption()
    53  	if err != nil {
    54  		return nil, err
    55  	}
    56  
    57  	conn, err := grpcclient.Dial(addr, grpcclient.FailFast(false), opt)
    58  	if err != nil {
    59  		return nil, err
    60  	}
    61  
    62  	return &gRPCVtctldClient{
    63  		cc: conn,
    64  		c:  vtctlservicepb.NewVtctldClient(conn),
    65  	}, nil
    66  }
    67  
    68  // NewWithDialOpts returns a vtctldclient.VtctldClient configured with the given
    69  // DialOptions. It is exported for use in vtadmin.
    70  func NewWithDialOpts(addr string, failFast grpcclient.FailFast, opts ...grpc.DialOption) (vtctldclient.VtctldClient, error) {
    71  	conn, err := grpcclient.Dial(addr, failFast, opts...)
    72  	if err != nil {
    73  		return nil, err
    74  	}
    75  
    76  	return &gRPCVtctldClient{
    77  		cc: conn,
    78  		c:  vtctlservicepb.NewVtctldClient(conn),
    79  	}, nil
    80  }
    81  
    82  // Close is part of the vtctldclient.VtctldClient interface.
    83  func (client *gRPCVtctldClient) Close() error {
    84  	err := client.cc.Close()
    85  	if err == nil {
    86  		client.c = nil
    87  	}
    88  
    89  	return err
    90  }
    91  
    92  // WaitForReady is part of the vtctldclient.VtctldClient interface.
    93  func (client *gRPCVtctldClient) WaitForReady(ctx context.Context) error {
    94  	// The gRPC implementation of WaitForReady uses the gRPC Connectivity API
    95  	// See https://github.com/grpc/grpc/blob/master/doc/connectivity-semantics-and-api.md
    96  	for {
    97  		select {
    98  		// A READY connection to the vtctld could not be established
    99  		// within the context timeout. The caller should close their
   100  		// existing connection and establish a new one.
   101  		case <-ctx.Done():
   102  			return ErrConnectionTimeout
   103  
   104  		// Wait to transition to READY state
   105  		default:
   106  			connState := client.cc.GetState()
   107  
   108  			switch connState {
   109  			case connectivity.Ready:
   110  				return nil
   111  
   112  			// Per https://github.com/grpc/grpc/blob/master/doc/connectivity-semantics-and-api.md,
   113  			// a client that enters the SHUTDOWN state never leaves this state, and all new RPCs should
   114  			// fail immediately. Further polling is futile, in other words, and so we
   115  			// return an error immediately to indicate that the caller can close the connection.
   116  			case connectivity.Shutdown:
   117  				return ErrConnectionShutdown
   118  
   119  			// If the connection is IDLE, CONNECTING, or in a TRANSIENT_FAILURE mode,
   120  			// then we wait to see if it will transition to a READY state.
   121  			default:
   122  				if !client.cc.WaitForStateChange(ctx, connState) {
   123  					// If the client has failed to transition, fail so that the caller can close the connection.
   124  					return fmt.Errorf("failed to transition from state %s", connState)
   125  				}
   126  			}
   127  		}
   128  	}
   129  }
   130  
   131  func init() {
   132  	vtctldclient.Register("grpc", gRPCVtctldClientFactory)
   133  }