Fix leak when reconnecting to back end without reload. Fixes eclipse-…

…theia#13215 (eclipse-theia#13250) Contributed on behalf of STMicroelectronics Signed-off-by: Thomas Mäder <t.s.maeder@gmail.com>
tsmaeder · Jan 16, 2024 · 0dbca07 · 0dbca07
1 parent 15caac5
commit 0dbca07
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 26 deletions.
diff --git a/packages/core/src/browser/messaging/ws-connection-source.ts b/packages/core/src/browser/messaging/ws-connection-source.ts
@@ -87,38 +87,58 @@ export class WebSocketConnectionSource implements ConnectionSource {
         this._socket.connect();
     }
 
-    protected handleSocketConnected(): void {
-        if (this.currentChannel) {
-            const reconnectListener = (hasConnection: boolean) => {
-                this._socket.off(ConnectionManagementMessages.RECONNECT, reconnectListener);
-                if (hasConnection) {
-                    this.writeBuffer.flush(this.socket);
+    protected negogiateReconnect(): void {
+        const reconnectListener = (hasConnection: boolean) => {
+            this._socket.off(ConnectionManagementMessages.RECONNECT, reconnectListener);
+            if (hasConnection) {
+                console.info(`reconnect succeeded on ${this.socket.id}`);
+                this.writeBuffer!.flush(this.socket);
+            } else {
+                if (FrontendApplicationConfigProvider.get().reloadOnReconnect) {
+                    window.location.reload(); // this might happen in the preload module, when we have no window service yet
                 } else {
-                    if (FrontendApplicationConfigProvider.get().reloadOnReconnect) {
-                        window.location.reload(); // this might happen in the preload module, when we have no window service yet
-                    } else {
-                        this.connectNewChannel();
-                    }
+                    console.info(`reconnect failed on ${this.socket.id}`);
+                    this.currentChannel.onCloseEmitter.fire({ reason: 'reconnecting channel' });
+                    this.currentChannel.close();
+                    this.writeBuffer.drain();
+                    this.socket.disconnect();
+                    this.socket.connect();
+                    this.negotiateInitialConnect();
                 }
-            };
-            this._socket.on(ConnectionManagementMessages.RECONNECT, reconnectListener);
-            this._socket.emit(ConnectionManagementMessages.RECONNECT, this.frontendIdProvider.getId());
+            }
+        };
+        this._socket.on(ConnectionManagementMessages.RECONNECT, reconnectListener);
+        console.info(`sending reconnect on ${this.socket.id}`);
+        this._socket.emit(ConnectionManagementMessages.RECONNECT, this.frontendIdProvider.getId());
+    }
+
+    protected negotiateInitialConnect(): void {
+        const initialConnectListener = () => {
+            console.info(`initial connect received on ${this.socket.id}`);
+
+            this._socket.off(ConnectionManagementMessages.INITIAL_CONNECT, initialConnectListener);
+            this.connectNewChannel();
+        };
+        this._socket.on(ConnectionManagementMessages.INITIAL_CONNECT, initialConnectListener);
+        console.info(`sending initial connect on ${this.socket.id}`);
+
+        this._socket.emit(ConnectionManagementMessages.INITIAL_CONNECT, this.frontendIdProvider.getId());
+    }
+
+    protected handleSocketConnected(): void {
+        if (this.currentChannel) {
+            this.negogiateReconnect();
         } else {
-            const initialConnectListener = () => {
-                this._socket.off(ConnectionManagementMessages.INITIAL_CONNECT, initialConnectListener);
-                this.connectNewChannel();
-            };
-            this._socket.on(ConnectionManagementMessages.INITIAL_CONNECT, initialConnectListener);
-            this._socket.emit(ConnectionManagementMessages.INITIAL_CONNECT, this.frontendIdProvider.getId());
+            this.negotiateInitialConnect();
         }
     }
 
     connectNewChannel(): void {
         if (this.currentChannel) {
-            this.writeBuffer.drain();
             this.currentChannel.close();
             this.currentChannel.onCloseEmitter.fire({ reason: 'reconnecting channel' });
         }
+        this.writeBuffer.drain();
         this.currentChannel = this.createChannel();
         this.onConnectionDidOpenEmitter.fire(this.currentChannel);
     }

diff --git a/packages/core/src/node/messaging/websocket-frontend-connection-service.ts b/packages/core/src/node/messaging/websocket-frontend-connection-service.ts
@@ -51,7 +51,9 @@ export class WebsocketFrontendConnectionService implements FrontendConnectionSer
             if (this.connectionsByFrontend.has(frontEndId)) {
                 this.closeConnection(frontEndId, 'reconnecting same front end');
             }
-            channelCreatedHandler(this.createConnection(socket, frontEndId));
+            const channel = this.createConnection(socket, frontEndId);
+            this.handleSocketDisconnect(socket, channel, frontEndId);
+            channelCreatedHandler(channel);
             socket.emit(ConnectionManagementMessages.INITIAL_CONNECT);
         };
 
@@ -63,6 +65,7 @@ export class WebsocketFrontendConnectionService implements FrontendConnectionSer
                 console.info(`Reconnecting to front end ${frontEndId}`);
                 socket.emit(ConnectionManagementMessages.RECONNECT, true);
                 channel.connect(socket);
+                this.handleSocketDisconnect(socket, channel, frontEndId);
                 const pendingTimeout = this.closeTimeouts.get(frontEndId);
                 clearTimeout(pendingTimeout);
                 this.closeTimeouts.delete(frontEndId);
@@ -89,11 +92,16 @@ export class WebsocketFrontendConnectionService implements FrontendConnectionSer
         connection.close();
     }
 
-    protected createConnection(socket: Socket, frontEndId: string): Channel {
+    protected createConnection(socket: Socket, frontEndId: string): ReconnectableSocketChannel {
         console.info(`creating connection for ${frontEndId}`);
         const channel = new ReconnectableSocketChannel();
         channel.connect(socket);
 
+        this.connectionsByFrontend.set(frontEndId, channel);
+        return channel;
+    }
+
+    handleSocketDisconnect(socket: Socket, channel: ReconnectableSocketChannel, frontEndId: string): void {
         socket.on('disconnect', evt => {
             console.info('socked closed');
             channel.disconnect();
@@ -111,9 +119,6 @@ export class WebsocketFrontendConnectionService implements FrontendConnectionSer
                 // timeout < 0: never close the back end
             }
         });
-
-        this.connectionsByFrontend.set(frontEndId, channel);
-        return channel;
     }
 
     markForClose(channelId: string): void {