From 196e43d1b653bd317a4abb5a9f4f1c84db6f18da Mon Sep 17 00:00:00 2001
From: wphan <william@drift.trade>
Date: Wed, 29 Nov 2023 16:12:08 -0800
Subject: [PATCH] update health check faster

---
 src/core/metrics.ts | 71 ++++++++++++++++-----------------------------
 src/index.ts        | 35 +++++++++++-----------
 2 files changed, 43 insertions(+), 63 deletions(-)

diff --git a/src/core/metrics.ts b/src/core/metrics.ts
index 99e9de0..6462da4 100644
--- a/src/core/metrics.ts
+++ b/src/core/metrics.ts
@@ -12,10 +12,9 @@ import {
 	commitHash,
 	driftEnv,
 	endpoint,
-	getSlotHealthCheckInfo,
 	wsEndpoint,
 } from '..';
-import { E_ALREADY_LOCKED, tryAcquire } from 'async-mutex';
+import { SlotSource } from '@drift-labs/sdk';
 
 /**
  * Creates {count} buckets of size {increment} starting from {start}. Each bucket stores the count of values within its "size".
@@ -129,34 +128,31 @@ let lastHealthCheckPerformed = Date.now() - healthCheckInterval;
  * We may be hit by multiple sources performing health checks on us, so this middleware will latch
  * to its health state and only update every `healthCheckInterval`.
  */
