handle server sunsubscribe

Author: leibale

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

@@ -3,7 +3,7 @@ import { AbortError, ErrorReply } from '../errors';
 import { RedisCommandArguments, RedisCommandRawReply } from '../commands';
 import RESP2Decoder from './RESP2/decoder';
 import encodeCommand from './RESP2/encoder';
-import { PubSub, PubSubListener, PubSubType } from './pub-sub';
+import { ChannelListeners, PubSub, PubSubCommand, PubSubListener, PubSubType } from './pub-sub';
 
 export interface QueueCommandOptions {
     asap?: boolean;
@@ -28,10 +28,10 @@ interface CommandWaitingForReply {
     returnBuffers?: boolean;
 }
 
-type PubSubCommand = ReturnType<typeof PubSub.prototype.subscribe | typeof PubSub.prototype.unsubscribe>;
-
 const PONG = Buffer.from('pong');
 
+type OnServerSUnsubscribe = (channel: string, listeners?: ChannelListeners) => void;
+
 export default class RedisCommandsQueue {
     static #flushQueue<T extends CommandWaitingForReply>(queue: LinkedList<T>, err: Error): void {
         while (queue.length) {
@@ -42,6 +42,7 @@ export default class RedisCommandsQueue {
     readonly #maxLength: number | null | undefined;
     readonly #waitingToBeSent = new LinkedList<CommandWaitingToBeSent>();
     readonly #waitingForReply = new LinkedList<CommandWaitingForReply>();
+    readonly #onServerSUnsubscribe: OnServerSUnsubscribe;
 
     readonly #pubSub = new PubSub();
 
@@ -55,7 +56,16 @@ export default class RedisCommandsQueue {
         onReply: reply => {
             if (this.#pubSub.isActive && Array.isArray(reply)) {
                 if (this.#pubSub.handleMessageReply(reply as Array<Buffer>)) return;
-                if (this.#pubSub.isStatusReply(reply as Array<Buffer>)) {
+                
+                const isShardedUnsubscribe = PubSub.isShardedUnsubscribe(reply as Array<Buffer>);
+                if (isShardedUnsubscribe && !this.#waitingForReply.length) {
+                    const channel = (reply[1] as Buffer).toString();
+                    this.#onServerSUnsubscribe(
+                        channel,
+                        this.#pubSub.removeShardedListeners(channel)
+                    );
+                    return;
+                } else if (isShardedUnsubscribe || PubSub.isStatusReply(reply as Array<Buffer>)) {
                     const head = this.#waitingForReply.head!.value;
                     if (
                         (Number.isNaN(head.channelsCounter!) && reply[2] === 0) ||
@@ -82,8 +92,12 @@ export default class RedisCommandsQueue {
         }
     });
 
-    constructor(maxLength: number | null | undefined) {
+    constructor(
+        maxLength: number | null | undefined,
+        onServerSUnsubscribe: OnServerSUnsubscribe
+    ) {
         this.#maxLength = maxLength;
+        this.#onServerSUnsubscribe = onServerSUnsubscribe;
     }
 
     addCommand<T = RedisCommandRawReply>(args: RedisCommandArguments, options?: QueueCommandOptions): Promise<T> {

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

@@ -8,6 +8,12 @@ import { defineScript } from '../lua-script';
 import { spy } from 'sinon';
 import { once } from 'events';
 import { ClientKillFilters } from '../commands/CLIENT_KILL';
+import { ClusterSlotStates } from '../commands/CLUSTER_SETSLOT';
+
+// We need to use 'require', because it's not possible with Typescript to import
+// function that are exported as 'module.exports = function`, without esModuleInterop
+// set to true.
+const calculateSlot = require('cluster-key-slot');
 
 export const SQUARE_SCRIPT = defineScript({
     SCRIPT: 'return ARGV[1] * ARGV[1];',
@@ -848,32 +854,65 @@ describe('Client', () => {
             disableClientSetup: true
         });
 
-        testUtils.testWithClient('shareded PubSub', async publisher => {
-            const subscriber = publisher.duplicate();
-
-            await subscriber.connect();
-
-            try {
-                const listener = spy();
-                await subscriber.sSubscribe('channel', listener);
-
-                await Promise.all([
-                    waitTillBeenCalled(listener),
-                    publisher.sPublish('channel', 'message')
-                ]);
+        describe('shareded PubSub', () => {
+            testUtils.isVersionGreaterThanHook([7]);
 
-                assert.ok(listener.calledOnceWithExactly('message', 'channel'));
+            testUtils.testWithClient('should be able to recive messages', async publisher => {
+                const subscriber = publisher.duplicate();
+    
+                await subscriber.connect();
+    
+                try {
+                    const listener = spy();
+                    await subscriber.sSubscribe('channel', listener);
+    
+                    await Promise.all([
+                        waitTillBeenCalled(listener),
+                        publisher.sPublish('channel', 'message')
+                    ]);
+    
+                    assert.ok(listener.calledOnceWithExactly('message', 'channel'));
+    
+                    await subscriber.sUnsubscribe();
+    
+                    // should be able to send commands
+                    await assert.doesNotReject(subscriber.ping());
+                } finally {
+                    await subscriber.disconnect();
+                }
+            }, {
+                ...GLOBAL.SERVERS.OPEN
+            });
 
-                await subscriber.sUnsubscribe();
+            testUtils.testWithClient('should handle sunsubscribe from server', async publisher => {
+                await publisher.clusterAddSlotsRange({ start: 0, end: 16383 });
 
-                // should be able to send commands
-                await assert.doesNotReject(subscriber.ping());
-            } finally {
-                await subscriber.disconnect();
-            }
-        }, {
-            ...GLOBAL.SERVERS.OPEN,
-            minimumDockerVersion: [7]
+                const subscriber = publisher.duplicate();
+    
+                await subscriber.connect();
+    
+                try {
+                    await subscriber.sSubscribe('channel', () => {});
+
+                    await Promise.all([
+                        publisher.clusterSetSlot(
+                            calculateSlot('channel'),
+                            ClusterSlotStates.NODE,
+                            await publisher.clusterMyId()
+                        ),
+                        once(subscriber, 'server-sunsubscribe')
+                    ]);
+                    
+                    assert.equal(
+                        await subscriber.ping(),
+                        'PONG'
+                    );
+                } finally {
+                    await subscriber.disconnect();
+                }
+            }, {
+                serverArguments: ['--cluster-enabled', 'yes']
+            });
         });
 
         testUtils.testWithClient('should handle errors in SUBSCRIBE', async publisher => {

packages/client/lib/client/index.ts

@@ -216,7 +216,10 @@ export default class RedisClient<
     }
 
     #initiateQueue(): RedisCommandsQueue {
-        return new RedisCommandsQueue(this.#options?.commandsQueueMaxLength);
+        return new RedisCommandsQueue(
+            this.#options?.commandsQueueMaxLength,
+            (channel, listeners) => this.emit('server-sunsubscribe', channel, listeners)
+        );
     }
 
     #initiateSocket(): RedisSocket {

packages/client/lib/client/pub-sub.ts

@@ -28,15 +28,31 @@ export type PubSubListener<
     RETURN_BUFFERS extends boolean = false
 > = <T extends RETURN_BUFFERS extends true ? Buffer : string>(message: T, channel: T) => unknown;
 
-interface ChannelListeners {
+export interface ChannelListeners {
     unsubscribing: boolean;
     buffers: Set<PubSubListener<true>>;
     strings: Set<PubSubListener<false>>;
 }
 
 type Listeners = Record<PubSubType, Map<string, ChannelListeners>>;
 
+export type PubSubCommand = ReturnType<typeof PubSub.prototype.subscribe | typeof PubSub.prototype.unsubscribe>;
+
 export class PubSub {
+    static isStatusReply(reply: Array<Buffer>): boolean {
+        return (
+            COMMANDS[PubSubType.CHANNELS].subscribe.equals(reply[0]) ||
+            COMMANDS[PubSubType.CHANNELS].unsubscribe.equals(reply[0]) ||
+            COMMANDS[PubSubType.PATTERNS].subscribe.equals(reply[0]) ||
+            COMMANDS[PubSubType.PATTERNS].unsubscribe.equals(reply[0]) ||
+            COMMANDS[PubSubType.SHARDED].subscribe.equals(reply[0])
+        );
+    }
+
+    static isShardedUnsubscribe(reply: Array<Buffer>): boolean {
+        return COMMANDS[PubSubType.SHARDED].unsubscribe.equals(reply[0]);
+    }
+    
     static #channelsArray(channels: string | Array<string>) {
         return (Array.isArray(channels) ? channels : [channels]);
     }
@@ -227,7 +243,7 @@ export class PubSub {
         this.#subscribing = 0;
     }
 
-    resubscribe() {
+    resubscribe(): Array<PubSubCommand> {
         const commands = [];
         for (const [type, listeners] of Object.entries(this.#listeners)) {
             if (!listeners.size) continue;
@@ -276,6 +292,15 @@ export class PubSub {
 
         return false;
     }
+
+    removeShardedListeners(channel: string): ChannelListeners | undefined {
+        const listeners = this.#listeners[PubSubType.SHARDED].get(channel);
+        if (listeners) {
+            this.#listeners[PubSubType.SHARDED].delete(channel);
+            this.#updateIsActive();
+            return listeners;
+        }
+    }
 
     #emitPubSubMessage(
         type: PubSubType,
@@ -304,15 +329,4 @@ export class PubSub {
             listener(messageString, channelString);
         }
     }
-
-    isStatusReply(reply: Array<Buffer>): boolean {
-        return (
-            COMMANDS[PubSubType.CHANNELS].subscribe.equals(reply[0]) ||
-            COMMANDS[PubSubType.CHANNELS].unsubscribe.equals(reply[0]) ||
-            COMMANDS[PubSubType.PATTERNS].subscribe.equals(reply[0]) ||
-            COMMANDS[PubSubType.PATTERNS].unsubscribe.equals(reply[0]) ||
-            COMMANDS[PubSubType.SHARDED].subscribe.equals(reply[0]) ||
-            COMMANDS[PubSubType.SHARDED].unsubscribe.equals(reply[0])
-        );
-    }
 }