gRPC Interceptor Rate Limiting

// ratelimit.go — unary and stream interceptors with per-method weights.
package ratelimit

import (
    "context"
    "fmt"
    "sync/atomic"

    "google.golang.org/grpc"
    "google.golang.org/grpc/codes"
    "google.golang.org/grpc/metadata"
    "google.golang.org/grpc/status"
)

// Cost per call, by full method name. Anything unlisted costs 1.
var methodCost = map[string]int{
    "/telemetry.Ingest/BulkImport": 25,
    "/billing.Invoices/Generate":   50,
}

type Limiter interface {
    // Consume returns allowed and the milliseconds until capacity returns.
    Consume(ctx context.Context, key string, cost int) (bool, int64, error)
}

func exhausted(retryMs int64) error {
    return status.Errorf(codes.ResourceExhausted,
        "rate limit exceeded; retry in %d ms", retryMs)
}

// Unary: one decision per call, weighted by method.
func UnaryInterceptor(l Limiter) grpc.UnaryServerInterceptor {
    return func(ctx context.Context, req any, info *grpc.UnaryServerInfo,
        handler grpc.UnaryHandler) (any, error) {

        key, err := identity(ctx)
        if err != nil {
            return nil, err
        }
        cost := 1
        if c, ok := methodCost[info.FullMethod]; ok {
            cost = c
        }
        ok, retryMs, err := l.Consume(ctx, key, cost)
        if err != nil {
            return handler(ctx, req) // fail open on limiter errors for read paths
        }
        if !ok {
            // Trailers carry the pacing hint; the code carries the semantic.
            _ = grpc.SetTrailer(ctx, metadata.Pairs("retry-after-ms", fmt.Sprint(retryMs)))
            return nil, exhausted(retryMs)
        }
        return handler(ctx, req)
    }
}

// Stream: charge the open, then charge every received message.
type countingStream struct {
    grpc.ServerStream
    ctx      context.Context
    key      string
    l        Limiter
    sample   int32
    interval int32
}

func (s *countingStream) RecvMsg(m any) error {
    if err := s.ServerStream.RecvMsg(m); err != nil {
        return err
    }
    // Sampled roll-up: charge `interval` units every `interval` messages.
    if atomic.AddInt32(&s.sample, 1)%s.interval == 0 {
        ok, retryMs, err := s.l.Consume(s.ctx, s.key, int(s.interval))
        if err == nil && !ok {
            return exhausted(retryMs) // ends the stream mid-flight, deliberately
        }
    }
    return nil
}

func StreamInterceptor(l Limiter) grpc.StreamServerInterceptor {
    return func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo,
        handler grpc.StreamHandler) error {

        key, err := identity(ss.Context())
        if err != nil {
            return err
        }
        if ok, retryMs, err := l.Consume(ss.Context(), key, 1); err == nil && !ok {
            return exhausted(retryMs) // reject at open: cheapest possible rejection
        }
        return handler(srv, &countingStream{
            ServerStream: ss, ctx: ss.Context(), key: key, l: l, interval: 25,
        })
    }
}

# 1. Unary: exceed the limit and confirm the code and the metadata hint.
grpcurl -H "x-api-key: grpc_test" -d '{}' api.example.com:443 telemetry.Ingest/Ping
# ERROR: Code: ResourceExhausted  Message: rate limit exceeded; retry in 850 ms
# Trailers: retry-after-ms: 850

# 2. Streaming: confirm the budget is enforced INSIDE the stream, not only at open.
ghz --insecure --proto telemetry.proto --call telemetry.Ingest/Ingest \
    --stream-interval 0 -n 100000 -c 1 api.example.com:443 | grep -A3 "Status code"
# expect a ResourceExhausted count > 0; a run with zero means the stream is uncharged

# 3. Health checks must never be limited.
for i in $(seq 1 5000); do grpcurl -plaintext localhost:50051 grpc.health.v1.Health/Check; done | \
  grep -c SERVING
# expect 5000