From 12281198e539632de5e97cac80eb69546b347103 Mon Sep 17 00:00:00 2001
From: wphan <william@drift.trade>
Date: Sat, 2 Dec 2023 14:02:44 -0800
Subject: [PATCH] add grace period to health check

---
 src/core/metrics.ts | 18 ++++++++++++++++--
 src/index.ts        |  7 +++++--
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/src/core/metrics.ts b/src/core/metrics.ts
index f06b981..05b30f5 100644
--- a/src/core/metrics.ts
+++ b/src/core/metrics.ts
@@ -135,14 +135,23 @@ const healthCheckInterval = 2 * (ORDERBOOK_UPDATE_INTERVAL ?? 1000); // ORDERBOO
 let lastHealthCheckSlot = -1;
 let lastHealthCheckState = true; // true = healthy, false = unhealthy
 let lastHealthCheckPerformed = Date.now() - healthCheckInterval;
+let lastTimeHealthy = Date.now() - healthCheckInterval;
+
 /**
  * Middleware that checks if we are in general healthy by checking that the bulk account loader slot
  * has changed recently.
  *
  * 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`.
+ *
+ * A grace period is also used to only report unhealthy if we have been unhealthy for a certain
+ * amount of time. This prevents reporting unhealthy even if we are just in the middle of a
+ * bulk account load.
  */
-const handleHealthCheck = (slotSource: SlotSource) => {
+const handleHealthCheck = (
+	healthCheckGracePeriod: number,
+	slotSource: SlotSource
+) => {
 	return async (_req, res, _next) => {
 		if (Date.now() < lastHealthCheckPerformed + healthCheckInterval) {
 			if (lastHealthCheckState) {
@@ -163,12 +172,17 @@ const handleHealthCheck = (slotSource: SlotSource) => {
 					Date.now() - lastHealthCheckPerformed
 				} ms`
 			);
+		} else {
+			lastTimeHealthy = Date.now();
 		}
 
 		lastHealthCheckSlot = lastSlotReceived;
 		lastHealthCheckPerformed = Date.now();
 
-		if (!lastHealthCheckState) {
+		if (
+			!lastHealthCheckState &&
+			Date.now() - lastTimeHealthy > healthCheckGracePeriod
+		) {
 			healthStatus = HEALTH_STATUS.UnhealthySlotSubscriber;
 
 			res.writeHead(500);
diff --git a/src/index.ts b/src/index.ts
index b8f14a9..896462f 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -331,9 +331,12 @@ const main = async () => {
 		}
 	};
 
-	app.get('/health', handleHealthCheck(slotSource));
+	app.get(
+		'/health',
+		handleHealthCheck(2 * WS_FALLBACK_FETCH_INTERVAL, slotSource)
+	);
 	app.get('/startup', handleStartup);
-	app.get('/', handleHealthCheck(slotSource));
+	app.get('/', handleHealthCheck(2 * WS_FALLBACK_FETCH_INTERVAL, slotSource));
 
 	if (FEATURE_FLAGS.ENABLE_ORDERS_ENDPOINTS) {
 		app.get('/orders/json/raw', async (_req, res, next) => {