socket.ts

import { EventEmitter, once } from 'node:events';
import net from 'node:net';
import tls from 'node:tls';
import { ConnectionTimeoutError, ClientClosedError, SocketClosedUnexpectedlyError, ReconnectStrategyError, SocketTimeoutError, SocketTimeoutDuringMaintenanceError } from '../errors';
import { setTimeout } from 'node:timers/promises';
import { RedisArgument } from '../RESP/types';
import { dbgMaintenance } from './enterprise-maintenance-manager';
import { publish, CHANNELS } from './tracing';
import { DEFAULT_KEEPALIVE_INITIAL_DELAY } from '../defaults';

type NetOptions = {
  tls?: false;
};

type ReconnectStrategyFunction = (retries: number, cause: Error) => false | Error | number;

type RedisSocketOptionsCommon = {
  /**
   * Connection timeout (in milliseconds)
   */
  connectTimeout?: number;
  /**
   * When the socket closes unexpectedly (without calling `.close()`/`.destroy()`), the client uses `reconnectStrategy` to decide what to do. The following values are supported:
   * 1. `false` -> do not reconnect, close the client and flush the command queue.
   * 2. `number` -> wait for `X` milliseconds before reconnecting.
   * 3. `(retries: number, cause: Error) => false | number | Error` -> `number` is the same as configuring a `number` directly, `Error` is the same as `false`, but with a custom error.
   */
  reconnectStrategy?: false | number | ReconnectStrategyFunction;
  /**
   * The timeout (in milliseconds) after which the socket will be closed. `undefined` means no timeout.
   */
  socketTimeout?: number;
}

type RedisTcpOptions = RedisSocketOptionsCommon & NetOptions & Omit<
  net.TcpNetConnectOpts,
  'timeout' | 'onread' | 'readable' | 'writable' | 'port'
> & {
  port?: number;
};

type RedisTlsOptions = RedisSocketOptionsCommon & tls.ConnectionOptions & {
  tls: true;
}

type RedisIpcOptions = RedisSocketOptionsCommon & Omit<
  net.IpcNetConnectOpts,
  'timeout' | 'onread' | 'readable' | 'writable'
> & {
  tls: false;
}

export type RedisTcpSocketOptions = RedisTcpOptions | RedisTlsOptions;

export type RedisSocketOptions = RedisTcpSocketOptions | RedisIpcOptions;

export type RedisSocketInitiator = () => void | Promise<unknown>;

export default class RedisSocket extends EventEmitter {
  readonly #initiator;
  readonly #connectTimeout;
  readonly #reconnectStrategy;
  readonly #socketFactory;
  readonly #socketTimeout;
  readonly #clientId: string;

  #maintenanceTimeout: number | undefined;

  #socket?: net.Socket | tls.TLSSocket;

  #isOpen = false;

  get isOpen() {
    return this.#isOpen;
  }

  #isReady = false;

  get isReady() {
    return this.#isReady;
  }

  #isSocketUnrefed = false;

  #socketEpoch = 0;

  get socketEpoch() {
    return this.#socketEpoch;
  }

  get host() {
    return this.#socket?.remoteAddress;
  }

  get port() {
    return this.#socket?.remotePort;
  }

  constructor(
    initiator: RedisSocketInitiator,
    clientId: string,
    options?: RedisSocketOptions,
  ) {
    super();

    this.#initiator = initiator;
    this.#connectTimeout = options?.connectTimeout ?? 5000;
    this.#reconnectStrategy = this.#createReconnectStrategy(options);
    this.#socketFactory = this.#createSocketFactory(options);
    this.#socketTimeout = options?.socketTimeout;
    this.#clientId = clientId;
  }

  #createReconnectStrategy(options?: RedisSocketOptions): ReconnectStrategyFunction {
    const strategy = options?.reconnectStrategy;
    if (strategy === false || typeof strategy === 'number') {
      return () => strategy;
    }

    if (strategy) {
      return (retries, cause) => {
        try {
const retryIn = strategy(retries, cause);
          if (retryIn !== false && !(retryIn instanceof Error) && typeof retryIn !== 'number') {
            throw new TypeError(`Reconnect strategy should return \`false | Error | number\`, got ${retryIn} instead`);
          }
          return retryIn;
        } catch (err) {
          publish(CHANNELS.ERROR, () => ({
            error: err as Error,
            origin: 'client',
            internal: false,
            clientId: this.#clientId
          }));
          this.emit('error', err);
          return this.defaultReconnectStrategy(retries, err);
        }
      };
    }

