refactor: rework timeouts

Author: nkaradzhov

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

@@ -3,9 +3,8 @@ 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, TimeoutDuringMaintanance, TimeoutError } from '../errors';
+import { AbortError, ErrorReply, CommandTimeoutDuringMaintananceError, TimeoutError } from '../errors';
 import { MonitorCallback } from '.';
-import assert from 'assert';
 
 export interface CommandOptions<T = TypeMapping> {
   chainId?: symbol;
@@ -71,10 +70,12 @@ export default class RedisCommandsQueue {
 
   #pushHandlers: PushHandler[] = [this.#onPush.bind(this)];
 
-  // If this value is set, we are in a maintenance mode.
-  // This means any existing commands should have their timeout
-  // overwritten to the new timeout. And all new commands should
-  // have their timeout set as the new timeout.
+  #inMaintenance = false;
+
+  set inMaintenance(value: boolean) {
+    this.#inMaintenance = value;
+  }
+
   #maintenanceCommandTimeout: number | undefined
 
   setMaintenanceCommandTimeout(ms: number | undefined) {
@@ -87,29 +88,24 @@ export default class RedisCommandsQueue {
     this.#toWrite.forEachNode(node => {
       const command = node.value;
 
-      // If the command didnt have a timeout, skip it
-      if (!command.timeout) return;
-
-      // Remove existing timeout listener
+      // Remove timeout listener if it exists
       RedisCommandsQueue.#removeTimeoutListener(command)
 
-      //TODO see if this is needed
-      // // Keep a flag to know if we were in maintenance at this point in time.
-      // // To be used in the timeout listener, which needs to know which exact error to use.
-      // const wasMaintenance = !!this.#maintenanceCommandTimeout
-
       // Determine newTimeout
       const newTimeout = this.#maintenanceCommandTimeout ?? command.timeout?.originalTimeout;
-      assert(newTimeout !== undefined, 'Trying to reset timeout to `undefined`')
+      // if no timeout is given and the command didnt have any timeout before, skip
+      if (!newTimeout)  return;
+
 
+      // Overwrite the command's timeout
       const signal = AbortSignal.timeout(newTimeout);
       command.timeout = {
         signal,
         listener: () => {
           this.#toWrite.remove(node);
-          command.reject(this.#maintenanceCommandTimeout ? new TimeoutDuringMaintanance(newTimeout) : new TimeoutError());
+          command.reject(this.#inMaintenance ? new CommandTimeoutDuringMaintananceError(newTimeout) : new TimeoutError());
         },
-        originalTimeout: command.timeout.originalTimeout
+        originalTimeout: command.timeout?.originalTimeout
       };
       signal.addEventListener('abort', command.timeout.listener, { once: true });
     });
@@ -219,23 +215,17 @@ export default class RedisCommandsQueue {
         typeMapping: options?.typeMapping
       };
 
-      // If #commandTimeout was explicitly set, this
-      // means we are in maintenance mode and should
+      // If #maintenanceCommandTimeout was explicitly set, we should
       // use it instead of the timeout provided by the command
       const timeout = this.#maintenanceCommandTimeout || options?.timeout
       if (timeout) {
 
-        //TODO see if this is needed
-        // // Keep a flag to know if we were in maintenance at this point in time.
-        // // To be used in the timeout listener, which needs to know which exact error to use.
-        // const wasMaintenance = !!this.#maintenanceCommandTimeout
-
         const signal = AbortSignal.timeout(timeout);
         value.timeout = {
           signal,
           listener: () => {
             this.#toWrite.remove(node);
-            value.reject(this.#maintenanceCommandTimeout ? new TimeoutDuringMaintanance(timeout) : new TimeoutError());
+            value.reject(this.#inMaintenance ? new CommandTimeoutDuringMaintananceError(timeout) : new TimeoutError());
           },
           originalTimeout: options?.timeout
         };

packages/client/lib/client/enterprise-maintenance-manager.ts

@@ -10,11 +10,16 @@ export const MAINTENANCE_EVENTS = {
 } as const;
 
 const PN = {
- MOVING: "MOVING",
- MIGRATING: "MIGRATING",
- MIGRATED: "MIGRATED",
- FAILING_OVER: "FAILING_OVER",
- FAILED_OVER: "FAILED_OVER",
+  MOVING: "MOVING",
+  MIGRATING: "MIGRATING",
+  MIGRATED: "MIGRATED",
+  FAILING_OVER: "FAILING_OVER",
+  FAILED_OVER: "FAILED_OVER",
+};
+
+export interface SocketTimeoutUpdate {
+  inMaintenance: boolean,
+  timeout?: number
 }
 
 export default class EnterpriseMaintenanceManager extends EventEmitter {
@@ -91,27 +96,25 @@ export default class EnterpriseMaintenanceManager extends EventEmitter {
   };
 
   #onMigrating = async () => {
-    this.#commandsQueue.setMaintenanceCommandTimeout(this.#getCommandTimeout());
-    this.emit(
-      MAINTENANCE_EVENTS.TIMEOUTS_UPDATE,
-      this.#getSocketTimeout(),
+    this.#commandsQueue.inMaintenance = true;
+    this.#commandsQueue.setMaintenanceCommandTimeout(
+      this.#options.gracefulMaintenance?.relaxedCommandTimeout,
     );
+
+    this.emit(MAINTENANCE_EVENTS.TIMEOUTS_UPDATE, {
+      inMaintenance: true,
+      timeout: this.#options.gracefulMaintenance?.relaxedSocketTimeout
+    } satisfies SocketTimeoutUpdate);
   };
 
   #onMigrated = async () => {
+    this.#commandsQueue.inMaintenance = false;
     this.#commandsQueue.setMaintenanceCommandTimeout(undefined);
-    this.emit(MAINTENANCE_EVENTS.TIMEOUTS_UPDATE, undefined);
-  };
 
-  #getSocketTimeout(): number | undefined {
-    return this.#options.gracefulMaintenance?.handleTimeouts === "error"
-      ? this.#options.socket?.socketTimeout
-      : this.#options.gracefulMaintenance?.handleTimeouts;
-  }
+    this.emit(MAINTENANCE_EVENTS.TIMEOUTS_UPDATE, {
+      inMaintenance: false,
+      timeout: undefined
+    } satisfies SocketTimeoutUpdate);
+  };
 
-  #getCommandTimeout(): number | undefined {
-    return this.#options.gracefulMaintenance?.handleTimeouts === "error"
-      ? this.#options.commandOptions?.timeout
-      : this.#options.gracefulMaintenance?.handleTimeouts;
-  }
-}
+};

packages/client/lib/client/index.ts

@@ -20,7 +20,7 @@ import { BasicClientSideCache, ClientSideCacheConfig, ClientSideCacheProvider }
 import { BasicCommandParser, CommandParser } from './parser';
 import SingleEntryCache from '../single-entry-cache';
 import { version } from '../../package.json'
-import EnterpriseMaintenanceManager, { MAINTENANCE_EVENTS } from './enterprise-maintenance-manager';
+import EnterpriseMaintenanceManager, { MAINTENANCE_EVENTS, SocketTimeoutUpdate } from './enterprise-maintenance-manager';
 
 export interface RedisClientOptions<
   M extends RedisModules = RedisModules,
@@ -177,11 +177,19 @@ export interface RedisClientOptions<
     /**
      * Designates how failed commands should be handled. A failed command is when the time isn’t sufficient to deal with the responses on the old connection before the server shuts it down
      */
-    handleFailedCommands: 'exception' | 'retry',
+    handleFailedCommands?: 'exception' | 'retry',
     /**
-     * Specify whether we should throw a TimeoutDuringMaintanance exception or provide more relaxed timeout, in order to minimize command timeouts during maintenance.
+     * Specifies a more relaxed timeout (in milliseconds) for commands during a maintenance window.
+     * This helps minimize command timeouts during maintenance. If not provided, the `commandOptions.timeout`
+     * will be used instead. Timeouts during maintenance period result in a `CommandTimeoutDuringMaintanance` error.
      */
-    handleTimeouts: 'error' | number,
+    relaxedCommandTimeout?: number,
+    /**
+     * Specifies a more relaxed timeout (in milliseconds) for the socket during a maintenance window.
+     * This helps minimize socket timeouts during maintenance. If not provided, the `socket.timeout`
+     * will be used instead. Timeouts during maintenance period result in a `SocketTimeoutDuringMaintanance` error.
+     */
+    relaxedSocketTimeout?: number
   }
 }
 
