fix(cluster): keep batch commands routed by slot during migration

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

@@ -29,6 +29,146 @@ describe('RedisCommandsQueue', () => {
     });
   });
 
+  describe('chainInExecution', () => {
+    it('is undefined before anything is written, and matches a chain once part of it is sent', () => {
+      const queue = createQueue();
+      const chainId = Symbol('MULTI Chain');
+
+      assert.strictEqual(queue.chainInExecution, undefined);
+
+      queue.addCommand(['MULTI'], { chainId }).catch(() => {});
+      queue.addCommand(['SET', 'k', 'v'], { chainId }).catch(() => {});
+      queue.addCommand(['EXEC'], { chainId }).catch(() => {});
+
+      // Nothing sent yet - the whole chain is still just sitting in #toWrite.
+      assert.strictEqual(queue.chainInExecution, undefined);
+
+      const writer = queue.commandsToWrite();
+      writer.next(); // "sends" MULTI
+      writer.next(); // "sends" SET
+
+      // Two of the three commands are now in #waitingForReply (out of view);
+      // chainInExecution should point at this chain, and the one command
+      // still in #toWrite (EXEC) should carry a matching chainId - that's
+      // the tail-of-an-in-flight-chain signal cluster-slots.ts relies on.
+      assert.strictEqual(queue.chainInExecution, chainId);
+
+      const remaining = queue.extractAllCommands();
+      assert.strictEqual(remaining.length, 1);
+      assert.strictEqual(remaining[0].chainId, queue.chainInExecution);
+    });
+
+    it('leaves a finished chain with no queued tail to relocate', () => {
+      const queue = createQueue();
+      const chainId = Symbol('MULTI Chain');
+
+      queue.addCommand(['MULTI'], { chainId }).catch(() => {});
+      queue.addCommand(['EXEC'], { chainId }).catch(() => {});
+
+      const writer = queue.commandsToWrite();
+      writer.next(); // "sends" MULTI
+      writer.next(); // "sends" EXEC - whole chain is now in #waitingForReply
+
+      // The whole chain was sent, so nothing of it remains in #toWrite -
+      // extractAllCommands has nothing left to (mis)classify as its tail.
+      assert.strictEqual(queue.chainInExecution, chainId);
+      assert.strictEqual(queue.extractAllCommands().length, 0);
+    });
+  });
+
+  describe('extractCommandsForSlots', () => {
+    it('leaves the queued tail of an in-flight chain and everything after it in place', () => {
+      const queue = createQueue();
+      const chainId = Symbol('MULTI Chain');
+
+      queue.addCommand(['MULTI'], { chainId, slotNumber: 1 }).catch(() => {});
+      queue.addCommand(['SET', 'k', 'v'], { chainId, slotNumber: 1 }).catch(() => {});
+      queue.addCommand(['EXEC'], { chainId, slotNumber: 1 }).catch(() => {});
+      // Queued after the chain, on the same connection. If this were
+      // relocated to another node while the transaction is still pending
+      // here, it could run before the transaction completes - reading 'k'
+      // before SET applies, even though it was queued after EXEC.
+      queue.addCommand(['GET', 'k'], { slotNumber: 1 }).catch(() => {});
+
+      const writer = queue.commandsToWrite();
+      writer.next(); // "sends" MULTI - SET, EXEC and GET are still queued behind it
+
+      const extracted = queue.extractCommandsForSlots(new Set([1]));
+
+      // Nothing is extracted: reaching the in-flight chain's tail stops the
+      // scan entirely, so GET stays behind it in queue order too.
+      assert.deepStrictEqual(extracted, []);
+      assert.strictEqual(queue.pendingCount, 4);
+    });
+
+    it('is unaffected by an already-fully-sent chain', () => {
+      const queue = createQueue();
+      const chainId = Symbol('MULTI Chain');
+
+      queue.addCommand(['MULTI'], { chainId, slotNumber: 1 }).catch(() => {});
+      queue.addCommand(['EXEC'], { chainId, slotNumber: 1 }).catch(() => {});
+      queue.addCommand(['GET', 'k'], { slotNumber: 1 }).catch(() => {});
+
+      const writer = queue.commandsToWrite();
+      writer.next(); // "sends" MULTI
+      writer.next(); // "sends" EXEC - the whole chain is now in #waitingForReply
+
+      const extracted = queue.extractCommandsForSlots(new Set([1]));
+
+      // chainInExecution still points at this chain (nothing has been sent
+      // since), but none of its commands remain in #toWrite to (mis)match -
+      // the trailing GET is extracted normally.
+      assert.strictEqual(queue.chainInExecution, chainId);
+      assert.deepStrictEqual(
+        extracted.map(command => command.args?.[0]),
+        ['GET'],
+      );
+    });
+
+    it('extracts a second, fully-queued chain atomically without mistaking it for the in-flight chain\'s tail', () => {
+      const queue = createQueue();
+      const chainA = Symbol('Chain A (in-flight, different slot)');
+      const chainB = Symbol('Chain B (fully queued, migrating slot)');
+
+      // Chain A is in flight on a slot that isn't migrating - its queued
+      // tail (SET, EXEC) carries a chainId that will equal chainInExecution.
+      queue.addCommand(['MULTI'], { chainId: chainA, slotNumber: 5 }).catch(() => {});
+      queue.addCommand(['SET', 'a', '1'], { chainId: chainA, slotNumber: 5 }).catch(() => {});
+      queue.addCommand(['EXEC'], { chainId: chainA, slotNumber: 5 }).catch(() => {});
+
+      // Chain B is queued entirely behind chain A, on the slot that IS
+      // migrating. None of its commands have been sent, so none of them
+      // carry chainInExecution's id - it's a distinct chain, not a
+      // continuation of chain A's tail.
+      queue.addCommand(['MULTI'], { chainId: chainB, slotNumber: 1 }).catch(() => {});
+      queue.addCommand(['SET', 'b', '2'], { chainId: chainB, slotNumber: 1 }).catch(() => {});
+      queue.addCommand(['EXEC'], { chainId: chainB, slotNumber: 1 }).catch(() => {});
+
+      const writer = queue.commandsToWrite();
+      writer.next(); // "sends" chain A's MULTI - chainInExecution now points at chain A
+
+      const extracted = queue.extractCommandsForSlots(new Set([1]));
+
+      // Chain A's tail lives on slot 5, so it's skipped over (not in the
+      // migrating slot set) without ever triggering the in-flight-tail
+      // guard. Chain B, entirely on slot 1, is then extracted as a whole -
+      // its distinct chainId never matches chainInExecution, so it's never
+      // mistaken for chain A's tail and relocates atomically, in order.
+      assert.deepStrictEqual(
+        extracted.map(command => command.args?.[0]),
+        ['MULTI', 'SET', 'EXEC'],
+      );
+      assert.ok(extracted.every(command => command.chainId === chainB));
+
+      // Chain A's tail stays behind, untouched, on its own slot.
+      const remaining = queue.extractAllCommands();
+      assert.deepStrictEqual(
+        remaining.map(command => command.args?.[0]),
+        ['SET', 'EXEC'],
+      );
+    });
+  });
+
   describe('addCommand', () => {
     it('does not keep a command if timeout listener setup fails', async () => {
       const queue = createQueue();
@@ -72,6 +212,49 @@ describe('RedisCommandsQueue', () => {
     });
   });
 
+  describe('full node loss (in-flight chain)', () => {
+    // Mirrors cluster-slots.ts's full-node-loss path: extractAllCommands()
+    // pulls the in-flight chain's queued tail out of #toWrite and rejects it
+    // explicitly, then the source client is destroyed, which calls
+    // flushAll() and rejects whatever's left in #waitingForReply - the
+    // chain's already-sent head. Both halves need to reject; otherwise the
+    // transaction half-commits: the server already applied the head against
+    // a connection this client is abandoning, while the tail is rejected as
+    // never sent.
+    it('rejects the already-sent head of an in-flight chain along with its queued tail', async () => {
+      const queue = createQueue();
+      const chainId = Symbol('MULTI Chain');
+
+      const multiPromise = queue.addCommand(['MULTI'], { chainId });
+      const setPromise = queue.addCommand(['SET', 'k', 'v'], { chainId });
+      const execPromise = queue.addCommand(['EXEC'], { chainId });
+      [multiPromise, setPromise, execPromise].forEach(promise => promise.catch(() => {}));
+
+      const writer = queue.commandsToWrite();
+      writer.next(); // "sends" MULTI - now in #waitingForReply
+      writer.next(); // "sends" SET - now in #waitingForReply
+
+      // EXEC is still the queued tail in #toWrite.
+      const remaining = queue.extractAllCommands();
+      assert.strictEqual(remaining.length, 1);
+      assert.strictEqual(remaining[0].chainId, queue.chainInExecution);
+
+      // cluster-slots.ts rejects the queued tail explicitly, since it can't
+      // be safely relocated to another node...
+      queue.rejectCommands(remaining, new DisconnectsClientError());
+
+      // ...then destroy()'s call to flushAll() rejects whatever's left in
+      // #waitingForReply - the already-sent MULTI and SET.
+      queue.flushAll(new DisconnectsClientError());
+
+      await Promise.all([
+        assert.rejects(multiPromise, DisconnectsClientError),
+        assert.rejects(setPromise, DisconnectsClientError),
+        assert.rejects(execPromise, DisconnectsClientError),
+      ]);
+    });
+  });
+
   describe('prependCommandsToWrite', () => {
 
     it('rebinds an abort listener so it removes the command from its new queue', async () => {

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

@@ -144,6 +144,18 @@ export default class RedisCommandsQueue {
     return this.#toWrite.length + this.#waitingForReply.length;
   }
 
+  /**
+   * The chainId of the most recently written command, if it belonged to a
+   * MULTI/pipeline chain. While this is set, any command still in `#toWrite`
+   * with a matching `chainId` is the tail of a chain that's already partially
+   * sent to the server - the rest of that chain is in `#waitingForReply` and
+   * out of view, so this tail must not be treated as a complete, relocatable
+   * chain (see `flushWaitingForReply`, which rejects it instead of resending it).
+   */
+  get chainInExecution(): symbol | undefined {
+    return this.#chainInExecution;
+  }
+
   constructor(
     respVersion: RespVersions,
     maxLength: number | null | undefined,
@@ -691,30 +703,45 @@ export default class RedisCommandsQueue {
    * Extracts commands for the given slots from the toWrite queue.
    * Some commands don't have "slotNumber", which means they are not designated to particular slot/node.
    * We ignore those.
+   *
+   * A command whose chainId matches `chainInExecution` is the queued tail of a
+   * chain (MULTI/pipeline) whose head was already written to this socket - the
+   * rest is in `#waitingForReply`, out of view. Relocating just the tail would
+   * split the transaction across two connections, and this connection isn't
+   * going away (unlike a full node loss), so extraction stops there entirely:
+   * that command and everything still queued behind it stay on this
+   * connection, preserving the order they were originally sent in, and get
+   * sent here as originally queued. If a key has by then moved to another
+   * node, that surfaces as a normal MOVED/ASK error through this method's
+   * caller - unlike single commands, MULTI/pipeline execution doesn't go
+   * through cluster's redirect-and-retry loop, so this isn't automatically
+   * retried, but that's an existing, separate limitation this method isn't
+   * responsible for preventing.
    */
   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)
-      ) {
-        const command = current.value;
-        if (command.abort) {
-          RedisCommandsQueue.#removeAbortListener(command);
-        }
-        if (command.timeout) {
-          RedisCommandsQueue.#removeTimeoutListener(command);
-        }
-        result.push(command);
-        const toRemove = current;
-        current = current.next;
-        this.#toWrite.remove(toRemove);
-      } else {
-        // Move to next node even if we don't extract this command
+      if (current.value.slotNumber === undefined || !slots.has(current.value.slotNumber)) {
         current = current.next;
+        continue;
+      }
+
+      if (this.#chainInExecution !== undefined && current.value.chainId === this.#chainInExecution) {
+        break;
       }
+
+      const command = current.value;
+      if (command.abort) {
+        RedisCommandsQueue.#removeAbortListener(command);
+      }
+      if (command.timeout) {
+        RedisCommandsQueue.#removeTimeoutListener(command);
+      }
+      result.push(command);
+      const toRemove = current;
+      current = current.next;
+      this.#toWrite.remove(toRemove);
     }
     return result;
   }

packages/client/lib/client/index.ts

@@ -1508,7 +1508,8 @@ export default class RedisClient<
    */
   async _executePipeline(
     commands: Array<RedisMultiQueuedCommand>,
-    selectedDB?: number
+    selectedDB?: number,
+    slotNumber?: number
   ) {
     if (!this._self.#socket.isOpen) {
       return Promise.reject(new ClientClosedError());
@@ -1524,7 +1525,8 @@ export default class RedisClient<
             const traced = trace(CHANNELS.TRACE_COMMAND,
               () => this._self.#queue.addCommand(args, {
                 chainId,
-                typeMapping: this._commandOptions?.typeMapping
+                typeMapping: this._commandOptions?.typeMapping,
+                slotNumber
               }),
               () => ({
                 ...this._self.#commandTraceContext(args),
@@ -1564,7 +1566,8 @@ export default class RedisClient<
    */
   async _executeMulti(
     commands: Array<RedisMultiQueuedCommand>,
-    selectedDB?: number
+    selectedDB?: number,
+    slotNumber?: number
   ) {
     const dirtyWatch = this._self.#dirtyWatch;
     this._self.#dirtyWatch = undefined;
@@ -1590,20 +1593,21 @@ export default class RedisClient<
         const typeMapping = this._commandOptions?.typeMapping;
         const chainId = Symbol('MULTI Chain');
         const promises: Array<Promise<unknown>> = [
-          this._self.#queue.addCommand(['MULTI'], { chainId }),
+          this._self.#queue.addCommand(['MULTI'], { chainId, slotNumber }),
         ];
 
         for (const { args } of commands) {
           promises.push(
             this._self.#queue.addCommand(args, {
               chainId,
-              typeMapping
+              typeMapping,
+              slotNumber
             })
           );
         }
 
         promises.push(
-          this._self.#queue.addCommand(['EXEC'], { chainId })
+          this._self.#queue.addCommand(['EXEC'], { chainId, slotNumber })
         );
 
         this._self.#scheduleWrite();

packages/client/lib/cluster/batch-routing.spec.ts

@@ -0,0 +1,60 @@
+import { strict as assert } from 'node:assert';
+import RedisCluster from '.';
+
+describe('Cluster batch routing', () => {
+  it('passes the selected slot to MULTI execution', async () => {
+    const cluster = RedisCluster.create({ rootNodes: [] });
+    let capturedSlotNumber: number | undefined;
+
+    const fakeClient = {
+      _executeMulti: async (_commands: unknown, _selectedDB?: number, slotNumber?: number) => {
+        capturedSlotNumber = slotNumber;
+        return ['OK'];
+      }
+    };
+
+    cluster._slots.getClientAndSlotNumber = async () => ({
+      client: fakeClient as never,
+      slotNumber: 123
+    });
+
+    assert.deepEqual(
+      await cluster.multi()
+        .sendCommand(['SET', 'key', 'value'], {
+          firstKey: 'key',
+          isReadonly: false
+        })
+        .exec(),
+      ['OK']
+    );
+    assert.equal(capturedSlotNumber, 123);
+  });
+
+  it('passes the selected slot to pipeline execution', async () => {
+    const cluster = RedisCluster.create({ rootNodes: [] });
+    let capturedSlotNumber: number | undefined;
+
+    const fakeClient = {
+      _executePipeline: async (_commands: unknown, _selectedDB?: number, slotNumber?: number) => {
+        capturedSlotNumber = slotNumber;
+        return ['OK'];
+      }
+    };
+
+    cluster._slots.getClientAndSlotNumber = async () => ({
+      client: fakeClient as never,
+      slotNumber: 456
+    });
+
+    assert.deepEqual(
+      await cluster.multi()
+        .sendCommand(['SET', 'key', 'value'], {
+          firstKey: 'key',
+          isReadonly: false
+        })
+        .execAsPipeline(),
+      ['OK']
+    );
+    assert.equal(capturedSlotNumber, 456);
+  });
+});