vitess.io/vitess@v0.16.2/go/vt/vttablet/tabletconn/tablet_conn.go (about)

     1  /*
     2  Copyright 2019 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 tabletconn
    18  
    19  import (
    20  	"sync"
    21  
    22  	"github.com/spf13/pflag"
    23  
    24  	"vitess.io/vitess/go/vt/grpcclient"
    25  	"vitess.io/vitess/go/vt/log"
    26  	"vitess.io/vitess/go/vt/servenv"
    27  	"vitess.io/vitess/go/vt/vterrors"
    28  	"vitess.io/vitess/go/vt/vttablet/queryservice"
    29  
    30  	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
    31  	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
    32  )
    33  
    34  var (
    35  	// ConnClosed is returned when the underlying connection was closed.
    36  	ConnClosed = vterrors.New(vtrpcpb.Code_UNAVAILABLE, "vttablet: Connection Closed")
    37  
    38  	tabletProtocol = "grpc"
    39  )
    40  
    41  // RegisterFlags registers the tabletconn flags on a given flagset. It is
    42  // exported for tests that need to inject a particular TabletProtocol.
    43  func RegisterFlags(fs *pflag.FlagSet) {
    44  	fs.StringVar(&tabletProtocol, "tablet_protocol", "grpc", "Protocol to use to make queryservice RPCs to vttablets.")
    45  }
    46  
    47  func init() {
    48  	for _, cmd := range []string{
    49  		"vtcombo",
    50  		"vtctl",
    51  		"vtctld",
    52  		"vtctldclient",
    53  		"vtgate",
    54  		"vttablet",
    55  	} {
    56  		servenv.OnParseFor(cmd, RegisterFlags)
    57  	}
    58  }
    59  
    60  // TabletDialer represents a function that will return a QueryService
    61  // object that can communicate with a tablet. Only the tablet's
    62  // HostName and PortMap should be used (and maybe the alias for debug
    63  // messages).
    64  //
    65  // timeout represents the connection timeout. If set to 0, this
    66  // connection should be established in the background and the
    67  // TabletDialer should return right away.
    68  type TabletDialer func(tablet *topodatapb.Tablet, failFast grpcclient.FailFast) (queryservice.QueryService, error)
    69  
    70  var dialers = make(map[string]TabletDialer)
    71  
    72  // mu This mutex helps us prevent data races when registering / getting dialers
    73  var mu sync.Mutex
    74  
    75  // RegisterDialer is meant to be used by TabletDialer implementations
    76  // to self register.
    77  func RegisterDialer(name string, dialer TabletDialer) {
    78  	mu.Lock()
    79  	defer mu.Unlock()
    80  	if _, ok := dialers[name]; ok {
    81  		log.Fatalf("Dialer %s already exists", name)
    82  	}
    83  	dialers[name] = dialer
    84  }
    85  
    86  // GetDialer returns the dialer to use, described by the command line flag
    87  func GetDialer() TabletDialer {
    88  	mu.Lock()
    89  	defer mu.Unlock()
    90  	td, ok := dialers[tabletProtocol]
    91  	if !ok {
    92  		log.Exitf("No dialer registered for tablet protocol %s", tabletProtocol)
    93  	}
    94  	return td
    95  }