Merge remote-tracking branch 'origin/main' into yuri/grpc-cluster

CHANGELOG.md

@@ -246,6 +246,7 @@
 * [ENHANCEMENT] Memberlist: Notifications can now be processed once per interval specified by `-memberlist.notify-interval` to reduce notify storms in large clusters. #592
 * [ENHANCEMENT] Memberlist: Implemented the `Delete` operation in the memberlist backed KV store. How frequently deleted entries are cleaned up is specified by the `-memberlist.obsolete-entries-timeout` flag. #612
 * [ENHANCEMENT] KV: Add `MockCountingClient`, which wraps the `kv.client` and can be used in order to count calls at specific functions of the interface. #618
+* [ENHANCEMENT] Server: Add interceptor support to `GrpcInflightMethodLimiter`. #643
 * [BUGFIX] spanlogger: Support multiple tenant IDs. #59
 * [BUGFIX] Memberlist: fixed corrupted packets when sending compound messages with more than 255 messages or messages bigger than 64KB. #85
 * [BUGFIX] Ring: `ring_member_ownership_percent` and `ring_tokens_owned` metrics are not updated on scale down. #109

server/limits.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"strings"
 
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/stats"
 	"google.golang.org/grpc/tap"
@@ -19,6 +20,12 @@ type GrpcInflightMethodLimiter interface {
 	// otherwise gRPC-server implementation-specific error will be returned to the client (codes.PermissionDenied in [email protected]).
 	RPCCallStarting(ctx context.Context, methodName string, md metadata.MD) (context.Context, error)
 
+	// RPCCallProcessing is called by a server interceptor, allowing request pre-processing or request blocking to be
+	// performed. The returned function will be applied after the request is handled, providing any error that occurred while
+	// handling the request.
+	RPCCallProcessing(ctx context.Context, methodName string) (func(error), error)
+
+	// RPCCallFinished is called when an RPC call is finished being handled.
 	RPCCallFinished(ctx context.Context)
 }
 
@@ -47,6 +54,31 @@ func (g *grpcInflightLimitCheck) TapHandle(ctx context.Context, info *tap.Info)
 	return g.methodLimiter.RPCCallStarting(ctx, info.FullMethodName, info.Header)
 }
 
+func (g *grpcInflightLimitCheck) UnaryServerInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+	finish, err := g.methodLimiter.RPCCallProcessing(ctx, info.FullMethod)
+	if err != nil {
+		return nil, err
+	}
+	result, err := handler(ctx, req)
+	if finish != nil {
+		finish(err)
+	}
+	return result, err
+
+}
+
+func (g *grpcInflightLimitCheck) StreamServerInterceptor(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+	finish, err := g.methodLimiter.RPCCallProcessing(ss.Context(), info.FullMethod)
+	if err != nil {
+		return err
+	}
+	err = handler(srv, ss)
+	if finish != nil {
+		finish(err)
+	}
+	return err
+}
+
 func (g *grpcInflightLimitCheck) TagRPC(ctx context.Context, _ *stats.RPCTagInfo) context.Context {
 	return ctx
 }

server/limits_test.go

@@ -322,6 +322,10 @@ func (m *methodLimiter) RPCCallStarting(ctx context.Context, methodName string,
 	return context.WithValue(ctx, ctxMethodName, methodName), nil
 }
 
+func (m *methodLimiter) RPCCallProcessing(_ context.Context, _ string) (func(error), error) {
+	return nil, nil
+}
+
 func (m *methodLimiter) RPCCallFinished(ctx context.Context) {
 	m.allInflight.Dec()
 

server/server.go

@@ -425,6 +425,13 @@ func newServer(cfg Config, metrics *Metrics) (*Server, error) {
 		PermitWithoutStream: cfg.GRPCServerPingWithoutStreamAllowed,
 	}
 
+	var grpcServerLimit *grpcInflightLimitCheck
+	if cfg.GrpcMethodLimiter != nil {
+		grpcServerLimit = newGrpcInflightLimitCheck(cfg.GrpcMethodLimiter)
+		grpcMiddleware = append(grpcMiddleware, grpcServerLimit.UnaryServerInterceptor)
+		grpcStreamMiddleware = append(grpcStreamMiddleware, grpcServerLimit.StreamServerInterceptor)
+	}
+
 	grpcOptions := []grpc.ServerOption{
 		grpc.ChainUnaryInterceptor(grpcMiddleware...),
 		grpc.ChainStreamInterceptor(grpcStreamMiddleware...),
@@ -436,9 +443,11 @@ func newServer(cfg Config, metrics *Metrics) (*Server, error) {
 		grpc.NumStreamWorkers(uint32(cfg.GRPCServerNumWorkers)),
 	}
 
-	if cfg.GrpcMethodLimiter != nil {
-		grpcServerLimit := newGrpcInflightLimitCheck(cfg.GrpcMethodLimiter)
-		grpcOptions = append(grpcOptions, grpc.InTapHandle(grpcServerLimit.TapHandle), grpc.StatsHandler(grpcServerLimit))
+	if grpcServerLimit != nil {
+		grpcOptions = append(grpcOptions,
+			grpc.StatsHandler(grpcServerLimit),
+			grpc.InTapHandle(grpcServerLimit.TapHandle),
+		)
 	}
 
 	if cfg.GRPCServerStatsTrackingEnabled {