diff --git a/README.md b/README.md index a801e337..ce018b31 100644 --- a/README.md +++ b/README.md @@ -960,6 +960,45 @@ By default the ratelimit gRPC server binds to `0.0.0.0:8081`. To change this set socket then set `GRPC_UDS`, e.g. `GRPC_UDS=//ratelimit.sock` and leave `GRPC_HOST` and `GRPC_PORT` unmodified. +`GRPC_MAX_CONCURRENT_STREAMS` limits active streams **per gRPC connection**. +It defaults to `0`, which leaves grpc-go's effectively unlimited default in +place. Valid caps are `1` through `4294967294`; grpc-go treats `4294967295` as +unlimited, so that value is rejected. A valid cap is advertised in HTTP/2 +SETTINGS. A client respecting that setting waits for stream capacity or its own +deadline; it does not receive an immediate `ResourceExhausted` response from +this limit. A peer that sends streams past the advertised cap can instead receive +HTTP/2 `REFUSED_STREAM`. +Size the cap using the number of upstream connections and the desired +per-instance memory budget. This setting does not limit HTTP `/json` or the +number of connections; use `MAX_CONCURRENT_REQUESTS` for a per-instance bound +on admitted calls. + +**Request admission and deadlines** + +The following settings can protect the service from accumulating synchronous request handlers when its cache is slow: + +1. `MAX_CONCURRENT_REQUESTS`: maximum admitted `ShouldRateLimit` calls per service instance, shared by gRPC and HTTP. Default: `0` (disabled). Excess calls fail immediately; there is no admission queue. +1. `REQUEST_TIMEOUT`: deadline budget passed to admitted calls and their cache operations. Default: `0` (use only the caller's deadline). An earlier caller deadline is preserved. This setting alone does not limit concurrency. + +Both settings must be non-negative and take effect at process startup. Reloading descriptor configuration does not reset the admission limit or release occupied slots. Admission occurs after the request has been decoded, before configuration lookup, tracing attributes, or cache work. It does not limit connection counts, request sizes, descriptor counts, or work a backend starts asynchronously. Set `GRPC_MAX_CONCURRENT_STREAMS` separately to limit how many gRPC streams each connection can have open before calls reach admission. + +A slot stays occupied until the synchronous cache call and handler processing return, including after caller cancellation. `REQUEST_TIMEOUT` supplies a cancellation signal; it is **not a guarantee of prompt backend cleanup**. In particular, Radix can continue draining a cancelled response after returning one call and can hold later calls behind that response even after their deadlines expire. Such later handlers retain their slots. Commands already sent to Redis may still execute. Recovery must be checked by observing completed calls and successful new requests after the backend recovers. + +Overload returns gRPC `ResourceExhausted`, not an `OVER_LIMIT` quota decision. Cancellation and expiry return `Canceled` and `DeadlineExceeded` after the synchronous work returns. Global and descriptor shadow modes do not override these service errors. Envoy handles them according to its separate `failure_mode_deny` setting; check that policy before enabling admission limits, and avoid immediate retry loops. + +HTTP `/json` now preserves its request context, including when both limits are disabled. Service error codes `ResourceExhausted`, `DeadlineExceeded`, and `Canceled` map to HTTP `503`, `504`, and `408` respectively. A successful quota rejection continues to use `429`. + +With either request limit enabled, the service records the following metrics without per-request or per-descriptor labels. The names below are exported by the default Prometheus mapper; custom mappers need equivalent entries. + +| Metric | Meaning | +| ---------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `ratelimit_service_request_admission_admitted_total` | Calls admitted for processing. | +| `ratelimit_service_request_admission_rejected_total` | Calls rejected because all admission slots were occupied; excludes already-cancelled callers and quota decisions. | +| `ratelimit_service_request_admission_in_flight` | Admitted handlers whose synchronous processing and completion-metric recording have not finished. This is not a count of all outstanding Redis operations or internally retained responses. | +| `ratelimit_service_request_admission_completed_duration_seconds` | Histogram of admitted synchronous processing time up to recording the completion metric, including time retained after caller cancellation. It excludes the recording call's own wait and has no sample while cache work is still blocked. | + +The corresponding StatsD prefix is `ratelimit.service.request_admission`, with suffixes `admitted`, `rejected`, `in_flight`, and `completed_duration`. The duration is emitted in milliseconds and converted to seconds by the default Prometheus mapper, independently of `PROMETHEUS_RESPONSE_TIME_AS_MILLISECONDS`. + # Request Fields For information on the fields of a Ratelimit gRPC request please read the information @@ -1343,10 +1382,10 @@ The deployment type can be specified with the `REDIS_TYPE` / `REDIS_PERSECOND_TY ### Connection Timeout -Controls the maximum duration for Redis connection establishment, read operations, and write operations. +Controls the timeout for Redis connection establishment, not command I/O. `REQUEST_TIMEOUT` supplies a request deadline subject to the cancellation and cleanup limitations described above. -1. `REDIS_TIMEOUT`: sets the timeout for Redis connection and I/O operations. Default: `10s` -1. `REDIS_PERSECOND_TIMEOUT`: sets the timeout for per-second Redis connection and I/O operations. Default: `10s` +1. `REDIS_TIMEOUT`: timeout for Redis connection establishment. Default: `10s` +1. `REDIS_PERSECOND_TIMEOUT`: timeout for per-second Redis connection establishment. Default: `10s` ### Pool On-Empty Behavior diff --git a/src/server/grpc_streams_test.go b/src/server/grpc_streams_test.go new file mode 100644 index 00000000..017ad03c --- /dev/null +++ b/src/server/grpc_streams_test.go @@ -0,0 +1,117 @@ +package server + +import ( + "context" + "io" + "net" + "testing" + "time" + + "github.com/stretchr/testify/require" + "golang.org/x/net/http2" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/health" + healthpb "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/status" + "google.golang.org/grpc/test/bufconn" + + "github.com/envoyproxy/ratelimit/src/settings" +) + +func TestGrpcMaxConcurrentStreamsAdvertised(t *testing.T) { + for _, tc := range []struct { + name string + limit uint32 + }{ + {name: "default unlimited"}, + {name: "configured limit", limit: 16}, + } { + t.Run(tc.name, func(t *testing.T) { + s := settings.Settings{GrpcMaxConcurrentStreams: tc.limit} + grpcServer := grpc.NewServer(grpcServerOptions(s)...) + listener := bufconn.Listen(1024 * 1024) + serveDone := make(chan error, 1) + go func() { serveDone <- grpcServer.Serve(listener) }() + t.Cleanup(func() { + grpcServer.Stop() + require.NoError(t, <-serveDone) + }) + + conn, err := listener.Dial() + require.NoError(t, err) + defer conn.Close() + require.NoError(t, conn.SetDeadline(time.Now().Add(5*time.Second))) + _, err = io.WriteString(conn, http2.ClientPreface) + require.NoError(t, err) + framer := http2.NewFramer(conn, conn) + require.NoError(t, framer.WriteSettings()) + frame, err := framer.ReadFrame() + require.NoError(t, err) + serverSettings, ok := frame.(*http2.SettingsFrame) + require.True(t, ok, "first server frame must be SETTINGS") + + var advertised uint32 + found := false + require.NoError(t, serverSettings.ForeachSetting(func(setting http2.Setting) error { + if setting.ID == http2.SettingMaxConcurrentStreams { + advertised = setting.Val + found = true + } + return nil + })) + if tc.limit == 0 { + require.False(t, found, "zero must preserve grpc-go's default") + } else { + require.True(t, found, "server must advertise the configured cap") + require.Equal(t, tc.limit, advertised) + } + }) + } +} + +func TestGrpcStreamLimitWaitsUntilCapacityOrCallerDeadline(t *testing.T) { + s := settings.Settings{GrpcMaxConcurrentStreams: 1} + grpcServer := grpc.NewServer(grpcServerOptions(s)...) + healthpb.RegisterHealthServer(grpcServer, health.NewServer()) + listener := bufconn.Listen(1024 * 1024) + serveDone := make(chan error, 1) + go func() { serveDone <- grpcServer.Serve(listener) }() + t.Cleanup(func() { + grpcServer.Stop() + require.NoError(t, <-serveDone) + }) + + clientConn, err := grpc.NewClient("passthrough:///bufnet", + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() }), + ) + require.NoError(t, err) + defer clientConn.Close() + client := healthpb.NewHealthClient(clientConn) + request := &healthpb.HealthCheckRequest{Service: "ratelimit"} + + firstCtx, cancelFirst := context.WithTimeout(context.Background(), 5*time.Second) + defer cancelFirst() + first, err := client.Watch(firstCtx, request) + require.NoError(t, err) + _, err = first.Recv() + require.NoError(t, err) // First Watch now holds the only active stream. + + secondCtx, cancelSecond := context.WithTimeout(context.Background(), 500*time.Millisecond) + defer cancelSecond() + second, err := client.Watch(secondCtx, request) + if err == nil { + _, err = second.Recv() + } + require.Equal(t, codes.DeadlineExceeded, status.Code(err)) + + cancelFirst() + thirdCtx, cancelThird := context.WithTimeout(context.Background(), 3*time.Second) + defer cancelThird() + third, err := client.Watch(thirdCtx, request) + require.NoError(t, err) + _, err = third.Recv() + require.NoError(t, err) +} diff --git a/src/server/server_impl.go b/src/server/server_impl.go index a82f2821..739b8b9d 100644 --- a/src/server/server_impl.go +++ b/src/server/server_impl.go @@ -26,8 +26,10 @@ import ( gostats "github.com/lyft/gostats" logger "github.com/sirupsen/logrus" "google.golang.org/grpc" + "google.golang.org/grpc/codes" "google.golang.org/grpc/health" healthpb "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/status" "github.com/envoyproxy/ratelimit/src/limiter" "github.com/envoyproxy/ratelimit/src/settings" @@ -85,7 +87,7 @@ func NewJsonHandler(svc pb.RateLimitServiceServer) func(http.ResponseWriter, *ht return func(writer http.ResponseWriter, request *http.Request) { var req pb.RateLimitRequest - ctx := context.Background() + ctx := request.Context() body, err := io.ReadAll(request.Body) if err != nil { @@ -103,7 +105,16 @@ func NewJsonHandler(svc pb.RateLimitServiceServer) func(http.ResponseWriter, *ht resp, err := svc.ShouldRateLimit(ctx, &req) if err != nil { logger.Warnf("error: %s", err.Error()) - writeHttpStatus(writer, http.StatusBadRequest) + httpStatus := http.StatusBadRequest + switch status.Code(err) { + case codes.ResourceExhausted: + httpStatus = http.StatusServiceUnavailable + case codes.DeadlineExceeded: + httpStatus = http.StatusGatewayTimeout + case codes.Canceled: + httpStatus = http.StatusRequestTimeout + } + writeHttpStatus(writer, httpStatus) return } @@ -246,18 +257,7 @@ func newServer(s settings.Settings, name string, statsManager stats.Manager, loc ret.store.AddStatGenerator(limiter.NewLocalCacheStats(localCache, ret.scope.Scope("localcache"))) } - keepaliveOpt := grpc.KeepaliveParams(keepalive.ServerParameters{ - MaxConnectionAge: s.GrpcMaxConnectionAge, - MaxConnectionAgeGrace: s.GrpcMaxConnectionAgeGrace, - }) - grpcOptions := []grpc.ServerOption{ - keepaliveOpt, - grpc.ChainUnaryInterceptor( - s.GrpcUnaryInterceptor, // chain otel interceptor after the input interceptor - otelgrpc.UnaryServerInterceptor(), - ), - grpc.StreamInterceptor(otelgrpc.StreamServerInterceptor()), - } + grpcOptions := grpcServerOptions(s) if s.GrpcServerUseTLS { grpcServerTlsConfig := s.GrpcServerTlsConfig ret.grpcCertProvider = provider.NewCertProvider(s, ret.store, s.GrpcServerTlsCert, s.GrpcServerTlsKey) @@ -349,6 +349,25 @@ func newServer(s settings.Settings, name string, statsManager stats.Manager, loc return ret } +func grpcServerOptions(s settings.Settings) []grpc.ServerOption { + keepaliveOpt := grpc.KeepaliveParams(keepalive.ServerParameters{ + MaxConnectionAge: s.GrpcMaxConnectionAge, + MaxConnectionAgeGrace: s.GrpcMaxConnectionAgeGrace, + }) + grpcOptions := []grpc.ServerOption{ + keepaliveOpt, + grpc.ChainUnaryInterceptor( + s.GrpcUnaryInterceptor, // chain otel interceptor after the input interceptor + otelgrpc.UnaryServerInterceptor(), + ), + grpc.StreamInterceptor(otelgrpc.StreamServerInterceptor()), + } + if s.GrpcMaxConcurrentStreams > 0 { + grpcOptions = append(grpcOptions, grpc.MaxConcurrentStreams(s.GrpcMaxConcurrentStreams)) + } + return grpcOptions +} + func (server *server) Stop() { server.grpcServer.GracefulStop() server.listenerMu.Lock() diff --git a/src/service/ratelimit.go b/src/service/ratelimit.go index 85dfa90d..a2420ff0 100644 --- a/src/service/ratelimit.go +++ b/src/service/ratelimit.go @@ -10,6 +10,7 @@ import ( "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/trace" + "google.golang.org/grpc/status" "google.golang.org/protobuf/types/known/structpb" "github.com/envoyproxy/ratelimit/src/settings" @@ -59,6 +60,7 @@ type service struct { globalQuotaMode bool responseDynamicMetadataEnabled bool useCalendarMonthRateLimit bool + requestLimits *requestLimits } func (this *service) SetConfig(updateEvent provider.ConfigUpdateEvent, healthyWithAtLeastOneConfigLoad bool) { @@ -443,6 +445,26 @@ func (this *service) ShouldRateLimit( ctx context.Context, request *pb.RateLimitRequest, ) (finalResponse *pb.RateLimitResponse, finalError error) { + if this.requestLimits != nil { + var release func() + ctx, release, finalError = this.requestLimits.acquire(ctx) + if finalError != nil { + return nil, finalError + } + defer release() + // Do not start cache work if cancellation raced with admission. Check + // again on return because a backend may finish after its deadline. + if err := ctx.Err(); err != nil { + return nil, status.FromContextError(err).Err() + } + defer func() { + if err := ctx.Err(); err != nil { + finalResponse = nil + finalError = status.FromContextError(err).Err() + } + }() + } + logger.Debugf("ShouldRateLimit: %+v", request) // Generate trace _, span := tracer.Start( @@ -493,6 +515,7 @@ func (this *service) GetCurrentConfig() (config.RateLimitConfig, bool, bool) { func NewService(cache limiter.RateLimitCache, configProvider provider.RateLimitConfigProvider, statsManager stats.Manager, health *server.HealthChecker, clock utils.TimeSource, shadowMode, forceStart bool, healthyWithAtLeastOneConfigLoad bool, + options ...ServiceOption, ) RateLimitServiceServer { newService := &service{ configLock: sync.RWMutex{}, @@ -505,6 +528,9 @@ func NewService(cache limiter.RateLimitCache, configProvider provider.RateLimitC globalQuotaMode: false, customHeaderClock: clock, } + for _, option := range options { + option(newService) + } if !forceStart { logger.Info("Waiting for initial ratelimit config update event") diff --git a/src/service/request_limits.go b/src/service/request_limits.go new file mode 100644 index 00000000..c75dd346 --- /dev/null +++ b/src/service/request_limits.go @@ -0,0 +1,79 @@ +package ratelimit + +import ( + "context" + "time" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/envoyproxy/ratelimit/src/stats" +) + +// ServiceOption configures a service before it starts accepting requests. +type ServiceOption func(*service) + +// WithRequestLimits limits concurrent ShouldRateLimit calls and supplies a +// deadline to their cache operations. Zero disables the corresponding limit. +// The deadline does not guarantee that a cache operation returns promptly. +func WithRequestLimits(maxConcurrentRequests int, timeout time.Duration) ServiceOption { + if maxConcurrentRequests < 0 { + panic("MAX_CONCURRENT_REQUESTS must be >= 0") + } + if timeout < 0 { + panic("REQUEST_TIMEOUT must be >= 0") + } + return func(s *service) { + if maxConcurrentRequests == 0 && timeout == 0 { + return + } + limits := &requestLimits{timeout: timeout, stats: s.stats.RequestAdmission} + if maxConcurrentRequests > 0 { + limits.active = make(chan struct{}, maxConcurrentRequests) + } + s.requestLimits = limits + } +} + +type requestLimits struct { + active chan struct{} + timeout time.Duration + stats stats.RequestAdmissionStats +} + +// acquire does not queue requests. The caller must defer release until all +// synchronous cache work has returned, even if the request context is cancelled. +func (l *requestLimits) acquire(ctx context.Context) (context.Context, func(), error) { + if err := ctx.Err(); err != nil { + return nil, nil, status.FromContextError(err).Err() + } + if l.active != nil { + select { + case l.active <- struct{}{}: + default: + l.stats.Rejected.Inc() + return nil, nil, status.Error(codes.ResourceExhausted, "maximum concurrent rate limit requests reached") + } + } + + cancel := func() {} + if l.timeout > 0 { + ctx, cancel = context.WithTimeout(ctx, l.timeout) + } + l.stats.Admitted.Inc() + l.stats.InFlight.Add(1) + start := time.Now() + return ctx, func() { + cancel() + defer func() { + l.stats.InFlight.Sub(1) + if l.active != nil { + <-l.active + } + }() + // Record processing time up to metric emission, not just until caller + // cancellation. Milliseconds are converted to seconds by the mapper. + // Keep the admission slot during synchronous metric export as well. + l.stats.CompletedDuration.AddValue(float64(time.Since(start)) / float64(time.Millisecond)) + }, nil +} diff --git a/src/service/request_limits_redis_test.go b/src/service/request_limits_redis_test.go new file mode 100644 index 00000000..c8014668 --- /dev/null +++ b/src/service/request_limits_redis_test.go @@ -0,0 +1,173 @@ +package ratelimit + +import ( + "bufio" + "context" + "math/rand" + "net" + "sync" + "testing" + "time" + + pb "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v3" + gostats "github.com/lyft/gostats" + "github.com/mediocregopher/radix/v4/resp" + "github.com/mediocregopher/radix/v4/resp/resp3" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/envoyproxy/ratelimit/src/redis" + "github.com/envoyproxy/ratelimit/src/settings" + "github.com/envoyproxy/ratelimit/src/stats" + "github.com/envoyproxy/ratelimit/src/utils" +) + +// stalledRedis receives real RESP commands over loopback, but withholds command +// replies until released. Reads continue while replies are stalled, so the test +// can prove that another request was actually submitted behind a cancelled read. +func stalledRedis(t *testing.T) (string, <-chan struct{}, func()) { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + commands := make(chan struct{}, 16) + unblock := make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(unblock) }) } + var workers sync.WaitGroup + var mu sync.Mutex + var conns []net.Conn + workers.Add(1) + go func() { + defer workers.Done() + for { + conn, err := listener.Accept() + if err != nil { + return + } + mu.Lock() + conns = append(conns, conn) + mu.Unlock() + workers.Add(1) + go func() { + defer workers.Done() + defer conn.Close() + type reply struct { + value string + stall bool + } + replies := make(chan reply, 16) + writerDone := make(chan struct{}) + go func() { + defer close(writerDone) + for response := range replies { + if response.stall { + <-unblock + } + if _, err := conn.Write([]byte(response.value)); err != nil { + return + } + } + }() + defer func() { close(replies); <-writerDone }() + reader := bufio.NewReader(conn) + opts := resp.NewOpts() + for { + var command []string + if err := resp3.Unmarshal(reader, &command, opts); err != nil { + return + } + response := reply{value: "-ERR unsupported test command\r\n"} + switch command[0] { + case "PING": + response.value = "+PONG\r\n" + case "INCRBY": + commands <- struct{}{} + response = reply{value: ":1\r\n", stall: true} + case "EXPIRE": + response.value = ":1\r\n" + } + select { + case replies <- response: + case <-writerDone: + return + } + } + }() + } + }() + t.Cleanup(func() { + release() + _ = listener.Close() + mu.Lock() + for _, conn := range conns { + _ = conn.Close() + } + mu.Unlock() + workers.Wait() + }) + return listener.Addr().String(), commands, release +} + +func TestRequestAdmissionBoundsWorkBehindStalledRedisCancellation(t *testing.T) { + addr, commands, releaseRedis := stalledRedis(t) + store := gostats.NewStore(gostats.NewNullSink(), false) + manager := stats.NewStatManager(store, settings.Settings{}) + client := redis.NewClientImpl(context.Background(), store, false, "", "tcp", "single", addr, + 1, 0, 0, nil, false, nil, time.Second, "WAIT", "", time.Millisecond, time.Millisecond, time.Second, false) + t.Cleanup(func() { + // Unblock the peer before closing the pool, including on assertion + // failure: pool shutdown may wait for active shared calls to return. + releaseRedis() + _ = client.Close() + }) + cache := redis.NewFixedRateLimitCacheImpl(client, nil, utils.NewTimeSourceImpl(), + rand.New(utils.NewLockedSource(1)), 0, nil, 0.8, "", manager, false, false) + s := newAdmissionTestService(t, cache, 1, 200*time.Millisecond) + + // The first deadline returns to its caller, but Radix can still be draining + // that response internally. This increment does not claim to stop Redis work. + firstDone := make(chan error, 1) + go func() { _, err := s.ShouldRateLimit(context.Background(), admissionRequest()); firstDone <- err }() + waitAdmissionSignal(t, commands) + require.Equal(t, codes.DeadlineExceeded, status.Code(waitAdmissionResult(t, firstDone))) + require.Zero(t, s.stats.RequestAdmission.InFlight.Value()) + + // A second request is written behind the unread response. Cancelling it + // cannot interrupt the earlier response's drain; its service slot must stay + // occupied until its own synchronous cache call has returned. + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + secondDone := make(chan error, 1) + go func() { _, err := s.ShouldRateLimit(ctx, admissionRequest()); secondDone <- err }() + waitAdmissionSignal(t, commands) + cancel() + select { + case <-secondDone: + t.Fatal("Radix unexpectedly returned while the earlier response was still withheld") + case <-time.After(50 * time.Millisecond): + } + require.Equal(t, uint64(1), s.stats.RequestAdmission.InFlight.Value()) + for range 100 { + response, err := s.ShouldRateLimit(context.Background(), admissionRequest()) + require.Nil(t, response) + require.Equal(t, codes.ResourceExhausted, status.Code(err)) + } + require.Equal(t, uint64(2), s.stats.RequestAdmission.Admitted.Value()) + require.Equal(t, uint64(100), s.stats.RequestAdmission.Rejected.Value()) + select { + case <-commands: + t.Fatal("a rejected request submitted another Redis increment") + default: + } + + // Recovery is established only after the server releases the withheld + // responses. A subsequent normal request can then reuse the freed capacity. + releaseRedis() + require.Equal(t, codes.Canceled, status.Code(waitAdmissionResult(t, secondDone))) + require.Zero(t, s.stats.RequestAdmission.InFlight.Value()) + response, err := s.ShouldRateLimit(context.Background(), admissionRequest()) + require.NoError(t, err) + require.Equal(t, pb.RateLimitResponse_OK, response.OverallCode) + require.Equal(t, uint64(3), s.stats.RequestAdmission.Admitted.Value()) +} diff --git a/src/service/request_limits_test.go b/src/service/request_limits_test.go new file mode 100644 index 00000000..d6e5d15b --- /dev/null +++ b/src/service/request_limits_test.go @@ -0,0 +1,336 @@ +package ratelimit + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + ratelimitv3 "github.com/envoyproxy/go-control-plane/envoy/extensions/common/ratelimit/v3" + pb "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v3" + gostats "github.com/lyft/gostats" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/envoyproxy/ratelimit/src/config" + "github.com/envoyproxy/ratelimit/src/limiter" + "github.com/envoyproxy/ratelimit/src/redis" + "github.com/envoyproxy/ratelimit/src/settings" + "github.com/envoyproxy/ratelimit/src/stats" +) + +type admissionTestCache func(context.Context) []*pb.RateLimitResponse_DescriptorStatus + +func (f admissionTestCache) DoLimit(ctx context.Context, _ *pb.RateLimitRequest, _ []*config.RateLimit) []*pb.RateLimitResponse_DescriptorStatus { + return f(ctx) +} + +func (admissionTestCache) Flush() {} + +func admissionOK() []*pb.RateLimitResponse_DescriptorStatus { + return []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OK}} +} + +func admissionRequest() *pb.RateLimitRequest { + return &pb.RateLimitRequest{ + Domain: "test", + Descriptors: []*ratelimitv3.RateLimitDescriptor{{ + Entries: []*ratelimitv3.RateLimitDescriptor_Entry{{Key: "key", Value: "value"}}, + }}, + } +} + +func newAdmissionTestService(t *testing.T, cache limiter.RateLimitCache, max int, timeout time.Duration) *service { + t.Helper() + store := gostats.NewStore(gostats.NewNullSink(), false) + manager := stats.NewStatManager(store, settings.Settings{}) + s := &service{ + cache: cache, + stats: manager.NewServiceStats(), + config: config.NewRateLimitConfigImpl([]config.RateLimitConfigToLoad{{ + Name: "test", + ConfigYaml: &config.YamlRoot{ + Domain: "test", + Descriptors: []config.YamlDescriptor{{ + Key: "key", Value: "value", + RateLimit: &config.YamlRateLimit{RequestsPerUnit: 100, Unit: "second"}, + }}, + }, + }}, manager, false), + } + WithRequestLimits(max, timeout)(s) + return s +} + +func waitAdmissionResult(t *testing.T, done <-chan error) error { + t.Helper() + select { + case err := <-done: + return err + case <-time.After(3 * time.Second): + t.Fatal("request did not return") + return nil + } +} + +func waitAdmissionSignal(t *testing.T, signal <-chan struct{}) { + t.Helper() + select { + case <-signal: + case <-time.After(3 * time.Second): + t.Fatal("cache operation did not start") + } +} + +func TestRequestAdmissionRejectsExcessWithoutQueueing(t *testing.T) { + const max = 3 + entered := make(chan struct{}, max) + unblock := make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(unblock) }) } + t.Cleanup(release) + var calls atomic.Int64 + s := newAdmissionTestService(t, admissionTestCache(func(context.Context) []*pb.RateLimitResponse_DescriptorStatus { + calls.Add(1) + entered <- struct{}{} + <-unblock + return admissionOK() + }), max, 0) + done := make(chan error, max) + for range max { + go func() { + _, err := s.ShouldRateLimit(context.Background(), admissionRequest()) + done <- err + }() + waitAdmissionSignal(t, entered) + } + require.Equal(t, uint64(max), s.stats.RequestAdmission.InFlight.Value()) + + // The overload response is independent of descriptor/global shadow mode: + // it is a service failure, never a successful OVER_LIMIT quota decision. + for _, shadow := range []bool{false, true} { + s.configLock.Lock() + s.globalShadowMode = shadow + s.configLock.Unlock() + for range 50 { + response, err := s.ShouldRateLimit(context.Background(), admissionRequest()) + require.Nil(t, response) + require.Equal(t, codes.ResourceExhausted, status.Code(err)) + } + } + require.Equal(t, int64(max), calls.Load()) + require.Equal(t, uint64(100), s.stats.RequestAdmission.Rejected.Value()) + require.Equal(t, uint64(max), s.stats.RequestAdmission.Admitted.Value()) + release() + for range max { + require.NoError(t, waitAdmissionResult(t, done)) + } + require.Zero(t, s.stats.RequestAdmission.InFlight.Value()) + response, err := s.ShouldRateLimit(context.Background(), admissionRequest()) + require.NoError(t, err) + require.Equal(t, pb.RateLimitResponse_OK, response.OverallCode) +} + +type admissionConfigUpdate struct{ config.RateLimitConfig } + +func (u admissionConfigUpdate) GetConfig() (config.RateLimitConfig, any) { + return u.RateLimitConfig, nil +} + +func TestRequestAdmissionCancellationAndConfigReloadKeepSlotUntilCacheReturns(t *testing.T) { + entered := make(chan struct{}, 1) + unblock := make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(unblock) }) } + t.Cleanup(release) + s := newAdmissionTestService(t, admissionTestCache(func(ctx context.Context) []*pb.RateLimitResponse_DescriptorStatus { + entered <- struct{}{} + <-ctx.Done() + <-unblock // Deliberately model a backend which has not finished cancellation. + return admissionOK() + }), 1, 0) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan error, 1) + go func() { _, err := s.ShouldRateLimit(ctx, admissionRequest()); done <- err }() + waitAdmissionSignal(t, entered) + cancel() + s.SetConfig(admissionConfigUpdate{s.config}, false) + + _, err := s.ShouldRateLimit(context.Background(), admissionRequest()) + require.Equal(t, codes.ResourceExhausted, status.Code(err)) + require.Equal(t, uint64(1), s.stats.RequestAdmission.InFlight.Value()) + select { + case <-done: + t.Fatal("cancelled caller released its slot before cache work returned") + default: + } + release() + require.Equal(t, codes.Canceled, status.Code(waitAdmissionResult(t, done))) + require.Zero(t, s.stats.RequestAdmission.InFlight.Value()) +} + +func TestRequestAdmissionDeadlinePreservesEarlierCallerDeadline(t *testing.T) { + for _, shorterCaller := range []bool{false, true} { + t.Run(map[bool]string{false: "server deadline", true: "earlier caller deadline"}[shorterCaller], func(t *testing.T) { + timeout := 100 * time.Millisecond + ctx := context.Background() + var expected time.Time + if shorterCaller { + timeout = time.Hour + } + var observed time.Time + s := newAdmissionTestService(t, admissionTestCache(func(ctx context.Context) []*pb.RateLimitResponse_DescriptorStatus { + observed, _ = ctx.Deadline() + <-ctx.Done() + return admissionOK() // A late success must not become a quota response. + }), 1, timeout) + if shorterCaller { + var cancel context.CancelFunc + ctx, cancel = context.WithDeadline(ctx, time.Now().Add(100*time.Millisecond)) + defer cancel() + expected, _ = ctx.Deadline() + } + before := time.Now() + response, err := s.ShouldRateLimit(ctx, admissionRequest()) + require.Nil(t, response) + require.Equal(t, codes.DeadlineExceeded, status.Code(err)) + if shorterCaller { + require.Equal(t, expected, observed) + } else { + require.False(t, observed.Before(before.Add(timeout))) + require.True(t, observed.Before(time.Now().Add(timeout))) + } + require.Zero(t, s.stats.RequestAdmission.InFlight.Value()) + }) + } +} + +func TestRequestAdmissionRejectsAlreadyCancelledContextBeforeCache(t *testing.T) { + s := newAdmissionTestService(t, admissionTestCache(func(context.Context) []*pb.RateLimitResponse_DescriptorStatus { + t.Fatal("cache should not be called") + return nil + }), 1, 0) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + response, err := s.ShouldRateLimit(ctx, admissionRequest()) + require.Nil(t, response) + require.Equal(t, codes.Canceled, status.Code(err)) + require.Zero(t, s.stats.RequestAdmission.Admitted.Value()) + require.Zero(t, s.stats.RequestAdmission.Rejected.Value()) + require.Zero(t, s.stats.RequestAdmission.InFlight.Value()) +} + +func TestRequestAdmissionPanicReleasesSlot(t *testing.T) { + for _, tc := range []struct { + name string + value any + recovered bool + }{ + {"Redis error", redis.RedisError("backend failed"), true}, + {"unexpected panic", "unexpected failure", false}, + } { + t.Run(tc.name, func(t *testing.T) { + calls := 0 + s := newAdmissionTestService(t, admissionTestCache(func(context.Context) []*pb.RateLimitResponse_DescriptorStatus { + calls++ + if calls == 1 { + panic(tc.value) + } + return admissionOK() + }), 1, 0) + if tc.recovered { + response, err := s.ShouldRateLimit(context.Background(), admissionRequest()) + require.Nil(t, response) + require.Equal(t, tc.value, err) + } else { + require.PanicsWithValue(t, tc.value, func() { + _, _ = s.ShouldRateLimit(context.Background(), admissionRequest()) + }) + } + require.Zero(t, s.stats.RequestAdmission.InFlight.Value()) + response, err := s.ShouldRateLimit(context.Background(), admissionRequest()) + require.NoError(t, err) + require.Equal(t, pb.RateLimitResponse_OK, response.OverallCode) + }) + } +} + +func TestRequestTimeoutAloneDoesNotLimitConcurrency(t *testing.T) { + entered := make(chan struct{}, 2) + s := newAdmissionTestService(t, admissionTestCache(func(ctx context.Context) []*pb.RateLimitResponse_DescriptorStatus { + entered <- struct{}{} + <-ctx.Done() + return admissionOK() + }), 0, time.Hour) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan error, 2) + for range 2 { + go func() { _, err := s.ShouldRateLimit(ctx, admissionRequest()); done <- err }() + waitAdmissionSignal(t, entered) + } + require.Nil(t, s.requestLimits.active) + require.Equal(t, uint64(2), s.stats.RequestAdmission.InFlight.Value()) + cancel() + for range 2 { + require.Equal(t, codes.Canceled, status.Code(waitAdmissionResult(t, done))) + } + require.Zero(t, s.stats.RequestAdmission.InFlight.Value()) +} + +type admissionTimerSink struct{ flush func() } + +func (admissionTimerSink) FlushCounter(string, uint64) {} +func (admissionTimerSink) FlushGauge(string, uint64) {} +func (s admissionTimerSink) FlushTimer(string, float64) { s.flush() } + +func TestRequestAdmissionRetainsSlotThroughMetricExport(t *testing.T) { + entered := make(chan struct{}, 1) + unblock := make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(unblock) }) } + t.Cleanup(release) + s := newAdmissionTestService(t, admissionTestCache(func(context.Context) []*pb.RateLimitResponse_DescriptorStatus { + return admissionOK() + }), 1, 0) + store := gostats.NewStore(admissionTimerSink{flush: func() { + entered <- struct{}{} + <-unblock + }}, false) + s.requestLimits.stats.CompletedDuration = store.NewMilliTimer("completed_duration") + done := make(chan error, 1) + go func() { _, err := s.ShouldRateLimit(context.Background(), admissionRequest()); done <- err }() + waitAdmissionSignal(t, entered) + require.Equal(t, uint64(1), s.stats.RequestAdmission.InFlight.Value()) + response, err := s.ShouldRateLimit(context.Background(), admissionRequest()) + require.Nil(t, response) + require.Equal(t, codes.ResourceExhausted, status.Code(err)) + release() + require.NoError(t, waitAdmissionResult(t, done)) + require.Zero(t, s.stats.RequestAdmission.InFlight.Value()) +} + +func TestRequestAdmissionMetricPanicReleasesSlot(t *testing.T) { + s := newAdmissionTestService(t, admissionTestCache(func(context.Context) []*pb.RateLimitResponse_DescriptorStatus { + return admissionOK() + }), 1, 0) + store := gostats.NewStore(admissionTimerSink{flush: func() { panic("metric failure") }}, false) + s.requestLimits.stats.CompletedDuration = store.NewMilliTimer("completed_duration") + require.PanicsWithValue(t, "metric failure", func() { + _, _ = s.ShouldRateLimit(context.Background(), admissionRequest()) + }) + require.Zero(t, s.stats.RequestAdmission.InFlight.Value()) + require.Empty(t, s.requestLimits.active) +} + +func TestRequestLimitsValidateAndDefaultToDisabled(t *testing.T) { + s := &service{} + WithRequestLimits(0, 0)(s) + assert.Nil(t, s.requestLimits) + assert.Panics(t, func() { WithRequestLimits(-1, 0) }) + assert.Panics(t, func() { WithRequestLimits(1, -time.Second) }) +} diff --git a/src/service_cmd/runner/runner.go b/src/service_cmd/runner/runner.go index 9ac6d09f..372b25c0 100644 --- a/src/service_cmd/runner/runner.go +++ b/src/service_cmd/runner/runner.go @@ -206,6 +206,7 @@ func (runner *Runner) Run() { s.GlobalShadowMode, s.ForceStartWithoutInitialConfig, s.HealthyWithAtLeastOneConfigLoaded, + ratelimit.WithRequestLimits(s.MaxConcurrentRequests, s.RequestTimeout), ) srv.AddDebugHttpEndpoint( diff --git a/src/settings/request_limits_test.go b/src/settings/request_limits_test.go new file mode 100644 index 00000000..07bb6e62 --- /dev/null +++ b/src/settings/request_limits_test.go @@ -0,0 +1,65 @@ +package settings + +import ( + "os" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestRequestLimitsSettings(t *testing.T) { + for _, tc := range []struct { + name, max, timeout string + wantMax int + wantTimeout time.Duration + invalid bool + }{ + {name: "disabled", max: "0", timeout: "0"}, + {name: "configured", max: "32", timeout: "250ms", wantMax: 32, wantTimeout: 250 * time.Millisecond}, + {name: "negative concurrency", max: "-1", timeout: "0", invalid: true}, + {name: "negative timeout", max: "0", timeout: "-1s", invalid: true}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv("MAX_CONCURRENT_REQUESTS", tc.max) + t.Setenv("REQUEST_TIMEOUT", tc.timeout) + if tc.invalid { + require.Panics(t, func() { NewSettings() }) + return + } + s := NewSettings() + require.Equal(t, tc.wantMax, s.MaxConcurrentRequests) + require.Equal(t, tc.wantTimeout, s.RequestTimeout) + }) + } +} + +func TestGrpcMaxConcurrentStreamsSettings(t *testing.T) { + t.Run("unset", func(t *testing.T) { + t.Setenv("GRPC_MAX_CONCURRENT_STREAMS", "") + require.NoError(t, os.Unsetenv("GRPC_MAX_CONCURRENT_STREAMS")) + s := NewSettings() + require.Zero(t, s.GrpcMaxConcurrentStreams) + }) + for _, tc := range []struct { + name, value string + want uint32 + invalid bool + }{ + {name: "disabled", value: "0"}, + {name: "configured", value: "16", want: 16}, + {name: "largest valid", value: "4294967294", want: 4294967294}, + {name: "grpc unlimited sentinel", value: "4294967295", invalid: true}, + {name: "negative", value: "-1", invalid: true}, + {name: "overflow", value: "4294967296", invalid: true}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv("GRPC_MAX_CONCURRENT_STREAMS", tc.value) + if tc.invalid { + require.Panics(t, func() { NewSettings() }) + return + } + require.Equal(t, tc.want, NewSettings().GrpcMaxConcurrentStreams) + }) + } +} diff --git a/src/settings/settings.go b/src/settings/settings.go index 460b0329..ecc66457 100644 --- a/src/settings/settings.go +++ b/src/settings/settings.go @@ -2,6 +2,7 @@ package settings import ( "crypto/tls" + "math" "time" "github.com/kelseyhightower/envconfig" @@ -26,6 +27,10 @@ type Settings struct { GrpcUds string `envconfig:"GRPC_UDS" default:""` GrpcHost string `envconfig:"GRPC_HOST" default:"0.0.0.0"` GrpcPort int `envconfig:"GRPC_PORT" default:"8081"` + // GrpcMaxConcurrentStreams limits active gRPC streams per connection. + // Zero leaves grpc-go's default (no practical stream limit) in place; + // MaxUint32 is rejected because grpc-go also treats it as unlimited. + GrpcMaxConcurrentStreams uint32 `envconfig:"GRPC_MAX_CONCURRENT_STREAMS" default:"0"` // GrpcServerTlsConfig configures grpc for the server GrpcServerTlsConfig *tls.Config // GrpcMaxConnectionAge is a duration for the maximum amount of time a connection may exist before it will be closed by sending a GoAway. @@ -54,6 +59,13 @@ type Settings struct { ConfigType string `envconfig:"CONFIG_TYPE" default:"FILE"` // ForceStartWithoutInitialConfig enables start the server without initial rate limit config event ForceStartWithoutInitialConfig bool `envconfig:"FORCE_START_WITHOUT_INITIAL_CONFIG" default:"false"` + // MaxConcurrentRequests limits admitted ShouldRateLimit calls across both + // gRPC and HTTP. Excess calls fail immediately; zero disables the limit. + MaxConcurrentRequests int `envconfig:"MAX_CONCURRENT_REQUESTS" default:"0"` + // RequestTimeout supplies a deadline to cache work without extending an + // earlier caller deadline. Zero uses only the caller's deadline. A backend + // may continue running after cancellation and retains its admission slot. + RequestTimeout time.Duration `envconfig:"REQUEST_TIMEOUT" default:"0"` // xDS rate limit configuration // ConfigGrpcXdsNodeId is the Node ID. xDS server should set snapshots to this Node ID @@ -281,6 +293,15 @@ func NewSettings() Settings { if err := envconfig.Process("", &s); err != nil { panic(err) } + if s.MaxConcurrentRequests < 0 { + panic("MAX_CONCURRENT_REQUESTS must be >= 0") + } + if s.RequestTimeout < 0 { + panic("REQUEST_TIMEOUT must be >= 0") + } + if s.GrpcMaxConcurrentStreams == math.MaxUint32 { + panic("GRPC_MAX_CONCURRENT_STREAMS must be less than 4294967295") + } // When we require TLS to connect to Redis, we check if we need to connect using the provided key-pair. RedisTlsConfig(s.RedisTls || s.RedisPerSecondTls)(&s) MemcacheTlsConfig(s.MemcacheTls)(&s) diff --git a/src/stats/manager.go b/src/stats/manager.go index 688cfb8f..faac97a9 100644 --- a/src/stats/manager.go +++ b/src/stats/manager.go @@ -47,6 +47,24 @@ type ServiceStats struct { ConfigLoadError gostats.Counter ShouldRateLimit ShouldRateLimitStats GlobalShadowMode gostats.Counter + RequestAdmission RequestAdmissionStats +} + +// Request admission metrics count handler calls, not Redis commands or quota hits. +type RequestAdmissionStats struct { + Admitted gostats.Counter + Rejected gostats.Counter + InFlight gostats.Gauge + CompletedDuration gostats.Timer +} + +func NewRequestAdmissionStats(scope gostats.Scope) RequestAdmissionStats { + return RequestAdmissionStats{ + Admitted: scope.NewCounter("admitted"), + Rejected: scope.NewCounter("rejected"), + InFlight: scope.NewGauge("in_flight"), + CompletedDuration: scope.NewMilliTimer("completed_duration"), + } } // Stats for an individual rate limit config entry. diff --git a/src/stats/manager_impl.go b/src/stats/manager_impl.go index aa195f88..9d433d6f 100644 --- a/src/stats/manager_impl.go +++ b/src/stats/manager_impl.go @@ -65,6 +65,7 @@ func (this *ManagerImpl) NewServiceStats() ServiceStats { ret.ConfigLoadError = this.serviceStatsScope.NewCounter("config_load_error") ret.ShouldRateLimit = this.NewShouldRateLimitStats() ret.GlobalShadowMode = this.serviceStatsScope.NewCounter("global_shadow_mode") + ret.RequestAdmission = NewRequestAdmissionStats(this.serviceStatsScope.Scope("request_admission")) return ret } diff --git a/src/stats/prom/default_mapper.yaml b/src/stats/prom/default_mapper.yaml index e03706e3..ac01fbaf 100644 --- a/src/stats/prom/default_mapper.yaml +++ b/src/stats/prom/default_mapper.yaml @@ -90,6 +90,20 @@ mappings: labels: grpc_method: "$1" + - match: "ratelimit.service.request_admission.admitted" + name: "ratelimit_service_request_admission_admitted_total" + match_metric_type: counter + - match: "ratelimit.service.request_admission.rejected" + name: "ratelimit_service_request_admission_rejected_total" + match_metric_type: counter + - match: "ratelimit.service.request_admission.in_flight" + name: "ratelimit_service_request_admission_in_flight" + match_metric_type: gauge + - match: "ratelimit.service.request_admission.completed_duration" + name: "ratelimit_service_request_admission_completed_duration_seconds" + timer_type: histogram + scale: 0.001 + - match: "ratelimit.service.config_load_success" name: "ratelimit_service_config_load_success" match_metric_type: counter diff --git a/src/stats/prom/request_admission_test.go b/src/stats/prom/request_admission_test.go new file mode 100644 index 00000000..ee8eee46 --- /dev/null +++ b/src/stats/prom/request_admission_test.go @@ -0,0 +1,86 @@ +package prom + +import ( + "fmt" + "io" + "net/http/httptest" + "testing" + "time" + + gostats "github.com/lyft/gostats" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promhttp" + dto "github.com/prometheus/client_model/go" + "github.com/stretchr/testify/require" + + "github.com/envoyproxy/ratelimit/src/settings" + "github.com/envoyproxy/ratelimit/src/stats" +) + +func TestRequestAdmissionMetricNamesTypesAndUnits(t *testing.T) { + for _, legacyResponseTime := range []bool{false, true} { + t.Run(fmt.Sprintf("legacy_response_time_%t", legacyResponseTime), func(t *testing.T) { + oldRegisterer, oldGatherer := prometheus.DefaultRegisterer, prometheus.DefaultGatherer + reg := prometheus.NewRegistry() + prometheus.DefaultRegisterer, prometheus.DefaultGatherer = reg, reg + t.Cleanup(func() { + prometheus.DefaultRegisterer, prometheus.DefaultGatherer = oldRegisterer, oldGatherer + }) + sink := NewPrometheusSink(WithAddr("127.0.0.1:0"), + WithPath(fmt.Sprintf("/metrics-admission-%t", legacyResponseTime)), + WithResponseTimeAsMilliseconds(legacyResponseTime)) + store := gostats.NewStore(sink, false) + admission := stats.NewStatManager(store, settings.Settings{}).NewServiceStats().RequestAdmission + admission.Admitted.Add(2) + admission.Rejected.Add(3) + admission.InFlight.Set(1) + admission.CompletedDuration.AllocateSpan().CompleteWithDuration(1250 * time.Millisecond) + store.Flush() + + var metrics map[string]*dto.MetricFamily + require.Eventually(t, func() bool { + families, err := reg.Gather() + if err != nil { + return false + } + metrics = make(map[string]*dto.MetricFamily) + for _, family := range families { + metrics[family.GetName()] = family + } + return metrics["ratelimit_service_request_admission_admitted_total"] != nil && + metrics["ratelimit_service_request_admission_rejected_total"] != nil && + metrics["ratelimit_service_request_admission_in_flight"] != nil && + metrics["ratelimit_service_request_admission_completed_duration_seconds"] != nil + }, time.Second, time.Millisecond) + + admitted := metrics["ratelimit_service_request_admission_admitted_total"] + require.Equal(t, dto.MetricType_COUNTER, admitted.GetType()) + require.Len(t, admitted.Metric, 1) + require.Empty(t, admitted.Metric[0].Label) + require.Equal(t, float64(2), admitted.Metric[0].Counter.GetValue()) + rejected := metrics["ratelimit_service_request_admission_rejected_total"] + require.Equal(t, dto.MetricType_COUNTER, rejected.GetType()) + require.Equal(t, float64(3), rejected.Metric[0].Counter.GetValue()) + inFlight := metrics["ratelimit_service_request_admission_in_flight"] + require.Equal(t, dto.MetricType_GAUGE, inFlight.GetType()) + require.Equal(t, float64(1), inFlight.Metric[0].Gauge.GetValue()) + duration := metrics["ratelimit_service_request_admission_completed_duration_seconds"] + require.Equal(t, dto.MetricType_HISTOGRAM, duration.GetType()) + require.Equal(t, uint64(1), duration.Metric[0].Histogram.GetSampleCount()) + require.Equal(t, 1.25, duration.Metric[0].Histogram.GetSampleSum()) + + endpoint := httptest.NewServer(promhttp.HandlerFor(reg, promhttp.HandlerOpts{})) + defer endpoint.Close() + response, err := endpoint.Client().Get(endpoint.URL + "/metrics") + require.NoError(t, err) + defer response.Body.Close() + body, err := io.ReadAll(response.Body) + require.NoError(t, err) + require.Contains(t, string(body), "ratelimit_service_request_admission_admitted_total 2\n") + require.Contains(t, string(body), "ratelimit_service_request_admission_rejected_total 3\n") + require.Contains(t, string(body), "ratelimit_service_request_admission_in_flight 1\n") + require.Contains(t, string(body), "ratelimit_service_request_admission_completed_duration_seconds_sum 1.25\n") + require.Contains(t, string(body), "ratelimit_service_request_admission_completed_duration_seconds_count 1\n") + }) + } +} diff --git a/test/mocks/stats/manager.go b/test/mocks/stats/manager.go index 82e7e382..1dd67768 100644 --- a/test/mocks/stats/manager.go +++ b/test/mocks/stats/manager.go @@ -35,6 +35,7 @@ func (m *MockStatManager) NewServiceStats() stats.ServiceStats { ret.ConfigLoadError = m.store.NewCounter("config_load_error") ret.ShouldRateLimit = m.NewShouldRateLimitStats() ret.GlobalShadowMode = m.store.NewCounter("global_shadow_mode") + ret.RequestAdmission = stats.NewRequestAdmissionStats(m.store.Scope("request_admission")) return ret } diff --git a/test/server/server_impl_test.go b/test/server/server_impl_test.go index 899f5640..37524822 100644 --- a/test/server/server_impl_test.go +++ b/test/server/server_impl_test.go @@ -10,6 +10,8 @@ import ( "testing" "github.com/stretchr/testify/mock" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "google.golang.org/protobuf/proto" pb "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v3" @@ -84,3 +86,26 @@ func TestJsonHandler(t *testing.T) { }, nil) assertHttpResponse(t, handler, `{"domain": "foo"}`, 429, "application/json", `{"overallCode":"OVER_LIMIT"}`) } + +func TestJsonHandlerPreservesRequestContextAndServiceFailures(t *testing.T) { + for _, tc := range []struct { + code codes.Code + httpStatus int + }{ + {codes.ResourceExhausted, http.StatusServiceUnavailable}, + {codes.DeadlineExceeded, http.StatusGatewayTimeout}, + {codes.Canceled, http.StatusRequestTimeout}, + } { + t.Run(tc.code.String(), func(t *testing.T) { + controller := gomock.NewController(t) + rls := mock_v3.NewMockRateLimitServiceServer(controller) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + rls.EXPECT().ShouldRateLimit(ctx, gomock.Any()).Return(nil, status.Error(tc.code, "service failure")) + req := httptest.NewRequest(http.MethodPost, "/json", strings.NewReader(`{"domain":"test"}`)).WithContext(ctx) + response := httptest.NewRecorder() + server.NewJsonHandler(rls)(response, req) + assert.Equal(t, tc.httpStatus, response.Code) + }) + } +}