From 12b1dbb3651346eee6b7eb892d839b8f505fc5a3 Mon Sep 17 00:00:00 2001
From: Nour Alharithi <nour.alharithi@gmail.com>
Date: Wed, 29 Nov 2023 11:42:10 -0800
Subject: [PATCH] fetch with ws

---
 src/dlobProvider.ts             |  7 +++
 src/index.ts                    |  6 +++
 src/publishers/dlobPublisher.ts | 83 ++++++++++++++++++++++++---------
 3 files changed, 74 insertions(+), 22 deletions(-)

diff --git a/src/dlobProvider.ts b/src/dlobProvider.ts
index 1470543..bf34a21 100644
--- a/src/dlobProvider.ts
+++ b/src/dlobProvider.ts
@@ -11,6 +11,7 @@ export type DLOBProvider = {
 	}>;
 	getUserAccount(publicKey: PublicKey): UserAccount | undefined;
 	size(): number;
+	fetch(): Promise<void>;
 };
 
 export function getDLOBProviderFromUserMap(userMap: UserMap): DLOBProvider {
@@ -38,6 +39,9 @@ export function getDLOBProviderFromUserMap(userMap: UserMap): DLOBProvider {
 		size: () => {
 			return userMap.size();
 		},
+		fetch: () => {
+			return userMap.sync();
+		},
 	};
 }
 
@@ -73,5 +77,8 @@ export function getDLOBProviderFromOrderSubscriber(
 		size(): number {
 			return orderSubscriber.usersAccounts.size;
 		},
+		fetch() {
+			return orderSubscriber.fetch();
+		},
 	};
 }
diff --git a/src/index.ts b/src/index.ts
index 888b6a1..9f44f72 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -62,6 +62,7 @@ const sdkConfig = initialize({ env: process.env.ENV });
 const stateCommitment: Commitment = 'processed';
 const serverPort = process.env.PORT || 6969;
 export const ORDERBOOK_UPDATE_INTERVAL = 1000;
+const WS_FALLBACK_FETCH_INTERVAL = ORDERBOOK_UPDATE_INTERVAL * 10;
 const useWebsocket = process.env.USE_WEBSOCKET?.toLowerCase() === 'true';
 const useOrderSubscriber =
 	process.env.USE_ORDER_SUBSCRIBER?.toLowerCase() === 'true';
@@ -282,6 +283,11 @@ const main = async () => {
 		`dlob provider initialized in ${Date.now() - initDLOBProviderStart} ms`
 	);
 	logger.info(`dlob provider size ${dlobProvider.size()}`);
+	if (useWebsocket) {
+		setInterval(async () => {
+			await dlobProvider.fetch();
+		}, WS_FALLBACK_FETCH_INTERVAL);
+	}
 
 	logger.info(`Initializing DLOBSubscriber...`);
 	const initDlobSubscriberStart = Date.now();
diff --git a/src/publishers/dlobPublisher.ts b/src/publishers/dlobPublisher.ts
index 9ad03d7..1863256 100644
--- a/src/publishers/dlobPublisher.ts
+++ b/src/publishers/dlobPublisher.ts
@@ -11,6 +11,8 @@ import {
 	BulkAccountLoader,
 	OrderSubscriber,
 	SlotSource,
+	DriftClientSubscriptionConfig,
+	SlotSubscriber,
 } from '@drift-labs/sdk';
 
 import { logger, setLogLevel } from '../utils/logger';
@@ -24,6 +26,9 @@ import {
 } from '../dlobProvider';
 
 require('dotenv').config();
+const stateCommitment: Commitment = 'processed';
+const ORDERBOOK_UPDATE_INTERVAL = 1000;
+const WS_FALLBACK_FETCH_INTERVAL = ORDERBOOK_UPDATE_INTERVAL * 10;
 const driftEnv = (process.env.ENV || 'devnet') as DriftEnv;
 const commitHash = process.env.COMMIT;
 const REDIS_HOST = process.env.REDIS_HOST || 'localhost';
