[perf]: More precise producer rate limiter (#989)

* perf: more accurate producer rate limiter * perf: more accurate producer rate limiter
apache · Mar 13, 2023 · bcbac9f · bcbac9f
1 parent 352c463
commit bcbac9f
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 11 deletions.
diff --git a/go.mod b/go.mod
@@ -26,6 +26,7 @@ require (
 	go.uber.org/atomic v1.7.0
 	golang.org/x/mod v0.5.1
 	golang.org/x/oauth2 v0.0.0-20210402161424-2e8d93401602
+	golang.org/x/time v0.0.0-20191024005414-555d28b269f0
 	google.golang.org/protobuf v1.26.0
 )
 

diff --git a/go.sum b/go.sum
@@ -427,6 +427,7 @@ golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
 golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
 golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
+golang.org/x/time v0.0.0-20191024005414-555d28b269f0 h1:/5xXl8Y5W96D+TtHSlonuFqGHIWVuyCkGJLwGh9JJFs=
 golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=

diff --git a/perf/perf-producer.go b/perf/perf-producer.go
@@ -22,6 +22,8 @@ import (
 	"encoding/json"
 	"time"
 
+	"golang.org/x/time/rate"
+
 	"github.com/bmizerany/perks/quantile"
 	"github.com/spf13/cobra"
 
@@ -104,16 +106,9 @@ func produce(produceArgs *ProduceArgs, stop <-chan struct{}) {
 	payload := make([]byte, produceArgs.MessageSize)
 
 	ch := make(chan float64)
-	rateLimitCh := make(chan time.Time, produceArgs.Rate)
-	go func(rateLimit int, interval time.Duration) {
-		if rateLimit <= 0 { // 0 as no limit enforced
-			return
-		}
-		for {
-			oldest := <-rateLimitCh
-			time.Sleep(interval - time.Since(oldest))
-		}
-	}(produceArgs.Rate, time.Second)
+
+	limit := rate.Every(time.Duration(float64(time.Second) / float64(produceArgs.Rate)))
+	rateLimiter := rate.NewLimiter(limit, produceArgs.Rate)
 
 	go func(stopCh <-chan struct{}) {
 		for {
@@ -125,7 +120,7 @@ func produce(produceArgs *ProduceArgs, stop <-chan struct{}) {
 
 			start := time.Now()
 			if produceArgs.Rate > 0 {
-				rateLimitCh <- start
+				_ = rateLimiter.Wait(context.TODO())
 			}
 
 			producer.SendAsync(ctx, &pulsar.ProducerMessage{