fix #1718 - add support for buffers in pubsub

Author: leibale

.github/README.md

@@ -150,6 +150,18 @@ Publish a message on a channel:
 await publisher.publish('channel', 'message');
 ```
 
+There is support for buffers as well:
+
+```typescript
+await subscriber.subscribe('channel', (message) => {
+  console.log(message); // <Buffer 6d 65 73 73 61 67 65>
+}, true);
+
+await subscriber.pSubscribe('channe*', (message, channel) => {
+  console.log(message, channel); // <Buffer 6d 65 73 73 61 67 65>, <Buffer 63 68 61 6e 6e 65 6c>
+}, true);
+```
+
 ### Scan Iterator
 
 [`SCAN`](https://redis.io/commands/scan) results can be looped over using [async iterators](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Symbol/asyncIterator):

packages/client/lib/client/commands-queue.ts

@@ -39,9 +39,29 @@ export enum PubSubUnsubscribeCommands {
     PUNSUBSCRIBE = 'PUNSUBSCRIBE'
 }
 
-export type PubSubListener = (message: string, channel: string) => unknown;
+type PubSubArgumentTypes = Buffer | string;
 
-export type PubSubListenersMap = Map<string, Set<PubSubListener>>;
+export type PubSubListener<
+    BUFFER_MODE extends boolean = false,
+    T = BUFFER_MODE extends true ? Buffer : string
+> = (message: T, channel: T) => unknown;
+
+interface PubSubListeners {
+    buffers: Set<PubSubListener<true>>;
+    strings: Set<PubSubListener<false>>;
+}
+
+type PubSubListenersMap = Map<string, PubSubListeners>;
+
+interface PubSubState {
+    subscribing: number;
+    subscribed: number;
+    unsubscribing: number;
+    listeners: {
+        channels: PubSubListenersMap;
+        patterns: PubSubListenersMap;
+    };
+}
 
 export default class RedisCommandsQueue {
     static #flushQueue<T extends CommandWaitingForReply>(queue: LinkedList<T>, err: Error): void {
@@ -50,10 +70,20 @@ export default class RedisCommandsQueue {
         }
     }
 
-    static #emitPubSubMessage(listeners: Set<PubSubListener>, message: string, channel: string): void {
-        for (const listener of listeners) {
+    static #emitPubSubMessage(listenersMap: PubSubListenersMap, message: Buffer, channel: Buffer, pattern?: Buffer): void {
+        const keyString = (pattern || channel).toString(),
+            listeners = listenersMap.get(keyString)!;
+        for (const listener of listeners.buffers) {
             listener(message, channel);
         }
+
+        if (!listeners.strings.size) return;
+
+        const messageString = message.toString(),
+            channelString = pattern ? channel.toString() : keyString;
+        for (const listener of listeners.strings) {
+            listener(messageString, channelString);
+        }
     }
 
     readonly #maxLength: number | null | undefined;
@@ -62,41 +92,43 @@ export default class RedisCommandsQueue {
 
     readonly #waitingForReply = new LinkedList<CommandWaitingForReply>();
 
-    readonly #pubSubState = {
-        subscribing: 0,
-        subscribed: 0,
-        unsubscribing: 0
-    };
+    #pubSubState: PubSubState | undefined;
 
-    readonly #pubSubListeners = {
-        channels: <PubSubListenersMap>new Map(),
-        patterns: <PubSubListenersMap>new Map()
+    static readonly #PUB_SUB_MESSAGES = {
+        message: Buffer.from('message'),
+        pMessage: Buffer.from('pmessage'),
+        subscribe: Buffer.from('subscribe'),
+        pSubscribe: Buffer.from('psubscribe'),
+        unsubscribe: Buffer.from('unsunscribe'),
+        pUnsubscribe: Buffer.from('punsubscribe')
     };
 
     readonly #parser = new RedisParser({
         returnReply: (reply: unknown) => {
-            if ((this.#pubSubState.subscribing || this.#pubSubState.subscribed) && Array.isArray(reply)) {
-                switch (reply[0]) {
-                    case 'message':
-                        return RedisCommandsQueue.#emitPubSubMessage(
-                            this.#pubSubListeners.channels.get(reply[1])!,
-                            reply[2],
-                            reply[1]
-                        );
-
-                    case 'pmessage':
-                        return RedisCommandsQueue.#emitPubSubMessage(
-                            this.#pubSubListeners.patterns.get(reply[1])!,
-                            reply[3],
-                            reply[2]
-                        );
-
-                    case 'subscribe':
-                    case 'psubscribe':
-                        if (--this.#waitingForReply.head!.value.channelsCounter! === 0) {
-                            this.#shiftWaitingForReply().resolve();
-                        }
-                        return;
+            if (this.#pubSubState && Array.isArray(reply)) {
+                if (RedisCommandsQueue.#PUB_SUB_MESSAGES.message.equals(reply[0])) {
+                    return RedisCommandsQueue.#emitPubSubMessage(
+                        this.#pubSubState.listeners.channels,
+                        reply[2],
+                        reply[1]
+                    );
+                } else if (RedisCommandsQueue.#PUB_SUB_MESSAGES.pMessage.equals(reply[0])) {
+                    return RedisCommandsQueue.#emitPubSubMessage(
+                        this.#pubSubState.listeners.patterns,
+                        reply[3],
+                        reply[2],
+                        reply[1]
+                    );
+                } else if (
+                    RedisCommandsQueue.#PUB_SUB_MESSAGES.subscribe.equals(reply[0]) ||
+                    RedisCommandsQueue.#PUB_SUB_MESSAGES.pSubscribe.equals(reply[0]) ||
+                    RedisCommandsQueue.#PUB_SUB_MESSAGES.unsubscribe.equals(reply[0]) ||
+                    RedisCommandsQueue.#PUB_SUB_MESSAGES.pUnsubscribe.equals(reply[0])
+                ) {
+                    if (--this.#waitingForReply.head!.value.channelsCounter! === 0) {
+                        this.#shiftWaitingForReply().resolve();
+                    }
+                    return;
                 }
             }
 
@@ -112,7 +144,7 @@ export default class RedisCommandsQueue {
     }
 
     addCommand<T = RedisCommandRawReply>(args: RedisCommandArguments, options?: QueueCommandOptions, bufferMode?: boolean): Promise<T> {
-        if (this.#pubSubState.subscribing || this.#pubSubState.subscribed) {
+        if (this.#pubSubState) {
             return Promise.reject(new Error('Cannot send commands in PubSub mode'));
         } else if (this.#maxLength && this.#waitingToBeSent.length + this.#waitingForReply.length >= this.#maxLength) {
             return Promise.reject(new Error('The queue is full'));
@@ -126,7 +158,7 @@ export default class RedisCommandsQueue {
                 chainId: options?.chainId,
                 bufferMode,
                 resolve,
-                reject,
+                reject
             });
 
             if (options?.signal) {
@@ -153,17 +185,41 @@ export default class RedisCommandsQueue {
         });
     }
 
-    subscribe(command: PubSubSubscribeCommands, channels: string | Array<string>, listener: PubSubListener): Promise<void> {
-        const channelsToSubscribe: Array<string> = [],
-            listeners = command === PubSubSubscribeCommands.SUBSCRIBE ? this.#pubSubListeners.channels : this.#pubSubListeners.patterns;
+    #initiatePubSubState(): PubSubState {
+        return this.#pubSubState ??= {
+            subscribed: 0,
+            subscribing: 0,
+            unsubscribing: 0,
+            listeners: {
+                channels: new Map(),
+                patterns: new Map()
+            }
+        };
+    }
+
+    subscribe<T extends boolean>(
+        command: PubSubSubscribeCommands,
+        channels: PubSubArgumentTypes | Array<PubSubArgumentTypes>,
+        listener: PubSubListener<T>,
+        bufferMode?: T
+    ): Promise<void> {
+        const pubSubState = this.#initiatePubSubState(),
+            channelsToSubscribe: Array<PubSubArgumentTypes> = [],
+            listenersMap = command === PubSubSubscribeCommands.SUBSCRIBE ? pubSubState.listeners.channels : pubSubState.listeners.patterns;
         for (const channel of (Array.isArray(channels) ? channels : [channels])) {
-            if (listeners.has(channel)) {
-                listeners.get(channel)!.add(listener);
-                continue;
+            const channelString = typeof channel === 'string' ? channel : channel.toString();
+            let listeners = listenersMap.get(channelString);
+            if (!listeners) {
+                listeners = {
+                    buffers: new Set(),
+                    strings: new Set()
+                };
+                listenersMap.set(channelString, listeners);
+                channelsToSubscribe.push(channel);
             }
 
-            listeners.set(channel, new Set([listener]));
-            channelsToSubscribe.push(channel);
+            // https://github.com/microsoft/TypeScript/issues/23132
+            (bufferMode ? listeners.buffers : listeners.strings).add(listener as any);
         }
 
         if (!channelsToSubscribe.length) {
@@ -173,8 +229,20 @@ export default class RedisCommandsQueue {
         return this.#pushPubSubCommand(command, channelsToSubscribe);
     }
 
-    unsubscribe(command: PubSubUnsubscribeCommands, channels?: string | Array<string>, listener?: PubSubListener): Promise<void> {
-        const listeners = command === PubSubUnsubscribeCommands.UNSUBSCRIBE ? this.#pubSubListeners.channels : this.#pubSubListeners.patterns;
+    unsubscribe<T extends boolean>(
+        command: PubSubUnsubscribeCommands,
+        channels?: string | Array<string>,
+        listener?: PubSubListener<T>,
+        bufferMode?: T
+    ): Promise<void> {
+        if (!this.#pubSubState) {
+            return Promise.resolve();
+        }
+
+        const listeners = command === PubSubUnsubscribeCommands.UNSUBSCRIBE ?
+            this.#pubSubState.listeners.channels :
+            this.#pubSubState.listeners.patterns;
+
         if (!channels) {
             const size = listeners.size;
             listeners.clear();
@@ -183,13 +251,16 @@ export default class RedisCommandsQueue {
 
         const channelsToUnsubscribe = [];
         for (const channel of (Array.isArray(channels) ? channels : [channels])) {
-            const set = listeners.get(channel);
-            if (!set) continue;
+            const sets = listeners.get(channel);
+            if (!sets) continue;
 
-            let shouldUnsubscribe = !listener;
+            let shouldUnsubscribe;
             if (listener) {
-                set.delete(listener);
-                shouldUnsubscribe = set.size === 0;
+                // https://github.com/microsoft/TypeScript/issues/23132
+                (bufferMode ? sets.buffers : sets.strings).delete(listener as any);
+                shouldUnsubscribe = !sets.buffers.size && !sets.strings.size;
+            } else {
+                shouldUnsubscribe = true;
             }
 
             if (shouldUnsubscribe) {
@@ -205,11 +276,12 @@ export default class RedisCommandsQueue {
         return this.#pushPubSubCommand(command, channelsToUnsubscribe);
     }
 
-    #pushPubSubCommand(command: PubSubSubscribeCommands | PubSubUnsubscribeCommands, channels: number | Array<string>): Promise<void> {
+    #pushPubSubCommand(command: PubSubSubscribeCommands | PubSubUnsubscribeCommands, channels: number | Array<PubSubArgumentTypes>): Promise<void> {
         return new Promise((resolve, reject) => {
-            const isSubscribe = command === PubSubSubscribeCommands.SUBSCRIBE || command === PubSubSubscribeCommands.PSUBSCRIBE,
+            const pubSubState = this.#initiatePubSubState(),
+                isSubscribe = command === PubSubSubscribeCommands.SUBSCRIBE || command === PubSubSubscribeCommands.PSUBSCRIBE,
                 inProgressKey = isSubscribe ? 'subscribing' : 'unsubscribing',
-                commandArgs: Array<string> = [command];
+                commandArgs: Array<PubSubArgumentTypes> = [command];
 
             let channelsCounter: number;
             if (typeof channels === 'number') { // unsubscribe only
@@ -219,35 +291,41 @@ export default class RedisCommandsQueue {
                 channelsCounter = channels.length;
             }
 
-            this.#pubSubState[inProgressKey] += channelsCounter;
+            pubSubState[inProgressKey] += channelsCounter;
 
             this.#waitingToBeSent.push({
                 args: commandArgs,
                 channelsCounter,
+                bufferMode: true,
                 resolve: () => {
-                    this.#pubSubState[inProgressKey] -= channelsCounter;
-                    this.#pubSubState.subscribed += channelsCounter * (isSubscribe ? 1 : -1);
+                    pubSubState[inProgressKey] -= channelsCounter;
+                    if (isSubscribe) {
+                        pubSubState.subscribed += channelsCounter;
+                    } else {
+                        pubSubState.subscribed -= channelsCounter;
+                        if (!pubSubState.subscribed && !pubSubState.subscribing && !pubSubState.subscribed) {
+                            this.#pubSubState = undefined;
+                        }
+                    }
                     resolve();
                 },
                 reject: () => {
-                    this.#pubSubState[inProgressKey] -= channelsCounter;
+                    pubSubState[inProgressKey] -= channelsCounter * (isSubscribe ? 1 : -1);
                     reject();
                 }
             });
         });
     }
 
     resubscribe(): Promise<any> | undefined {
-        if (!this.#pubSubState.subscribed && !this.#pubSubState.subscribing) {
+        if (!this.#pubSubState) {
             return;
         }
 
-        this.#pubSubState.subscribed = this.#pubSubState.subscribing = 0;
-
         // TODO: acl error on one channel/pattern will reject the whole command
         return Promise.all([
-            this.#pushPubSubCommand(PubSubSubscribeCommands.SUBSCRIBE, [...this.#pubSubListeners.channels.keys()]),
-            this.#pushPubSubCommand(PubSubSubscribeCommands.PSUBSCRIBE, [...this.#pubSubListeners.patterns.keys()])
+            this.#pushPubSubCommand(PubSubSubscribeCommands.SUBSCRIBE, [...this.#pubSubState.listeners.channels.keys()]),
+            this.#pushPubSubCommand(PubSubSubscribeCommands.PSUBSCRIBE, [...this.#pubSubState.listeners.patterns.keys()])
         ]);
     }
 
@@ -269,7 +347,10 @@ export default class RedisCommandsQueue {
     }
 
     parseResponse(data: Buffer): void {
-        this.#parser.setReturnBuffers(!!this.#waitingForReply.head?.value.bufferMode);
+        this.#parser.setReturnBuffers(
+            !!this.#waitingForReply.head?.value.bufferMode ||
+            !!this.#pubSubState?.subscribed
+        );
         this.#parser.execute(data);
     }
 

packages/client/lib/client/index.spec.ts

@@ -561,17 +561,27 @@ describe('Client', () => {
     }, GLOBAL.SERVERS.OPEN);
 
     testUtils.testWithClient('PubSub', async publisher => {
+        function assertStringListener(message: string, channel: string) {
+            assert.ok(typeof message === 'string');
+            assert.ok(typeof channel === 'string');
+        }
+
+        function assertBufferListener(message: Buffer, channel: Buffer) {
+            assert.ok(Buffer.isBuffer(message));
+            assert.ok(Buffer.isBuffer(channel));
+        }
+
         const subscriber = publisher.duplicate();
 
         await subscriber.connect();
 
         try {
-            const channelListener1 = spy(),
-                channelListener2 = spy(),
-                patternListener = spy();
+            const channelListener1 = spy(assertBufferListener),
+                channelListener2 = spy(assertStringListener),
+                patternListener = spy(assertStringListener);
 
             await Promise.all([
-                subscriber.subscribe('channel', channelListener1),
+                subscriber.subscribe('channel', channelListener1, true),
                 subscriber.subscribe('channel', channelListener2),
                 subscriber.pSubscribe('channel*', patternListener)
             ]);
@@ -580,14 +590,14 @@ describe('Client', () => {
                 waitTillBeenCalled(channelListener1),
                 waitTillBeenCalled(channelListener2),
                 waitTillBeenCalled(patternListener),
-                publisher.publish('channel', 'message')
+                publisher.publish(Buffer.from('channel'), Buffer.from('message'))
             ]);
 
-            assert.ok(channelListener1.calledOnceWithExactly('message', 'channel'));
+            assert.ok(channelListener1.calledOnceWithExactly(Buffer.from('message'), Buffer.from('channel')));
             assert.ok(channelListener2.calledOnceWithExactly('message', 'channel'));
             assert.ok(patternListener.calledOnceWithExactly('message', 'channel'));
 
-            await subscriber.unsubscribe('channel', channelListener1);
+            await subscriber.unsubscribe('channel', channelListener1, true);
             await Promise.all([
                 waitTillBeenCalled(channelListener2),
                 waitTillBeenCalled(patternListener),