Merge pull request #93 from drift-labs/master

attempted memory leak fix

package.json

@@ -6,7 +6,7 @@
   "license": "Apache-2.0",
   "dependencies": {
     "@coral-xyz/anchor": "^0.29.0",
-    "@drift-labs/sdk": "2.65.0-beta.1",
+    "@drift-labs/sdk": "2.65.0-beta.6",
     "@opentelemetry/api": "^1.1.0",
     "@opentelemetry/auto-instrumentations-node": "^0.31.1",
     "@opentelemetry/exporter-prometheus": "^0.31.0",

src/wsConnectionManager.ts

@@ -36,6 +36,7 @@ const WS_PORT = process.env.WS_PORT || '3000';
 console.log(`WS LISTENER PORT : ${WS_PORT}`);
 
 const REDIS_PASSWORD = process.env.REDIS_PASSWORD;
+const MAX_BUFFERED_AMOUNT = 300000;
 
 const safeGetRawChannelFromMessage = (message: any): string => {
 	return message?.channel;
@@ -103,7 +104,10 @@ async function main() {
 		const subscribers = channelSubscribers.get(subscribedChannel);
 		if (subscribers) {
 			subscribers.forEach((ws) => {
-				if (ws.readyState === WebSocket.OPEN && ws.bufferedAmount < 300000)
+				if (
+					ws.readyState === WebSocket.OPEN &&
+					ws.bufferedAmount < MAX_BUFFERED_AMOUNT
+				)
 					ws.send(
 						JSON.stringify({ channel: subscribedChannel, data: message })
 					);
@@ -245,10 +249,31 @@ async function main() {
 			}
 		});
 
+		// Set interval to send heartbeat every 5 seconds
+		const heartbeatInterval = setInterval(() => {
+			if (ws.readyState === WebSocket.OPEN) {
+				ws.send(JSON.stringify({ channel: 'heartbeat' }));
+			} else {
+				clearInterval(heartbeatInterval);
+			}
+		}, 5000);
+
+		// Buffer overflow check interval
+		const bufferInterval = setInterval(() => {
+			if (ws.bufferedAmount > MAX_BUFFERED_AMOUNT) {
+				if (ws.readyState === WebSocket.OPEN) {
+					ws.close(1008, 'Buffer overflow');
+				}
+				clearInterval(bufferInterval);
+			}
+		}, 10000);
+
 		// Handle disconnection
 		ws.on('close', () => {
 			console.log('Client disconnected');
-			// Clear any existing intervals and timeouts
+			// Clear any existing intervals
+			clearInterval(heartbeatInterval);
+			clearInterval(bufferInterval);
 			channelSubscribers.forEach((subscribers, channel) => {
 				if (subscribers.delete(ws) && subscribers.size === 0) {
 					redisClient.client.unsubscribe(channel);
@@ -262,15 +287,6 @@ async function main() {
 		ws.on('error', (error) => {
 			console.error('Socket error:', error);
 		});
-
-		// Set interval to send heartbeat every 5 seconds
-		setInterval(() => {
-			ws.send(
-				JSON.stringify({
-					channel: 'heartbeat',
-				})
-			);
-		}, 5000);
 	});
 
 	server.listen(WS_PORT, () => {
@@ -285,19 +301,6 @@ async function main() {
 	server.on('error', (error) => {
 		console.error('Server error:', error);
 	});
-
-	// Periodic check for bad clients and disconnect them to alleviate backpressure
-	setInterval(() => {
-		let set = new Set<WebSocket>();
-		for (const [_channel, wsSet] of channelSubscribers) {
-			set = new Set([...set, ...wsSet]);
-		}
-		for (const ws of set) {
-			if (ws.bufferedAmount > 350000) {
-				ws.close();
-			}
-		}
-	}, 5000);
 }
 
 async function recursiveTryCatch(f: () => void) {

yarn.lock

@@ -115,10 +115,10 @@
     enabled "2.0.x"
     kuler "^2.0.0"
 
-"@drift-labs/sdk@2.65.0-beta.1":
-  version "2.65.0-beta.1"
-  resolved "https://registry.yarnpkg.com/@drift-labs/sdk/-/sdk-2.65.0-beta.1.tgz#559e831691007e600863fb0b0358e4157bec3e6b"
-  integrity sha512-yjQ8oxtgIMjY4BHoCebm0cz3/8FHWge0PwT3xgnUJ3d7mBQJ4/vBv4kDXMbraD/a7NO8LUOdnThLcmreUFnJNw==
+"@drift-labs/sdk@2.65.0-beta.6":
+  version "2.65.0-beta.6"
+  resolved "https://registry.yarnpkg.com/@drift-labs/sdk/-/sdk-2.65.0-beta.6.tgz#7da1434ce04142f4adca573baf44dc3539303a0f"
+  integrity sha512-ENnkk1o1wahXoT3F1svjiHzVdSrxlVZ7NfUXmrtyQEqbcZtbUyoS6CqL9f7aC273ydw7xC0a3jrxEH0yMOrcVQ==
   dependencies:
     "@coral-xyz/anchor" "0.28.1-beta.2"
     "@ellipsis-labs/phoenix-sdk" "^1.4.2"