Source file src/github.com/grpc-ecosystem/go-grpc-middleware/ratelimit/ratelimit.go

Documentation: github.com/grpc-ecosystem/go-grpc-middleware/ratelimit

     1  package ratelimit
     2  
     3  import (
     4  	"context"
     5  
     6  	"google.golang.org/grpc"
     7  	"google.golang.org/grpc/codes"
     8  	"google.golang.org/grpc/status"
     9  )
    10  
    11  // Limiter defines the interface to perform request rate limiting.
    12  // If Limit function return true, the request will be rejected.
    13  // Otherwise, the request will pass.
    14  type Limiter interface {
    15  	Limit() bool
    16  }
    17  
    18  // UnaryServerInterceptor returns a new unary server interceptors that performs request rate limiting.
    19  func UnaryServerInterceptor(limiter Limiter) grpc.UnaryServerInterceptor {
    20  	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
    21  		if limiter.Limit() {
    22  			return nil, status.Errorf(codes.ResourceExhausted, "%s is rejected by grpc_ratelimit middleware, please retry later.", info.FullMethod)
    23  		}
    24  		return handler(ctx, req)
    25  	}
    26  }
    27  
    28  // StreamServerInterceptor returns a new stream server interceptor that performs rate limiting on the request.
    29  func StreamServerInterceptor(limiter Limiter) grpc.StreamServerInterceptor {
    30  	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    31  		if limiter.Limit() {
    32  			return status.Errorf(codes.ResourceExhausted, "%s is rejected by grpc_ratelimit middleware, please retry later.", info.FullMethod)
    33  		}
    34  		return handler(srv, stream)
    35  	}
    36  }
    37  

View as plain text