diff --git a/subscriptionHandler.ts b/subscriptionHandler.ts new file mode 100644 index 0000000..adc254b --- /dev/null +++ b/subscriptionHandler.ts @@ -0,0 +1,62 @@ +import { WS } from "./index"; + +export interface Subscriber { + send: (message: string) => void; +} + +export class SubscriptionMap { + channelsPerSubscriber = new Map>(); + subscribersPerChannel = new Map>(); + + // prepare + private prepareChannelSet(subscriber: Subscriber): void { + if (this.channelsPerSubscriber.has(subscriber)) return; + this.channelsPerSubscriber.set(subscriber, new Set()); + } + + private prepareSubscriberSet(channel: string): void { + if (this.subscribersPerChannel.has(channel)) return; + this.subscribersPerChannel.set(channel, new Set()); + } + + // utility + getChannelList(subscriber: Subscriber): Set | undefined { + return this.channelsPerSubscriber.get(subscriber); + } + + getClientList(channel: string): Set | undefined { + return this.subscribersPerChannel.get(channel); + } + + // modifications + set(subscriber: Subscriber, channel: string): void { + this.prepareChannelSet(subscriber); + this.prepareSubscriberSet(channel); + + this.getChannelList(subscriber)?.add(channel); + this.getClientList(channel)?.add(subscriber); + } + + delete(subscriber: Subscriber, channel: string): void { + const channelList = this.getChannelList(subscriber); + const clientList = this.getClientList(channel); + + channelList?.delete(channel); + if (channelList?.size == 0) { + this.channelsPerSubscriber.delete(subscriber); + } + + clientList?.delete(subscriber); + if (clientList?.size == 0) { + this.subscribersPerChannel.delete(channel); + } + } + + // messages + forwardMessage(channels: string[], messageString: string): void { + channels.forEach((channel) => { + const subscribers = this.getClientList(channel); + subscribers?.forEach((subscriber) => subscriber.send(messageString)); + }); + } +} diff --git a/subscriptionMap.ts b/subscriptionMap.ts deleted file mode 100644 index 6b77e1e..0000000 --- a/subscriptionMap.ts +++ /dev/null @@ -1,50 +0,0 @@ -import { WS } from "./index"; - -export class SubscriptionMap { - channelsPerClient = new Map>(); - clientsPerChannel = new Map>(); - - // prepare - private prepareChannelSet(ws: WS) { - if (this.channelsPerClient.has(ws)) return; - this.channelsPerClient.set(ws, new Set()); - } - - private prepareClientSet(channel: string) { - if (this.clientsPerChannel.has(channel)) return; - this.clientsPerChannel.set(channel, new Set()); - } - - // utility - getChannelList(ws: WS) { - return this.channelsPerClient.get(ws); - } - - getClientList(channel: string) { - return this.clientsPerChannel.get(channel); - } - - // modifications - set(ws: WS, channel: string) { - this.prepareChannelSet(ws); - this.prepareClientSet(channel); - - this.getChannelList(ws)?.add(channel); - this.getClientList(channel)?.add(ws); - } - - delete(ws: WS, channel: string) { - const channelList = this.getChannelList(ws); - const clientList = this.getClientList(channel); - - channelList?.delete(channel); - if (channelList?.size == 0) { - this.channelsPerClient.delete(ws); - } - - clientList?.delete(ws); - if (clientList?.size == 0) { - this.clientsPerChannel.delete(channel); - } - } -} diff --git a/websocket-handler.ts b/websocket-handler.ts index a4052be..f827e13 100644 --- a/websocket-handler.ts +++ b/websocket-handler.ts @@ -1,3 +1,7 @@ +import { + Subscriber, + SubscriptionMap, +} from "./subscriptionHandler"; import { doIfIsString, parseMessage, @@ -7,7 +11,6 @@ import { } from "./utility"; import Crypto from "crypto"; -import { SubscriptionMap } from "./subscriptionMap"; import { WS } from "./index"; import { defaultConfig } from "./config-handler"; @@ -35,7 +38,9 @@ export interface WebSocketMessage { const knownMessageIds = new Set(); export const subscriptionMap = new SubscriptionMap(); + export const clientsConnected = new Set(); + export const serversConnectedAsClient = new Set(); export const serversConnected = new Set(); export const serversDisconnected = new Set(); @@ -58,14 +63,20 @@ export function forgetConnection(ws: WS): void { }); } -export function trackSubscription(ws: WS, channel: string): void { - subscriptionMap.set(ws, channel); - confirmSubscription(ws, channel); +export function trackSubscription( + subscriber: Subscriber, + channel: string +): void { + subscriptionMap.set(subscriber, channel); + confirmSubscription(subscriber, channel); } -export function forgetSubscription(ws: WS, channel: string): void { - subscriptionMap.delete(ws, channel); - confirmSubscription(ws, channel); +export function forgetSubscription( + subscriber: Subscriber, + channel: string +): void { + subscriptionMap.delete(subscriber, channel); + confirmSubscription(subscriber, channel); } // OTHER SERVERS @@ -131,7 +142,7 @@ export function subscribeChannel(server: WS, channel: string): void { // AUDIT export function getWebSocketStats(): [string, number][] { return [ - ["channels", subscriptionMap.clientsPerChannel.size], + ["channels", subscriptionMap.subscribersPerChannel.size], ["clients connected", clientsConnected.size], ["servers connected", getServerCount()], ["servers disconnected", serversDisconnected.size], @@ -190,10 +201,7 @@ export function sendMessage( const messageObject = { uuid, messageChannel, messageBody }; const messageString = stringifyMessage(messageObject); - channels.forEach((channel) => { - const subscribers = subscriptionMap.getClientList(channel); - subscribers?.forEach((ws) => ws.send(messageString)); - }); + subscriptionMap.forwardMessage(channels, messageString); } function requestServerConnection(ws: WS): void { @@ -204,11 +212,16 @@ function requestServerConnection(ws: WS): void { ws.send(messageString); } -function confirmSubscription(ws: WS, messageChannel: string): void { +function confirmSubscription( + subscriber: Subscriber, + messageChannel: string +): void { const messageObject = { messageChannel, - subscribed: subscriptionMap.getChannelList(ws) ? true : false, + subscribed: subscriptionMap.getChannelList(subscriber)?.has(messageChannel) + ? true + : false, }; const messageString = stringifyMessage(messageObject); - ws.send(messageString); + subscriber.send(messageString); }