diff --git a/src/models/server.ts b/src/models/server.ts index ef10163..65f14ea 100644 --- a/src/models/server.ts +++ b/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 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 { + 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({ client_id: String, ip: { @@ -20,7 +64,11 @@ const ServerSchema = new Schema({ 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 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 }; }); diff --git a/src/provisioning.ts b/src/provisioning.ts index c4b37c6..1bfe597 100644 --- a/src/provisioning.ts +++ b/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 { 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) { diff --git a/src/types/mongoose/server.ts b/src/types/mongoose/server.ts index 2b19850..9f21429 100644 --- a/src/types/mongoose/server.ts +++ b/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 {