feat: add subscribeBatchSize option to split subscribe packets for AWS IoT Core (#1995)

Author: Sheile

README.md

@@ -475,6 +475,8 @@ The arguments are:
       For ws/wss protocols only. Can be used to implement a custom websocket subprotocol or implementation.
   - `resubscribe` : if connection is broken and reconnects,
     subscribed topics are automatically subscribed again (default `true`)
+  - `subscribeBatchSize` : optional `number`
+    Maximum number of topics per SUBSCRIBE packet. When the number of topics to subscribe exceeds this value, the client will automatically split them into multiple SUBSCRIBE packets of this size.
   - `messageIdProvider`: custom messageId provider. when `new UniqueMessageIdProvider()` is set, then non conflict messageId is provided.
   - `log`: custom log function. Default uses [debug](https://www.npmjs.com/package/debug) package.
   - `manualConnect`: prevents the constructor to call `connect`. In this case after the `mqtt.connect` is called you should call `client.connect` manually.

src/lib/client.ts

@@ -29,6 +29,7 @@ import { type ClientRequestArgs } from 'http'
 import {
 	type DoneCallback,
 	type ErrorWithReasonCode,
+	ErrorWithSubackPacket,
 	type GenericCallback,
 	type IStream,
 	MQTTJS_VERSION,
@@ -61,6 +62,7 @@ const defaultConnectOptions: IClientOptions = {
 	connectTimeout: 30 * 1000,
 	clean: true,
 	resubscribe: true,
+	subscribeBatchSize: null,
 	writeCache: true,
 	timerVariant: 'auto',
 }
@@ -220,6 +222,14 @@ export interface IClientOptions extends ISecureClientOptions {
 	 */
 	resubscribe?: boolean
 
+	/**
+	 * Maximum number of topics to include in a single SUBSCRIBE packet.
+	 * When subscribing to more topics than this, the client will automatically
+	 * split them into batches of this size.
+	 * This is useful on AWS IoT Core, which limits each SUBSCRIBE packet to 8 topics.
+	 */
+	subscribeBatchSize?: number
+
 	/** when defined this function will be called to transform the url string generated by MqttClient from provided options */
 	transformWsUrl?: (
 		url: string,
@@ -1209,16 +1219,13 @@ export default class MqttClient extends TypedEventEmitter<MqttClientEventCallbac
 			return this
 		}
 
-		const subscribeProc = () => {
-			const messageId = this._nextId()
-			if (messageId === null) {
-				this.log('No messageId left')
-				return false
-			}
-
+		const subscribeChunkedSubs = (
+			chunkedSubs: ISubscriptionRequest[],
+			messageId: number,
+		): Promise<ISubackPacket> => {
 			const packet: ISubscribePacket = {
 				cmd: 'subscribe',
-				subscriptions: subs,
+				subscriptions: chunkedSubs,
 				// qos: 1,
 				// retain: false,
 				// dup: false,
@@ -1233,7 +1240,7 @@ export default class MqttClient extends TypedEventEmitter<MqttClientEventCallbac
 			if (this.options.resubscribe) {
 				this.log('subscribe :: resubscribe true')
 				const topics = []
-				subs.forEach((sub) => {
+				chunkedSubs.forEach((sub) => {
 					if (this.options.reconnectPeriod > 0) {
 						const topic: IClientSubscribeOptions = { qos: sub.qos }
 						if (version === 5) {
@@ -1249,21 +1256,61 @@ export default class MqttClient extends TypedEventEmitter<MqttClientEventCallbac
 				this.messageIdToTopic[packet.messageId] = topics
 			}
 
-			this.outgoing[packet.messageId] = {
-				volatile: true,
-				cb(err, packet2: ISubackPacket) {
-					if (!err) {
-						const { granted } = packet2
-						for (let i = 0; i < granted.length; i += 1) {
-							subs[i].qos = granted[i] as QoS
+			const promise = new Promise<ISubackPacket>((resolve, reject) => {
+				this.outgoing[packet.messageId] = {
+					volatile: true,
+					cb(err, packet2: ISubackPacket) {
+						if (!err) {
+							const { granted } = packet2
+							for (
+								let grantedI = 0;
+								grantedI < granted.length;
+								grantedI += 1
+							) {
+								chunkedSubs[grantedI].qos = granted[
+									grantedI
+								] as QoS
+							}
 						}
-					}
 
-					callback(err, subs, packet2)
-				},
-			}
+						if (!err) {
+							resolve(packet2)
+						} else {
+							reject(
+								new ErrorWithSubackPacket(err.message, packet2),
+							)
+						}
+					},
+				}
+			})
 			this.log('subscribe :: call _sendPacket')
 			this._sendPacket(packet)
+			return promise
+		}
+
+		const subscribeProc = () => {
+			const batchSize = this.options.subscribeBatchSize ?? subs.length
+			const subscribePromises: Promise<ISubackPacket>[] = []
+
+			for (let i = 0; i < subs.length; i += batchSize) {
+				const chunkedSubs = subs.slice(i, i + batchSize)
+				const messageId = this._nextId()
+				if (messageId === null) {
+					this.log('No messageId left')
+					return false
+				}
+				subscribePromises.push(
+					subscribeChunkedSubs(chunkedSubs, messageId),
+				)
+			}
+			Promise.all(subscribePromises)
+				.then((packets) => {
+					callback(null, subs, packets.at(-1))
+				})
+				.catch((err: ErrorWithSubackPacket) => {
+					callback(err, subs, err.packet)
+				})
+
 			return true
 		}
 

src/lib/shared.ts

@@ -1,4 +1,4 @@
-import type { Packet } from 'mqtt-packet'
+import type { Packet, ISubackPacket } from 'mqtt-packet'
 import type { Duplex } from 'stream'
 import type MqttClient from './client'
 import type { IClientOptions } from './client'
@@ -42,6 +42,19 @@ export class ErrorWithReasonCode extends Error {
 	}
 }
 
+export class ErrorWithSubackPacket extends Error {
+	public packet: ISubackPacket
+
+	public constructor(message: string, packet: ISubackPacket) {
+		super(message)
+		this.packet = packet
+
+		// We need to set the prototype explicitly
+		Object.setPrototypeOf(this, ErrorWithSubackPacket.prototype)
+		Object.getPrototypeOf(this).name = 'ErrorWithSubackPacket'
+	}
+}
+
 // eslint-disable-next-line @typescript-eslint/ban-types
 export type Constructor<T = {}> = new (...args: any[]) => T
 

test/node/abstract_client.ts

@@ -2619,6 +2619,48 @@ export default function abstractTest(server, config, ports) {
 				})
 			})
 		})
+
+		it('should send multiple subscribe packets when topic count exceeds batchSize', function _test(t, done) {
+			const client = connect({ subscribeBatchSize: 2 })
+			const subs = ['test1', 'test2', 'test3']
+			client.once('connect', () => {
+				client.subscribe(subs)
+			})
+
+			const spy = sinon.spy()
+			server.once('client', (serverClient) => {
+				serverClient.on('subscribe', spy)
+			})
+			client.on('end', () => {
+				assert.strictEqual(spy.callCount, 2)
+				for (let i = 0; i < 2; i++) {
+					// i.e. [{topic: 'a', qos: 0}, {topic: 'b', qos: 0}]
+					const expected = subs
+						.slice(i * 2, i * 2 + 2)
+						.map((topic) => {
+							const result: ISubscriptionRequest = {
+								topic,
+								qos: 0,
+							}
+							if (version === 5) {
+								result.nl = false
+								result.rap = false
+								result.rh = 0
+							}
+							return result
+						})
+
+					assert.deepStrictEqual(
+						spy.getCall(i).args[0].subscriptions,
+						expected,
+					)
+				}
+				done()
+			})
+			setTimeout(() => {
+				client.end()
+			}, 300)
+		})
 	})
 
 	describe('receiving messages', () => {