Fix concurrent sub + quick unsub race condition

Author: mlaflamm

src/redis-pubsub.ts

@@ -80,6 +80,7 @@ export class RedisPubSub implements PubSubEngine {
 
     this.subscriptionMap = {};
     this.subsRefsMap = new Map<string, Set<number>>();
+    this.subsPendingRefsMap = new Map<string, { refs: number[], pending: Promise<number> }>();
     this.currentSubscriptionId = 0;
   }
 
@@ -102,22 +103,44 @@ export class RedisPubSub implements PubSubEngine {
     }
 
     const refs = this.subsRefsMap.get(triggerName);
-    if (refs.size > 0) {
+
+    const pendingRefs = this.subsPendingRefsMap.get(triggerName)
+    if (pendingRefs != null) {
+      // A pending remote subscribe call is currently in flight, piggyback on it
+      pendingRefs.refs.push(id)
+      return pendingRefs.pending.then(() => id)
+    } else if (refs.size > 0) {
+      // Already actively subscribed to redis
       refs.add(id);
       return Promise.resolve(id);
     } else {
-      return new Promise<number>((resolve, reject) => {
+      // New subscription.
+      // Keep a pending state until the remote subscribe call is completed
+      const pending = new Deferred()
+      const subsPendingRefsMap = this.subsPendingRefsMap
+      subsPendingRefsMap.set(triggerName, { refs: [], pending });
+
+      const sub = new Promise<number>((resolve, reject) => {
         const subscribeFn = options['pattern'] ? this.redisSubscriber.psubscribe : this.redisSubscriber.subscribe;
 
         subscribeFn.call(this.redisSubscriber, triggerName, err => {
           if (err) {
+            subsPendingRefsMap.delete(triggerName)
             reject(err);
           } else {
+            // Add ids of subscribe calls initiated when waiting for the remote call response
+            const pendingRefs = subsPendingRefsMap.get(triggerName)
+            pendingRefs.refs.forEach((id) => refs.add(id))
+            subsPendingRefsMap.delete(triggerName)
+
             refs.add(id);
             resolve(id);
           }
         });
       });
+      // Ensure waiting subscribe will complete
+      sub.then(pending.resolve).catch(pending.reject)
+      return sub;
     }
   }
 
@@ -167,6 +190,7 @@ export class RedisPubSub implements PubSubEngine {
 
   private readonly subscriptionMap: { [subId: number]: [string, OnMessage<unknown>] };
   private readonly subsRefsMap: Map<string, Set<number>>;
+  private readonly subsPendingRefsMap: Map<string, { refs: number[], pending: Promise<number> }>;
   private currentSubscriptionId: number;
 
   private onMessage(pattern: string, channel: string, message: string) {
@@ -191,6 +215,19 @@ export class RedisPubSub implements PubSubEngine {
   }
 }
 
+// Unexported deferrable promise used to complete waiting subscribe calls
+function Deferred() {
+  const p = this.promise = new Promise((resolve, reject) => {
+    this.resolve = resolve;
+    this.reject = reject;
+  });
+  this.then = p.then.bind(p);
+  this.catch = p.catch.bind(p);
+  if (p.finally) {
+    this.finally = p.finally.bind(p);
+  }
+}
+
 export type Path = Array<string | number>;
 export type Trigger = string | Path;
 export type TriggerTransform = (

src/test/tests.ts

@@ -140,6 +140,60 @@ describe('RedisPubSub', () => {
       });
   });
 
+  it('concurrent subscribe, unsubscribe first sub before second sub complete', done => {
+    const subs = {
+      first: null as Promise<number>,
+      second: null as Promise<number>,
+    }
+
+    let firstCb, secondCb
+    const redisSubCallback = (channel, cb) => {
+      process.nextTick(() => {
+        if (!firstCb) {
+          firstCb = () => cb(null, channel)
+          // Handling first call, init second sub
+          subs.second = pubSub.subscribe('Posts', () => null)
+          // Continue first sub callback
+          firstCb()
+        } else {
+          secondCb = () => cb(null, channel)
+        }
+      })
+    }
+    const subscribeStub = stub().callFn(redisSubCallback);
+    const mockRedisClientWithSubStub = {...mockRedisClient, ...{subscribe: subscribeStub}};
+    const mockOptionsWithSubStub = {...mockOptions, ...{subscriber: (mockRedisClientWithSubStub as any)}}
+    const pubSub = new RedisPubSub(mockOptionsWithSubStub);
+
+    // First leg of the test, init first sub and immediately unsubscribe. The second sub is triggered in the redis cb
+    // before the first promise sub complete
+    subs.first = pubSub.subscribe('Posts', () => null)
+      .then(subId => {
+        // This assertion is done against a private member, if you change the internals, you may want to change that
+        expect((pubSub as any).subscriptionMap[subId]).not.to.be.an('undefined');
+        pubSub.unsubscribe(subId);
+
+        // Continue second sub callback
+        subs.first.then(() => secondCb())
+        return subId;
+      });
+
+    // Second leg of the test, here we have unsubscribed from the first sub. We try unsubbing from the second sub
+    // as soon it is ready
+    subs.first
+      .then((subId) => {
+        // This assertion is done against a private member, if you change the internals, you may want to change that
+        expect((pubSub as any).subscriptionMap[subId]).to.be.an('undefined');
+        expect(() => pubSub.unsubscribe(subId)).to.throw(`There is no subscription of id "${subId}"`);
+
+        return subs.second.then(secondSubId => {
+          pubSub.unsubscribe(secondSubId);
+        })
+      .then(done)
+      .catch(done)
+    });
+  });
+
   it('will not unsubscribe from the redis channel if there is another subscriber on it\'s subscriber list', done => {
     const pubSub = new RedisPubSub(mockOptions);
     const subscriptionPromises = [