@@ -32,10 +37,6 @@ const REDIS_PASSWORD = process.env.REDIS_PASSWORD;
 
 //@ts-ignore
 const sdkConfig = initialize({ env: process.env.ENV });
-
-const stateCommitment: Commitment = 'processed';
-const ORDERBOOK_UPDATE_INTERVAL = 1000;
-
 let driftClient: DriftClient;
 
 const opts = program.opts();
@@ -45,6 +46,7 @@ const endpoint = process.env.ENDPOINT;
 const wsEndpoint = process.env.WS_ENDPOINT;
 const useOrderSubscriber =
 	process.env.USE_ORDER_SUBSCRIBER?.toLowerCase() === 'true';
+const useWebsocket = process.env.USE_WEBSOCKET?.toLowerCase() === 'true';
 
 logger.info(`RPC endpoint: ${endpoint}`);
 logger.info(`WS endpoint:  ${wsEndpoint}`);
@@ -60,20 +62,48 @@ const main = async () => {
 		commitment: stateCommitment,
 	});
 
-	const bulkAccountLoader = new BulkAccountLoader(
-		connection,
-		stateCommitment,
-		ORDERBOOK_UPDATE_INTERVAL
-	);
+	// only set when polling
+	let bulkAccountLoader: BulkAccountLoader | undefined;
+
+	// only set when using websockets
+	let slotSubscriber: SlotSubscriber | undefined;
+
+	let accountSubscription: DriftClientSubscriptionConfig;
+	let slotSource: SlotSource;
+
+	if (!useWebsocket) {
+		bulkAccountLoader = new BulkAccountLoader(
+			connection,
+			stateCommitment,
+			ORDERBOOK_UPDATE_INTERVAL
+		);
+
+		accountSubscription = {
+			type: 'polling',
+			accountLoader: bulkAccountLoader,
+		};
+
+		slotSource = {
+			getSlot: () => bulkAccountLoader!.getSlot(),
+		};
+	} else {
+		accountSubscription = {
+			type: 'websocket',
+			commitment: stateCommitment,
+		};
+		slotSubscriber = new SlotSubscriber(connection);
+		await slotSubscriber.subscribe();
+
+		slotSource = {
+			getSlot: () => slotSubscriber!.getSlot(),
+		};
+	}
 
 	driftClient = new DriftClient({
 		connection,
 		wallet,
 		programID: clearingHousePublicKey,
-		accountSubscription: {
-			type: 'polling',
-			accountLoader: bulkAccountLoader,
-		},
+		accountSubscription,
 		env: driftEnv,
 	});
 
@@ -90,15 +120,23 @@ const main = async () => {
 		logger.error(e);
 	});
 
-	let slotSource: SlotSource;
 	let dlobProvider: DLOBProvider;
 	if (useOrderSubscriber) {
-		const orderSubscriber = new OrderSubscriber({
-			driftClient,
-			subscriptionConfig: {
+		let subscriptionConfig;
+		if (useWebsocket) {
+			subscriptionConfig = {
+				type: 'websocket',
+			};
+		} else {
+			subscriptionConfig = {
 				type: 'polling',
 				frequency: ORDERBOOK_UPDATE_INTERVAL,
-			},
+			};
+		}
+
+		const orderSubscriber = new OrderSubscriber({
+			driftClient,
+			subscriptionConfig,
 		});
 
 		dlobProvider = getDLOBProviderFromOrderSubscriber(orderSubscriber);
@@ -114,10 +152,6 @@ const main = async () => {
 		);
 
 		dlobProvider = getDLOBProviderFromUserMap(userMap);
-
-		slotSource = {
-			getSlot: () => bulkAccountLoader.getSlot(),
-		};
 	}
 
 	await dlobProvider.subscribe();
@@ -133,6 +167,11 @@ const main = async () => {
 		redisClient,
 	});
 	await dlobSubscriber.subscribe();
+	if (useWebsocket) {
+		setInterval(async () => {
+			await dlobProvider.fetch();
+		}, WS_FALLBACK_FETCH_INTERVAL);
+	}
 
 	console.log('DLOBSubscriber Publishing Messages');
 };