Hitless upgrades

.gitignore

@@ -8,3 +8,4 @@ node_modules/
 dump.rdb
 documentation/
 tsconfig.tsbuildinfo
+*.log

packages/client/lib/RESP/decoder.ts

@@ -98,7 +98,7 @@ export class Decoder {
     this.#next = undefined;
     return this.#decodeTypeValue(type, chunk);
   }
-    
+
   #decodeTypeValue(type, chunk) {
     switch (type) {
       case RESP_TYPES.NULL:
@@ -128,7 +128,7 @@ export class Decoder {
             chunk
           )
         );
-      
+
       case RESP_TYPES.DOUBLE:
         return this.#handleDecodedValue(
           this.onReply,
@@ -137,7 +137,7 @@ export class Decoder {
             chunk
           )
         );
-      
+
       case RESP_TYPES.SIMPLE_STRING:
         return this.#handleDecodedValue(
           this.onReply,
@@ -146,7 +146,7 @@ export class Decoder {
             chunk
           )
         );
-      
+
       case RESP_TYPES.BLOB_STRING:
         return this.#handleDecodedValue(
           this.onReply,
@@ -170,7 +170,7 @@ export class Decoder {
           this.onErrorReply,
           this.#decodeSimpleError(chunk)
         );
-      
+
       case RESP_TYPES.BLOB_ERROR:
         return this.#handleDecodedValue(
           this.onErrorReply,
@@ -188,7 +188,7 @@ export class Decoder {
           this.onReply,
           this.#decodeSet(this.getTypeMapping(), chunk)
         );
-      
+
       case RESP_TYPES.MAP:
         return this.#handleDecodedValue(
           this.onReply,
@@ -421,17 +421,17 @@ export class Decoder {
           return this.#cursor === chunk.length ?
             this.#decodeDoubleExponent.bind(this, d) :
             this.#decodeDoubleExponent(d, chunk);
-        
+
         case ASCII['\r']:
           this.#cursor = cursor + 2; // skip \r\n
           return isNegative ? -double : double;
       }
-      
+
       if (decimalIndex < Decoder.#DOUBLE_DECIMAL_MULTIPLIERS.length) {
         double += (byte - ASCII['0']) * Decoder.#DOUBLE_DECIMAL_MULTIPLIERS[decimalIndex++];
       }
     } while (++cursor < chunk.length);
-    
+
     this.#cursor = cursor;
     return this.#decodeDoubleDecimal.bind(this, isNegative, decimalIndex, double);
   }
@@ -613,7 +613,7 @@ export class Decoder {
   }
 
   #decodeVerbatimStringFormat(stringLength, chunk) {
-    const formatCb = this.#decodeStringWithLength.bind(this, 3, 1, String); 
+    const formatCb = this.#decodeStringWithLength.bind(this, 3, 1, String);
     return this.#cursor >= chunk.length ?
       this.#continueDecodeVerbatimStringFormat.bind(this, stringLength, formatCb) :
       this.#continueDecodeVerbatimStringFormat(stringLength, formatCb, chunk);
