fix(microservices) add subscriptions counter to mqtt and redis

Author: kamilmysliwiec

integration/microservices/e2e/sum-mqtt.spec.ts

@@ -49,10 +49,7 @@ describe('MQTT transport', () => {
       .expect(200, '15');
   });
 
-  /**
-   * Figure out race conditions here (flaky tests)
   it(`/POST (concurrent)`, function() {
-    this.retries(10);
     return request(server)
       .post('/concurrent')
       .send([
@@ -68,7 +65,7 @@ describe('MQTT transport', () => {
         Array.from({ length: 10 }, (v, k) => k + 91),
       ])
       .expect(200, 'true');
-  }).timeout(5000); */
+  }).timeout(5000);
 
   it(`/POST (streaming)`, () => {
     return request(server)

packages/microservices/client/client-mqtt.ts

@@ -17,6 +17,7 @@ let mqttPackage: any = {};
 
 export class ClientMqtt extends ClientProxy {
   protected readonly logger = new Logger(ClientProxy.name);
+  protected readonly subscriptionsCount = new Map<string, number>();
   protected readonly url: string;
   protected mqttClient: MqttClient;
   protected connection: Promise<any>;
@@ -121,19 +122,30 @@ export class ClientMqtt extends ClientProxy {
       const pattern = this.normalizePattern(partialPacket.pattern);
       const serializedPacket = this.serializer.serialize(packet);
       const responseChannel = this.getResPatternName(pattern);
+      let subscriptionsCount =
+        this.subscriptionsCount.get(responseChannel) || 0;
 
-      this.mqttClient.subscribe(responseChannel, (err: any) => {
-        if (err) {
-          return;
-        }
+      const publishPacket = () => {
+        subscriptionsCount = this.subscriptionsCount.get(responseChannel) || 0;
+        this.subscriptionsCount.set(responseChannel, subscriptionsCount + 1);
         this.routingMap.set(packet.id, callback);
         this.mqttClient.publish(
           this.getAckPatternName(pattern),
           JSON.stringify(serializedPacket),
         );
-      });
+      };
+
+      if (subscriptionsCount <= 0) {
+        this.mqttClient.subscribe(
+          responseChannel,
+          (err: any) => !err && publishPacket(),
+        );
+      } else {
+        publishPacket();
+      }
+
       return () => {
-        this.mqttClient.unsubscribe(responseChannel);
+        this.unsubscribeFromChannel(responseChannel);
         this.routingMap.delete(packet.id);
       };
     } catch (err) {
@@ -151,4 +163,13 @@ export class ClientMqtt extends ClientProxy {
       ),
     );
   }
+
+  protected unsubscribeFromChannel(channel: string) {
+    const subscriptionCount = this.subscriptionsCount.get(channel);
+    this.subscriptionsCount.set(channel, subscriptionCount - 1);
+
+    if (subscriptionCount - 1 <= 0) {
+      this.mqttClient.unsubscribe(channel);
+    }
+  }
 }

packages/microservices/client/client-redis.ts

@@ -21,6 +21,7 @@ let redisPackage: any = {};
 
 export class ClientRedis extends ClientProxy {
   protected readonly logger = new Logger(ClientProxy.name);
+  protected readonly subscriptionsCount = new Map<string, number>();
   protected readonly url: string;
   protected pubClient: RedisClient;
   protected subClient: RedisClient;
@@ -152,20 +153,30 @@ export class ClientRedis extends ClientProxy {
       const pattern = this.normalizePattern(partialPacket.pattern);
       const serializedPacket = this.serializer.serialize(packet);
       const responseChannel = this.getResPatternName(pattern);
+      let subscriptionsCount =
+        this.subscriptionsCount.get(responseChannel) || 0;
 
-      this.routingMap.set(packet.id, callback);
-      this.subClient.subscribe(responseChannel, (err: any) => {
-        if (err) {
-          return;
-        }
+      const publishPacket = () => {
+        subscriptionsCount = this.subscriptionsCount.get(responseChannel) || 0;
+        this.subscriptionsCount.set(responseChannel, subscriptionsCount + 1);
+        this.routingMap.set(packet.id, callback);
         this.pubClient.publish(
           this.getAckPatternName(pattern),
           JSON.stringify(serializedPacket),
         );
-      });
+      };
+
+      if (subscriptionsCount <= 0) {
+        this.subClient.subscribe(
+          responseChannel,
+          (err: any) => !err && publishPacket(),
+        );
+      } else {
+        publishPacket();
+      }
 
       return () => {
-        this.subClient.unsubscribe(responseChannel);
+        this.unsubscribeFromChannel(responseChannel);
         this.routingMap.delete(packet.id);
       };
     } catch (err) {
@@ -183,4 +194,13 @@ export class ClientRedis extends ClientProxy {
       ),
     );
   }
+
+  protected unsubscribeFromChannel(channel: string) {
+    const subscriptionCount = this.subscriptionsCount.get(channel);
+    this.subscriptionsCount.set(channel, subscriptionCount - 1);
+
+    if (subscriptionCount - 1 <= 0) {
+      this.subClient.unsubscribe(channel);
+    }
+  }
 }