Skip to content

Commit

Permalink
Set connection deadlines
Browse files Browse the repository at this point in the history
Set dial/read/write deadlines.
  • Loading branch information
klauspost committed Jul 11, 2024
1 parent 9e01ed3 commit 6e37762
Showing 1 changed file with 13 additions and 1 deletion.
14 changes: 13 additions & 1 deletion cli/benchserver.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@ import (
"context"
"errors"
"fmt"
"net"
"net/http"
"net/url"
"os"
"strconv"
Expand Down Expand Up @@ -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)
Expand All @@ -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 {
Expand All @@ -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]
Expand All @@ -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
}
Expand Down

0 comments on commit 6e37762

Please sign in to comment.