From e579a5d69e7bb79fb885b5baca427599e1b514bc Mon Sep 17 00:00:00 2001 From: duncandu Date: Tue, 24 May 2022 12:42:36 +0800 Subject: [PATCH] Add '--rate' to limit the frequency of requests --- README.md | 5 ++-- go.mod | 1 + go.sum | 2 ++ main.go | 70 ++++++++++++++++++++++++++++++++++++++++++++++++++-- requester.go | 17 ++++++++++++- 5 files changed, 90 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index b920219..22fc17f 100644 --- a/README.md +++ b/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 diff --git a/go.mod b/go.mod index ca4cb6b..91bbce0 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 698e61e..f380d07 100644 --- a/go.sum +++ b/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= diff --git a/main.go b/main.go index 39571a4..94ec1f6 100644 --- a/main.go +++ b/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 diff --git a/requester.go b/requester.go index fd27dae..32bc47d 100644 --- a/requester.go +++ b/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