fix alibaba#83; fix consumer interceptor

gorexlv · Mar 9, 2020 · 8bb7e34 · 8bb7e34
1 parent f39e051
commit 8bb7e34
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 9 deletions.
diff --git a/adapter/rocketmq/consumer.go b/adapter/rocketmq/consumer.go
@@ -8,20 +8,23 @@ import (
 	"github.com/apache/rocketmq-client-go/v2/primitive"
 )
 
+// SentinelConsumerInterceptor returns interceptor for consumer
 func SentinelConsumerInterceptor(opts ...Option) primitive.Interceptor {
 	options := evaluateOptions(opts)
 	return func(ctx context.Context, req, reply interface{}, next primitive.Invoker) error {
-		producerCtx := primitive.GetProducerCtx(ctx)
-		resourceName := producerCtx.Message.Topic
+		consumerCtx, _ := primitive.GetConsumerCtx(ctx)
+		resourceName := consumerCtx.MQ.Topic
 
-		if options.resourceExtract != nil {
-			resourceName = options.resourceExtract(producerCtx)
+		if options.consumerResourceExtract != nil {
+			resourceName = options.consumerResourceExtract(consumerCtx)
 		}
 
 		entry, err := sentinel.Entry(
 			resourceName,
 			sentinel.WithResourceType(base.ResTypeMQ),
 			sentinel.WithTrafficType(base.Inbound),
+			sentinel.WithArgs("topic", consumerCtx.MQ.Topic),
+			sentinel.WithArgs("broker", consumerCtx.MQ.BrokerName),
 		)
 
 		if err != nil {

diff --git a/adapter/rocketmq/option.go b/adapter/rocketmq/option.go
@@ -7,17 +7,27 @@ import (
 type (
 	Option func(*options)
 	options struct {
-		resourceExtract func(*primitive.ProducerCtx) string
+		consumerResourceExtract func(*primitive.ConsumeMessageContext) string
+		providerResourceExtract func(*primitive.ProducerCtx) string
 		blockFallback primitive.Interceptor
 	}
 )
 
-func WithResourceExtract(fn func(ctx *primitive.ProducerCtx) string) Option {
+// WithConsumerResourceExtract set consumerResourceExtract
+func WithConsumerResourceExtract(fn func(ctx *primitive.ConsumeMessageContext) string) Option {
 	return func(options *options) {
-		options.resourceExtract = fn
+		options.consumerResourceExtract = fn
 	}
 }
 
+// WithProviderResourceExtract set providerResourceExtract
+func WithProviderResourceExtract(fn func(ctx *primitive.ProducerCtx) string) Option {
+	return func(options *options) {
+		options.providerResourceExtract = fn
+	}
+}
+
+// WithBlockFallback set blockFallback
 func WithBlockFallback(fn primitive.Interceptor) Option {
 	return func(options *options) {
 		options.blockFallback = fn

diff --git a/adapter/rocketmq/producer.go b/adapter/rocketmq/producer.go
@@ -8,20 +8,23 @@ import (
 	"github.com/apache/rocketmq-client-go/v2/primitive"
 )
 
+// SentinelProviderInterceptor returns interceptor for producer
 func SentinelProviderInterceptor(opts ...Option) primitive.Interceptor {
 	options := evaluateOptions(opts)
 	return func(ctx context.Context, req, reply interface{}, next primitive.Invoker) error {
 		producerCtx := primitive.GetProducerCtx(ctx)
 		resourceName := producerCtx.Message.Topic
 
-		if options.resourceExtract != nil {
-			resourceName = options.resourceExtract(producerCtx)
+		if options.providerResourceExtract != nil {
+			resourceName = options.providerResourceExtract(producerCtx)
 		}
 
 		entry, err := sentinel.Entry(
 			resourceName,
 			sentinel.WithResourceType(base.ResTypeMQ),
 			sentinel.WithTrafficType(base.Outbound),
+			sentinel.WithArgs("topic", producerCtx.Message.Topic),
+			sentinel.WithArgs("broker", producerCtx.MQ.BrokerName),
 		)
 
 		if err != nil {