Add websocket TCP write timeouts (#18988)

Add 3 second write timeout to writes.

This will make dead TCP connections terminate in a reasonable time.

Fixes writes blocking for reconnection.
This commit is contained in:
Klaus Post 2024-02-06 13:34:46 -08:00 committed by GitHub
parent ebc6c9b498
commit 22687c1f50
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -179,6 +179,7 @@ const (
writeBufferSize = 16 << 10 writeBufferSize = 16 << 10
defaultDialTimeout = 2 * time.Second defaultDialTimeout = 2 * time.Second
connPingInterval = 10 * time.Second connPingInterval = 10 * time.Second
connWriteTimeout = 3 * time.Second
) )
type connectionParams struct { type connectionParams struct {
@ -600,6 +601,10 @@ func (c *Connection) sendMsg(conn net.Conn, msg message, payload msgp.MarshalSiz
if c.outgoingBytes != nil { if c.outgoingBytes != nil {
c.outgoingBytes(int64(len(dst))) c.outgoingBytes(int64(len(dst)))
} }
err = conn.SetWriteDeadline(time.Now().Add(connWriteTimeout))
if err != nil {
return err
}
return wsutil.WriteMessage(conn, c.side, ws.OpBinary, dst) return wsutil.WriteMessage(conn, c.side, ws.OpBinary, dst)
} }
@ -1104,6 +1109,11 @@ func (c *Connection) handleMessages(ctx context.Context, conn net.Conn) {
return return
} }
PutByteBuffer(toSend) PutByteBuffer(toSend)
err = conn.SetWriteDeadline(time.Now().Add(connWriteTimeout))
if err != nil {
logger.LogIf(ctx, fmt.Errorf("conn.SetWriteDeadline: %w", err))
return
}
_, err = buf.WriteTo(conn) _, err = buf.WriteTo(conn)
if err != nil { if err != nil {
logger.LogIf(ctx, fmt.Errorf("ws write: %w", err)) logger.LogIf(ctx, fmt.Errorf("ws write: %w", err))
@ -1143,6 +1153,11 @@ func (c *Connection) handleMessages(ctx context.Context, conn net.Conn) {
return return
} }
// buf is our local buffer, so we can reuse it. // buf is our local buffer, so we can reuse it.
err = conn.SetWriteDeadline(time.Now().Add(connWriteTimeout))
if err != nil {
logger.LogIf(ctx, fmt.Errorf("conn.SetWriteDeadline: %w", err))
return
}
_, err = buf.WriteTo(conn) _, err = buf.WriteTo(conn)
if err != nil { if err != nil {
logger.LogIf(ctx, fmt.Errorf("ws write: %w", err)) logger.LogIf(ctx, fmt.Errorf("ws write: %w", err))