add onDisconnect callback

Author: Mahad-10

lib/session.ts

@@ -43,6 +43,7 @@ export class Session {
     private _baseSession: IBaseSession;
     private _wampSession: WAMPSession;
     private _idGen: SessionScopeIDGenerator = new SessionScopeIDGenerator();
+    private _disconnectCallbacks: Array<() => Promise<void>> = [];
 
     private _callRequests: Map<number, {
         resolve: (value: Result) => void,
@@ -76,6 +77,7 @@ export class Session {
     constructor(baseSession: IBaseSession) {
         this._baseSession = baseSession;
         this._wampSession = new WAMPSession(baseSession.serializer());
+        this._baseSession.onDisconnect(async () => { await this.markDisconnected();});
 
         (async () => {
             for (; ;) {
@@ -85,6 +87,10 @@ export class Session {
         })();
     }
 
+    onDisconnect(callback: () => Promise<void>): void {
+        this._disconnectCallbacks.push(callback);
+    }
+
     private get _nextID(): number {
         return this._idGen.next();
     }
@@ -234,15 +240,15 @@ export class Session {
                     throw new ProtocolError(wampErrorString(message));
             }
         } else if (message instanceof Goodbye) {
-            this.markDisconnected()
+            await this.markDisconnected()
         } else {
             throw new ProtocolError(`Unexpected message type ${typeof message}`);
         }
     }
 
-    private markDisconnected() {
-        if (!this.isConnected()) {
-            return
+    private async markDisconnected() {
+        if (this._disconnectCallbacks.length > 0) {
+            await Promise.all(this._disconnectCallbacks.map(cb => cb()));
         }
 
         if (this._goodbyeRequest && !this._goodbyeRequest.isCompleted) {

lib/types.ts

@@ -48,13 +48,18 @@ export abstract class IBaseSession {
     isConnected(): boolean {
         throw new Error("UnimplementedError");
     }
+
+    onDisconnect(callback: () => Promise<void>): void {
+        throw new Error("UnimplementedError");
+    }
 }
 
 export class BaseSession extends IBaseSession {
     private readonly _ws: WebSocket;
     private readonly _wsMessageHandler: any;
     private readonly sessionDetails: SessionDetails;
     private readonly _serializer: Serializer;
+    private _disconnectCallbacks: Array<() => Promise<void>> = [];
 
     constructor(
         ws: WebSocket,
@@ -70,6 +75,9 @@ export class BaseSession extends IBaseSession {
 
         // close cleanly on abrupt client disconnect
         this._ws.addEventListener("close", async () => {
+            if (this._disconnectCallbacks.length > 0) {
+                await Promise.all(this._disconnectCallbacks.map(cb => cb()));
+            }
             await this.close();
         });
     }
@@ -134,6 +142,10 @@ export class BaseSession extends IBaseSession {
     isConnected(): boolean {
         return this._ws.readyState === WebSocket.OPEN;
     }
+
+    onDisconnect(callback: () => Promise<void>): void {
+        this._disconnectCallbacks.push(callback);
+    }
 }
 
 export class Result {