vitess.io/vitess@v0.16.2/go/vt/throttler/throttlerclient/throttlerclient.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 throttlerclient defines the generic RPC client interface for the
    18  // throttler service. It has to be implemented for the different RPC frameworks
    19  // e.g. gRPC.
    20  package throttlerclient
    21  
    22  import (
    23  	"fmt"
    24  	"log"
    25  
    26  	"github.com/spf13/pflag"
    27  
    28  	"vitess.io/vitess/go/vt/servenv"
    29  
    30  	"context"
    31  
    32  	throttlerdatapb "vitess.io/vitess/go/vt/proto/throttlerdata"
    33  )
    34  
    35  // protocol specifics which RPC client implementation should be used.
    36  var protocol = "grpc"
    37  
    38  func init() {
    39  	servenv.OnParseFor("vttablet", registerFlags)
    40  }
    41  
    42  func registerFlags(fs *pflag.FlagSet) {
    43  	fs.StringVar(&protocol, "throttler_client_protocol", protocol, "the protocol to use to talk to the integrated throttler service")
    44  }
    45  
    46  // Client defines the generic RPC interface for the throttler service.
    47  type Client interface {
    48  	// MaxRates returns the current max rate for each throttler of the process.
    49  	MaxRates(ctx context.Context) (map[string]int64, error)
    50  
    51  	// SetMaxRate allows to change the current max rate for all throttlers
    52  	// of the process.
    53  	// It returns the names of the updated throttlers.
    54  	SetMaxRate(ctx context.Context, rate int64) ([]string, error)
    55  
    56  	// GetConfiguration returns the configuration of the MaxReplicationlag module
    57  	// for the given throttler or all throttlers if "throttlerName" is empty.
    58  	GetConfiguration(ctx context.Context, throttlerName string) (map[string]*throttlerdatapb.Configuration, error)
    59  
    60  	// UpdateConfiguration (partially) updates the configuration of the
    61  	// MaxReplicationlag module for the given throttler or all throttlers if
    62  	// "throttlerName" is empty.
    63  	// If "copyZeroValues" is true, fields with zero values will be copied
    64  	// as well.
    65  	// The function returns the names of the updated throttlers.
    66  	UpdateConfiguration(ctx context.Context, throttlerName string, configuration *throttlerdatapb.Configuration, copyZeroValues bool) ([]string, error)
    67  
    68  	// ResetConfiguration resets the configuration of the MaxReplicationlag module
    69  	// to the initial configuration for the given throttler or all throttlers if
    70  	// "throttlerName" is empty.
    71  	// The function returns the names of the updated throttlers.
    72  	ResetConfiguration(ctx context.Context, throttlerName string) ([]string, error)
    73  
    74  	// Close will terminate the connection and free resources.
    75  	Close()
    76  }
    77  
    78  // Factory has to be implemented and must create a new RPC client for a given
    79  // "addr".
    80  type Factory func(addr string) (Client, error)
    81  
    82  var factories = make(map[string]Factory)
    83  
    84  // RegisterFactory allows a client implementation to register itself.
    85  func RegisterFactory(name string, factory Factory) {
    86  	if _, ok := factories[name]; ok {
    87  		log.Fatalf("RegisterFactory: %s already exists", name)
    88  	}
    89  	factories[name] = factory
    90  }
    91  
    92  // New will return a client for the selected RPC implementation.
    93  func New(addr string) (Client, error) {
    94  	factory, ok := factories[protocol]
    95  	if !ok {
    96  		return nil, fmt.Errorf("unknown throttler client protocol: %v", protocol)
    97  	}
    98  	return factory(addr)
    99  }