    return this.defaultReconnectStrategy;
  }

  #createSocketFactory(options?: RedisSocketOptions) {
    // TLS
    if (options?.tls === true) {
      const withDefaults: tls.ConnectionOptions = {
        ...options,
        port: options?.port ?? 6379,
        // https://nodejs.org/api/tls.html#tlsconnectoptions-callback "Any socket.connect() option not already listed"
        // @types/node is... incorrect...
        // @ts-expect-error - @types/node omits socket.connect noDelay.
        noDelay: options?.noDelay ?? true,
        // @ts-expect-error - @types/node omits socket.connect keepAlive.
        keepAlive: options?.keepAlive ?? true,
        // @ts-expect-error - @types/node omits socket.connect keepAliveInitialDelay.
        keepAliveInitialDelay: options?.keepAliveInitialDelay ?? DEFAULT_KEEPALIVE_INITIAL_DELAY,
        timeout: undefined,
        onread: undefined,
        readable: true,
        writable: true
      };
      return {
        create() {
          return tls.connect(withDefaults);
        },
        event: 'secureConnect'
      };
    }

    // IPC
    if (options && 'path' in options) {
      const withDefaults: net.IpcNetConnectOpts = {
        ...options,
        timeout: undefined,
        onread: undefined,
        readable: true,
        writable: true
      };
      return {
        create() {
          return net.createConnection(withDefaults);
        },
        event: 'connect'
      };
    }

    // TCP
    const withDefaults: net.TcpNetConnectOpts = {
      ...options,
      port: options?.port ?? 6379,
      noDelay: options?.noDelay ?? true,
      keepAlive: options?.keepAlive ?? true,
      keepAliveInitialDelay: options?.keepAliveInitialDelay ?? DEFAULT_KEEPALIVE_INITIAL_DELAY,
      timeout: undefined,
      onread: undefined,
      readable: true,
      writable: true
    };
    return {
      create() {
        return net.createConnection(withDefaults);
      },
      event: 'connect'
    };
  }

  /**
   * The single choke point where `reconnectStrategy` giving up (`false` or an
   * `Error`) is handled: closes the socket for good and emits `'terminated'`
   * so a caller reacting only to `'error'` — which also fires on every
   * *retried* disconnect — can tell "still retrying" apart from "reconnection
   * has permanently stopped, the client is unusable from here on".
   */
  #shouldReconnect(retries: number, cause: Error) {
    const retryIn = this.#reconnectStrategy(retries, cause);
    if (retryIn === false) {
      this.#isOpen = false;
      publish(CHANNELS.ERROR, () => ({
        error: cause,
        origin: 'client',
        internal: false,
        clientId: this.#clientId
      }));
      this.emit('terminated', cause);
      this.emit('error', cause);
      return cause;
    } else if (retryIn instanceof Error) {
      this.#isOpen = false;
      publish(CHANNELS.ERROR, () => ({
        error: cause,
        origin: 'client',
        internal: false,
        clientId: this.#clientId
      }));
      const terminatedBy = new ReconnectStrategyError(retryIn, cause);
      this.emit('terminated', terminatedBy);
      this.emit('error', cause);
      return terminatedBy;
    }

    return retryIn;
  }

  async connect(): Promise<void> {
    if (this.#isOpen) {
      throw new Error('Socket already opened');
    }

    this.#isOpen = true;
    return this.#connect();
  }

  async #connect(): Promise<void> {
    let retries = 0;
    do {
      try {
        const connectStartTime = performance.now();
        const socket = this.#socket = await this.#createSocket();
        this.emit('connect');

        try {
          await this.#initiateWhileSocketAlive(socket);

          // Check if socket was closed/destroyed during initiator execution
          if (!this.#socket || this.#socket.destroyed || !this.#socket.readable || !this.#socket.writable) {
            const retryIn = this.#shouldReconnect(retries++, new SocketClosedUnexpectedlyError());
            if (typeof retryIn !== 'number') { throw retryIn; }
            await setTimeout(retryIn);
            this.emit('reconnecting');
            continue;
          }
        } catch (err) {
          // #socket may already be undefined if the client was destroyed while
          // the initiator was suspended (destroySocket cleared it).
          this.#socket?.destroy();
          this.#socket = undefined;
          throw err;
        }
        this.#isReady = true;
        this.#socketEpoch++;
        publish(CHANNELS.CONNECTION_READY, () => ({
          clientId: this.#clientId,
          serverAddress: this.host,
          serverPort: this.port,
          createTimeMs: performance.now() - connectStartTime,
        }));
        this.emit('ready');
      } catch (err) {
        // The client was closed while connecting (e.g. destroy()/quit() raced
        // an async initiator). Abort the attempt without emitting error/
        // reconnecting or scheduling a retry — the shutdown is intentional.
        if (!this.#isOpen) throw err;

        const retryIn = this.#shouldReconnect(retries++, err as Error);
        if (typeof retryIn !== 'number') {
          throw retryIn;
        }

        publish(CHANNELS.ERROR, () => ({
          error: err as Error,
          origin: 'client',
          internal: false,
          clientId: this.#clientId
        }));
        this.emit('error', err);
        await setTimeout(retryIn);
        this.emit('reconnecting');
      }
    } while (this.#isOpen && !this.#isReady);
  }

  /**
   * Awaits the initiator, rejecting as soon as the socket errors or closes.
   * If the socket dies while the initiator is suspended (e.g. on DNS
   * resolution or an async credentials provider), commands it enqueues
   * afterwards are never written nor flushed — without this guard `#connect`
   * would stay suspended forever without emitting a terminal event.
   */
  #initiateWhileSocketAlive(socket: net.Socket | tls.TLSSocket) {
    let onSocketDied!: (err?: unknown) => void;
    const socketDied = new Promise<never>((_, reject) => {
      onSocketDied = (err?: unknown) => {
        reject(err instanceof Error ? err : new SocketClosedUnexpectedlyError());
      };
      socket.once('error', onSocketDied);
      socket.once('close', onSocketDied);
    });

    // Defer into a microtask so a synchronous throw from the initiator becomes
    // a rejection routed through the race below, instead of escaping before the
    // socketDied listeners are registered and cleaned up (which would leak them
    // and surface the raced rejection as an unhandled rejection).
    const initiated = Promise.resolve().then(() => this.#initiator());
    // an abandoned initiator can still reject later (its stranded commands get
    // flushed on a subsequent failure) — the raced error already drove the retry
    initiated.catch(() => {});

    return Promise.race([initiated, socketDied]).finally(() => {
      socket.removeListener('error', onSocketDied);
      socket.removeListener('close', onSocketDied);
    });
  }

  setMaintenanceTimeout(ms?: number) {
    dbgMaintenance(`Set socket timeout to ${ms}`);
    if (this.#maintenanceTimeout === ms) {
      dbgMaintenance(`Socket already set maintenanceCommandTimeout to ${ms}, skipping`);
      return;
    };

    this.#maintenanceTimeout = ms;

    if(ms !== undefined) {
      this.#socket?.setTimeout(ms);
      publish(CHANNELS.CONNECTION_RELAXED_TIMEOUT, () => ({ clientId: this.#clientId, value: 1 }));
    } else {
      this.#socket?.setTimeout(this.#socketTimeout ?? 0);
      publish(CHANNELS.CONNECTION_RELAXED_TIMEOUT, () => ({ clientId: this.#clientId, value: -1 }));
    }
  }

  async #createSocket(): Promise<net.Socket | tls.TLSSocket> {
    const socket = this.#socketFactory.create();

    let onTimeout;
    if (this.#connectTimeout !== undefined) {
      onTimeout = () => socket.destroy(new ConnectionTimeoutError());
      socket.once('timeout', onTimeout);
      socket.setTimeout(this.#connectTimeout);
    }

    if (this.#isSocketUnrefed) {
      socket.unref();
    }

    await once(socket, this.#socketFactory.event);

    if (onTimeout) {
      socket.removeListener('timeout', onTimeout);
    }

    if (this.#socketTimeout) {
      socket.once('timeout', () => {
        const error = this.#maintenanceTimeout
          ? new SocketTimeoutDuringMaintenanceError(this.#maintenanceTimeout)
          : new SocketTimeoutError(this.#socketTimeout!)
        socket.destroy(error);
      });
      socket.setTimeout(this.#socketTimeout);
    }

    socket
      .once('error', err => this.#onSocketError(err))
      .once('close', hadError => {
        if (hadError || !this.#isOpen || this.#socket !== socket) return;
        this.#onSocketError(new SocketClosedUnexpectedlyError());
      })
      .on('drain', () => this.emit('drain'))
      .on('data', data => this.emit('data', data));

    return socket;
  }

  #onSocketError(err: Error): void {
    const wasReady = this.#isReady;
    this.#isReady = false;
    const socket = this.#socket;
    this.#socket = undefined;
    socket?.removeAllListeners('data');
    socket?.destroy();

    if (wasReady) {
      publish(CHANNELS.CONNECTION_CLOSED, () => ({ clientId: this.#clientId, reason: 'error', wasConnected: true }));
    }

    if (!wasReady) {
      if (!this.#isOpen) {
        publish(CHANNELS.ERROR, () => ({
          error: err,
          origin: 'client',
          internal: false,
          clientId: this.#clientId
        }));
        this.emit('error', err);
      }
      return;
    }

    if (!this.#isOpen) {
      publish(CHANNELS.ERROR, () => ({
        error: err,
        origin: 'client',
        internal: false,
        clientId: this.#clientId
      }));
      this.emit('error', err);
      return;
    }

    const retryIn = this.#shouldReconnect(0, err);
    if (typeof retryIn !== 'number') return;

    publish(CHANNELS.ERROR, () => ({
      error: err,
      origin: 'client',
      internal: false,
      clientId: this.#clientId
    }));
    this.emit('error', err);
    if (!this.#isOpen) return;

    this.emit('reconnecting');
    this.#connect().catch(() => {
      // the error was already emitted, silently ignore it
    });
  }


  write(iterable: Iterable<ReadonlyArray<RedisArgument>>) {
    if (!this.#socket || !this.#socket.writable) return;

    this.#socket.cork();
    try {
      for (const args of iterable) {
        for (const toWrite of args) {
          this.#socket.write(toWrite);
        }

        if (this.#socket.writableNeedDrain) break;
      }
    } catch (err) {
      // net.Socket.write can throw synchronously on a half-closed socket
      // (writeAfterFIN -> EPIPE) before the 'close' event fires. The pending
      // command has already been moved to #waitingForReply by the queue's
      // generator, so the close handler will reject it on reconnect.
      if (!err || (err as NodeJS.ErrnoException).code !== 'EPIPE') {
        throw err;
      }
    } finally {
      this.#socket.uncork();
    }
  }

  async quit<T>(fn: () => Promise<T>): Promise<T> {
    if (!this.#isOpen) {
      throw new ClientClosedError();
    }

    this.#isOpen = false;
    const reply = await fn();
    this.destroySocket();
    return reply;
  }

  close() {
    if (!this.#isOpen) {
      throw new ClientClosedError();
    }

    this.#isOpen = false;
  }

  destroy() {
    // Idempotent: return instead of throwing when already closed. A terminal
    // connect failure (reconnectStrategy gave up) leaves #isOpen === false, and
    // the owning client still needs to dispose itself (unregister metrics,
    // dispose credentials) — throwing here would abort that cleanup. Returning
    // also means a repeated destroy() won't re-run destroySocket() and
    // republish CONNECTION_CLOSED / re-emit 'end'.
    if (!this.#isOpen) return;
    this.#isOpen = false;
    this.destroySocket();
  }

  destroySocket() {
    const wasReady = this.#isReady;
    this.#isReady = false;

    if (this.#socket) {
      this.#socket.destroy();
      this.#socket = undefined;
    }

    publish(CHANNELS.CONNECTION_CLOSED, () => ({ clientId: this.#clientId, reason: 'application_close', wasConnected: wasReady }));
    this.emit('end');
  }

  ref() {
    this.#isSocketUnrefed = false;
    this.#socket?.ref();
  }

  unref() {
    this.#isSocketUnrefed = true;
    this.#socket?.unref();
  }

  defaultReconnectStrategy(retries: number, cause: unknown) {
    // By default, do not reconnect on socket timeout.
    if (cause instanceof SocketTimeoutError) {
      return false;
    }

    // Generate a random jitter between 0 – 200 ms:
    const jitter = Math.floor(Math.random() * 200);
    // Delay is an exponential back off, (times^2) * 50 ms, with a maximum value of 2000 ms:
    const delay = Math.min(Math.pow(2, retries) * 50, 2000);

    return delay + jitter;
  }
}