From 49a95abe5ce492049087d3a365b6ff4b559366ea Mon Sep 17 00:00:00 2001
From: Dima Grossman <dima@grossman.io>
Date: Tue, 19 Dec 2023 14:17:00 +0200
Subject: [PATCH] feat(ws): add logging for socket handling

---
 .../src/socket/services/web-socket.worker.ts  |  5 +---
 .../external-services-route.usecase.ts        |  6 +++--
 apps/ws/src/socket/ws.gateway.ts              | 23 +++++++++++++++++++
 3 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/apps/ws/src/socket/services/web-socket.worker.ts b/apps/ws/src/socket/services/web-socket.worker.ts
index ff3f532fa26..bea5f86a34a 100644
--- a/apps/ws/src/socket/services/web-socket.worker.ts
+++ b/apps/ws/src/socket/services/web-socket.worker.ts
@@ -31,10 +31,7 @@ export class WebSocketWorker extends WebSocketsWorkerService {
         // eslint-disable-next-line @typescript-eslint/no-this-alias
         const _this = this;
 
-        Logger.verbose(
-          `Job ${job.id} / ${job.data.event} is being processed in the MemoryDB instance WebSocketWorker`,
-          LOG_CONTEXT
-        );
+        Logger.log(`Job ${job.id} / ${job.data.event} is being processed WebSocketWorker`, LOG_CONTEXT);
 
         nr.startBackgroundTransaction(
           ObservabilityBackgroundTransactionEnum.WS_SOCKET_QUEUE,
diff --git a/apps/ws/src/socket/usecases/external-services-route/external-services-route.usecase.ts b/apps/ws/src/socket/usecases/external-services-route/external-services-route.usecase.ts
index 901c42b0080..71f386f51f1 100644
--- a/apps/ws/src/socket/usecases/external-services-route/external-services-route.usecase.ts
+++ b/apps/ws/src/socket/usecases/external-services-route/external-services-route.usecase.ts
@@ -17,6 +17,8 @@ export class ExternalServicesRoute {
     const isOnline = await this.connectionExist(command);
 
     if (!isOnline) {
+      Logger.log(`Connection does not exist, ignoring command for ${command.userId}`, LOG_CONTEXT);
+
       return;
     }
 
@@ -37,10 +39,10 @@ export class ExternalServicesRoute {
     const { message, messageId } = command.payload || {};
     // TODO: Retro-compatibility for a bit just in case stalled messages
     if (message) {
-      Logger.verbose('Sending full message in the payload', LOG_CONTEXT);
+      Logger.log('Sending full message in the payload', LOG_CONTEXT);
       await this.wsGateway.sendMessage(command.userId, command.event, command.payload);
     } else if (messageId) {
-      Logger.verbose('Sending messageId in the payload, we need to retrieve the full message', LOG_CONTEXT);
+      Logger.log(`Sending messageId: ${messageId} in the payload, we need to retrieve the full message`, LOG_CONTEXT);
       const storedMessage = await this.messageRepository.findOne({
         _id: messageId,
         _environmentId: command._environmentId,
diff --git a/apps/ws/src/socket/ws.gateway.ts b/apps/ws/src/socket/ws.gateway.ts
index 7307096529b..8d93b2b9e68 100644
--- a/apps/ws/src/socket/ws.gateway.ts
+++ b/apps/ws/src/socket/ws.gateway.ts
@@ -21,6 +21,8 @@ export class WSGateway implements OnGatewayConnection, OnGatewayDisconnect, IDes
   server: Server | null;
 
   async handleDisconnect(connection: Socket) {
+    Logger.log(`New disconnect received from ${connection.id}`, LOG_CONTEXT);
+
     // eslint-disable-next-line @typescript-eslint/no-this-alias
     const _this = this;
 
@@ -44,6 +46,8 @@ export class WSGateway implements OnGatewayConnection, OnGatewayDisconnect, IDes
   }
 
   async handleConnection(connection: Socket) {
+    Logger.log(`New connection received from ${connection.id}`, LOG_CONTEXT);
+
     // eslint-disable-next-line @typescript-eslint/no-this-alias
     const _this = this;
 
@@ -93,6 +97,8 @@ export class WSGateway implements OnGatewayConnection, OnGatewayDisconnect, IDes
   private async processDisconnectionRequest(connection: Socket) {
     if (!this.isShutdown) {
       await this.handlerSubscriberDisconnection(connection);
+    } else {
+      Logger.log(`Skipped disconnect due to shutdown flag for connection ${connection.id}`, LOG_CONTEXT);
     }
   }
 
@@ -109,6 +115,11 @@ export class WSGateway implements OnGatewayConnection, OnGatewayDisconnect, IDes
     }
 
     const activeConnections = await this.getActiveConnections(connection, subscriber._id);
+
+    Logger.log(
+      `Disconnect request received from ${subscriber._id}. Active connections: ${activeConnections}`,
+      LOG_CONTEXT
+    );
     await this.subscriberOnlineService.handleDisconnection(subscriber, activeConnections);
   }
 
@@ -122,15 +133,27 @@ export class WSGateway implements OnGatewayConnection, OnGatewayDisconnect, IDes
     const token = this.extractToken(connection);
 
     if (!token || token === 'null') {
+      Logger.warn(`No token was found during counnection process for ${connection.id}`, LOG_CONTEXT);
+
       return this.disconnect(connection);
     }
 
     const subscriber = await this.getSubscriber(token);
     if (!subscriber) {
+      Logger.warn(`No subscriber was found for specified token ${connection.id}`, LOG_CONTEXT);
+
       return this.disconnect(connection);
     }
 
+    Logger.log(
+      `Connection request received from ${subscriber._id} external id: ${subscriber.subscriberId}`,
+      LOG_CONTEXT
+    );
+
     await connection.join(subscriber._id);
+
+    Logger.log(`Connection request accepted for ${subscriber._id}`, LOG_CONTEXT);
+
     await this.subscriberOnlineService.handleConnection(subscriber);
   }