add timeout for each DoH and DoQ query

xxnuo · Jul 8, 2023 · a00e89d · a00e89d
1 parent f4f5c68
commit a00e89d
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 5 deletions.
diff --git a/cmd/benchmark.go b/cmd/benchmark.go
@@ -152,7 +152,13 @@ func (b *Benchmark) Run(ctx context.Context) ([]*ResultStats, error) {
 
 	var query queryFunc
 	if b.useDoH {
-		query, network = b.getDoHClient()
+		var dohQuery queryFunc
+		dohQuery, network = b.getDoHClient()
+		query = func(ctx context.Context, s string, msg *dns.Msg) (*dns.Msg, error) {
+			reqTimeoutCtx, cancel := context.WithTimeout(ctx, dnsTimeout)
+			defer cancel()
+			return dohQuery(reqTimeoutCtx, s, msg)
+		}
 	}
 
 	if b.useQuic {
@@ -163,7 +169,9 @@ func (b *Benchmark) Run(ctx context.Context) ([]*ResultStats, error) {
 			return nil, err
 		}
 		query = func(ctx context.Context, _ string, msg *dns.Msg) (*dns.Msg, error) {
-			return quicClient.Send(ctx, msg)
+			reqTimeoutCtx, cancel := context.WithTimeout(ctx, dnsTimeout)
+			defer cancel()
+			return quicClient.Send(reqTimeoutCtx, msg)
 		}
 		network = "quic"
 	}

diff --git a/go.mod b/go.mod
@@ -12,7 +12,7 @@ require (
 	github.com/quic-go/quic-go v0.36.1
 	github.com/stretchr/testify v1.8.4
 	github.com/tantalor93/doh-go v0.1.0
-	github.com/tantalor93/doq-go v0.3.0
+	github.com/tantalor93/doq-go v0.4.0
 	go-hep.org/x/hep v0.33.0
 	go.uber.org/ratelimit v0.2.0
 	golang.org/x/net v0.11.0

diff --git a/go.sum b/go.sum
@@ -101,8 +101,8 @@ github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcU
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 github.com/tantalor93/doh-go v0.1.0 h1:GpC7qH2ciRFPIpfiFJ0FMmp1Q4dR1zNIhV16+yByciA=
 github.com/tantalor93/doh-go v0.1.0/go.mod h1:GvJN2FOEuFGO0STTCbggTgGsAYkL8b04xQY+vf72IV0=
-github.com/tantalor93/doq-go v0.3.0 h1:b5meyME5mAPIespTnoveWRMIbUW8ehQQvqiGpqXguNs=
-github.com/tantalor93/doq-go v0.3.0/go.mod h1:La3/FOfPtW62nZS96wsOoFjKuKf6DE5DokefmzdEXw8=
+github.com/tantalor93/doq-go v0.4.0 h1:VFp9HS8sqZoGfoHo2X+BceYnU03OGRmY+1zM3N2ZxlE=
+github.com/tantalor93/doq-go v0.4.0/go.mod h1:La3/FOfPtW62nZS96wsOoFjKuKf6DE5DokefmzdEXw8=
 github.com/xhit/go-str2duration/v2 v2.1.0 h1:lxklc02Drh6ynqX+DdPyp5pCKLUQpRT8bp8Ydu2Bstc=
 github.com/xhit/go-str2duration/v2 v2.1.0/go.mod h1:ohY8p+0f07DiV6Em5LKB0s2YpLtXVyJfNt1+BlmyAsU=
 github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=