feat: mq push consume supports ratelimit (#948)

Co-authored-by: hulongjun <[email protected]>
douyu · Aug 2, 2023 · 95b46f4 · 95b46f4
1 parent 5393a4d
commit 95b46f4
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 1 deletion.
diff --git a/pkg/client/rocketmq/push_consumer.go b/pkg/client/rocketmq/push_consumer.go
@@ -48,11 +48,17 @@ func (conf *PushConsumerConfig) Build() *PushConsumer {
 
 	xlog.Jupiter().Debug("rocketmq's config: ", xlog.String("name", name), xlog.Any("conf", conf))
 
+	var bucket *ratelimit.Bucket
+	if conf.Rate > 0 && conf.Capacity > 0 {
+		bucket = ratelimit.NewBucketWithRate(conf.Rate, conf.Capacity)
+	}
+
 	cc := &PushConsumer{
 		name:               name,
 		PushConsumerConfig: *conf,
 		subscribers:        make(map[string]func(context.Context, ...*primitive.MessageExt) (consumer.ConsumeResult, error)),
 		interceptors:       []primitive.Interceptor{},
+		bucket:             bucket,
 	}
 	cc.interceptors = append(cc.interceptors,
 		consumerMetricInterceptor(),

diff --git a/test/e2e/go.mod b/test/e2e/go.mod
@@ -108,7 +108,7 @@ require (
 	go.opentelemetry.io/otel/trace v1.16.0 // indirect
 	go.opentelemetry.io/proto/otlp v0.19.0 // indirect
 	go.uber.org/atomic v1.10.0 // indirect
-	go.uber.org/automaxprocs v1.5.2 // indirect
+	go.uber.org/automaxprocs v1.5.3 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 	go.uber.org/zap v1.24.0 // indirect
 	golang.org/x/arch v0.3.0 // indirect

diff --git a/test/e2e/go.sum b/test/e2e/go.sum
@@ -623,6 +623,7 @@ go.uber.org/atomic v1.10.0 h1:9qC72Qh0+3MqyJbAn8YU5xVq1frD8bn3JtD2oXtafVQ=
 go.uber.org/atomic v1.10.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
 go.uber.org/automaxprocs v1.5.2 h1:2LxUOGiR3O6tw8ui5sZa2LAaHnsviZdVOUZw4fvbnME=
 go.uber.org/automaxprocs v1.5.2/go.mod h1:eRbA25aqJrxAbsLO0xy5jVwPt7FQnRgjW+efnwa1WM0=
+go.uber.org/automaxprocs v1.5.3/go.mod h1:eRbA25aqJrxAbsLO0xy5jVwPt7FQnRgjW+efnwa1WM0=
 go.uber.org/goleak v1.2.1 h1:NBol2c7O1ZokfZ0LEU9K6Whx/KnwvepVetCUhtKja4A=
 go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
 go.uber.org/multierr v1.3.0/go.mod h1:VgVr7evmIr6uPjLBxg28wmKNXyqE9akIJ5XnfpiKl+4=