add support for extracting commands from queue

Author: nkaradzhov

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

@@ -349,6 +349,10 @@ export default class RedisCommandsQueue {
     return this.#pubSub.removeAllListeners();
   }
 
+  removePubSubListenersForSlots(slots: Set<number>) {
+    return this.#pubSub.removePubSubListenersForSlots(slots);
+  }
+
   resubscribe(chainId?: symbol) {
     const commands = this.#pubSub.resubscribe();
     if (!commands.length) return;
@@ -548,4 +552,46 @@ export default class RedisCommandsQueue {
       this.#waitingForReply.length === 0
     );
   }
+
+  /**
+   *
+   * Extracts commands for the given slots from the toWrite queue.
+   * Some commands dont have "slotNumber", which means they are not designated to particular slot/node.
+   * We ignore those.
+   */
+  extractCommandsForSlots(slots: Set<number>): CommandToWrite[] {
+    const result: CommandToWrite[] = [];
+    let current = this.#toWrite.head;
+    while(current !== undefined) {
+      if(current.value.slotNumber !== undefined && slots.has(current.value.slotNumber)) {
+        result.push(current.value);
+        const toRemove = current;
+        current = current.next;
+        this.#toWrite.remove(toRemove);
+      }
+    }
+    return result;
+  }
+
+  /**
+  * Gets all commands from the write queue without removing them.
+  */
+  getAllCommands(): CommandToWrite[] {
+    const result: CommandToWrite[] = [];
+    let current = this.#toWrite.head;
+    while(current) {
+      result.push(current.value);
+      current = current.next;
+    }
+    return result;
+  }
+
+  /**
+   * Prepends commands to the write queue in reverse.
+   */
+  prependCommandsToWrite(commands: CommandToWrite[]) {
+    for (let i = commands.length - 1; i <= 0; i--) {
+      this.#toWrite.unshift(commands[i]);
+    }
+  }
 }

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

@@ -1,5 +1,6 @@
 import { RedisArgument } from '../RESP/types';
 import { CommandToWrite } from './commands-queue';
+import calculateSlot from 'cluster-key-slot';
 
 export const PUBSUB_TYPE = {
   CHANNELS: 'CHANNELS',
@@ -420,6 +421,31 @@ export class PubSub {
     return result;
   }
 
+  removePubSubListenersForSlots(slots: Set<number>) {
+    const channels = new Map<string, ChannelListeners>();
+    for (const [channel, value] of this.listeners[PUBSUB_TYPE.CHANNELS]) {
+      if (slots.has(calculateSlot(channel))) {
+        channels.set(channel, value);
+        this.listeners[PUBSUB_TYPE.CHANNELS].delete(channel);
+      }
+    }
+
+    const sharded = new Map<string, ChannelListeners>();
+    for (const [chanel, value] of this.listeners[PUBSUB_TYPE.SHARDED]) {
+      if (slots.has(calculateSlot(chanel))) {
+        sharded.set(chanel, value);
+        this.listeners[PUBSUB_TYPE.SHARDED].delete(chanel);
+      }
+    }
+
+    this.#updateIsActive();
+
+    return {
+      [PUBSUB_TYPE.CHANNELS]: channels,
+      [PUBSUB_TYPE.SHARDED]: sharded
+    };
+  }
+
   #emitPubSubMessage(
     type: PubSubType,
     message: Buffer,