@@ -689,13 +689,13 @@ export class Decoder {
 
       case RESP_TYPES.BIG_NUMBER:
         return this.#decodeBigNumber(typeMapping[RESP_TYPES.BIG_NUMBER], chunk);
-      
+
       case RESP_TYPES.DOUBLE:
         return this.#decodeDouble(typeMapping[RESP_TYPES.DOUBLE], chunk);
-      
+
       case RESP_TYPES.SIMPLE_STRING:
         return this.#decodeSimpleString(typeMapping[RESP_TYPES.SIMPLE_STRING], chunk);
-      
+
       case RESP_TYPES.BLOB_STRING:
         return this.#decodeBlobString(typeMapping[RESP_TYPES.BLOB_STRING], chunk);
 
@@ -704,7 +704,7 @@ export class Decoder {
 
       case RESP_TYPES.SIMPLE_ERROR:
         return this.#decodeSimpleError(chunk);
-      
+
       case RESP_TYPES.BLOB_ERROR:
         return this.#decodeBlobError(chunk);
 
@@ -713,7 +713,7 @@ export class Decoder {
 
       case RESP_TYPES.SET:
         return this.#decodeSet(typeMapping, chunk);
-      
+
       case RESP_TYPES.MAP:
         return this.#decodeMap(typeMapping, chunk);
 
@@ -997,7 +997,7 @@ export class Decoder {
       // decode simple string map key as string (and not as buffer)
       case RESP_TYPES.SIMPLE_STRING:
         return this.#decodeSimpleString(String, chunk);
-      
+
       // decode blob string map key as string (and not as buffer)
       case RESP_TYPES.BLOB_STRING:
         return this.#decodeBlobString(String, chunk);
@@ -1028,7 +1028,7 @@ export class Decoder {
         this.#decodeNestedType.bind(this, typeMapping),
         typeMapping
       );
-    }      
+    }
 
     const value = this.#decodeNestedType(typeMapping, chunk);
     if (typeof value === 'function') {

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

@@ -19,6 +19,11 @@ export interface CommandOptions<T = TypeMapping> {
    * Timeout for the command in milliseconds
    */
   timeout?: number;
+  /**
+   * @internal
+   * The slot the command is targeted to (if any)
+   */
+  slotNumber?: number;
 }
 
 export interface CommandToWrite extends CommandWaitingForReply {
@@ -33,6 +38,7 @@ export interface CommandToWrite extends CommandWaitingForReply {
     listener: () => unknown;
     originalTimeout: number | undefined;
   } | undefined;
+  slotNumber?: number
 }
 
 interface CommandWaitingForReply {
@@ -186,14 +192,34 @@ export default class RedisCommandsQueue {
     this.#pushHandlers.push(handler);
   }
 
-  async waitForInflightCommandsToComplete(): Promise<void> {
+  async waitForInflightCommandsToComplete(options?: { timeoutMs?: number, flushOnTimeout?: boolean }): Promise<void> {
     // In-flight commands already completed
     if(this.#waitingForReply.length === 0) {
       return
     };
     // Otherwise wait for in-flight commands to fire `empty` event
     return new Promise(resolve => {
-      this.#waitingForReply.events.on('empty', resolve)
+      const onEmpty = () => {
+        if (timeoutId) clearTimeout(timeoutId);
+        resolve();
+      };
+
+      let timeoutId: ReturnType<typeof setTimeout> | undefined;
+      const timeoutMs = options?.timeoutMs;
+      if (timeoutMs !== undefined && timeoutMs > 0) {
+        timeoutId = setTimeout(() => {
+          this.#waitingForReply.events.off('empty', onEmpty);
+          const pendingCount = this.#waitingForReply.length;
+          dbgMaintenance(`waitForInflightCommandsToComplete timed out after ${timeoutMs}ms with ${pendingCount} commands still waiting`);
+          if (options?.flushOnTimeout && pendingCount > 0) {
+            dbgMaintenance(`Flushing ${pendingCount} commands that timed out waiting for reply`);
+            this.#flushWaitingForReply(new TimeoutError());
+          }
+          resolve(); // Resolve instead of reject - we don't want to fail the migration
+        }, timeoutMs);
+      }
+
+      this.#waitingForReply.events.once('empty', onEmpty);
     });
   }
 
@@ -219,6 +245,7 @@ export default class RedisCommandsQueue {
         channelsCounter: undefined,
         typeMapping: options?.typeMapping
       };
+      value.slotNumber = options?.slotNumber
 
       // If #maintenanceCommandTimeout was explicitly set, we should
       // use it instead of the timeout provided by the command
@@ -283,7 +310,8 @@ export default class RedisCommandsQueue {
       if (Array.isArray(reply)) {
         if (this.#onPush(reply)) return;
 
-        if (PONG.equals(reply[0] as Buffer)) {
+        const firstElement = typeof reply[0] === 'string' ? Buffer.from(reply[0]) : reply[0];
+        if (PONG.equals(firstElement as Buffer)) {
           const { resolve, typeMapping } = this.#waitingForReply.shift()!,
             buffer = ((reply[1] as Buffer).length === 0 ? reply[0] : reply[1]) as Buffer;
           resolve(typeMapping?.[RESP_TYPES.SIMPLE_STRING] === Buffer ? buffer : buffer.toString());
@@ -342,6 +370,10 @@ export default class RedisCommandsQueue {
     return this.#pubSub.removeAllListeners();
   }
 
+  removeShardedPubSubListenersForSlots(slots: Set<number>) {
+    return this.#pubSub.removeShardedPubSubListenersForSlots(slots);
+  }
+
   resubscribe(chainId?: symbol) {
     const commands = this.#pubSub.resubscribe();
     if (!commands.length) return;
@@ -541,4 +573,54 @@ 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);
+      } else {
+        // Move to next node even if we don't extract this command
+        current = current.next;
+      }
+    }
+    return result;
+  }
+
+  /**
+  * Gets all commands from the write queue without removing them.
+  */
+  extractAllCommands(): CommandToWrite[] {
+    const result: CommandToWrite[] = [];
+    let current = this.#toWrite.head;
+    while(current) {
+      result.push(current.value);
+      this.#toWrite.remove(current);
+      current = current.next;
+    }
+    return result;
+  }
+
+  /**
+   * Prepends commands to the write queue in reverse.
+   */
+  prependCommandsToWrite(commands: CommandToWrite[]) {
+    if (!commands.length) {
+      return;
+    }
+
+    for (let i = commands.length - 1; i >= 0; i--) {
+      this.#toWrite.unshift(commands[i]);
+    }
+  }
 }