implement queue drain mechanism

Author: nkaradzhov

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

@@ -1,10 +1,11 @@
-import { SinglyLinkedList, DoublyLinkedNode, DoublyLinkedList } from './linked-list';
+import { SinglyLinkedList, DoublyLinkedNode, DoublyLinkedList, makeEmptyAware } from './linked-list';
 import encodeCommand from '../RESP/encoder';
 import { Decoder, PUSH_TYPE_MAPPING, RESP_TYPES } from '../RESP/decoder';
 import { TypeMapping, ReplyUnion, RespVersions, RedisArgument } from '../RESP/types';
 import { ChannelListeners, PubSub, PubSubCommand, PubSubListener, PubSubType, PubSubTypeListeners } from './pub-sub';
 import { AbortError, ErrorReply, TimeoutError } from '../errors';
 import { MonitorCallback } from '.';
+import EventEmitter from 'events';
 
 export interface CommandOptions<T = TypeMapping> {
   chainId?: symbol;
@@ -54,18 +55,18 @@ export default class RedisCommandsQueue {
   readonly #respVersion;
   readonly #maxLength;
   readonly #toWrite = new DoublyLinkedList<CommandToWrite>();
-  readonly #waitingForReply = new SinglyLinkedList<CommandWaitingForReply>();
+  readonly #waitingForReply: SinglyLinkedList<CommandWaitingForReply>;
   readonly #onShardedChannelMoved;
   #chainInExecution: symbol | undefined;
   readonly decoder;
   readonly #pubSub = new PubSub();
+  readonly events = new EventEmitter();
 
   get isPubSubActive() {
     return this.#pubSub.isActive;
   }
 
   #invalidateCallback?: (key: RedisArgument | null) => unknown;
-  #movingCallback?: (afterMs: number, host: string, port: number) => void;
 
   constructor(
     respVersion: RespVersions,
@@ -76,6 +77,9 @@ export default class RedisCommandsQueue {
     this.#maxLength = maxLength;
     this.#onShardedChannelMoved = onShardedChannelMoved;
     this.decoder = this.#initiateDecoder();
+    const [waitingForReply, emptyEmitter] = makeEmptyAware(new SinglyLinkedList<CommandWaitingForReply>())
+    this.#waitingForReply = waitingForReply;
+    emptyEmitter.on('empty', this.events.on.bind(this.events, 'waitingForReplyEmpty'))
   }
 
   #onReply(reply: ReplyUnion) {
@@ -137,16 +141,10 @@ export default class RedisCommandsQueue {
               break;
             }
             case 'MOVING': {
-              if (this.#movingCallback) {
-                console.log('received moving', push)
-                const [_, afterMs, url] = push;
-                let [host, port] = url.toString().split(':');
-                //['18.200.246.58'] - for some reason the server sends the host this way
-                if(host.includes('[')) {
-                  host = host.slice(2, -2);
-                }
-                this.#movingCallback(afterMs, host, Number(port));
-              }
+              console.log('received moving', push)
+              const [_, afterMs, url] = push;
+              const [host, port] = url.toString().split(':');
+              this.events.emit('moving', afterMs, host, Number(port))
               break;
             }
           }
@@ -160,8 +158,8 @@ export default class RedisCommandsQueue {
     this.#invalidateCallback = callback;
   }
 
-  setMovingCallback(callback?: (afterMs: number, host: string, port: number) => void) {
-    this.#movingCallback = callback;
+  isWaitingForReply(): boolean {
+    return this.#waitingForReply.length > 0;
   }
 
   addCommand<T>(

packages/client/lib/client/index.ts

@@ -417,11 +417,16 @@ export default class RedisClient<
   #watchEpoch?: number;
   #clientSideCache?: ClientSideCacheProvider;
   #credentialsSubscription: Disposable | null = null;
+  // Flag used to pause writing to the socket during maintenance windows.
+  // When true, prevents new commands from being written while waiting for:
+  // 1. New socket to be ready after maintenance redirect
+  // 2. In-flight commands on the old socket to complete
+  #paused = false;
+
   get clientSideCache() {
     return this._self.#clientSideCache;
   }
 
-
   get options(): RedisClientOptions<M, F, S, RESP> | undefined {
     return this._self.#options;
   }
@@ -470,24 +475,60 @@ export default class RedisClient<
     this.#options = this.#initiateOptions(options);
     this.#queue = this.#initiateQueue();
     this.#socket = this.#initiateSocket(this.#options);
-
-    this.#queue.setMovingCallback(async (afterMs: number, host: string, port: number) => {
+    //  Queue
+    //     toWrite [ C D E ]
+    //     waitingForReply [ A B ]
+    //
+    //  time: ---1-2---3-4-5-6---------------------------
+    //
+    //  1. [EVENT] MOVING PN received
+    //  2. [ACTION] Pause writing ( we need to wait for new socket to connect and for all in-flight commands to complete )
+    //  3. [EVENT] New sock connected
+    //  4. [EVENT] In-flight commands completed
+    //  5. [ACTION] Unpause writing -> we are going to write to the new socket from now on
+    //  6. [ACTION] Destroy old socket
+    this.options?.gracefulMaintenance && this.#queue.events.on('moving', async (afterMs: number, host: string, port: number) => {
+      // 1
       console.log(`Moving to ${host}:${port} before ${afterMs}ms`);
+
+      // 2
+      this.#paused = true;
+
       const oldSocket = this.#socket;
-      const newSocket = this.#initiateSocket({
+      this.#socket = this.#initiateSocket({
         ...this.#options,
         socket: {
           ...this.#options?.socket,
           host,
           port
         }
       });
-      newSocket.on('ready', () => {
-        console.log(`Connected to ${host}:${port}, destroying old socket`);
-        oldSocket.destroy()
-        this.#socket = newSocket
+
+      // 3
+      this.#socket.once('ready', () => {
+        //TODO handshake...???
+        console.log(`Connected to ${host}:${port}`);
+
+        // 4
+        if(!this.#queue.isWaitingForReply()) {
+          // 5 and 6
+          oldSocket.destroy();
+          this.#paused = false;
+        }
+      });
+
+      // 4
+      this.#queue.events.once('waitingForReplyEmpty', () => {
+
+        // 3
+        if(this.#socket.isReady) {
+          // 5 and 6
+          oldSocket.destroy();
+          this.#paused = false;
+        }
       });
