mirror of
https://github.com/status-im/consul.git
synced 2025-01-09 05:23:04 +00:00
d4c435856b
Adds automation for generating the map of `gRPC Method Name → Rate Limit Type` used by the middleware introduced in #15550, and will ensure we don't forget to add new endpoints. Engineers must annotate their RPCs in the proto file like so: ``` rpc Foo(FooRequest) returns (FooResponse) { option (consul.internal.ratelimit.spec) = { operation_type: READ, }; } ``` When they run `make proto` a protoc plugin `protoc-gen-consul-rate-limit` will be installed that writes rate-limit specs as a JSON array to a file called `.ratelimit.tmp` (one per protobuf package/directory). After running Buf, `make proto` will execute a post-process script that will ingest all of the `.ratelimit.tmp` files and generate a Go file containing the mappings in the `agent/grpc-middleware` package. In the enterprise repository, it will write an additional file with the enterprise-only endpoints. If an engineer forgets to add the annotation to a new RPC, the plugin will return an error like so: ``` RPC Foo is missing rate-limit specification, fix it with: import "proto-public/annotations/ratelimit/ratelimit.proto"; service Bar { rpc Foo(...) returns (...) { option (hashicorp.consul.internal.ratelimit.spec) = { operation_type: OPERATION_READ | OPERATION_WRITE | OPERATION_EXEMPT, }; } } ``` In the future, this annotation can be extended to support rate-limit category (e.g. KV vs Catalog) and to determine the retry policy.
65 lines
1.9 KiB
Go
65 lines
1.9 KiB
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/peer"
|
|
"google.golang.org/grpc/status"
|
|
"google.golang.org/grpc/tap"
|
|
|
|
recovery "github.com/grpc-ecosystem/go-grpc-middleware/recovery"
|
|
|
|
"github.com/hashicorp/consul/agent/consul/rate"
|
|
)
|
|
|
|
// ServerRateLimiterMiddleware implements a ServerInHandle function to perform
|
|
// RPC rate limiting at the cheapest possible point (before the full request has
|
|
// been decoded).
|
|
func ServerRateLimiterMiddleware(limiter rate.RequestLimitsHandler, panicHandler recovery.RecoveryHandlerFunc, logger Logger) tap.ServerInHandle {
|
|
return func(ctx context.Context, info *tap.Info) (_ context.Context, retErr error) {
|
|
// This function is called before unary and stream RPC interceptors, so we
|
|
// must handle our own panics here.
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
retErr = panicHandler(r)
|
|
}
|
|
}()
|
|
|
|
// Do not rate-limit the xDS service, it handles its own limiting.
|
|
if info.FullMethodName == "/envoy.service.discovery.v3.AggregatedDiscoveryService/DeltaAggregatedResources" {
|
|
return ctx, nil
|
|
}
|
|
|
|
peer, ok := peer.FromContext(ctx)
|
|
if !ok {
|
|
// This should never happen!
|
|
return ctx, status.Error(codes.Internal, "gRPC rate limit middleware unable to read peer")
|
|
}
|
|
|
|
operationType, ok := rpcRateLimitSpecs[info.FullMethodName]
|
|
if !ok {
|
|
logger.Warn("failed to determine which rate limit to apply to RPC", "rpc", info.FullMethodName)
|
|
return ctx, nil
|
|
}
|
|
|
|
err := limiter.Allow(rate.Operation{
|
|
Name: info.FullMethodName,
|
|
SourceAddr: peer.Addr,
|
|
Type: operationType,
|
|
})
|
|
|
|
switch {
|
|
case err == nil:
|
|
return ctx, nil
|
|
case errors.Is(err, rate.ErrRetryElsewhere):
|
|
return ctx, status.Error(codes.ResourceExhausted, err.Error())
|
|
case errors.Is(err, rate.ErrRetryLater):
|
|
return ctx, status.Error(codes.Unavailable, err.Error())
|
|
default:
|
|
return ctx, status.Error(codes.Internal, err.Error())
|
|
}
|
|
}
|
|
}
|