vitess.io/vitess@v0.16.2/go/vt/throttler/grpcthrottlerclient/grpcthrottlerclient.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 grpcthrottlerclient contains the gRPC version of the throttler client protocol.
    18  package grpcthrottlerclient
    19  
    20  import (
    21  	"flag"
    22  
    23  	"context"
    24  
    25  	"google.golang.org/grpc"
    26  
    27  	"vitess.io/vitess/go/vt/grpcclient"
    28  	"vitess.io/vitess/go/vt/throttler/throttlerclient"
    29  	"vitess.io/vitess/go/vt/vterrors"
    30  
    31  	throttlerdatapb "vitess.io/vitess/go/vt/proto/throttlerdata"
    32  	throttlerservicepb "vitess.io/vitess/go/vt/proto/throttlerservice"
    33  )
    34  
    35  var (
    36  	cert = flag.String("throttler_client_grpc_cert", "", "the cert to use to connect")
    37  	key  = flag.String("throttler_client_grpc_key", "", "the key to use to connect")
    38  	ca   = flag.String("throttler_client_grpc_ca", "", "the server ca to use to validate servers when connecting")
    39  	crl  = flag.String("throttler_client_grpc_crl", "", "the server crl to use to validate server certificates when connecting")
    40  	name = flag.String("throttler_client_grpc_server_name", "", "the server name to use to validate server certificate")
    41  )
    42  
    43  type client struct {
    44  	conn       *grpc.ClientConn
    45  	gRPCClient throttlerservicepb.ThrottlerClient
    46  }
    47  
    48  func factory(addr string) (throttlerclient.Client, error) {
    49  	opt, err := grpcclient.SecureDialOption(*cert, *key, *ca, *crl, *name)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	conn, err := grpcclient.Dial(addr, grpcclient.FailFast(false), opt)
    54  	if err != nil {
    55  		return nil, err
    56  	}
    57  	gRPCClient := throttlerservicepb.NewThrottlerClient(conn)
    58  
    59  	return &client{conn, gRPCClient}, nil
    60  }
    61  
    62  // MaxRates is part of the throttlerclient.Client interface and returns the
    63  // current max rate for each throttler of the process.
    64  func (c *client) MaxRates(ctx context.Context) (map[string]int64, error) {
    65  	response, err := c.gRPCClient.MaxRates(ctx, &throttlerdatapb.MaxRatesRequest{})
    66  	if err != nil {
    67  		return nil, vterrors.FromGRPC(err)
    68  	}
    69  	return response.Rates, nil
    70  }
    71  
    72  // SetMaxRate is part of the throttlerclient.Client interface and sets the rate
    73  // on all throttlers of the server.
    74  func (c *client) SetMaxRate(ctx context.Context, rate int64) ([]string, error) {
    75  	request := &throttlerdatapb.SetMaxRateRequest{
    76  		Rate: rate,
    77  	}
    78  
    79  	response, err := c.gRPCClient.SetMaxRate(ctx, request)
    80  	if err != nil {
    81  		return nil, vterrors.FromGRPC(err)
    82  	}
    83  	return response.Names, nil
    84  }
    85  
    86  // GetConfiguration is part of the throttlerclient.Client interface.
    87  func (c *client) GetConfiguration(ctx context.Context, throttlerName string) (map[string]*throttlerdatapb.Configuration, error) {
    88  	response, err := c.gRPCClient.GetConfiguration(ctx, &throttlerdatapb.GetConfigurationRequest{
    89  		ThrottlerName: throttlerName,
    90  	})
    91  	if err != nil {
    92  		return nil, vterrors.FromGRPC(err)
    93  	}
    94  	return response.Configurations, nil
    95  }
    96  
    97  // UpdateConfiguration is part of the throttlerclient.Client interface.
    98  func (c *client) UpdateConfiguration(ctx context.Context, throttlerName string, configuration *throttlerdatapb.Configuration, copyZeroValues bool) ([]string, error) {
    99  	response, err := c.gRPCClient.UpdateConfiguration(ctx, &throttlerdatapb.UpdateConfigurationRequest{
   100  		ThrottlerName:  throttlerName,
   101  		Configuration:  configuration,
   102  		CopyZeroValues: copyZeroValues,
   103  	})
   104  	if err != nil {
   105  		return nil, vterrors.FromGRPC(err)
   106  	}
   107  	return response.Names, nil
   108  }
   109  
   110  // ResetConfiguration is part of the throttlerclient.Client interface.
   111  func (c *client) ResetConfiguration(ctx context.Context, throttlerName string) ([]string, error) {
   112  	response, err := c.gRPCClient.ResetConfiguration(ctx, &throttlerdatapb.ResetConfigurationRequest{
   113  		ThrottlerName: throttlerName,
   114  	})
   115  	if err != nil {
   116  		return nil, vterrors.FromGRPC(err)
   117  	}
   118  	return response.Names, nil
   119  }
   120  
   121  // Close is part of the throttlerclient.Client interface.
   122  func (c *client) Close() {
   123  	c.conn.Close()
   124  }
   125  
   126  func init() {
   127  	throttlerclient.RegisterFactory("grpc", factory)
   128  }