gRPC: Rate limiting an API on a per-RPC basis

Viewed 2948

I am looking for a way to rate-limit RPCs separately with high granularity, and to my dismay, there are not many options available for this issue. I am trying to replace a REST API with gRPC, and one of the most important features for me was the ability to add middleware for each route. Unfortunately, go-grpc-middleware only applies middleware to an entire server.

In my imagination, an ideal rate-limiting middleware for gRPC would use similar tricks as go-proto-validators, where the proto file would contain configurations for the ratelimiting itself.

1 Answers

Figured I could post a snippet for reference of how this would look like in practice, using go-grpc-middleware WithUnaryServerChain and a unary interceptor.

The idea is to add a grpc.UnaryInterceptor to the server, which will be invoked with an instance of *grpc.UnaryServerInfo. This object exports the field FullMethod, which holds the qualified name of the RPC method being called.

In the interceptor you can then implement arbitrary code before actually calling the RPC handler, including RPC-specific rate limiting logic.

// import grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
// import "google.golang.org/grpc"

    grpcServer := grpc.NewServer(
        // WithUnaryServerChain is a grpc.Server config option that accepts multiple unary interceptors.
        grpc_middleware.WithUnaryServerChain(
            // UnaryServerInterceptor provides a hook to intercept the execution of a unary RPC on the server. info
            // contains all the information of this RPC the interceptor can operate on. And handler is the wrapper
            // of the service method implementation. It is the responsibility of the interceptor to invoke handler
            // to complete the RPC.
            grpc.UnaryInterceptor(func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
                // FullMethod is the full RPC method string, i.e., /package.service/method.
                switch info.FullMethod {
                case "/mypackage.someservice/DoThings":
                    // ... rate limiting code
                    // if all is good, then call the handler
                    return handler(ctx, req)
                }
            }),
            // other `grpc.ServerOption` opts
        ),
    )
Related