add rate limit

biz/emitter.go

@@ -12,12 +12,14 @@ type Emitter struct {
 	sync.WaitGroup
 	plugins     *InOutPlugins
 	filterChain filter.Filter
+	limiter     Limiter
 }
 
 // NewEmitter creates and initializes new Emitter object.
-func NewEmitter(f filter.Filter) *Emitter {
+func NewEmitter(f filter.Filter, lim Limiter) *Emitter {
 	var e Emitter
 	e.filterChain = f
+	e.limiter = lim
 	return &e
 }
 
@@ -58,11 +60,17 @@ func (e *Emitter) CopyMulty(src PluginReader, writers ...PluginWriter) error {
 			continue
 		}
 		msg, ok := e.filterChain.Filter(msg)
-		if ok {
-			for _, dst := range writers {
-				if err = dst.Write(msg); err != nil {
-					slog.Error("dst.Write:%v", err)
-				}
+		if !ok {
+			continue
+		}
+
+		if e.limiter != nil && !e.limiter.Allow() {
+			continue
+		}
+
+		for _, dst := range writers {
+			if err = dst.Write(msg); err != nil {
+				slog.Error("dst.Write:%v", err)
 			}
 		}
 	}

biz/itf.go

@@ -16,3 +16,7 @@ type PluginWriter interface {
 	io.Closer
 	Write(msg *protocol.Message) (err error)
 }
+
+type Limiter interface {
+	Allow() bool
+}

biz/ratelimit.go

@@ -0,0 +1,14 @@
+package biz
+
+import (
+	"github.com/vearne/grpcreplay/config"
+	"golang.org/x/time/rate"
+)
+
+func NewRateLimit(settings *config.AppSettings) Limiter {
+	if settings.RateLimitQPS > 0 {
+		value := settings.RateLimitQPS
+		return rate.NewLimiter(rate.Limit(value), value)
+	}
+	return nil
+}

config/settings.go

@@ -92,6 +92,11 @@ type AppSettings struct {
 
 	// --- filter ---
 	IncludeFilterMethodMatch string `json:"include-filter-method-match"`
+
+	// --- rate limit ---
+	// Query per second
+	RateLimitQPS int `json:"rate-limit-qps"`
+
 	// --- other ---
 	Codec string `json:"codec"`
 }

go.mod

@@ -43,6 +43,7 @@ require (
 	golang.org/x/sync v0.7.0 // indirect
 	golang.org/x/sys v0.22.0 // indirect
 	golang.org/x/text v0.16.0 // indirect
+	golang.org/x/time v0.6.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240528184218-531527333157 // indirect
 	stathat.com/c/consistent v1.0.0 // indirect
 )

go.sum

@@ -218,6 +218,8 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4=
 golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI=
+golang.org/x/time v0.6.0 h1:eTDhh4ZXt5Qf0augr54TN6suAUudPcawVZeIAPU7D4U=
+golang.org/x/time v0.6.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 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=
 golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=

main.go

@@ -94,6 +94,10 @@ func init() {
 	flag.StringVar(&settings.IncludeFilterMethodMatch, "include-filter-method-match", "",
 		`filter requests when the method matches the specified regular expression`)
 
+	// rate limit
+	flag.IntVar(&settings.RateLimitQPS, "rate-limit-qps", -1,
+		`the capture rate per second limit for Query`)
+
 	// rocketmq
 	flag.Var(&config.MultiStringOption{Params: &settings.OutputRocketMQNameServer},
 		"output-rocketmq-name-server",
@@ -129,7 +133,8 @@ func main() {
 	if err != nil {
 		slog.Fatal("create FilterChain error:%v", err)
 	}
-	emitter := biz.NewEmitter(filterChain)
+	limiter := biz.NewRateLimit(&settings)
+	emitter := biz.NewEmitter(filterChain, limiter)
 	plugins := biz.NewPlugins(&settings)
 
 	slog.Info("plugins:%v", plugins)