Add support for OnDisconnect when WS connection is closed

v4 only

Author: StephenSorriaux

Diff for: engineio/server.v4.go

@@ -111,6 +111,7 @@ func (v4 *serverV4) serveTransport(w http.ResponseWriter, r *http.Request) (tran
 	ctx = v4.sessions.WithCancel(ctx)
 	ctx = v4.sessions.WithInterval(ctx, v4.pingInterval)
 	ctx = v4.sessions.WithTimeout(ctx, v4.pingTimeout)
+	ctx = v4.sessions.WithDisconnectOnClose(ctx)
 
 	go func() {
 		v4.transportRunError <- upgrade.transport.Run(w, r.WithContext(ctx), append(v4.eto, opts...)...)

Diff for: engineio/session/manage.go

@@ -9,8 +9,9 @@ const (
 	SessionIntervalKey      sessionCtxKey = "interval"
 	SessionExtendTimeoutKey sessionCtxKey = "timeout-extend"
 
-	SessionCloseChannelKey  sessionCtxKey = "cancel-channel"
-	SessionCloseFunctionKey sessionCtxKey = "cancel-function"
+	SessionCloseChannelKey       sessionCtxKey = "cancel-channel"
+	SessionCloseFunctionKey      sessionCtxKey = "cancel-function"
+	SessionDisconnectFunctionKey sessionCtxKey = "cancel-disconnect"
 )
 
 type (

Diff for: engineio/sessions.go

@@ -25,6 +25,7 @@ type transportSessions interface {
 	WithCancel(ctx context.Context) context.Context
 	WithTimeout(ctx context.Context, d time.Duration) context.Context
 	WithInterval(ctx context.Context, d time.Duration) context.Context
+	WithDisconnectOnClose(ctx context.Context) context.Context
 }
 
 type sessions struct {
@@ -88,6 +89,23 @@ type lifecycle struct {
 	removeTransport func(SessionID)
 }
 
+func (c *lifecycle) WithDisconnectOnClose(ctx context.Context) context.Context {
+	sessionID, ok := ctx.Value(ctxSessionID).(SessionID)
+	if !ok {
+		// there is no session to attach the context to
+		return ctx
+	}
+	ctx = context.WithValue(ctx, eios.SessionDisconnectFunctionKey, func() func() {
+		return func() {
+			c.removeSession(sessionID)
+			if c.removeTransport != nil {
+				c.removeTransport(sessionID)
+			}
+		}
+	})
+	return ctx
+}
+
 func (c *lifecycle) WithCancel(ctx context.Context) context.Context {
 	sessionID, ok := ctx.Value(ctxSessionID).(SessionID)
 	if !ok {

Diff for: engineio/transport/transport.websocket.go

@@ -66,6 +66,16 @@ func (t *WebsocketTransport) With(opts ...Option) {
 
 func (t *WebsocketTransport) InnerTransport() *Transport { return t.Transport }
 
+func (t *WebsocketTransport) ConnClose(code ws.StatusCode, reason string, ctx context.Context) error {
+	err := t.conn.Close(code, reason)
+	if done, ok := ctx.Value(eios.SessionDisconnectFunctionKey).(func() func()); ok {
+		if cleanup := done(); cleanup != nil {
+			cleanup()
+		}
+	}
+	return err
+}
+
 func (t *WebsocketTransport) Run(w http.ResponseWriter, r *http.Request, opts ...Option) (err error) {
 	t.With(opts...)
 
@@ -95,7 +105,7 @@ func (t *WebsocketTransport) Run(w http.ResponseWriter, r *http.Request, opts ..
 	grp.Go(func() error { return t.outgoing(r.WithContext(ctx)) })
 
 	err = grp.Wait()
-	t.conn.Close(ws.StatusNormalClosure, "done")
+	t.ConnClose(ws.StatusNormalClosure, "done", ctx)
 	return err
 }
 
@@ -161,7 +171,7 @@ func (t *WebsocketTransport) incoming(ctx context.Context) (err error) {
 
 	var done func()
 	var reason string
-	defer func() { t.conn.Close(ws.StatusNormalClosure, reason) }()
+	defer func() { t.ConnClose(ws.StatusNormalClosure, reason, ctx) }()
 
 	var start = time.Now()
 Write:
@@ -259,7 +269,7 @@ func (t *WebsocketTransport) outgoing(r *http.Request) (err error) {
 	}
 
 	var unbuffered = new(sync.WaitGroup)
-	defer t.conn.Close(ws.StatusNormalClosure, "read")
+	defer t.ConnClose(ws.StatusNormalClosure, "read", ctx)
 
 	for {
 		if !t.buffered {
@@ -309,7 +319,7 @@ func (t *WebsocketTransport) outgoing(r *http.Request) (err error) {
 				}
 			}
 			t.conn.CloseRead(ctx)
-			t.conn.Close(ws.StatusNormalClosure, "cross origin WebSocket accepted")
+			t.ConnClose(ws.StatusNormalClosure, "cross origin WebSocket accepted", ctx)
 			return nil
 		case eiop.PingPacket:
 			cw, err := t.conn.Writer(ctx, ws.MessageText)