From 313d8958222725ca8dc810e6e1d48459c5d54c3c Mon Sep 17 00:00:00 2001
From: Artem Stoianov <zarianec91@gmail.com>
Date: Tue, 31 May 2022 23:52:42 +0200
Subject: [PATCH] fix: fix infinite loop in graceful shutdown handler

---
 src/index.ts | 34 ++++++++++++++++++----------------
 1 file changed, 18 insertions(+), 16 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 4b5b9d5b..27fc63ec 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -37,8 +37,8 @@ logger.info(`Start probe version ${VERSION} in a ${process.env['NODE_ENV'] ?? 'p
 
 function connect() {
 	const worker = {
-		jobs: new Map<string, number>(),
 		active: false,
+		jobs: 0,
 	};
 
 	const socket = io(`${getConfValue<string>('api.host')}/probes`, {
@@ -57,18 +57,19 @@ function connect() {
 			logger.debug('connection to API established');
 		})
 		.on('disconnect', (reason: string): void => {
-			logger.debug(`disconnected from API. (${reason})`);
+			logger.debug(`disconnected from API: ${reason}`);
 			if (reason === 'io server disconnect') {
 				socket.connect();
 			}
 		})
 		.on('connect_error', error => {
-			logger.error('connection to API failed', error);
-
+			logger.error('connection to API failed:', error);
 			const isFatalError = fatalConnectErrors.some(fatalError => error.message.startsWith(fatalError));
 
-			if (!isFatalError) {
-				socket.connect();
+			if (isFatalError) {
+				// At that stage socket.connected=false already,
+				// but we want to stop reconnections for fatal errors
+				socket.disconnect();
 			}
 		})
 		.on('api:error', apiErrorHandler)
@@ -81,8 +82,6 @@ function connect() {
 			const {id: measurementId, measurement} = data;
 			const testId = cryptoRandomString({length: 16, type: 'alphanumeric'});
 
-			worker.jobs.set(measurementId, Date.now());
-
 			logger.debug(`${measurement.type} request ${data.id} received`, data);
 			socket.emit('probe:measurement:ack', {id: testId, measurementId}, async () => {
 				const handler = handlersMap.get(measurement.type);
@@ -91,11 +90,13 @@ function connect() {
 				}
 
 				try {
+					worker.jobs++;
 					await handler.run(socket, measurementId, testId, measurement);
-					worker.jobs.delete(measurementId);
 				} catch (error: unknown) {
 					// Todo: maybe we should notify api as well
 					logger.error('failed to run the measurement.', error);
+				} finally {
+					worker.jobs--;
 				}
 			});
 		});
@@ -106,14 +107,15 @@ function connect() {
 
 		logger.debug('SIGTERM received');
 
-		const closeInterval = setInterval(() => {
-			if (worker.jobs.size === 0) {
-				clearInterval(closeInterval);
+		if (worker.jobs <= 0) {
+			logger.debug('exit immediately');
+			process.exit();
+		}
 
-				logger.debug('closing process');
-				process.exit(0);
-			}
-		}, 100);
+		setTimeout(() => {
+			logger.debug('exit after graceful shutdown period');
+			process.exit();
+		}, 30_000);
 	});
 }