Merge pull request #317 from PretendoNetwork/feat/basic-health-check

Add basic health check before returning server address

src/models/server.ts

@@ -1,7 +1,51 @@
+import dgram from 'node:dgram';
+import crypto from 'node:crypto';
 import { Schema, model } from 'mongoose';
 import uniqueValidator from 'mongoose-unique-validator';
+import { LOG_WARN } from '@/logger';
 import type { IServer, IServerConnectInfo, IServerMethods, ServerModel } from '@/types/mongoose/server';
 
+// * Kinda ugly to slap this in with the Mongoose stuff but it's fine for now
+// TODO - Maybe move this one day?
+const socket = dgram.createSocket('udp4');
+const pendingHealthCheckRequests = new Map<string, () => void>();
+
+socket.on('message', (msg: Buffer, _rinfo: dgram.RemoteInfo) => {
+	const uuid = msg.toString();
+	const resolve = pendingHealthCheckRequests.get(uuid);
+
+	if (resolve) {
+		resolve();
+	}
+});
+
+socket.bind();
+
+function healthCheck(target: { host: string; port: number }): Promise<string> {
+	return new Promise((resolve, reject) => {
+		const uuid = crypto.randomUUID();
+
+		const timeout = setTimeout(() => {
+			pendingHealthCheckRequests.delete(uuid);
+			reject(new Error('No valid response received'));
+		}, 2 * 1000); // TODO - Make this configurable? 2 seconds seems fine for now
+
+		pendingHealthCheckRequests.set(uuid, () => {
+			clearTimeout(timeout);
+			pendingHealthCheckRequests.delete(uuid);
+			resolve(target.host);
+		});
+
+		socket.send(Buffer.from(uuid), target.port, target.host, (error) => {
+			if (error) {
+				clearTimeout(timeout);
+				pendingHealthCheckRequests.delete(uuid);
+				reject(error);
+			}
+		});
+	});
+}
+
 const ServerSchema = new Schema<IServer, ServerModel, IServerMethods>({
 	client_id: String,
 	ip: {
@@ -20,7 +64,11 @@ const ServerSchema = new Schema<IServer, ServerModel, IServerMethods>({
 	access_mode: String,
 	maintenance_mode: Boolean,
 	device: Number,
-	aes_key: String
+	aes_key: String,
+	health_check_port: {
+		type: Number,
+		required: false
+	}
 });
 
 ServerSchema.plugin(uniqueValidator, { message: '{PATH} already in use.' });
@@ -31,9 +79,43 @@ ServerSchema.method('getServerConnectInfo', async function (): Promise<IServerCo
 		throw new Error(`No IP configured for server ${this._id}`);
 	}
 
-	const randomIp = ipList[Math.floor(Math.random() * ipList.length)];
+	const randomIP = ipList[Math.floor(Math.random() * ipList.length)];
+
+	if (!this.health_check_port) {
+		return {
+			ip: randomIP,
+			port: this.port
+		};
+	}
+
+	// * Remove the random IP from the race pool to remove the duplicate health check
+	const healthCheckTargets = ipList.filter(ip => ip !== randomIP).map(ip => ({
+		host: ip,
+		port: this.health_check_port!
+	}));
+
+	// * Default to the random IP in case nothing responded in time
+	// * and just Hope For The Best:tm:
+	let target = randomIP;
+
+	// * Check the random IP and start the race at the same time, preferring
+	// * the result of the random IP should it succeed. Worst case scenario
+	// * this takes 2 seconds to complete
+	const [randomResult, raceResult] = await Promise.allSettled([
+		healthCheck({ host: randomIP, port: this.health_check_port! }),
+		Promise.race(healthCheckTargets.map(target => healthCheck(target)))
+	]);
+
+	if (randomResult.status === 'rejected') {
+		if (raceResult.status === 'fulfilled') {
+			target = raceResult.value;
+		} else {
+			LOG_WARN(`Server ${this.service_name} failed to find healthy NEX server. Using the randomly selected IP ${target}`);
+		}
+	}
+
 	return {
-		ip: randomIp,
+		ip: target,
 		port: this.port
 	};
 });

src/provisioning.ts

@@ -15,7 +15,8 @@ const serverProvisioningSchema = z.object({
 		name: z.string(),
 		ip: z.string().optional(),
 		ipList: z.array(z.string()).optional(),
-		port: z.coerce.number()
+		port: z.coerce.number(),
+		health_check_port: z.coerce.number().optional()
 	}))
 });
 
@@ -43,7 +44,8 @@ export async function handleServerProvisioning(): Promise<void> {
 				service_name: server.name,
 				ipList: server.ipList,
 				ip: server.ip,
-				port: server.port
+				port: server.port,
+				health_check_port: server.health_check_port
 			}
 		});
 		if (!result) {

src/types/mongoose/server.ts

@@ -13,6 +13,7 @@ export interface IServer {
 	maintenance_mode: boolean;
 	device: number;
 	aes_key: string;
+	health_check_port?: number;
 }
 
 export interface IServerConnectInfo {