Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
Prevent leaking goroutines from websocket connections (#253)
The write to the CloseNotify() channel was blocking, preventing the
Read() call from exiting, and preventing ServeHTTP from returning. This
meant that requests tended to stay around forever in a deadlocked state.

Using the context package prevent this, as the cancellation is
non-blocking. It depends on Go 1.7.

Note that the writer still needs to support the CloseNotifier interface,
otherwise other things break, but it does not seem to be necessary to
signal using it.
  • Loading branch information
amerry authored and johanbrandhorst committed Oct 10, 2018
1 parent b89ec63 commit 1464bd9
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 4 deletions.
7 changes: 5 additions & 2 deletions go/grpcweb/websocket_wrapper.go
Expand Up @@ -3,6 +3,7 @@ package grpcweb
import (
"bufio"
"bytes"
"context"
"encoding/binary"
"errors"
"io"
Expand Down Expand Up @@ -110,6 +111,7 @@ type webSocketWrappedReader struct {
respWriter *webSocketResponseWriter
remainingBuffer []byte
remainingError error
cancel context.CancelFunc
}

func (w *webSocketWrappedReader) Close() error {
Expand Down Expand Up @@ -152,7 +154,7 @@ func (w *webSocketWrappedReader) Read(p []byte) (int, error) {
messageType, framePayload, err := w.wsConn.ReadMessage()
if err == io.EOF || messageType == -1 {
// The client has closed the connection. Indicate to the response writer that it should close
w.respWriter.closeNotifyChan <- true
w.cancel()
return 0, io.EOF
}

Expand Down Expand Up @@ -193,12 +195,13 @@ func (w *webSocketWrappedReader) Read(p []byte) (int, error) {
return len(p), nil
}

func newWebsocketWrappedReader(wsConn *websocket.Conn, respWriter *webSocketResponseWriter) *webSocketWrappedReader {
func newWebsocketWrappedReader(wsConn *websocket.Conn, respWriter *webSocketResponseWriter, cancel context.CancelFunc) *webSocketWrappedReader {
return &webSocketWrappedReader{
wsConn: wsConn,
respWriter: respWriter,
remainingBuffer: nil,
remainingError: nil,
cancel: cancel,
}
}

Expand Down
8 changes: 6 additions & 2 deletions go/grpcweb/wrapper.go
Expand Up @@ -7,6 +7,7 @@ import (
"net/http"
"net/url"

"context"
"strings"
"time"

Expand Down Expand Up @@ -148,14 +149,17 @@ func (w *WrappedGrpcServer) handleWebSocket(wsConn *websocket.Conn, req *http.Re
return
}

ctx, cancelFunc := context.WithCancel(req.Context())
defer cancelFunc()

respWriter := newWebSocketResponseWriter(wsConn)
wrappedReader := newWebsocketWrappedReader(wsConn, respWriter)
wrappedReader := newWebsocketWrappedReader(wsConn, respWriter, cancelFunc)

req.Body = wrappedReader
req.Method = http.MethodPost
req.Header = headers

w.server.ServeHTTP(respWriter, hackIntoNormalGrpcRequest(req))
w.server.ServeHTTP(respWriter, hackIntoNormalGrpcRequest(req.WithContext(ctx)))
}

// IsGrpcWebRequest determines if a request is a gRPC-Web request by checking that the "content-type" is
Expand Down

0 comments on commit 1464bd9

Please sign in to comment.