From f3d64ff988f4ea2bf56b480f048203f307e06be9 Mon Sep 17 00:00:00 2001
From: Nour Alharithi <nour.alharithi@gmail.com>
Date: Sun, 29 Oct 2023 15:55:44 -0700
Subject: [PATCH] *prettify + new example

---
 example/{newClient.ts => wsClient.ts} | 15 +++++++++----
 src/wsConnectionManager.ts            | 32 +++++++++++++++++----------
 2 files changed, 31 insertions(+), 16 deletions(-)
 rename example/{newClient.ts => wsClient.ts} (50%)

diff --git a/example/newClient.ts b/example/wsClient.ts
similarity index 50%
rename from example/newClient.ts
rename to example/wsClient.ts
index 5dba733..efb101b 100644
--- a/example/newClient.ts
+++ b/example/wsClient.ts
@@ -1,17 +1,24 @@
 import WebSocket from 'ws';
+// const ws = new WebSocket('wss://master.dlob.drift.trade/ws');
 const ws = new WebSocket('ws://localhost:3000/ws');
+import { sleep } from '../src/utils/utils';
 
-ws.on('open', () => {
+ws.on('open', async () => {
   console.log('Connected to the server');
   ws.send(JSON.stringify({ type: 'subscribe', channel: 'SOL-PERP' }));
+  ws.send(JSON.stringify({ type: 'subscribe', channel: 'LINK-PERP' }));
+  ws.send(JSON.stringify({ type: 'subscribe', channel: 'INJ-PERP' }));
+  await sleep(5000);
+  
+  ws.send(JSON.stringify({ type: 'unsubscribe', channel: 'SOL-PERP' }));
+  console.log("####################");
 });
 
 ws.on('message', (data: WebSocket.Data) => {
   try {
     const message = JSON.parse(data.toString());
-    if (message.channel === 'SOL-PERP') {
-      console.log('Received data:', message.data);
-    }
+    console.log(`Received data from market ${message.channel}`);
+    // book data is in message.data
   } catch (e) {
     console.error('Invalid message:', data);
   }
diff --git a/src/wsConnectionManager.ts b/src/wsConnectionManager.ts
index 41cdc1a..867a5fe 100644
--- a/src/wsConnectionManager.ts
+++ b/src/wsConnectionManager.ts
@@ -12,7 +12,7 @@ app.use(compression());
 app.set('trust proxy', 1);
 
 const server = http.createServer(app);
-const wss = new WebSocketServer({server, path: '/ws'});
+const wss = new WebSocketServer({ server, path: '/ws' });
 
 const REDIS_HOST = process.env.REDIS_HOST || 'localhost';
 const REDIS_PORT = process.env.REDIS_PORT || '6379';
@@ -28,7 +28,7 @@ async function main() {
 	redisClient.client.on('message', (subscribedChannel, message) => {
 		const subscribers = channelSubscribers.get(subscribedChannel);
 		subscribers.forEach((ws) => {
-			ws.send(JSON.stringify({channel: subscribedChannel, data: message }));
+			ws.send(JSON.stringify({ channel: subscribedChannel, data: message }));
 		});
 	});
 
@@ -43,14 +43,22 @@ async function main() {
 					const channel = parsedMessage.channel;
 					if (!subscribedChannels.has(channel)) {
 						console.log('Subscribing to channel', channel);
-						redisClient.client.subscribe(channel).then(() => {
-							subscribedChannels.add(channel);
-						}).catch(() => {
-							ws.send(JSON.stringify({channel, error: `Invalid channel: ${channel}`}));
-							return;
-						});
-					}	
-				
+						redisClient.client
+							.subscribe(channel)
+							.then(() => {
+								subscribedChannels.add(channel);
+							})
+							.catch(() => {
+								ws.send(
+									JSON.stringify({
+										channel,
+										error: `Invalid channel: ${channel}`,
+									})
+								);
+								return;
+							});
+					}
+
 					if (!channelSubscribers.get(channel)) {
 						const subscribers = new Set<WebSocket>();
 						channelSubscribers.set(channel, subscribers);
@@ -65,7 +73,7 @@ async function main() {
 						channelSubscribers.get(channel).delete(ws);
 					}
 					break;
-				}	
+				}
 				default:
 					break;
 			}
@@ -81,7 +89,7 @@ async function main() {
 					console.log('Disconnecting because of ping/pong timeout');
 					ws.terminate();
 				}
-			}, 5000);  // 5 seconds to wait for a pong
+			}, 5000); // 5 seconds to wait for a pong
 			ws.ping();
 		}, 30000);