diff --git a/packages/controller/src/BaseConnection.ts b/packages/controller/src/BaseConnection.ts index 91f2a5ab..980b0a14 100644 --- a/packages/controller/src/BaseConnection.ts +++ b/packages/controller/src/BaseConnection.ts @@ -33,6 +33,10 @@ export default class BaseConnection extends lib.Link { this.handle(lib.ModPackGetRequest, this.handleModPackGetRequest.bind(this)); this.handle(lib.ModPackGetDefaultRequest, this.handleModPackGetDefaultRequest.bind(this)); this.handle(lib.ModDownloadRequest, this.handleModDownloadRequest.bind(this)); + + this.connector.on("close", () => { + this._controller.subscriptions.unsubscribe(this); + }); } async disconnect(code: number, reason: string) { diff --git a/packages/lib/src/subscriptions.ts b/packages/lib/src/subscriptions.ts index 9ecd493d..411449f4 100644 --- a/packages/lib/src/subscriptions.ts +++ b/packages/lib/src/subscriptions.ts @@ -101,7 +101,7 @@ export class SubscriptionRequest { type EventData = { subscriptionUpdate?: SubscriptionRequestHandler, - subscriptions: Map void }>, + subscriptions: Set, }; /** @@ -134,7 +134,7 @@ export class SubscriptionController { } this._events.set(entry.name, { subscriptionUpdate: subscriptionUpdate, - subscriptions: new Map(), + subscriptions: new Set(), }); } @@ -151,15 +151,26 @@ export class SubscriptionController { if (!eventData) { throw new Error(`Event ${entry.name} is not a registered as subscribable`); } - for (let [link, subscription] of eventData.subscriptions) { + for (let link of eventData.subscriptions) { if ((link.connector as WebSocketBaseConnector).closing) { eventData.subscriptions.delete(link); - } else if (subscription.all) { + } else { link.send(event); } } } + /** + * Unsubscribe from all events of a given link. + * Used when a link is closed to stop all active subscriptions. + * @param link - Link to stop sending events to. + */ + unsubscribe(link: Link) { + for (let eventData of this._events.values()) { + eventData.subscriptions.delete(link); + } + } + /** * Handle incoming event subscription requests * @param event - incomming event. @@ -179,23 +190,12 @@ export class SubscriptionController { let eventReplay: Event | null = null; const link: Link = this.controller.wsServer.controlConnections.get(src.id); if (event.subscribe === false) { - let onceClose = eventData.subscriptions.get(link)?.onceClose; - if (onceClose) { - link.connector.off("close", onceClose); - eventData.subscriptions.delete(link); - } + eventData.subscriptions.delete(link); } else { + eventData.subscriptions.add(link); if (eventData.subscriptionUpdate) { eventReplay = await eventData.subscriptionUpdate(event, src, dst); } - let onceClose = eventData.subscriptions.get(link)?.onceClose; - if (!onceClose) { - onceClose = () => eventData.subscriptions.delete(link); - link.connector.once("close", onceClose); - } - eventData.subscriptions.set( - link, { all: event.subscribe, onceClose: onceClose }, - ); } return new SubscriptionResponse(eventReplay); } diff --git a/test/lib/subscriptions.js b/test/lib/subscriptions.js index 3405fcfa..14abba32 100644 --- a/test/lib/subscriptions.js +++ b/test/lib/subscriptions.js @@ -243,6 +243,24 @@ describe("lib/subscriptions", function() { }); }); + describe("unsubscribe()", function() { + beforeEach(function() { + subscriptions.handle(RegisteredEvent); + for (let connectorData of connectorSetupDate) { + subscriptions._handleRequest(connectorData.request, connectorData.src, connectorData.dst); + } + }); + + it("should remove an active subscription", async function() { + const link = mockController.wsServer.controlConnections.get(0); + subscriptions.unsubscribe(link); + subscriptions.broadcast(new RegisteredEvent()); + await onceConnectorSend(1); + assertNoEvent(0); + assertLastEvent(1, RegisteredEvent); + }); + }); + describe("_handleRequest()", function() { beforeEach(function() { subscriptions.handle(RegisteredEvent);