Merge pull request #122 from drift-labs/publisher-health-and-cleanup

graceful dlob publisher exit + cleanup

src/core/metrics.ts

@@ -45,6 +45,7 @@ export enum HEALTH_STATUS {
 	StaleBulkAccountLoader,
 	UnhealthySlotSubscriber,
 	LivenessTesting,
+	Restart,
 }
 
 const metricsPort =
@@ -169,6 +170,10 @@ let lastHealthCheckState = true; // true = healthy, false = unhealthy
 let lastHealthCheckPerformed = Date.now() - healthCheckInterval;
 let lastTimeHealthy = Date.now() - healthCheckInterval;
 
+export const setHealthStatus = (status: HEALTH_STATUS): void => {
+	healthStatus = status;
+};
+
 /**
  * Middleware that checks if we are in general healthy by checking that the bulk account loader slot
  * has changed recently.
@@ -185,6 +190,12 @@ const handleHealthCheck = (
 	slotSource: SlotSource
 ) => {
 	return async (_req, res, _next) => {
+		if (healthStatus === HEALTH_STATUS.Restart) {
+			res.writeHead(500);
+			res.end(`NOK`);
+			return;
+		}
+
 		if (Date.now() < lastHealthCheckPerformed + healthCheckInterval) {
 			if (lastHealthCheckState) {
 				res.writeHead(200);

src/dlob-subscriber/DLOBSubscriberIO.ts

@@ -16,6 +16,7 @@ import {
 	addOracletoResponse,
 	l2WithBNToStrings,
 } from '../utils/utils';
+import { setHealthStatus, HEALTH_STATUS } from '../core/metrics';
 
 type wsMarketArgs = {
 	marketIndex: number;
@@ -62,7 +63,7 @@ export class DLOBSubscriberIO extends DLOBSubscriber {
 
 	constructor(
 		config: DLOBSubscriptionConfig & {
-			env: DriftEnv,
+			env: DriftEnv;
 			redisClient: RedisClient;
 			perpMarketInfos: wsMarketInfo[];
 			spotMarketInfos: wsMarketInfo[];
@@ -209,11 +210,11 @@ export class DLOBSubscriberIO extends DLOBSubscriber {
 				this.killSwitchSlotDiffThreshold &&
 			!skipSlotCheck
 		) {
-			console.log(`Killing process due to slot diffs for market ${marketName}: 
+			console.log(`Unhealthy process due to slot diffs for market ${marketName}: 
 				dlobProvider slot: ${slot}
 				oracle slot: ${l2Formatted['oracleData']['slot']}
 			`);
-			process.exit(1);
+			setHealthStatus(HEALTH_STATUS.Restart);
 		}
 
 		// Check if times and slots are too different for market
@@ -227,11 +228,11 @@ export class DLOBSubscriberIO extends DLOBSubscriber {
 			Date.now() - lastMarketSlotAndTime.ts > MAKRET_STALENESS_THRESHOLD &&
 			!skipSlotCheck
 		) {
-			console.log(`Killing process due to same slot for market ${marketName} after > ${MAKRET_STALENESS_THRESHOLD}ms: 
+			console.log(`Unhealthy process due to same slot for market ${marketName} after > ${MAKRET_STALENESS_THRESHOLD}ms: 
 				dlobProvider slot: ${slot}
 				market slot: ${l2Formatted['marketSlot']}
 			`);
-			process.exit(1);
+			setHealthStatus(HEALTH_STATUS.Restart);
 		} else if (
 			lastMarketSlotAndTime &&
 			l2Formatted['marketSlot'] !== lastMarketSlotAndTime.slot

src/index.ts

@@ -2,7 +2,6 @@ import { program } from 'commander';
 import compression from 'compression';
 import cors from 'cors';
 import express from 'express';
-import rateLimit from 'express-rate-limit';
 import morgan from 'morgan';
 
 import { Commitment, Connection, Keypair, PublicKey } from '@solana/web3.js';
@@ -99,10 +98,6 @@ const SLOT_STALENESS_TOLERANCE =
 	parseInt(process.env.SLOT_STALENESS_TOLERANCE) || 20;
 const ROTATION_COOLDOWN = parseInt(process.env.ROTATION_COOLDOWN) || 5000;
 const useWebsocket = process.env.USE_WEBSOCKET?.toLowerCase() === 'true';
-const rateLimitCallsPerSecond = process.env.RATE_LIMIT_CALLS_PER_SECOND
-	? parseInt(process.env.RATE_LIMIT_CALLS_PER_SECOND)
-	: 1;
-const loadTestAllowed = process.env.ALLOW_LOAD_TEST?.toLowerCase() === 'true';
 const useRedis = process.env.USE_REDIS?.toLowerCase() === 'true';
 
 const logFormat =
@@ -119,23 +114,6 @@ app.use(compression());
 app.set('trust proxy', 1);
 app.use(logHttp);
 app.use(handleResponseTime);
-if (!FEATURE_FLAGS.DISABLE_RATE_LIMIT) {
-	app.use(
-		rateLimit({
-			windowMs: 1000, // 1 second
-			max: rateLimitCallsPerSecond,
-			standardHeaders: false, // Return rate limit info in the `RateLimit-*` headers
-			legacyHeaders: false, // Disable the `X-RateLimit-*` headers
-			skip: (req, _res) => {
-				if (!loadTestAllowed) {
-					return false;
-				}
-
-				return req.headers['user-agent'].includes('k6');
-			},
-		})
-	);
-}
 
 // strip off /dlob, if the request comes from exchange history server LB
 app.use((req, _res, next) => {

src/publishers/dlobPublisher.ts

@@ -40,6 +40,8 @@ import {
 } from '../dlobProvider';
 import FEATURE_FLAGS from '../utils/featureFlags';
 import { GeyserOrderSubscriber } from '../grpc/OrderSubscriberGRPC';
+import express from 'express';
+import { handleHealthCheck } from '../core/metrics';
 
 require('dotenv').config();
 const stateCommitment: Commitment = 'confirmed';
@@ -49,6 +51,9 @@ const REDIS_HOST = process.env.REDIS_HOST || 'localhost';
 const REDIS_PORT = process.env.REDIS_PORT || '6379';
 const REDIS_PASSWORD = process.env.REDIS_PASSWORD;
 
+// Set up express for health checks
+const app = express();
+
 //@ts-ignore
 const sdkConfig = initialize({ env: process.env.ENV });
 let driftClient: DriftClient;
@@ -420,6 +425,31 @@ const main = async () => {
 		recursiveFetch();
 	}
 
+	const handleStartup = async (_req, res, _next) => {
+		if (driftClient.isSubscribed && dlobProvider.size() > 0) {
+			res.writeHead(200);
+			res.end('OK');
+		} else {
+			res.writeHead(500);
+			res.end('Not ready');
+		}
+	};
+
+	app.get(
+		'/health',
+		handleHealthCheck(2 * WS_FALLBACK_FETCH_INTERVAL, dlobProvider)
+	);
+	app.get('/startup', handleStartup);
+	app.get('/', handleHealthCheck(2 * WS_FALLBACK_FETCH_INTERVAL, dlobProvider));
+	const server = app.listen(8080);
+
+	// Default keepalive is 5s, since the AWS ALB timeout is 60 seconds, clients
+	// sometimes get 502s.
+	// https://shuheikagawa.com/blog/2019/04/25/keep-alive-timeout/
+	// https://stackoverflow.com/a/68922692
+	server.keepAliveTimeout = 61 * 1000;
+	server.headersTimeout = 65 * 1000;
+
 	console.log('DLOBSubscriber Publishing Messages');
 };