graceful shutdown for websocket

[dozer111]

Hi, i use your library in this way

create handler

func generateQueryHandler(){
h := handler.New(schema)

h.AddTransport(transport.Websocket{
	Upgrader: websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
		EnableCompression: true,
	},
	KeepAlivePingInterval: cfg.Connection.WebsocketKeepAliveInterval,
})
h.AddTransport(transport.Options{})
h.AddTransport(transport.GET{})
h.AddTransport(transport.POST{})
h.AddTransport(transport.MultipartForm{})
h.SetQueryCache(lru.New(1000))
....
....
}

and then use it in server.

After i got SIGTERM, i close server by s.Shutdown - https://pkg.go.dev/net/http#Server.Shutdown

But default go http server cannot close websocket connections:

// Shutdown does not attempt to close nor wait for hijacked
// connections such as WebSockets. The caller of Shutdown should
// separately notify such long-lived connections of shutdown and wait
// for them to close, if desired. See RegisterOnShutdown for a way to
// register shutdown notification functions.
//
// Once Shutdown has been called on a server, it may not be reused;
// future calls to methods such as Serve will return ErrServerClosed.
func (srv *Server) Shutdown(ctx context.Context) error {

so, how to close websocket connections correct?

[dozer111]

the answer is InitFunc

Outside the code i has context with cancel(parentContext)

InitFunc: func(ctx context.Context, _ transport.InitPayload) (context.Context, *transport.InitPayload, error) {
	newCtx, initFuncCancel := context.WithCancel(ctx)

	go func() {
		<-parentContext.Done()
		initFuncCancel()
	}()

	return newCtx, nil, nil
},

---

Also, to ensure that websocket connections are init and close correctly i used this code

...
var liveWebsocketConnections atomic.Int64

h.AddTransport(transport.Websocket{
	...
	KeepAlivePingInterval: cfg.Connection.WebsocketKeepAliveInterval,
	InitFunc: func(ctx context.Context, _ transport.InitPayload) (context.Context, *transport.InitPayload, error) {
		liveWebsocketConnections.Add(1)
		logger.Warn(fmt.Sprintf("websocket open. Active connections %d", liveWebsocketConnections.Load()))
		newCtx, initFuncCancel := context.WithCancel(ctx)

		go func() {
			<-parentContext.Done()
			initFuncCancel()
		}()

		return newCtx, nil, nil
	},
	CloseFunc: func(ctx context.Context, closeCode int) {
		liveWebsocketConnections.Add(-1)
		logger.Warn(fmt.Sprintf("websocket closed. Active connections %d", liveWebsocketConnections.Load()))
	},
	...
})