add option for rate limitting concurrent workers

cmd/benchmark.go

@@ -36,8 +36,9 @@ type Benchmark struct {
 	Count       int64
 	Concurrency uint32
 
-	Rate     int
-	QperConn int64
+	Rate            int
+	RateLimitWorker int
+	QperConn        int64
 
 	Recurse bool
 
@@ -180,7 +181,14 @@ func (b *Benchmark) Run(ctx context.Context) ([]*ResultStats, error) {
 	var limit ratelimit.Limiter
 	if b.Rate > 0 {
 		limit = ratelimit.New(b.Rate)
-		limits = fmt.Sprintf("(limited to %s QPS)", highlightStr(b.Rate))
+		if b.RateLimitWorker == 0 {
+			limits = fmt.Sprintf("(limited to %s QPS overall)", highlightStr(b.Rate))
+		} else {
+			limits = fmt.Sprintf("(limited to %s QPS overall and %s QPS per concurrent worker)", highlightStr(b.Rate), highlightStr(b.RateLimitWorker))
+		}
+	}
+	if b.Rate == 0 && b.RateLimitWorker > 0 {
+		limits = fmt.Sprintf("(limited to %s QPS per concurrent worker)", highlightStr(b.RateLimitWorker))
 	}
 
 	if !b.Silent && !b.JSON {
@@ -215,21 +223,36 @@ func (b *Benchmark) Run(ctx context.Context) ([]*ResultStats, error) {
 			// nolint:gosec
 			rando := rand.New(rand.NewSource(time.Now().Unix()))
 
+			var workerLimit ratelimit.Limiter
+			if b.RateLimitWorker > 0 {
+				workerLimit = ratelimit.New(b.RateLimitWorker)
+			}
+
 			var i int64
 			for i = 0; i < b.Count || b.Duration != 0; i++ {
 				for _, qt := range qTypes {
 					for _, q := range questions {
+						if ctx.Err() != nil {
+							return
+						}
 						if rando.Float64() > b.Probability {
 							continue
 						}
+						if limit != nil {
+							if err := checkLimit(ctx, limit); err != nil {
+								return
+							}
+						}
+						if workerLimit != nil {
+							if err := checkLimit(ctx, workerLimit); err != nil {
+								return
+							}
+						}
 						var r *dns.Msg
 						m := dns.Msg{}
 						m.RecursionDesired = b.Recurse
 						m.Question = make([]dns.Question, 1)
 						question := dns.Question{Qtype: qt, Qclass: dns.ClassINET}
-						if ctx.Err() != nil {
-							return
-						}
 						st.Counters.Total++
 
 						// instead of setting the question, do this manually for lower overhead and lock free access to id
@@ -240,9 +263,6 @@ func (b *Benchmark) Run(ctx context.Context) ([]*ResultStats, error) {
 							m.Id = uint16(rando.Uint32())
 						}
 						m.Question[0] = question
-						if limit != nil {
-							limit.Take()
-						}
 
 						start := time.Now()
 						if b.useQuic || b.useDoH {
@@ -369,3 +389,18 @@ func (b *Benchmark) getDoHClient() (queryFunc, string) {
 		return dohClient.SendViaPost, network
 	}
 }
+
+func checkLimit(ctx context.Context, limiter ratelimit.Limiter) error {
+	done := make(chan struct{})
+	go func() {
+		limiter.Take()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}

cmd/root.go

@@ -43,10 +43,13 @@ func init() {
 	pApp.Flag("rate-limit", "Apply a global questions / second rate limit.").
 		Short('l').Default("0").IntVar(&benchmark.Rate)
 
+	pApp.Flag("rate-limit-worker", "Apply a questions / second rate limit for each concurrent worker specified by --concurrency option.").
+		Default("0").IntVar(&benchmark.RateLimitWorker)
+
 	pApp.Flag("query-per-conn", "Queries on a connection before creating a new one. 0: unlimited. Applicable for plain DNS and DoT, this option is not considered for DoH or DoQ.").
 		Default("0").Int64Var(&benchmark.QperConn)
 
-	pApp.Flag("recurse", "Allow DNS recursion. Enabled by default. DNS recursion can be disabled by --no-recurse.").
+	pApp.Flag("recurse", "Allow DNS recursion. Enabled by default.").
 		Short('r').Default("true").BoolVar(&benchmark.Recurse)
 
 	pApp.Flag("probability", "Each provided hostname will be used with provided probability. Value 1 and above means that each hostname will be used by each concurrent benchmark goroutine. Useful for randomizing queries across benchmark goroutines.").
@@ -65,7 +68,7 @@ func init() {
 
 	pApp.Flag("read", "DNS read timeout.").Default(dnsTimeout.String()).DurationVar(&benchmark.ReadTimeout)
 
-	pApp.Flag("codes", "Enable counting DNS return codes. Enabled by default. Specifying --no-codes disables code counting.").
+	pApp.Flag("codes", "Enable counting DNS return codes. Enabled by default.").
 		Default("true").BoolVar(&benchmark.Rcodes)
 
 	pApp.Flag("min", "Minimum value for timing histogram.").
@@ -77,7 +80,7 @@ func init() {
 	pApp.Flag("precision", "Significant figure for histogram precision.").
 		Default("1").PlaceHolder("[1-5]").IntVar(&benchmark.HistPre)
 
-	pApp.Flag("distribution", "Display distribution histogram of timings to stdout. Enabled by default. Specifying --no-distribution disables histogram display.").
+	pApp.Flag("distribution", "Display distribution histogram of timings to stdout. Enabled by default.").
 		Default("true").BoolVar(&benchmark.HistDisplay)
 
 	pApp.Flag("csv", "Export distribution to CSV.").
@@ -87,7 +90,7 @@ func init() {
 
 	pApp.Flag("silent", "Disable stdout.").Default("false").BoolVar(&benchmark.Silent)
 
-	pApp.Flag("color", "ANSI Color output. Enabled by default. By specifying --no-color disables coloring.").
+	pApp.Flag("color", "ANSI Color output. Enabled by default.").
 		Default("true").BoolVar(&benchmark.Color)
 
 	pApp.Flag("plot", "Plot benchmark results and export them to the directory.").

docs/Gemfile.lock

@@ -71,6 +71,7 @@ GEM
 PLATFORMS
   arm64-darwin-21
   universal-darwin-22
+  universal-darwin-23
   x86_64-darwin-19
   x86_64-linux
 

docs/index.md

@@ -27,41 +27,43 @@ usage: dnspyre [<flags>] <queries>...
 
 A high QPS DNS benchmark.
 
+
 Flags:
-      --help                   Show context-sensitive help (also try --help-long and --help-man).
+      --[no-]help              Show context-sensitive help (also try --help-long and --help-man).
   -s, --server="127.0.0.1"     DNS server IP:port to test. IPv6 is also supported, for example '[fddd:dddd::]:53'. DoH (DNS over HTTPS) servers are supported such as `https://1.1.1.1/dns-query`, when such server is provided, the benchmark automatically
                                switches to the use of DoH. Note that path on which the DoH server handles requests (like `/dns-query`) has to be provided as well. DoQ (DNS over QUIC) servers are also supported, such as `quic://dns.adguard-dns.com`,
                                when such server is provided the benchmark switches to the use of DoQ.
   -t, --type=A ...             Query type. Repeatable flag. If multiple query types are specified then each query will be duplicated for each type.
   -n, --number=NUMBER          How many times the provided queries are repeated. Note that the total number of queries issued = types*number*concurrency*len(queries).
   -c, --concurrency=1          Number of concurrent queries to issue.
   -l, --rate-limit=0           Apply a global questions / second rate limit.
+      --rate-limit-worker=0    Apply a questions / second rate limit for each concurrent worker specified by --concurrency option.
       --query-per-conn=0       Queries on a connection before creating a new one. 0: unlimited. Applicable for plain DNS and DoT, this option is not considered for DoH or DoQ.
-  -r, --recurse                Allow DNS recursion. Enabled by default. DNS recursion can be disabled by --no-recurse.
+  -r, --[no-]recurse           Allow DNS recursion. Enabled by default.
       --probability=1          Each provided hostname will be used with provided probability. Value 1 and above means that each hostname will be used by each concurrent benchmark goroutine. Useful for randomizing queries across benchmark goroutines.
       --edns0=0                Enable EDNS0 with specified size.
       --ednsopt=""             code[:value], Specify EDNS option with code point code and optionally payload of value as a hexadecimal string. code must be an arbitrary numeric value.
-      --tcp                    Use TCP for DNS requests.
-      --dot                    Use DoT (DNS over TLS) for DNS requests.
+      --[no-]tcp               Use TCP for DNS requests.
+      --[no-]dot               Use DoT (DNS over TLS) for DNS requests.
       --write=1s               DNS write timeout.
       --read=4s                DNS read timeout.
-      --codes                  Enable counting DNS return codes. Enabled by default. Specifying --no-codes disables code counting.
+      --[no-]codes             Enable counting DNS return codes. Enabled by default.
       --min=400µs              Minimum value for timing histogram.
       --max=4s                 Maximum value for timing histogram.
       --precision=[1-5]        Significant figure for histogram precision.
-      --distribution           Display distribution histogram of timings to stdout. Enabled by default. Specifying --no-distribution disables histogram display.
+      --[no-]distribution      Display distribution histogram of timings to stdout. Enabled by default.
       --csv=/path/to/file.csv  Export distribution to CSV.
-      --json                   Report benchmark results as JSON.
-      --silent                 Disable stdout.
-      --color                  ANSI Color output. Enabled by default. By specifying --no-color disables coloring.
+      --[no-]json              Report benchmark results as JSON.
+      --[no-]silent            Disable stdout.
+      --[no-]color             ANSI Color output. Enabled by default.
       --plot=/path/to/folder   Plot benchmark results and export them to the directory.
       --plotf=png              Format of graphs. Supported formats: png, jpg.
       --doh-method=post        HTTP method to use for DoH requests. Supported values: get, post.
       --doh-protocol=1.1       HTTP protocol to use for DoH requests. Supported values: 1.1, 2 and 3.
-      --insecure               Disables server TLS certificate validation. Applicable for DoT, DoH and DoQ.
+      --[no-]insecure          Disables server TLS certificate validation. Applicable for DoT, DoH and DoQ.
   -d, --duration=1m            Specifies for how long the benchmark should be executing, the benchmark will run for the specified time while sending DNS requests in an infinite loop based on the data source. After running for the specified duration,
                                the benchmark is canceled. This option is exclusive with --number option. The duration is specified in GO duration format e.g. 10s, 15m, 1h.
-      --version                Show application version.
+      --[no-]version           Show application version.
 
 Args:
   <queries>  Queries to issue. It can be a local file referenced using @<file-path>, for example @data/2-domains. It can also be resource accessible using HTTP, like https://raw.githubusercontent.com/Tantalor93/dnspyre/master/data/1000-domains, in that

docs/ratelimit.md

@@ -0,0 +1,24 @@
+---
+title: Rate limiting
+layout: default
+parent: Examples
+---
+
+# Rate limiting
+dnspyre supports rate limiting number of queries send per second, for that you can use `--rate-limit` and `--rate-limit-worker` flags
+
+`--rate-limit` is used for setting a global rate limit, meaning that all concurrent workers spawned based on `--concurrency` flag will share this limit.
+It might happen, that some workers will be starving and the load generated by the `dnspyre` will not be evenly generated from all workers.
+
+For example this will generate load for 10 seconds using 10 concurrent workers and limit the load to 1 query per second
+```
+dnspyre --duration 10s -c 10 --rate-limit 1 --server '8.8.8.8' google.com
+```
+
+\
+`--rate-limit-worker` is used for setting a rate limit **separately** for each concurrent worker spawned based on `--concurrency` flag.
+
+For example this will generate load for 10 seconds using 10 concurrent workers and limit the load generated by each worker to 1 query per second
+```
+dnspyre --duration 10s -c 10 --rate-limit-worker 1 --server '8.8.8.8' google.com
+```