perf: gracefully shutdown the infinite loop

Author: wa0x6e

src/graphql/operations/options.ts

@@ -1,6 +1,5 @@
-import { capture } from '@snapshot-labs/snapshot-sentry';
-import snapshot from '@snapshot-labs/snapshot.js';
 import log from '../../helpers/log';
+import { createLoop } from '../../helpers/loop';
 import db from '../../helpers/mysql';
 
 const RUN_INTERVAL = 120e3;
@@ -11,23 +10,11 @@ export default async function () {
   return options;
 }
 
-async function loadOptions() {
-  options = await db.queryAsync('SELECT s.* FROM options s');
-}
-
-async function run() {
-  while (true) {
-    try {
-      log.info('[options] Start options refresh');
-      await loadOptions();
-      log.info(`[options] ${options.length} options reloaded`);
-      log.info('[options] End options refresh');
-    } catch (e: any) {
-      capture(e);
-      log.error(`[options] failed to refresh options, ${JSON.stringify(e)}`);
-    }
-    await snapshot.utils.sleep(RUN_INTERVAL);
+createLoop({
+  name: 'options',
+  interval: RUN_INTERVAL,
+  task: async () => {
+    options = await db.queryAsync('SELECT s.* FROM options s');
+    log.info(`[options] ${options.length} options reloaded`);
   }
-}
-
-run();
+});

src/helpers/loop.ts

@@ -0,0 +1,56 @@
+import { capture } from '@snapshot-labs/snapshot-sentry';
+import snapshot from '@snapshot-labs/snapshot.js';
+import log from './log';
+import {
+  getShutdownStatus,
+  registerStopFunction,
+  ShutdownFunction
+} from './shutdown';
+
+interface LoopOptions {
+  name: string;
+  interval: number;
+  task: () => Promise<void>;
+  maxConsecutiveFailsBeforeCapture?: number;
+}
+
+export const createLoop = (options: LoopOptions): ShutdownFunction => {
+  const {
+    name,
+    interval,
+    task,
+    maxConsecutiveFailsBeforeCapture = 1
+  } = options;
+  let isRunning = false;
+  let consecutiveFailsCount = 0;
+
+  const run = async () => {
+    isRunning = true;
+    while (!getShutdownStatus() && isRunning) {
+      try {
+        log.info(`[${name}] Start ${name} refresh`);
+        await task();
+        consecutiveFailsCount = 0; // Reset on success
+        log.info(`[${name}] End ${name} refresh`);
+      } catch (e: any) {
+        consecutiveFailsCount++;
+
+        if (consecutiveFailsCount >= maxConsecutiveFailsBeforeCapture) {
+          capture(e);
+        }
+        log.error(`[${name}] failed to refresh, ${JSON.stringify(e)}`);
+      }
+      await snapshot.utils.sleep(interval);
+    }
+  };
+
+  const stop = async () => {
+    isRunning = false;
+  };
+
+  // Register stop function and start the loop
+  registerStopFunction(stop);
+  run();
+
+  return stop;
+};

src/helpers/mysql.ts

@@ -4,6 +4,7 @@ import mysql from 'mysql';
 import Connection from 'mysql/lib/Connection';
 import Pool from 'mysql/lib/Pool';
 import log from './log';
+import { registerStopFunction } from './shutdown';
 
 const connectionLimit = parseInt(process.env.CONNECTION_LIMIT || '25');
 log.info(`[mysql] connection limit ${connectionLimit}`);
@@ -37,15 +38,24 @@ const sequencerDB = mysql.createPool(sequencerConfig);
 
 bluebird.promisifyAll([Pool, Connection]);
 
-export const closeDatabase = (): Promise<void> => {
+const createCloseFunction = (pool: Pool, name: string) => (): Promise<void> => {
   return new Promise(resolve => {
-    hubDB.end(() => {
-      sequencerDB.end(() => {
-        console.log('[mysql] Database connection pools closed');
-        resolve();
-      });
+    pool.end((err: Error | null) => {
+      if (err) {
+        log.error(`[mysql] Error closing ${name}:`, err);
+      } else {
+        log.info(`[mysql] ${name} connection pool closed`);
+      }
+      resolve();
     });
   });
 };
 
+const closeHubDB = createCloseFunction(hubDB, 'hubDB');
+const closeSequencerDB = createCloseFunction(sequencerDB, 'sequencerDB');
+
+// Register each database shutdown individually with the shutdown helper
+registerStopFunction(closeHubDB);
+registerStopFunction(closeSequencerDB);
+
 export { hubDB as default, sequencerDB };

src/helpers/shutdown.ts

@@ -0,0 +1,26 @@
+import log from './log';
+
+export type ShutdownFunction = () => Promise<void>;
+
+let isShuttingDown = false;
+const stopFunctions: ShutdownFunction[] = [];
+
+export const getShutdownStatus = () => isShuttingDown;
+
+export const registerStopFunction = (fn: ShutdownFunction) => {
+  stopFunctions.push(fn);
+};
+
+export const initiateShutdown = async () => {
+  isShuttingDown = true;
+
+  // Execute all stop functions with individual error handling
+  const results = await Promise.allSettled(stopFunctions.map(fn => fn()));
+
+  // Log any failures but don't throw
+  results.forEach((result, index) => {
+    if (result.status === 'rejected') {
+      log.error(`Stop function ${index} failed:`, result.reason);
+    }
+  });
+};

src/helpers/spaces.ts

@@ -1,7 +1,6 @@
-import { capture } from '@snapshot-labs/snapshot-sentry';
-import snapshot from '@snapshot-labs/snapshot.js';
 import networks from '@snapshot-labs/snapshot.js/src/networks.json';
 import { uniq } from 'lodash';
+import { createLoop } from './loop';
 import log from './log';
 import db from './mysql';
 
@@ -329,20 +328,12 @@ export async function getSpace(id: string) {
   };
 }
 
-export default async function run() {
-  while (true) {
-    try {
-      log.info('[spaces] Start spaces refresh');
-
-      await loadSpaces();
-      await loadSpacesMetrics();
-
-      sortSpaces();
-      log.info('[spaces] End spaces refresh');
-    } catch (e: any) {
-      capture(e);
-      log.error(`[spaces] failed to load spaces, ${JSON.stringify(e)}`);
-    }
-    await snapshot.utils.sleep(RUN_INTERVAL);
+export default createLoop({
+  name: 'spaces',
+  interval: RUN_INTERVAL,
+  task: async () => {
+    await loadSpaces();
+    await loadSpacesMetrics();
+    sortSpaces();
   }
-}
+});

src/helpers/strategies.ts

@@ -1,7 +1,7 @@
 import { URL } from 'url';
 import { capture } from '@snapshot-labs/snapshot-sentry';
 import snapshot from '@snapshot-labs/snapshot.js';
-import log from './log';
+import { createLoop } from './loop';
 import { spacesMetadata } from './spaces';
 
 export let strategies: any[] = [];
@@ -14,8 +14,6 @@ const scoreApiURL: URL = new URL(
 scoreApiURL.pathname = '/api/strategies';
 const uri = scoreApiURL.toString();
 
-let consecutiveFailsCount = 0;
-
 async function loadStrategies() {
   const res = await snapshot.utils.getJSON(uri);
 
@@ -57,23 +55,11 @@ async function loadStrategies() {
   );
 }
 
-async function run() {
-  while (true) {
-    try {
-      log.info('[strategies] Start strategies refresh');
-      await loadStrategies();
-      consecutiveFailsCount = 0;
-      log.info('[strategies] End strategies refresh');
-    } catch (e: any) {
-      consecutiveFailsCount++;
-
-      if (consecutiveFailsCount >= 3) {
-        capture(e);
-      }
-      log.error(`[strategies] failed to load ${JSON.stringify(e)}`);
-    }
-    await snapshot.utils.sleep(RUN_INTERVAL);
-  }
-}
-
-run();
+createLoop({
+  name: 'strategies',
+  interval: RUN_INTERVAL,
+  task: async () => {
+    await loadStrategies();
+  },
+  maxConsecutiveFailsBeforeCapture: 3
+});

src/index.ts

@@ -8,8 +8,8 @@ import graphql from './graphql';
 import { checkKeycard } from './helpers/keycard';
 import log from './helpers/log';
 import initMetrics from './helpers/metrics';
-import { closeDatabase } from './helpers/mysql';
 import rateLimit from './helpers/rateLimit';
+import { initiateShutdown } from './helpers/shutdown';
 import refreshSpacesCache from './helpers/spaces';
 import './helpers/strategies';
 
@@ -38,20 +38,24 @@ const server = app.listen(PORT, () =>
 );
 
 const gracefulShutdown = async (signal: string) => {
-  console.log(`Received ${signal}. Starting graceful shutdown...`);
-
-  server.close(async () => {
-    console.log('Express server closed.');
-
-    try {
-      await closeDatabase();
-      console.log('Graceful shutdown completed.');
-      process.exit(0);
-    } catch (error) {
-      console.error('Error during shutdown:', error);
-      process.exit(1);
-    }
+  log.info(`Received ${signal}. Starting graceful shutdown...`);
+
+  // Stop all background processes (loops + database)
+  await initiateShutdown();
+  log.info('Background processes stopped.');
+
+  // Close Express server
+  server.close(() => {
+    log.info('Express server closed.');
+    log.info('Graceful shutdown completed.');
+    process.exit(0);
   });
+
+  // Fallback timeout for the entire shutdown process
+  setTimeout(() => {
+    log.error('Graceful shutdown timeout exceeded, forcing exit');
+    process.exit(1);
+  }, 15000); // 15 seconds total shutdown timeout
 };
 
 process.on('SIGINT', () => gracefulShutdown('SIGINT'));