From 6e37762a4163904b0b82edda1d8628a2bf2cc9ef Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Thu, 11 Jul 2024 10:39:07 +0200
Subject: [PATCH] Set connection deadlines

Set dial/read/write deadlines.
---
 cli/benchserver.go | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/cli/benchserver.go b/cli/benchserver.go
index a1129c7..9e15327 100644
--- a/cli/benchserver.go
+++ b/cli/benchserver.go
@@ -21,6 +21,8 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net"
+	"net/http"
 	"net/url"
 	"os"
 	"strconv"
@@ -321,6 +323,7 @@ func (c *connections) roundTrip(i int, req serverRequest) (*clientReply, error)
 	for {
 		req.ClientIdx = i
 		conn := c.ws[i]
+		conn.SetWriteDeadline(time.Now().Add(2 * time.Second))
 		err := conn.WriteJSON(req)
 		if err != nil {
 			c.errLn(err)
@@ -329,6 +332,8 @@ func (c *connections) roundTrip(i int, req serverRequest) (*clientReply, error)
 			}
 			return nil, err
 		}
+		// Replies can be bigger, use longer deadline.
+		conn.SetReadDeadline(time.Now().Add(10 * time.Second))
 		var resp clientReply
 		err = conn.ReadJSON(&resp)
 		if err != nil {
@@ -345,6 +350,13 @@ func (c *connections) roundTrip(i int, req serverRequest) (*clientReply, error)
 // connect to a client.
 func (c *connections) connect(i int) error {
 	tries := 0
+	dialer := &websocket.Dialer{
+		NetDial: func(network, addr string) (net.Conn, error) {
+			return net.DialTimeout(network, addr, time.Second)
+		},
+		Proxy:            http.ProxyFromEnvironment,
+		HandshakeTimeout: 2 * time.Second,
+	}
 	for {
 		err := func() error {
 			host := c.hosts[i]
@@ -354,7 +366,7 @@ func (c *connections) connect(i int) error {
 			u := url.URL{Scheme: "ws", Host: host, Path: "/ws"}
 			c.info("Connecting to ", u.String())
 			var err error
-			c.ws[i], _, err = websocket.DefaultDialer.Dial(u.String(), nil)
+			c.ws[i], _, err = dialer.Dial(u.String(), nil)
 			if err != nil {
 				return err
 			}