feat(connection): added handlers on open and close connections

packages/core/src/com/logger-service.ts

@@ -97,7 +97,7 @@ function getValue(message: LogValue): LogValueData {
     return logValue;
 }
 
-function logToConsole({ message, metadata = {}, level }: LogMessage): void {
+export function logToConsole({ message, metadata = {}, level }: LogMessage): void {
     switch (level) {
         case LogLevel.DEBUG:
             console.log(message, metadata);

packages/runtime-node/src/node-env-manager.ts

@@ -10,7 +10,7 @@ import { IDisposable, SetMultiMap } from '@dazl/patterns';
 import { fileURLToPath } from 'node:url';
 import { parseArgs } from 'node:util';
 import { extname } from 'node:path';
-import { WsServerHost } from './ws-node-host.js';
+import { ConnectionHandlers, WsServerHost } from './ws-node-host.js';
 import { ILaunchHttpServerOptions, launchEngineHttpServer } from './launch-http-server.js';
 import { workerThreadInitializer2 } from './worker-thread-initializer2.js';
 import { bindMetricsListener, type PerformanceMetrics } from './metrics-utils.js';
@@ -49,7 +49,12 @@ export class NodeEnvManager implements IDisposable {
     ) {}
     public async autoLaunch(
         runtimeOptions: Map<string, string | boolean | undefined>,
-        serverOptions: ILaunchHttpServerOptions = {},
+        {
+            connectionHandlers,
+            ...serverOptions
+        }: ILaunchHttpServerOptions & {
+            connectionHandlers?: ConnectionHandlers;
+        } = {},
     ) {
         process.env.ENGINE_FLOW_V2_DIST_URL = this.importMeta.url;
         const disposeMetricsListener = bindMetricsListener(() => this.collectMetricsFromAllOpenEnvironments());
@@ -60,6 +65,20 @@ export class NodeEnvManager implements IDisposable {
         runtimeOptions.set('enginePort', port.toString());
 
         const clientsHost = new WsServerHost(socketServer);
+        const disposeOnConnectionOpen = connectionHandlers?.onConnectionOpen
+            ? clientsHost.registerConnectionHandler(connectionHandlers.onConnectionOpen)
+            : undefined;
+        const disposeOnConnectionClose = connectionHandlers?.onConnectionClose
+            ? clientsHost.registerDisconnectionHandler(connectionHandlers.onConnectionClose)
+            : undefined;
+        const disposeOnReconnection = connectionHandlers?.onConnectionReconnect
+            ? clientsHost.registerReconnectionHandler(connectionHandlers.onConnectionReconnect)
+            : undefined;
+        const disposeConnectionHandlers = () => {
+            disposeOnConnectionOpen?.();
+            disposeOnConnectionClose?.();
+            disposeOnReconnection?.();
+        };
         clientsHost.addEventListener('message', handleRegistrationOnMessage);
         const forwardingCom = new Communication(clientsHost, 'clients-host-com');
         function handleRegistrationOnMessage({ data }: { data: Message }) {
@@ -84,6 +103,7 @@ export class NodeEnvManager implements IDisposable {
             disposeMetricsListener();
             await this.closeAll();
             clientsHost.removeEventListener('message', handleRegistrationOnMessage);
+            disposeConnectionHandlers();
             await clientsHost.dispose();
             await close();
         };

packages/runtime-node/src/ws-node-host.ts

@@ -1,7 +1,19 @@
 import type io from 'socket.io';
-import { BaseHost, type Message } from '@dazl/engine-core';
+import { BaseHost, LogLevel, logToConsole, type Message } from '@dazl/engine-core';
 import { SafeDisposable, type IDisposable } from '@dazl/patterns';
 
+export interface IConnectionEvent {
+    clientId: string;
+    socket: io.Socket;
+}
+export type IConnectionHandler = (event: IConnectionEvent) => void;
+
+export interface ConnectionHandlers {
+    onConnectionOpen?: IConnectionHandler;
+    onConnectionClose?: IConnectionHandler;
+    onConnectionReconnect?: IConnectionHandler;
+}
+
 export class WsHost extends BaseHost {
     constructor(private socket: io.Socket) {
         super();
@@ -15,6 +27,9 @@ export class WsHost extends BaseHost {
 }
 
 export class WsServerHost extends BaseHost implements IDisposable {
+    private connectionHandlers = new Set<IConnectionHandler>();
+    private disconnectionHandlers = new Set<IConnectionHandler>();
+    private reconnectionHandlers = new Set<IConnectionHandler>();
     private socketToEnvId = new Map<string, { socket: io.Socket; clientID: string }>();
     private clientIdToSocket = new Map<string, io.Socket>();
     private disposables = new SafeDisposable(WsServerHost.name);
@@ -26,6 +41,30 @@ export class WsServerHost extends BaseHost implements IDisposable {
         this.server.on('connection', this.onConnection);
         this.disposables.add('connection', () => this.server.off('connection', this.onConnection));
         this.disposables.add('clear handlers', () => this.handlers.clear());
+        this.disposables.add('clear connection handlers', () => this.connectionHandlers.clear());
+        this.disposables.add('clear disconnection handlers', () => this.disconnectionHandlers.clear());
+        this.disposables.add('clear reconnection handlers', () => this.reconnectionHandlers.clear());
+    }
+
+    public registerConnectionHandler(handler: IConnectionHandler) {
+        this.connectionHandlers.add(handler);
+        return () => {
+            this.connectionHandlers.delete(handler);
+        };
+    }
+
+    public registerDisconnectionHandler(handler: IConnectionHandler) {
+        this.disconnectionHandlers.add(handler);
+        return () => {
+            this.disconnectionHandlers.delete(handler);
+        };
+    }
+
+    public registerReconnectionHandler(handler: IConnectionHandler) {
+        this.reconnectionHandlers.add(handler);
+        return () => {
+            this.reconnectionHandlers.delete(handler);
+        };
     }
 
     public postMessage(data: Message) {
@@ -42,17 +81,40 @@ export class WsServerHost extends BaseHost implements IDisposable {
         }
     }
 
+    private callWithErrorHandling(callback: () => void): void {
+        try {
+            callback();
+        } catch (error) {
+            logToConsole({
+                message: error instanceof Error ? `${error.message} ${error.stack}` : String(error),
+                level: LogLevel.ERROR,
+                timestamp: Date.now(),
+            });
+        }
+    }
+
     private onConnection = (socket: io.Socket): void => {
         const clientId = socket.handshake.auth?.clientId || socket.id;
-
-        // disconnect previous connection
         const existingSocket = this.clientIdToSocket.get(clientId);
-        if (existingSocket && existingSocket.connected) {
-            existingSocket.disconnect(true);
-        }
 
         this.clientIdToSocket.set(clientId, socket);
 
+        const connectionEvent: IConnectionEvent = {
+            clientId,
+            socket,
+        };
+
+        if (existingSocket && existingSocket.connected) {
+            existingSocket.disconnect(true); // disconnect previous connection
+            for (const handler of this.reconnectionHandlers) {
+                this.callWithErrorHandling(() => handler(connectionEvent));
+            }
+        } else {
+            for (const handler of this.connectionHandlers) {
+                this.callWithErrorHandling(() => handler(connectionEvent));
+            }
+        }
+
         const nameSpace = (original: string) => `${clientId}/${original}`;
         const onMessage = (message: Message): void => {
             // this mapping should not be here because of forwarding of messages
@@ -87,6 +149,9 @@ export class WsServerHost extends BaseHost implements IDisposable {
             }
             if (this.clientIdToSocket.get(clientId) === socket) {
                 this.clientIdToSocket.delete(clientId);
+                for (const handler of this.disconnectionHandlers) {
+                    this.callWithErrorHandling(() => handler(connectionEvent));
+                }
             }
         });
     };

packages/runtime-node/test/node-env.manager.unit.ts

@@ -1,6 +1,8 @@
 import { expect } from 'chai';
+import sinon from 'sinon';
 import { BaseHost, COM, Communication, WsClientHost } from '@dazl/engine-core';
 import {
+    IConnectionHandler,
     launchEngineHttpServer,
     NodeEnvManager,
     type NodeEnvsFeatureMapping,
@@ -153,6 +155,100 @@ describe('NodeEnvManager', () => {
         });
     });
 
+    describe('NodeEnvManager with connection handlers', () => {
+        it('should call connection, reconnection and disconnection handlers', async () => {
+            const connectionHandler = sinon.spy();
+            const disconnectionHandler = sinon.spy();
+            const reconnectionHandler = sinon.spy();
+
+            const featureEnvironmentsMapping: NodeEnvsFeatureMapping = {
+                featureToEnvironments: {
+                    'test-feature': [aEnv.env],
+                },
+                availableEnvironments: {
+                    a: {
+                        env: aEnv.env,
+                        endpointType: 'single',
+                        envType: 'node',
+                    },
+                },
+            };
+
+            const manager = disposeAfterTest(new NodeEnvManager(meta, featureEnvironmentsMapping));
+            const { port } = await manager.autoLaunch(new Map([['feature', 'test-feature']]), {
+                connectionHandlers: {
+                    onConnectionOpen: connectionHandler,
+                    onConnectionClose: disconnectionHandler,
+                    onConnectionReconnect: reconnectionHandler,
+                },
+            });
+
+            // Create a client connection 1
+            const initialClientId = 'test-client-id';
+            const client1 = new WsClientHost('http://localhost:' + port, {
+                auth: {
+                    clientId: initialClientId,
+                },
+            });
+            const communication1 = new Communication(new BaseHost(), testCommunicationId);
+            communication1.registerEnv(aEnv.env, client1);
+            communication1.registerMessageHandler(client1);
+
+            await client1.connected;
+
+            // Verify connection handler was called
+            expect(connectionHandler.callCount).to.equal(1);
+            const [args1] = connectionHandler.firstCall.args as Parameters<IConnectionHandler>;
+            expect(args1.clientId).to.equal(initialClientId);
+            expect(args1.socket).to.have.property('id');
+
+            // Replace a client connection
+            const waitDisconnectFirstSocket = new Promise<void>((resolve) => {
+                args1.socket.on('disconnect', () => {
+                    resolve();
+                });
+            });
+            const client2 = new WsClientHost('http://localhost:' + port, {
+                auth: {
+                    clientId: initialClientId,
+                },
+            });
+            const communication2 = new Communication(new BaseHost(), testCommunicationId);
+            communication2.registerEnv(aEnv.env, client2);
+            communication2.registerMessageHandler(client2);
+
+            await client2.connected;
+
+            // Verify reconnection handler was called
+            expect(reconnectionHandler.callCount).to.equal(1);
+            const [args2] = reconnectionHandler.firstCall.args as Parameters<IConnectionHandler>;
+            expect(args2.clientId).to.equal(initialClientId);
+            expect(args2.socket).to.have.property('id');
+            expect(args2.socket.id).to.not.equal(args1.socket.id);
+
+            // Verify disconnection handler was not called after the client connection was replaced
+            await waitDisconnectFirstSocket;
+            expect(disconnectionHandler.callCount).to.equal(0);
+
+            // Disconnect an active socket
+            await new Promise<void>((resolve) => {
+                args2.socket.on('disconnect', () => {
+                    resolve();
+                });
+                client2.disconnectSocket();
+            });
+
+            // Verify disconnection handler for active connection was called
+            expect(disconnectionHandler.callCount).to.equal(1);
+            const [disconnectArgs1] = disconnectionHandler.firstCall.args as Parameters<IConnectionHandler>;
+            expect(disconnectArgs1.clientId).to.equal(initialClientId);
+            expect(disconnectArgs1.socket.id).to.equal(args2.socket.id);
+
+            await communication1.dispose();
+            await communication2.dispose();
+        });
+    });
+
     function getClientCom(port: number) {
         const host = new WsClientHost('http://localhost:' + port, {});
         const com = new Communication(new BaseHost(), testCommunicationId);