-const handleHealthCheck = async (req, res, next) => {
-	if (Date.now() < lastHealthCheckPerformed + healthCheckInterval) {
-		if (lastHealthCheckState) {
-			res.writeHead(200);
-			res.end('OK');
-			lastHealthCheckPerformed = Date.now();
-			return;
-		}
-		// always check if last check was unhealthy (give it another chance to recover)
-	}
-
-	const { lastSlotReceived, lastSlotReceivedMutex } = getSlotHealthCheckInfo();
-
-	try {
-		await tryAcquire(lastSlotReceivedMutex).runExclusive(async () => {
-			// healthy if slot has advanced since the last check
-			lastHealthCheckState = lastSlotReceived > lastHealthCheckSlot;
-			if (!lastHealthCheckState) {
-				logger.error(
-					`Unhealthy: lastSlot: ${lastSlotReceived}, lastHealthCheckSlot: ${lastHealthCheckSlot}, timeSinceLastCheck: ${
-						Date.now() - lastHealthCheckPerformed
-					} ms`
-				);
+const handleHealthCheck = (slotSource: SlotSource) => {
+	return async (_req, res, _next) => {
+		if (Date.now() < lastHealthCheckPerformed + healthCheckInterval) {
+			if (lastHealthCheckState) {
+				res.writeHead(200);
+				res.end('OK');
+				lastHealthCheckPerformed = Date.now();
+				return;
 			}
+			// always check if last check was unhealthy (give it another chance to recover)
+		}
 
-			lastHealthCheckSlot = lastSlotReceived;
-			lastHealthCheckPerformed = Date.now();
-		});
+		// healthy if slot has advanced since the last check
+		const lastSlotReceived = slotSource.getSlot();
+		lastHealthCheckState = lastSlotReceived > lastHealthCheckSlot;
+		if (!lastHealthCheckState) {
+			logger.error(
+				`Unhealthy: lastSlot: ${lastSlotReceived}, lastHealthCheckSlot: ${lastHealthCheckSlot}, timeSinceLastCheck: ${
+					Date.now() - lastHealthCheckPerformed
+				} ms`
+			);
+		}
+
+		lastHealthCheckSlot = lastSlotReceived;
+		lastHealthCheckPerformed = Date.now();
 
 		if (!lastHealthCheckState) {
 			healthStatus = HEALTH_STATUS.UnhealthySlotSubscriber;
@@ -169,24 +165,7 @@ const handleHealthCheck = async (req, res, next) => {
 		healthStatus = HEALTH_STATUS.Ok;
 		res.writeHead(200);
 		res.end('OK');
-	} catch (e) {
-		if (e === E_ALREADY_LOCKED) {
-			// someone else is updating the slot, just return the latched state
-			if (lastHealthCheckState) {
-				logger.error(`DlobServer is healthy (busy mutex)`);
-				res.writeHead(200);
-				res.end('OK');
-			} else {
-				logger.error(`DlobServer is not healthy (busy mutex)`);
-				res.writeHead(500);
-				res.end(`NOK`);
-			}
-		} else {
-			logger.error(`Error in health check: ${e.message}`);
-			console.error(e);
-			next(e);
-		}
-	}
+	};
 };
 
 export {
diff --git a/src/index.ts b/src/index.ts
index 9f44f72..5c6a6ed 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -32,7 +32,6 @@ import {
 
 import { logger, setLogLevel } from './utils/logger';
 
-import { Mutex } from 'async-mutex';
 import * as http from 'http';
 import { handleHealthCheck } from './core/metrics';
 import { handleResponseTime } from './core/middleware';
@@ -136,13 +135,6 @@ logger.info(`Commit:       ${commitHash}`);
 
 let MARKET_SUBSCRIBERS: SubscriberLookup = {};
 
-const lastSlotReceivedMutex = new Mutex();
-let lastSlotReceived: number;
-
-export const getSlotHealthCheckInfo = () => {
-	return { lastSlotReceived, lastSlotReceivedMutex };
-};
-
 const initializeAllMarketSubscribers = async (driftClient: DriftClient) => {
 	const markets: SubscriberLookup = {};
 
@@ -272,10 +264,6 @@ const main = async () => {
 		logger.error(e);
 	});
 
-	setInterval(async () => {
-		lastSlotReceived = slotSource.getSlot();
-	}, ORDERBOOK_UPDATE_INTERVAL);
-
 	logger.info(`Initializing DLOB Provider...`);
 	const initDLOBProviderStart = Date.now();
 	await dlobProvider.subscribe();
@@ -284,9 +272,15 @@ const main = async () => {
 	);
 	logger.info(`dlob provider size ${dlobProvider.size()}`);
 	if (useWebsocket) {
-		setInterval(async () => {
-			await dlobProvider.fetch();
-		}, WS_FALLBACK_FETCH_INTERVAL);
+		const recursiveFetch = (delay = WS_FALLBACK_FETCH_INTERVAL) => {
+			setTimeout(() => {
+				dlobProvider.fetch().then(() => {
+					// eslint-disable-next-line @typescript-eslint/no-unused-vars
+					recursiveFetch();
+				});
+			}, delay);
+		};
+		recursiveFetch();
 	}
 
 	logger.info(`Initializing DLOBSubscriber...`);
@@ -302,7 +296,14 @@ const main = async () => {
 		`DLOBSubscriber initialized in ${Date.now() - initDlobSubscriberStart} ms`
 	);
 
+	logger.info(`Initializing all market subscribers...`);
+	const initAllMarketSubscribersStart = Date.now();
 	MARKET_SUBSCRIBERS = await initializeAllMarketSubscribers(driftClient);
+	logger.info(
+		`All market subscribers initialized in ${
+			Date.now() - initAllMarketSubscribersStart
+		} ms`
+	);
 
 	const handleStartup = async (_req, res, _next) => {
 		if (driftClient.isSubscribed && dlobProvider.size() > 0) {
@@ -314,9 +315,9 @@ const main = async () => {
 		}
 	};
 
-	app.get('/health', handleHealthCheck);
+	app.get('/health', handleHealthCheck(slotSource));
 	app.get('/startup', handleStartup);
-	app.get('/', handleHealthCheck);
+	app.get('/', handleHealthCheck(slotSource));
 
 	app.get('/orders/json/raw', async (_req, res, next) => {
 		try {