Merge d673e23 into 98dda60

grpc/grpc.go

@@ -2,9 +2,9 @@ package grpc
 
 import (
 	"context"
-	"fmt"
 
 	golangGrpc "google.golang.org/grpc"
+	"google.golang.org/grpc/status"
 )
 
 // UnaryServerInterceptor will trace requests to the given grpc server.
@@ -17,7 +17,8 @@ func UnaryServerInterceptor(opts ...InterceptorOption) golangGrpc.UnaryServerInt
 	return func(ctx context.Context, req interface{}, info *golangGrpc.UnaryServerInfo, handler golangGrpc.UnaryHandler) (interface{}, error) {
 		token, ok := cfg.limiter.Acquire(ctx)
 		if !ok {
-			return nil, fmt.Errorf("limit exceeded for limiter=%v", cfg.limiter)
+			errResp, errCode, err := cfg.limitExceededResponseClassifier(ctx, info.FullMethod, req, cfg.limiter)
+			return errResp, status.Error(errCode, err.Error())
 		}
 		resp, err := handler(ctx, req)
 		respType := cfg.serverResponseClassifer(ctx, req, info, resp, err)
@@ -43,7 +44,8 @@ func UnaryClientInterceptor(opts ...InterceptorOption) golangGrpc.UnaryClientInt
 	return func(ctx context.Context, method string, req, reply interface{}, cc *golangGrpc.ClientConn, invoker golangGrpc.UnaryInvoker, opts ...golangGrpc.CallOption) error {
 		token, ok := cfg.limiter.Acquire(ctx)
 		if !ok {
-			return fmt.Errorf("limit exceeded for limiter=%v", cfg.limiter)
+			_, errCode, err := cfg.limitExceededResponseClassifier(ctx, method, req, cfg.limiter)
+			return status.Error(errCode, err.Error())
 		}
 		err := invoker(ctx, method, req, reply, cc, opts...)
 		respType := cfg.clientResponseClassifer(ctx, method, req, reply, err)

grpc/option.go

@@ -2,8 +2,10 @@ package grpc
 
 import (
 	"context"
+	"fmt"
 
 	golangGrpc "google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 
 	"github.com/platinummonkey/go-concurrency-limits/core"
 	"github.com/platinummonkey/go-concurrency-limits/limit"
@@ -23,6 +25,10 @@ const (
 	ResponseTypeDropped
 )
 
+// LimitExceededResponseClassifier is a method definition for defining the error response type when the limit is exceeded
+// and a token is not able to be acquired. By default the RESOURCE_EXHUASTED type is returend.
+type LimitExceededResponseClassifier func(ctx context.Context, method string, req interface{}, l core.Limiter) (interface{}, codes.Code, error)
+
 // ClientResponseClassifier is a method definition for defining custom response types to the limiter algorithm to
 // correctly handle certain types of errors or embedded data.
 type ClientResponseClassifier func(ctx context.Context, method string, req, reply interface{}, err error) ResponseType
@@ -33,6 +39,15 @@ type ServerResponseClassifier func(
 	ctx context.Context, req interface{}, info *golangGrpc.UnaryServerInfo, resp interface{}, err error,
 ) ResponseType
 
+func defaultLimitExceededResponseClassifier(
+	ctx context.Context,
+	method string,
+	req interface{},
+	l core.Limiter,
+) (interface{}, codes.Code, error) {
+	return nil, codes.ResourceExhausted, fmt.Errorf("limit exceeded for limiter=%v", l)
+}
+
 func defaultClientResponseClassifier(
 	ctx context.Context,
 	method string,
@@ -63,6 +78,7 @@ type interceptorConfig struct {
 	name                    string
 	tags                    []string
 	limiter                 core.Limiter
+	limitExceededResponseClassifier LimitExceededResponseClassifier
 	serverResponseClassifer ServerResponseClassifier
 	clientResponseClassifer ClientResponseClassifier
 }
@@ -87,6 +103,7 @@ func defaults(cfg *interceptorConfig) {
 		core.EmptyMetricRegistryInstance,
 		tags...,
 	)
+	cfg.limitExceededResponseClassifier = defaultLimitExceededResponseClassifier
 	cfg.clientResponseClassifer = defaultClientResponseClassifier
 	cfg.serverResponseClassifer = defaultServerResponseClassifier
 }
@@ -112,6 +129,13 @@ func WithLimiter(limiter core.Limiter) InterceptorOption {
 	}
 }
 
+// WithClientResponseTypeClassifier sets the response classifier for the intercepted client
+func WithLimitExceededResponseClassifier(classifier LimitExceededResponseClassifier) InterceptorOption {
+	return func(cfg *interceptorConfig) {
+		cfg.limitExceededResponseClassifier = classifier
+	}
+}
+
 // WithClientResponseTypeClassifier sets the response classifier for the intercepted client
 func WithClientResponseTypeClassifier(classifier ClientResponseClassifier) InterceptorOption {
 	return func(cfg *interceptorConfig) {