Close Websocket Connection on Context close/cancel

graphql/handler/transport/websocket.go

@@ -170,6 +170,10 @@ func (c *wsConnection) run() {
 		go c.ping(ctx)
 	}
 
+	// Close the connection when the context is cancelled.
+	// Will optionally send a "close reason" that is retrieved from the context.
+	go c.closeOnCancel(ctx)
+
 	for {
 		start := graphql.Now()
 		m, err := c.me.NextMessage()
@@ -227,6 +231,15 @@ func (c *wsConnection) ping(ctx context.Context) {
 	}
 }
 
+func (c *wsConnection) closeOnCancel(ctx context.Context) {
+	<-ctx.Done()
+
+	if r := closeReasonForContext(ctx); r != "" {
+		c.sendConnectionError(r)
+	}
+	c.close(websocket.CloseNormalClosure, "terminated")
+}
+
 func (c *wsConnection) subscribe(start time.Time, msg *message) {
 	ctx := graphql.StartOperationTrace(c.ctx)
 	var params *graphql.RawParams

graphql/handler/transport/websocket_close_reason.go

@@ -0,0 +1,22 @@
+package transport
+
+import (
+	"context"
+)
+
+// A private key for context that only this package can access. This is important
+// to prevent collisions between different context uses
+var closeReasonCtxKey = &wsCloseReasonContextKey{"close-reason"}
+
+type wsCloseReasonContextKey struct {
+	name string
+}
+
+func AppendCloseReason(ctx context.Context, reason string) context.Context {
+	return context.WithValue(ctx, closeReasonCtxKey, reason)
+}
+
+func closeReasonForContext(ctx context.Context) string {
+	reason, _ := ctx.Value(closeReasonCtxKey).(string)
+	return reason
+}

graphql/handler/transport/websocket_graphql_transport_ws.go

@@ -21,18 +21,16 @@ const (
 	graphqltransportwsPongMsg           = graphqltransportwsMessageType("pong")
 )
 
-var (
-	allGraphqltransportwsMessageTypes = []graphqltransportwsMessageType{
-		graphqltransportwsConnectionInitMsg,
-		graphqltransportwsConnectionAckMsg,
-		graphqltransportwsSubscribeMsg,
-		graphqltransportwsNextMsg,
-		graphqltransportwsErrorMsg,
-		graphqltransportwsCompleteMsg,
-		graphqltransportwsPingMsg,
-		graphqltransportwsPongMsg,
-	}
-)
+var allGraphqltransportwsMessageTypes = []graphqltransportwsMessageType{
+	graphqltransportwsConnectionInitMsg,
+	graphqltransportwsConnectionAckMsg,
+	graphqltransportwsSubscribeMsg,
+	graphqltransportwsNextMsg,
+	graphqltransportwsErrorMsg,
+	graphqltransportwsCompleteMsg,
+	graphqltransportwsPingMsg,
+	graphqltransportwsPongMsg,
+}
 
 type (
 	graphqltransportwsMessageExchanger struct {

graphql/handler/transport/websocket_test.go

@@ -265,6 +265,32 @@ func TestWebsocketInitFunc(t *testing.T) {
 		require.NoError(t, err)
 		assert.Equal(t, "ok", resp.Empty)
 	})
+
+	t.Run("can set a deadline on a websocket connection and close it with a reason", func(t *testing.T) {
+		h := testserver.New()
+		var cancel func()
+		h.AddTransport(transport.Websocket{
+			InitFunc: func(ctx context.Context, _ transport.InitPayload) (newCtx context.Context, _ error) {
+				newCtx, cancel = context.WithTimeout(transport.AppendCloseReason(ctx, "beep boop"), time.Millisecond*5)
+				return
+			},
+		})
+		srv := httptest.NewServer(h)
+		defer srv.Close()
+
+		c := wsConnect(srv.URL)
+		require.NoError(t, c.WriteJSON(&operationMessage{Type: connectionInitMsg}))
+		assert.Equal(t, connectionAckMsg, readOp(c).Type)
+		assert.Equal(t, connectionKeepAliveMsg, readOp(c).Type)
+
+		// Cancel should contain an actual value now, so let's call it when we exit this scope (to make the linter happy)
+		defer cancel()
+
+		time.Sleep(time.Millisecond * 10)
+		m := readOp(c)
+		assert.Equal(t, m.Type, connectionErrorMsg)
+		assert.Equal(t, string(m.Payload), `{"message":"beep boop"}`)
+	})
 }
 
 func TestWebsocketGraphqltransportwsSubprotocol(t *testing.T) {