Skip to content

Commit

Permalink
fix: WebSocketConnection: Timeout on Read shuts down writing connection
Browse files Browse the repository at this point in the history
webSocketConnection Read and Write had separate
timeouts which, when expired, mutually closed the
websocket. In case only one party of the connection
was sending, the websocket was closed after the
timeout despite the fact that the connection was alive.
This comes from misunderstanding the impact of the Read/Write context parameter. See coder/websocket#242
  • Loading branch information
philippseith committed Oct 21, 2021
1 parent 0b9f47b commit 85843c5
Showing 1 changed file with 64 additions and 16 deletions.
80 changes: 64 additions & 16 deletions websocketconnection.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,44 +4,41 @@ import (
"bytes"
"context"
"fmt"

"github.com/teivah/onecontext"
"nhooyr.io/websocket"
"time"
)

type webSocketConnection struct {
ConnectionBase
conn *websocket.Conn
transferMode TransferMode
watchDogChan chan dogFood
}

func newWebSocketConnection(parentContext context.Context, requestContext context.Context, connectionID string, conn *websocket.Conn) *webSocketConnection {
ctx, _ := onecontext.Merge(parentContext, requestContext)
w := &webSocketConnection{
conn: conn,
conn: conn,
watchDogChan: make(chan dogFood, 1),
ConnectionBase: ConnectionBase{
ctx: ctx,
connectionID: connectionID,
},
}
go w.watchDog(ctx)
return w
}

func (w *webSocketConnection) Write(p []byte) (n int, err error) {
if err := w.Context().Err(); err != nil {
return 0, fmt.Errorf("webSocketConnection canceled: %w", w.ctx.Err())
}
ctx := w.ctx
if w.timeout > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(w.ctx, w.Timeout())
defer cancel() // has no effect because timeoutCtx is either done or not used anymore after websocket returns. But it keeps lint quiet
}
messageType := websocket.MessageText
if w.transferMode == BinaryTransferMode {
messageType = websocket.MessageBinary
}
err = w.conn.Write(ctx, messageType, p)
err = w.conn.Write(w.resetWatchDog(), messageType, p)
if err != nil {
return 0, err
}
Expand All @@ -52,19 +49,70 @@ func (w *webSocketConnection) Read(p []byte) (n int, err error) {
if err := w.Context().Err(); err != nil {
return 0, fmt.Errorf("webSocketConnection canceled: %w", w.ctx.Err())
}
ctx := w.ctx
if w.timeout > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(w.ctx, w.Timeout())
defer cancel() // has no effect because timeoutCtx is either done or not used anymore after websocket returns. But it keeps lint quiet
}
_, data, err := w.conn.Read(ctx)
_, data, err := w.conn.Read(w.resetWatchDog())
if err != nil {
return 0, err
}
return bytes.NewReader(data).Read(p)
}

// resetWatchDog resets the common watchDog for Read and Write.
// the watchDog will stop waiting for the last set timeout and wait for the new timeout.
func (w *webSocketConnection) resetWatchDog() context.Context {
ctx := w.ctx
food := dogFood{timeout: w.timeout}
if w.timeout > 0 {
ctx, food.bark = context.WithCancel(w.ctx)
}
w.watchDogChan <- food
return ctx
}

// dogFood is used to reset the watchDog
type dogFood struct {
// After this, the dog will bark
timeout time.Duration
bark context.CancelFunc
}

// watchDog is the common watchDog for Read and Write. It stops the connection (aka closes the Websocket)
// when the last timeout has elapsed. If resetWatchDog is called before the last timeout has elapsed,
// the watchDog will restart waiting for the new timeout. If timeout is set to 0, it will not wait at all.
func (w *webSocketConnection) watchDog(ctx context.Context) {
var timer *time.Timer
var cancelTimeoutChan chan struct{}
for {
select {
case <-ctx.Done():
return
case food := <-w.watchDogChan:
if timer != nil {
if !timer.Stop() {
go func() {
<-timer.C
}()
}
go func() {
cancelTimeoutChan <- struct{}{}
}()
}
if food.timeout != 0 {
timer = time.NewTimer(food.timeout)
cancelTimeoutChan = make(chan struct{}, 1)
go func() {
select {
case <-cancelTimeoutChan:
case <-timer.C:
food.bark()
}
}()
} else {
timer = nil
}
}
}
}

func (w *webSocketConnection) TransferMode() TransferMode {
return w.transferMode
}
Expand Down

0 comments on commit 85843c5

Please sign in to comment.