Add '--rate' to limit the frequency of requests

README.md

@@ -5,7 +5,7 @@
 [![GitHub license](https://img.shields.io/github/license/six-ddc/plow.svg)](https://github.com/six-ddc/plow/blob/main/LICENSE)
 [![made-with-Go](https://img.shields.io/badge/Made%20with-Go-1f425f.svg)](http://golang.org)
 
-Plow is a HTTP(S) benchmarking tool, written in Golang. It uses
+Plow is an HTTP(S) benchmarking tool, written in Golang. It uses
 excellent [fasthttp](https://github.com/valyala/fasthttp#http-client-comparison-with-nethttp) instead of Go's default
 net/http due to its lightning fast performance.
 
@@ -104,6 +104,7 @@ Examples:
 Flags:
       --help                   Show context-sensitive help.
   -c, --concurrency=1          Number of connections to run concurrently
+      --rate=infinity          Number of requests per time unit, examples: --rate 50 --rate 10/ms
   -n, --requests=-1            Number of requests to run
   -d, --duration=DURATION      Duration of test, examples: -d 10s -d 3m
   -i, --interval=200ms         Print snapshot result every interval, use 0 to print once at the end
@@ -124,8 +125,8 @@ Flags:
       --resp-timeout=DURATION  Timeout for full response reading
       --socks5=ip:port         Socks5 proxy
       --auto-open-browser      Specify whether auto open browser to show Web charts
-      --[no-]summary           Only print the summary without realtime reports
       --[no-]clean             Clean the histogram bar once its finished. Default is true
+      --[no-]summary           Only print the summary without realtime reports
       --version                Show application version.
 
   Flags default values also read from env PLOW_SOME_FLAG, such as PLOW_TIMEOUT=5s equals to --timeout=5s

go.mod

@@ -13,6 +13,7 @@ require (
 	github.com/nicksnyder/go-i18n v1.10.1 // indirect
 	github.com/valyala/fasthttp v1.33.0
 	go.uber.org/automaxprocs v1.4.0
+	golang.org/x/time v0.0.0-20220411224347-583f2d630306
 	gopkg.in/alecthomas/kingpin.v3-unstable v3.0.0-20191105091915-95d230a53780
 	gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect

go.sum

@@ -57,6 +57,8 @@ golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9sn
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 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-20220411224347-583f2d630306 h1:+gHMid33q6pen7kv9xvT+JRinntgeXO2AeZVd0AWD3w=
+golang.org/x/time v0.0.0-20220411224347-583f2d630306/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 gopkg.in/alecthomas/kingpin.v3-unstable v3.0.0-20191105091915-95d230a53780 h1:CEBpW6C191eozfEuWdUmIAHn7lwlLxJ7HVdr2e2Tsrw=
 gopkg.in/alecthomas/kingpin.v3-unstable v3.0.0-20191105091915-95d230a53780/go.mod h1:3HH7i1SgMqlzxCcBmUHW657sD4Kvv9sC3HpL3YukzwA=

main.go

@@ -2,16 +2,20 @@ package main
 
 import (
 	"fmt"
+	"golang.org/x/time/rate"
 	"io/ioutil"
 	"net"
 	"os"
+	"strconv"
 	"strings"
+	"time"
 
 	"gopkg.in/alecthomas/kingpin.v3-unstable"
 )
 
 var (
 	concurrency = kingpin.Flag("concurrency", "Number of connections to run concurrently").Short('c').Default("1").Int()
+	reqRate     = rateFlag(kingpin.Flag("rate", "Number of requests per time unit, examples: --rate 50 --rate 10/ms").Default("infinity"))
 	requests    = kingpin.Flag("requests", "Number of requests to run").Short('n').Default("-1").Int64()
 	duration    = kingpin.Flag("duration", "Duration of test, examples: -d 10s -d 3m").Short('d').PlaceHolder("DURATION").Duration()
 	interval    = kingpin.Flag("interval", "Print snapshot result every interval, use 0 to print once at the end").Short('i').Default("200ms").Duration()
@@ -99,9 +103,71 @@ Examples:
 {{end -}}
 `
 
+type rateFlagValue struct {
+	infinity bool
+	limit    rate.Limit
+	v        string
+}
+
+func (f *rateFlagValue) Set(v string) error {
+	if v == "infinity" {
+		f.infinity = true
+		return nil
+	}
+
+	retErr := fmt.Errorf("--rate format %q doesn't match the \"freq/duration\" (i.e. 50/1s)", v)
+	ps := strings.SplitN(v, "/", 2)
+	switch len(ps) {
+	case 1:
+		ps = append(ps, "1s")
+	case 0:
+		return retErr
+	}
+
+	freq, err := strconv.Atoi(ps[0])
+	if err != nil {
+		return retErr
+	}
+	if freq == 0 {
+		f.infinity = true
+		return nil
+	}
+
+	switch ps[1] {
+	case "ns", "us", "µs", "ms", "s", "m", "h":
+		ps[1] = "1" + ps[1]
+	}
+
+	per, err := time.ParseDuration(ps[1])
+	if err != nil {
+		return retErr
+	}
+
+	f.limit = rate.Limit(float64(freq) / per.Seconds())
+	f.v = v
+	return nil
+}
+
+func (f *rateFlagValue) Limit() *rate.Limit {
+	if f.infinity {
+		return nil
+	}
+	return &f.limit
+}
+
+func (f *rateFlagValue) String() string {
+	return f.v
+}
+
+func rateFlag(c *kingpin.Clause) (target *rateFlagValue) {
+	target = new(rateFlagValue)
+	c.SetValue(target)
+	return
+}
+
 func main() {
 	kingpin.UsageTemplate(CompactUsageTemplate).
-		Version("1.1.0").
+		Version("1.2.0").
 		Author("six-ddc@github").
 		Resolver(kingpin.PrefixedEnvarResolver("PLOW_", ";")).
 		Help = `A high-performance HTTP benchmarking tool with real-time web UI and terminal displaying`
@@ -160,7 +226,7 @@ func main() {
 		host:        *host,
 	}
 
-	requester, err := NewRequester(*concurrency, *requests, *duration, &clientOpt)
+	requester, err := NewRequester(*concurrency, *requests, *duration, reqRate.Limit(), &clientOpt)
 	if err != nil {
 		errAndExit(err.Error())
 		return

requester.go

@@ -7,6 +7,7 @@ import (
 	"github.com/valyala/fasthttp"
 	"github.com/valyala/fasthttp/fasthttpproxy"
 	"go.uber.org/automaxprocs/maxprocs"
+	"golang.org/x/time/rate"
 	"io/ioutil"
 	"net"
 	url2 "net/url"
@@ -90,6 +91,7 @@ func ThroughputInterceptorDial(dial fasthttp.DialFunc, r *int64, w *int64) fasth
 
 type Requester struct {
 	concurrency int
+	reqRate     *rate.Limit
 	requests    int64
 	duration    time.Duration
 	clientOpt   *ClientOpt
@@ -128,13 +130,14 @@ type ClientOpt struct {
 	host        string
 }
 
-func NewRequester(concurrency int, requests int64, duration time.Duration, clientOpt *ClientOpt) (*Requester, error) {
+func NewRequester(concurrency int, requests int64, duration time.Duration, reqRate *rate.Limit, clientOpt *ClientOpt) (*Requester, error) {
 	maxResult := concurrency * 100
 	if maxResult > 8192 {
 		maxResult = 8192
 	}
 	r := &Requester{
 		concurrency: concurrency,
+		reqRate:     reqRate,
 		requests:    requests,
 		duration:    duration,
 		clientOpt:   clientOpt,
@@ -304,6 +307,11 @@ func (r *Requester) Run() {
 		})
 	}
 
+	var limiter *rate.Limiter
+	if r.reqRate != nil {
+		limiter = rate.NewLimiter(*r.reqRate, 1)
+	}
+
 	semaphore := r.requests
 	for i := 0; i < r.concurrency; i++ {
 		r.wg.Add(1)
@@ -330,6 +338,13 @@ func (r *Requester) Run() {
 				default:
 				}
 
+				if limiter != nil {
+					err := limiter.Wait(ctx)
+					if err != nil {
+						continue
+					}
+				}
+
 				if r.requests > 0 && atomic.AddInt64(&semaphore, -1) < 0 {
 					cancelFunc()
 					return