vitess.io/vitess@v0.16.2/go/vt/throttler/grpcthrottlerserver/grpcthrottlerserver.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 grpcthrottlerserver contains the gRPC implementation of the server
    18  // side of the throttler service.
    19  package grpcthrottlerserver
    20  
    21  import (
    22  	"context"
    23  
    24  	"google.golang.org/grpc"
    25  
    26  	"vitess.io/vitess/go/vt/servenv"
    27  	"vitess.io/vitess/go/vt/throttler"
    28  
    29  	throttlerdatapb "vitess.io/vitess/go/vt/proto/throttlerdata"
    30  	throttlerservicepb "vitess.io/vitess/go/vt/proto/throttlerservice"
    31  )
    32  
    33  // Server is the gRPC server implementation of the Throttler service.
    34  type Server struct {
    35  	throttlerservicepb.UnimplementedThrottlerServer
    36  	manager throttler.Manager
    37  }
    38  
    39  // NewServer creates a new RPC server for a given throttler manager.
    40  func NewServer(m throttler.Manager) *Server {
    41  	return &Server{manager: m}
    42  }
    43  
    44  // MaxRates implements the gRPC server interface. It returns the current max
    45  // rate for each throttler of the process.
    46  func (s *Server) MaxRates(_ context.Context, request *throttlerdatapb.MaxRatesRequest) (_ *throttlerdatapb.MaxRatesResponse, err error) {
    47  	defer servenv.HandlePanic("throttler", &err)
    48  
    49  	rates := s.manager.MaxRates()
    50  	return &throttlerdatapb.MaxRatesResponse{
    51  		Rates: rates,
    52  	}, nil
    53  }
    54  
    55  // SetMaxRate implements the gRPC server interface. It sets the rate on all
    56  // throttlers controlled by the manager.
    57  func (s *Server) SetMaxRate(_ context.Context, request *throttlerdatapb.SetMaxRateRequest) (_ *throttlerdatapb.SetMaxRateResponse, err error) {
    58  	defer servenv.HandlePanic("throttler", &err)
    59  
    60  	names := s.manager.SetMaxRate(request.Rate)
    61  	return &throttlerdatapb.SetMaxRateResponse{
    62  		Names: names,
    63  	}, nil
    64  }
    65  
    66  // GetConfiguration implements the gRPC server interface.
    67  func (s *Server) GetConfiguration(_ context.Context, request *throttlerdatapb.GetConfigurationRequest) (_ *throttlerdatapb.GetConfigurationResponse, err error) {
    68  	defer servenv.HandlePanic("throttler", &err)
    69  
    70  	configurations, err := s.manager.GetConfiguration(request.ThrottlerName)
    71  	if err != nil {
    72  		return nil, err
    73  	}
    74  	return &throttlerdatapb.GetConfigurationResponse{
    75  		Configurations: configurations,
    76  	}, nil
    77  }
    78  
    79  // UpdateConfiguration implements the gRPC server interface.
    80  func (s *Server) UpdateConfiguration(_ context.Context, request *throttlerdatapb.UpdateConfigurationRequest) (_ *throttlerdatapb.UpdateConfigurationResponse, err error) {
    81  	defer servenv.HandlePanic("throttler", &err)
    82  
    83  	names, err := s.manager.UpdateConfiguration(request.ThrottlerName, request.Configuration, request.CopyZeroValues)
    84  	if err != nil {
    85  		return nil, err
    86  	}
    87  	return &throttlerdatapb.UpdateConfigurationResponse{
    88  		Names: names,
    89  	}, nil
    90  }
    91  
    92  // ResetConfiguration implements the gRPC server interface.
    93  func (s *Server) ResetConfiguration(_ context.Context, request *throttlerdatapb.ResetConfigurationRequest) (_ *throttlerdatapb.ResetConfigurationResponse, err error) {
    94  	defer servenv.HandlePanic("throttler", &err)
    95  
    96  	names, err := s.manager.ResetConfiguration(request.ThrottlerName)
    97  	if err != nil {
    98  		return nil, err
    99  	}
   100  	return &throttlerdatapb.ResetConfigurationResponse{
   101  		Names: names,
   102  	}, nil
   103  }
   104  
   105  // RegisterServer registers a new throttler server instance with the gRPC server.
   106  func RegisterServer(s *grpc.Server, m throttler.Manager) {
   107  	throttlerservicepb.RegisterThrottlerServer(s, NewServer(m))
   108  }
   109  
   110  func init() {
   111  	servenv.OnRun(func() {
   112  		if servenv.GRPCCheckServiceMap("throttler") {
   113  			RegisterServer(servenv.GRPCServer, throttler.GlobalManager)
   114  		}
   115  	})
   116  }