feat: support multiple handlers for a single subscription (#66)

Author: jonahsnider

src/connections/abstract_connection.ts

@@ -9,7 +9,6 @@
 
 import Emittery from 'emittery'
 import type { Redis, Cluster } from 'ioredis'
-import * as errors from '../errors.js'
 import type {
   PubSubOptions,
   ConnectionEvents,
@@ -38,8 +37,8 @@ export abstract class AbstractConnection<
   /**
    * A list of active subscriptions and pattern subscription
    */
-  protected subscriptions: Map<string, PubSubChannelHandler> = new Map()
-  protected psubscriptions: Map<string, PubSubPatternHandler> = new Map()
+  protected subscriptions: Map<string, Set<PubSubChannelHandler>> = new Map()
+  protected psubscriptions: Map<string, Set<PubSubPatternHandler>> = new Map()
 
   /**
    * The last error emitted by the `error` event. We set it to `null` after
@@ -230,19 +229,23 @@ export abstract class AbstractConnection<
      * Listen for messages
      */
     this.ioSubscriberConnection!.on('message', (channel, message) => {
-      const handler = this.subscriptions.get(channel)
-      if (handler) {
-        handler(message)
+      const handlers = this.subscriptions.get(channel)
+      if (handlers) {
+        for (const handler of handlers) {
+          handler(message)
+        }
       }
     })
 
     /**
      * Listen for pattern messages
      */
     this.ioSubscriberConnection!.on('pmessage', (pattern, channel, message) => {
-      const handler = this.psubscriptions.get(pattern)
-      if (handler) {
-        handler(channel, message)
+      const handlers = this.psubscriptions.get(pattern)
+      if (handlers) {
+        for (const handler of handlers) {
+          handler(channel, message)
+        }
       }
     })
   }
@@ -278,13 +281,6 @@ export abstract class AbstractConnection<
      */
     this.setupSubscriberConnection()
 
-    /**
-     * Disallow multiple subscriptions to a single channel
-     */
-    if (this.subscriptions.has(channel)) {
-      throw new errors.E_MULTIPLE_REDIS_SUBSCRIPTIONS([channel])
-    }
-
     /**
      * If the subscriptions map is empty, it means we have no active subscriptions
      * on the given channel, hence we should make one subscription and also set
@@ -296,7 +292,12 @@ export abstract class AbstractConnection<
           options?.onSubscription(count as number)
         }
         this.emit('subscription:ready', { count: count as number, connection: this })
-        this.subscriptions.set(channel, handler)
+        const subscriptions = this.subscriptions.get(channel)
+        if (subscriptions) {
+          subscriptions.add(handler)
+        } else {
+          this.subscriptions.set(channel, new Set([handler]))
+        }
       })
       .catch((error) => {
         if (options?.onError) {
@@ -309,8 +310,19 @@ export abstract class AbstractConnection<
   /**
    * Unsubscribe from a channel
    */
-  unsubscribe(channel: string) {
-    this.subscriptions.delete(channel)
+  unsubscribe(channel: string, handler?: PubSubChannelHandler) {
+    if (handler) {
+      const subscriptions = this.subscriptions.get(channel)
+      if (subscriptions) {
+        subscriptions.delete(handler)
+      }
+
+      if (subscriptions && subscriptions.size !== 0) {
+        return Promise.resolve()
+      }
+    } else {
+      this.subscriptions.delete(channel)
+    }
     return this.ioSubscriberConnection!.unsubscribe(channel)
   }
 
@@ -324,13 +336,6 @@ export abstract class AbstractConnection<
      */
     this.setupSubscriberConnection()
 
-    /**
-     * Disallow multiple subscriptions to a single channel
-     */
-    if (this.psubscriptions.has(pattern)) {
-      throw new errors.E_MULTIPLE_REDIS_PSUBSCRIPTIONS([pattern])
-    }
-
     /**
      * If the subscriptions map is empty, it means we have no active subscriptions
      * on the given channel, hence we should make one subscription and also set
@@ -342,7 +347,12 @@ export abstract class AbstractConnection<
           options?.onSubscription(count as number)
         }
         this.emit('psubscription:ready', { count: count as number, connection: this })
-        this.psubscriptions.set(pattern, handler)
+        const psubscriptions = this.psubscriptions.get(pattern)
+        if (psubscriptions) {
+          psubscriptions.add(handler)
+        } else {
+          this.psubscriptions.set(pattern, new Set([handler]))
+        }
       })
       .catch((error) => {
         if (options?.onError) {
@@ -355,8 +365,20 @@ export abstract class AbstractConnection<
   /**
    * Unsubscribe from a given pattern
    */
-  punsubscribe(pattern: string) {
-    this.psubscriptions.delete(pattern)
+  punsubscribe(pattern: string, handler?: PubSubPatternHandler) {
+    if (handler) {
+      const psubscriptions = this.psubscriptions.get(pattern)
+      if (psubscriptions) {
+        psubscriptions.delete(handler)
+      }
+
+      if (psubscriptions && psubscriptions.size !== 0) {
+        return Promise.resolve()
+      }
+    } else {
+      this.psubscriptions.delete(pattern)
+    }
+
     return this.ioSubscriberConnection!.punsubscribe(pattern)
   }
 

src/errors.ts

@@ -9,12 +9,14 @@
 
 import { createError } from '@poppinss/utils'
 
+/** @deprecated */
 export const E_MULTIPLE_REDIS_SUBSCRIPTIONS = createError<[string]>(
   'Cannot subscribe to "%s" channel. Channel already has an active subscription',
   'E_MULTIPLE_REDIS_SUBSCRIPTIONS',
   500
 )
 
+/** @deprecated */
 export const E_MULTIPLE_REDIS_PSUBSCRIPTIONS = createError<[string]>(
   'Cannot subscribe to "%s" pattern. Pattern already has an active subscription',
   'E_MULTIPLE_REDIS_PSUBSCRIPTIONS',

tests/redis_connection.spec.ts

@@ -138,7 +138,7 @@ test.group('Redis connection', () => {
     assert.equal(message, JSON.stringify({ username: 'virk' }))
   })
 
-  test('throw error when subscribing to a channel twice', async ({ cleanup }) => {
+  test('call all handlers subscribed to a channel', async ({ assert, cleanup }) => {
     const connection = new RedisConnection('main', {
       host: process.env.REDIS_HOST,
       port: Number(process.env.REDIS_PORT),
@@ -147,10 +147,32 @@ test.group('Redis connection', () => {
 
     await pEvent(connection, 'ready')
 
-    connection.subscribe('new:user', () => {})
-    await pEvent(connection, 'subscription:ready')
-    connection.subscribe('new:user', () => {})
-  }).throws('Cannot subscribe to "new:user" channel. Channel already has an active subscription')
+    const [message1] = await Promise.all([
+      new Promise<any>((resolve) => {
+        connection.subscribe('new:user', (message) => {
+          resolve(message)
+        })
+      }),
+      pEvent(connection, 'subscription:ready').then(() => {
+        connection.publish('new:user', JSON.stringify({ username: 'virk' }))
+      }),
+    ])
+
+    assert.equal(message1, JSON.stringify({ username: 'virk' }))
+
+    const [message2] = await Promise.all([
+      new Promise<any>((resolve) => {
+        connection.subscribe('new:user', (message) => {
+          resolve(message)
+        })
+      }),
+      pEvent(connection, 'subscription:ready').then(() => {
+        connection.publish('new:user', JSON.stringify({ username: 'virk' }))
+      }),
+    ])
+
+    assert.equal(message2, JSON.stringify({ username: 'virk' }))
+  })
 
   test('subscribe to a pattern and listen for messages', async ({ assert, cleanup }) => {
     const connection = new RedisConnection('main', {
@@ -174,7 +196,7 @@ test.group('Redis connection', () => {
     assert.equal(message.data, JSON.stringify({ username: 'virk' }))
   })
 
-  test('throw error when subscribing to a pattern twice', async ({ cleanup }) => {
+  test('call all handlers subscribed to a pattern', async ({ assert, cleanup }) => {
     const connection = new RedisConnection('main', {
       host: process.env.REDIS_HOST,
       port: Number(process.env.REDIS_PORT),
@@ -183,11 +205,34 @@ test.group('Redis connection', () => {
 
     await pEvent(connection, 'ready')
 
-    connection.psubscribe('user:*', () => {})
-    await pEvent(connection, 'psubscription:ready')
+    const [message1] = await Promise.all([
+      new Promise<any>((resolve) => {
+        connection.psubscribe('user:*', (channel, message) => {
+          resolve({ channel, message })
+        })
+      }),
+      pEvent(connection, 'psubscription:ready').then(() => {
+        connection.publish('user:add', JSON.stringify({ username: 'virk' }))
+      }),
+    ])
+
+    assert.equal(message1.channel, 'user:add')
+    assert.equal(message1.message, JSON.stringify({ username: 'virk' }))
 
-    connection.psubscribe('user:*', () => {})
-  }).throws('Cannot subscribe to "user:*" pattern. Pattern already has an active subscription')
+    const [message2] = await Promise.all([
+      new Promise<any>((resolve) => {
+        connection.psubscribe('user:*', (channel, message) => {
+          resolve({ channel, message })
+        })
+      }),
+      pEvent(connection, 'psubscription:ready').then(() => {
+        connection.publish('user:add', JSON.stringify({ username: 'virk' }))
+      }),
+    ])
+
+    assert.equal(message2.channel, 'user:add')
+    assert.equal(message2.message, JSON.stringify({ username: 'virk' }))
+  })
 
   test('unsubscribe from a channel', async ({ cleanup }) => {
     const connection = new RedisConnection('main', {