@@ -490,7 +498,10 @@ export default class RedisClient<
       new EnterpriseMaintenanceManager(this.#queue, this.#options!)
         .on(MAINTENANCE_EVENTS.PAUSE_WRITING, () => this._self.#pausedForMaintenance = true )
         .on(MAINTENANCE_EVENTS.RESUME_WRITING, this.#resumeFromMaintenance.bind(this))
-        .on(MAINTENANCE_EVENTS.TIMEOUTS_UPDATE, (mtm: number | undefined) => this._self.#socket.setMaintenanceTimeout(mtm))
+        .on(MAINTENANCE_EVENTS.TIMEOUTS_UPDATE, (value: SocketTimeoutUpdate) => {
+          this._self.#socket.inMaintenance = value.inMaintenance;
+          this._self.#socket.setMaintenanceTimeout(value.timeout);
+        })
     }
 
     if (options?.clientSideCache) {

packages/client/lib/client/socket.ts

@@ -1,7 +1,7 @@
 import { EventEmitter, once } from 'node:events';
 import net from 'node:net';
 import tls from 'node:tls';
-import { ConnectionTimeoutError, ClientClosedError, SocketClosedUnexpectedlyError, ReconnectStrategyError, SocketTimeoutError, TimeoutDuringMaintanance } from '../errors';
+import { ConnectionTimeoutError, ClientClosedError, SocketClosedUnexpectedlyError, ReconnectStrategyError, SocketTimeoutError, SocketTimeoutDuringMaintananceError } from '../errors';
 import { setTimeout } from 'node:timers/promises';
 import { RedisArgument } from '../RESP/types';
 
@@ -81,6 +81,12 @@ export default class RedisSocket extends EventEmitter {
     return this.#socketEpoch;
   }
 
+  #inMaintenance = false;
+
+  set inMaintenance(value: boolean) {
+    this.#inMaintenance = value;
+  }
+
   constructor(options?: RedisSocketOptions) {
     super();
 
@@ -237,12 +243,14 @@ export default class RedisSocket extends EventEmitter {
   }
 
   setMaintenanceTimeout(ms?: number) {
+    if (this.#maintenanceTimeout === ms) return;
+
     this.#maintenanceTimeout = ms;
 
     if(ms !== undefined) {
       this.#socket?.setTimeout(ms);
-    } else if (this.#socketTimeout !== undefined) {
-      this.#socket?.setTimeout(this.#socketTimeout);
+    } else {
+      this.#socket?.setTimeout(this.#socketTimeout ?? 0);
     }
   }
 
@@ -268,8 +276,8 @@ export default class RedisSocket extends EventEmitter {
 
     if (this.#socketTimeout) {
       socket.once('timeout', () => {
-        const error = this.#maintenanceTimeout
-          ? new TimeoutDuringMaintanance(this.#socketTimeout!)
+        const error = this.#inMaintenance
+          ? new SocketTimeoutDuringMaintananceError(this.#socketTimeout!)
           : new SocketTimeoutError(this.#socketTimeout!)
         socket.destroy(error);
       });

packages/client/lib/errors.ts

@@ -76,12 +76,17 @@ export class BlobError extends ErrorReply {}
 
 export class TimeoutError extends Error {}
 
-export class TimeoutDuringMaintanance extends Error {
+export class SocketTimeoutDuringMaintananceError extends TimeoutError {
   constructor(timeout: number) {
     super(`Socket timeout during maintenance. Expecting data, but didn't receive any in ${timeout}ms.`);
   }
 }
 
+export class CommandTimeoutDuringMaintananceError extends TimeoutError {
+  constructor(timeout: number) {
+    super(`Command timeout during maintenance. Waited to write command for more than ${timeout}ms.`);
+  }
+}
 
 export class MultiErrorReply extends ErrorReply {
   replies: Array<ErrorReply>;