58 lines
1.5 KiB
JavaScript
58 lines
1.5 KiB
JavaScript
const { WebSocketServer } = require('ws');
|
|
const log = require('./logger');
|
|
|
|
let wss = null;
|
|
// Map jobId → Set of connected clients
|
|
const subscribers = new Map();
|
|
|
|
function attach(server) {
|
|
wss = new WebSocketServer({ server, path: '/ws' });
|
|
|
|
wss.on('connection', (ws) => {
|
|
ws._jobId = null;
|
|
|
|
ws.on('message', (data) => {
|
|
try {
|
|
const msg = JSON.parse(data);
|
|
if (msg.type === 'subscribe' && msg.jobId) {
|
|
// Unsubscribe from previous
|
|
if (ws._jobId) {
|
|
const prev = subscribers.get(ws._jobId);
|
|
if (prev) { prev.delete(ws); if (prev.size === 0) subscribers.delete(ws._jobId); }
|
|
}
|
|
// Subscribe to new job
|
|
ws._jobId = msg.jobId;
|
|
if (!subscribers.has(msg.jobId)) subscribers.set(msg.jobId, new Set());
|
|
subscribers.get(msg.jobId).add(ws);
|
|
}
|
|
} catch {}
|
|
});
|
|
|
|
ws.on('close', () => {
|
|
if (ws._jobId) {
|
|
const subs = subscribers.get(ws._jobId);
|
|
if (subs) { subs.delete(ws); if (subs.size === 0) subscribers.delete(ws._jobId); }
|
|
}
|
|
});
|
|
});
|
|
|
|
log.info('WebSocket server attached at /ws');
|
|
}
|
|
|
|
function notify(jobId, data) {
|
|
const subs = subscribers.get(jobId);
|
|
if (!subs || subs.size === 0) return;
|
|
const msg = JSON.stringify({ jobId, ...data });
|
|
for (const ws of subs) {
|
|
if (ws.readyState === 1) { // OPEN
|
|
ws.send(msg);
|
|
}
|
|
}
|
|
}
|
|
|
|
function getConnectionCount() {
|
|
return wss ? wss.clients.size : 0;
|
|
}
|
|
|
|
module.exports = { attach, notify, getConnectionCount };
|