-      await newSocket.connect()
+
+      await this.#socket.connect()
     });
 
     if (options?.clientSideCache) {
@@ -1120,6 +1161,9 @@ export default class RedisClient<
   }
 
   #write() {
+    if(this.#paused) {
+      return
+    }
     this.#socket.write(this.#queue.commandsToWrite());
   }
 

packages/client/lib/client/linked-list.ts

@@ -1,3 +1,5 @@
+import EventEmitter from "events";
+
 export interface DoublyLinkedNode<T> {
   value: T;
   previous: DoublyLinkedNode<T> | undefined;
@@ -32,7 +34,7 @@ export class DoublyLinkedList<T> {
         next: undefined,
         value
       };
-    } 
+    }
 
     return this.#tail = this.#tail.next = {
       previous: this.#tail,
@@ -93,7 +95,7 @@ export class DoublyLinkedList<T> {
       node.previous!.next = node.next;
       node.previous = undefined;
     }
-    
+
     node.next = undefined;
   }
 
@@ -111,6 +113,33 @@ export class DoublyLinkedList<T> {
   }
 }
 
+// This function takes an object that has a `length` property
+// and returns both a proxy and an event emitter.
+// The proxy will act the same as the original object.
+// And the event emitter will emit an `empty` event whenever the `length` becomes zero.
+export const makeEmptyAware = <T extends { length: number }>(obj: T): [T, EventEmitter] => {
+  const eventEmitter = new EventEmitter();
+  const proxy = new Proxy(obj, {
+    get(target, prop, receiver) {
+      const original = Reflect.get(target, prop, receiver);
+      if (typeof original === 'function') {
+        return function (...args: any[]) {
+          const oldLength = target.length;
+          const ret = original.apply(target, args);
+          const newLength = target.length;
+          if(oldLength !== newLength && newLength === 0) {
+            eventEmitter.emit('empty')
+          }
+          return ret
+        };
+      } else {
+        return original;
+      };
+    },
+  });
+  return [ proxy, eventEmitter ];
+}
+
 export interface SinglyLinkedNode<T> {
   value: T;
   next: SinglyLinkedNode<T> | undefined;