LupiNex-Connect/servers/websocket/connections.js

/*
    ! /servers/websocket/connections.js – Central WebSocket Connection Manager
    ? This module handles all websocket connections as a pool and automatically sends heartbeats
*/

// % import nodejs dependencies
const path = require('path');

// % import external dependencies
const rootDir = path.resolve(__dirname, '..', '..');
const { generatePrefixed } = require(path.resolve(rootDir, 'services', 'idgen.js'));

// % create a new connection mapping
let connections = new Map();
let registry = new Map();

// % set heartbeat settings
let heartbeatInterval = null;
const heartbeatTime = 5000; // 5 seconds

// $ add a connection
function add(socket, codename = null) {
    // % generate a unique id for websocket connections
    const id = generatePrefixed('ws');

    // ~ add connection to connection pool
    const conn = { socket, codename, lastHeartbeat: Date.now() };
    connections.set(id, conn);

    // ~ add registry keys if codename is provided
    if (codename) registry.set(codename, id);

    // ~ listen for 'pong' events from client to update heartbeat timestamp
    socket.on('pong', () => conn.lastHeartbeat = Date.now());

    // ~ error handling
    socket.on('error', (err) => remove(id));

    // ~ close connection
    socket.on('close', () => remove(id));

    // ~ return client id
    return id;
}

// $ remove a connection by id
function remove(id) {
    // % find id in connection map
    const conn = connections.get(id);

    // ~ close and remove connection
    if (conn) {
        // ~ cleanup registry if codename exists
        if (conn.codename) registry.delete(conn.codename);

        // ~ close connection
        if (conn.socket.readyState === conn.socket.OPEN) conn.socket.close();

        // ~ remove connection from pool
        connections.delete(id);
    }
}

// $ remove all connections
function removeAll() {
    // ~ loop through connection mapping
    for (const id of connections.keys()) {
        remove(id);
    }
}

// $ get a connection by id
function get(id) {
    return connections.get(id);
}

// $ get a connection by codename
function getByCodename(codename) {
    // % get id from codename
    const id = registry.get(codename);

    // ~ return connection
    return id ? connections.get(id) : null;
}

// $ get all connections
function getAll() {
    // % create a reference for all connections
    const all = {};

    // ~ loop through connections and fill "all" reference
    for (const [id, data] of connections.entries()) {
        all[id] = {
            codename: data.codename,
            lastHeartbeat: data.lastHeartbeat
        };
    }

    // ~ return the connection reference
    return all;
}

// $ send a message to a specific connection by id
function send(id, message) {
    // % get id
    const conn = connections.get(id);

    // ~ check if connection exists and is open
    if (conn && conn.socket.readyState === conn.socket.OPEN) {
        conn.socket.send(JSON.stringify(message));
        return true;
    }

    // ~ return immediately after a failure
    return false;
}

// $ send a message to a specific connection by codename
function sendTo(codename, message) {
    // % get id
    const id = registry.get(codename);

    // ~ send a message if id exists
    if (id) return send(id, message);

    // ~ immediately return if id not exists
    return false;
}

// $ broadcast to all connections
function broadcast(message) {
    // % serialize payload
    const payload = JSON.stringify(message);

    // ~ loop through all established connections
    for (const conn of connections.values()) {
        if (conn.socket.readyState === conn.socket.OPEN) {
            conn.socket.send(payload);
        }
    }
}

// $ start the heartbeat system
function startBeat() {
    if (heartbeatInterval) return;

    heartbeatInterval = setInterval(() => {
        const now = Date.now();

        for (const [id, conn] of connections.entries()) {
            // ~ send ping to active connections
            if (conn.socket.readyState === conn.socket.OPEN) {
                conn.socket.ping();
            }

            // ~ remove stale connections (no pong in 2 heartbeat intervals)
            if (now - conn.lastHeartbeat > heartbeatTime * 2) {
                remove(id);
            }
        }
    }, heartbeatTime);
}

// $ stop the heartbeat system
function stopBeat() {
    if (heartbeatInterval) {
        clearInterval(heartbeatInterval);
        heartbeatInterval = null;
    }
}

// § export the websocket connection module
module.exports = { 
    add,
    remove,
    removeAll,
    get,
    getByCodename,
    getAll,
    send,
    sendTo,
    broadcast,
    startBeat,
    stopBeat
};