Filtered subscriptions (#788)

* Add filtered subscriptions

* Add tests for subscription filters

* Stricter assertions

* Un-nest if statements

packages/lib/src/subscriptions.ts

@@ -39,30 +39,45 @@ export class SubscriptionRequest {
 		public eventName: string,
 		public subscribe: boolean,
 		public lastRequestTimeMs: number = 0,
+		public filters: string | string[] | undefined = undefined,
 	) {
 		if (!Link._eventsByName.has(eventName)) {
 			throw new Error(`Unregistered Event class ${eventName}`);
 		}
 	}
 
-	static jsonSchema = Type.Tuple([
-		Type.String(),
-		Type.Boolean(),
-		Type.Number(),
+	static jsonSchema = Type.Union([
+		Type.Tuple([
+			Type.String(),
+			Type.Boolean(),
+			Type.Number(),
+		]),
+		Type.Tuple([
+			Type.String(),
+			Type.Boolean(),
+			Type.Number(),
+			Type.Union([Type.String(), Type.Array(Type.String())]),
+		]),
 	]);
 
 	toJSON() {
+		if (this.filters !== undefined) {
+			return [this.eventName, this.subscribe, this.lastRequestTimeMs, this.filters];
+		}
 		return [this.eventName, this.subscribe, this.lastRequestTimeMs];
 	}
 
 	static fromJSON(json: Static<typeof SubscriptionRequest.jsonSchema>): SubscriptionRequest {
-		return new this(...json);
+		const [eventName, subscribe, lastRequestTimeMs, filters] = json;
+		return new this(eventName, subscribe, lastRequestTimeMs, filters);
 	}
 }
 
 type Subscriber = {
 	link: Link,
 	dst: Address,
+	/** Optional filters this subscriber is interested in. Undefined means all events */
+	filters?: string[],
 }
 
 type EventData = {
@@ -77,6 +92,13 @@ type EventData = {
 export class SubscriptionController {
 	_events = new Map<string, EventData>();
 
+	private static normalizeFilters(value?: string | string[]): string[] | undefined {
+		if (value === undefined) {
+			return undefined;
+		}
+		return Array.isArray(value) ? value : [value];
+	}
+
 	/**
 	 * Allow clients to subscribe to an event by telling the subscription controller to accept them
 	 * Has an optional subscription update handler which is called when a client subscribes
@@ -102,7 +124,7 @@ export class SubscriptionController {
 	 * Broadcast an event to all subscribers of that event
 	 * @param event - Event to broadcast.
 	 */
-	broadcast<T>(event: Event<T>) {
+	broadcast<T>(event: Event<T>, filter?: string | string[]) {
 		const entry = Link._eventsByClass.get(event.constructor);
 		if (!entry) {
 			throw new Error(`Unregistered Event class ${event.constructor.name}`);
@@ -112,10 +134,19 @@ export class SubscriptionController {
 			throw new Error(`Event ${entry.name} is not a registered as subscribable`);
 		}
 		const toRemove: Subscriber[] = [];
+		const filtersToMatch = SubscriptionController.normalizeFilters(filter);
+
 		for (const subscriber of eventData.subscriptions) {
 			if ((subscriber.link.connector as WebSocketBaseConnector).closing) {
 				toRemove.push(subscriber);
-			} else {
+				continue;
+			}
+
+			if (
+				filtersToMatch === undefined
+				|| subscriber.filters === undefined
+				|| subscriber.filters.some(f => filtersToMatch.includes(f))
+			) {
 				subscriber.link.sendTo(subscriber.dst, event);
 			}
 		}
@@ -164,13 +195,37 @@ export class SubscriptionController {
 			throw new Error(`Event ${event.eventName} is not a registered as subscribable`);
 		}
 		if (event.subscribe === false) {
-			const index = eventData.subscriptions.findLastIndex(subscriber => subscriber.dst.addressedTo(src));
-			if (index >= 0) {
-				eventData.subscriptions.splice(index, 1);
+			// Unsubscribe logic
+			const subscriber = eventData.subscriptions.find(sub => sub.dst.addressedTo(src));
+			if (!subscriber) {
+				return false;
+			}
+			if (event.filters === undefined || subscriber.filters === undefined) {
+				// Remove entire subscription for this address
+				eventData.subscriptions = eventData.subscriptions.filter(sub => !sub.dst.addressedTo(src));
+			} else {
+				const filtersToRemove = Array.isArray(event.filters) ? event.filters : [event.filters];
+				subscriber.filters = subscriber.filters.filter(f => !filtersToRemove.includes(f));
+				if (subscriber.filters.length === 0) {
+					// No filters left, remove subscription
+					eventData.subscriptions = eventData.subscriptions.filter(sub => sub !== subscriber);
+				}
 			}
 		} else {
-			if (!eventData.subscriptions.some(subscriber => subscriber.dst.addressedTo(src))) {
-				eventData.subscriptions.push({ link: link, dst: src });
+			// Subscribe logic
+			let subscriber = eventData.subscriptions.find(sub => sub.dst.addressedTo(src));
+			const newFilters = SubscriptionController.normalizeFilters(event.filters);
+			if (!subscriber) {
+				eventData.subscriptions.push({ link: link, dst: src, filters: newFilters });
+			} else if (subscriber.filters === undefined || newFilters === undefined) {
+				// Either was already subscribed to all or new request wants all
+				subscriber.filters = undefined;
+			} else {
+				for (const f of newFilters) {
+					if (!subscriber.filters.includes(f)) {
+						subscriber.filters.push(f);
+					}
+				}
 			}
 			if (eventData.subscriptionUpdate) {
 				const eventReplay = await eventData.subscriptionUpdate(event, src, dst);

test/lib/subscriptions.js

@@ -88,6 +88,22 @@ describe("lib/subscriptions", function() {
 			assert.deepEqual(reconstructed, request);
 		});
 
+		it("should be round trip json serialisable with a single filter", function() {
+			const request = new lib.SubscriptionRequest(RegisteredEvent.name, true, 456, "foo");
+			const json = JSON.stringify(request);
+			assert.equal(typeof json, "string");
+			const reconstructed = lib.SubscriptionRequest.fromJSON(JSON.parse(json));
+			assert.deepEqual(reconstructed, request);
+		});
+
+		it("should be round trip json serialisable with multiple filters", function() {
+			const request = new lib.SubscriptionRequest(RegisteredEvent.name, true, 789, ["foo", "bar"]);
+			const json = JSON.stringify(request);
+			assert.equal(typeof json, "string");
+			const reconstructed = lib.SubscriptionRequest.fromJSON(JSON.parse(json));
+			assert.deepEqual(reconstructed, request);
+		});
+
 		it("should throw when given an unregistered event", function() {
 			assert.throws(
 				() => new lib.SubscriptionRequest(UnregisteredEvent.name, true),
@@ -314,6 +330,171 @@ describe("lib/subscriptions", function() {
 				await new Promise(resolve => setImmediate(resolve));
 				assertNoEvent(0);
 			});
+
+			describe("filters()", function() {
+				it("should notify only subscribers matching the broadcast filter", async function() {
+					const reqFoo = new lib.SubscriptionRequest(RegisteredEvent.name, true, 0, "foo");
+					await subscriptions.handleRequest(
+						getLink(0), reqFoo, connectorSetupDate[0].dst, connectorSetupDate[0].src
+					);
+					const reqBar = new lib.SubscriptionRequest(RegisteredEvent.name, true, 0, "bar");
+					await subscriptions.handleRequest(
+						getLink(1), reqBar, connectorSetupDate[1].dst, connectorSetupDate[1].src
+					);
+
+					const before0 = getLink(0).connector.sentMessages.length;
+					const before1 = getLink(1).connector.sentMessages.length;
+
+					subscriptions.broadcast(new RegisteredEvent(), "foo");
+					await onceConnectorSend(0);
+					assert.equal(getLink(0).connector.sentMessages.length, before0 + 1);
+					assert.equal(getLink(1).connector.sentMessages.length, before1);
+
+					subscriptions.broadcast(new RegisteredEvent(), "bar");
+					await onceConnectorSend(1);
+					assert.equal(getLink(0).connector.sentMessages.length, before0 + 1);
+					assert.equal(getLink(1).connector.sentMessages.length, before1 + 1);
+
+					const beforeBaz0 = getLink(0).connector.sentMessages.length;
+					const beforeBaz1 = getLink(1).connector.sentMessages.length;
+					subscriptions.broadcast(new RegisteredEvent(), "baz");
+					await new Promise(resolve => setImmediate(resolve));
+					assert.equal(getLink(0).connector.sentMessages.length, beforeBaz0);
+					assert.equal(getLink(1).connector.sentMessages.length, beforeBaz1);
+				});
+
+				it("should notify subscribers with no filters for any broadcast filter", async function() {
+					const reqAll = new lib.SubscriptionRequest(RegisteredEvent.name, true);
+					await subscriptions.handleRequest(
+						getLink(0), reqAll, connectorSetupDate[0].dst, connectorSetupDate[0].src
+					);
+					const reqBaz = new lib.SubscriptionRequest(RegisteredEvent.name, true, 0, "baz");
+					await subscriptions.handleRequest(
+						getLink(1), reqBaz, connectorSetupDate[1].dst, connectorSetupDate[1].src
+					);
+
+					const before0 = getLink(0).connector.sentMessages.length;
+					const before1 = getLink(1).connector.sentMessages.length;
+
+					subscriptions.broadcast(new RegisteredEvent(), "foo");
+					await onceConnectorSend(0);
+					assert.equal(getLink(0).connector.sentMessages.length, before0 + 1);
+					assert.equal(getLink(1).connector.sentMessages.length, before1);
+
+					subscriptions.broadcast(new RegisteredEvent(), "baz");
+					await Promise.all([onceConnectorSend(0), onceConnectorSend(1)]);
+					assert.equal(getLink(0).connector.sentMessages.length, before0 + 2);
+					assert.equal(getLink(1).connector.sentMessages.length, before1 + 1);
+				});
+
+				it("should merge multiple filter subscriptions for the same subscriber", async function() {
+					const reqFoo = new lib.SubscriptionRequest(RegisteredEvent.name, true, 0, "alpha");
+					await subscriptions.handleRequest(
+						getLink(0), reqFoo, connectorSetupDate[0].dst, connectorSetupDate[0].src
+					);
+					const reqBar = new lib.SubscriptionRequest(RegisteredEvent.name, true, 0, "beta");
+					await subscriptions.handleRequest(
+						getLink(0), reqBar, connectorSetupDate[0].dst, connectorSetupDate[0].src
+					);
+
+					const before0 = getLink(0).connector.sentMessages.length;
+
+					subscriptions.broadcast(new RegisteredEvent(), "alpha");
+					await onceConnectorSend(0);
+					assert.equal(getLink(0).connector.sentMessages.length, before0 + 1);
+
+					subscriptions.broadcast(new RegisteredEvent(), "beta");
+					await onceConnectorSend(0);
+					assert.equal(getLink(0).connector.sentMessages.length, before0 + 2);
+
+					subscriptions.broadcast(new RegisteredEvent(), "gamma");
+					await new Promise(resolve => setImmediate(resolve));
+					assert.equal(getLink(0).connector.sentMessages.length, before0 + 2);
+				});
+
+				it("should reset to all when subscribing with no filters after having filters", async function() {
+					const reqFoo = new lib.SubscriptionRequest(RegisteredEvent.name, true, 0, "foo");
+					await subscriptions.handleRequest(
+						getLink(0), reqFoo, connectorSetupDate[0].dst, connectorSetupDate[0].src
+					);
+					const reqAll = new lib.SubscriptionRequest(RegisteredEvent.name, true);
+					await subscriptions.handleRequest(
+						getLink(0), reqAll, connectorSetupDate[0].dst, connectorSetupDate[0].src
+					);
+
+					subscriptions.broadcast(new RegisteredEvent(), "baz");
+					await onceConnectorSend(0);
+				});
+
+				it("should remove only specified filters on unsubscribe with filters", async function() {
+					const reqFilters = new lib.SubscriptionRequest(RegisteredEvent.name, true, 0, ["foo", "bar"]);
+					await subscriptions.handleRequest(
+						getLink(0), reqFilters, connectorSetupDate[0].dst, connectorSetupDate[0].src
+					);
+
+					const unsubFoo = new lib.SubscriptionRequest(RegisteredEvent.name, false, 0, "foo");
+					await subscriptions.handleRequest(
+						getLink(0), unsubFoo, connectorSetupDate[0].dst, connectorSetupDate[0].src
+					);
+
+					const before0 = getLink(0).connector.sentMessages.length;
+
+					subscriptions.broadcast(new RegisteredEvent(), "foo");
+					await new Promise(resolve => setImmediate(resolve));
+					assert.equal(getLink(0).connector.sentMessages.length, before0);
+
+					subscriptions.broadcast(new RegisteredEvent(), "bar");
+					await onceConnectorSend(0);
+					assert.equal(getLink(0).connector.sentMessages.length, before0 + 1);
+				});
+
+				it("should remove entire subscription when last filter is unsubscribed", async function() {
+					const reqFoo = new lib.SubscriptionRequest(RegisteredEvent.name, true, 0, "foo");
+					await subscriptions.handleRequest(
+						getLink(0), reqFoo, connectorSetupDate[0].dst, connectorSetupDate[0].src
+					);
+
+					const unsubFoo = new lib.SubscriptionRequest(RegisteredEvent.name, false, 0, "foo");
+					await subscriptions.handleRequest(
+						getLink(0), unsubFoo, connectorSetupDate[0].dst, connectorSetupDate[0].src
+					);
+
+					subscriptions.broadcast(new RegisteredEvent(), "foo");
+					await new Promise(resolve => setImmediate(resolve));
+					assertNoEvent(0);
+				});
+
+				// eslint-disable-next-line max-len
+				it("should remove entire subscription when unsubscribing any filter from an all-subscription", async function() {
+					const reqAll = new lib.SubscriptionRequest(RegisteredEvent.name, true);
+					await subscriptions.handleRequest(
+						getLink(0), reqAll, connectorSetupDate[0].dst, connectorSetupDate[0].src
+					);
+					const unsubAny = new lib.SubscriptionRequest(RegisteredEvent.name, false, 0, "foo");
+					await subscriptions.handleRequest(
+						getLink(0), unsubAny, connectorSetupDate[0].dst, connectorSetupDate[0].src
+					);
+
+					subscriptions.broadcast(new RegisteredEvent());
+					await new Promise(resolve => setImmediate(resolve));
+					assertNoEvent(0);
+				});
+
+				it("should match any of multiple broadcast filters", async function() {
+					const reqFoo = new lib.SubscriptionRequest(RegisteredEvent.name, true, 0, "foo");
+					await subscriptions.handleRequest(
+						getLink(0), reqFoo, connectorSetupDate[0].dst, connectorSetupDate[0].src
+					);
+					const reqBar = new lib.SubscriptionRequest(RegisteredEvent.name, true, 0, "bar");
+					await subscriptions.handleRequest(
+						getLink(1), reqBar, connectorSetupDate[1].dst, connectorSetupDate[1].src
+					);
+
+					subscriptions.broadcast(new RegisteredEvent(), ["bar", "baz"]);
+					await onceConnectorSend(1);
+					assertNoEvent(0);
+				});
+			});
 		});
 	});