nest/packages/microservices/server/server-rmq.ts

/* eslint-disable @typescript-eslint/no-redundant-type-constituents */
import {
  isNil,
  isString,
  isUndefined,
} from '@nestjs/common/utils/shared.utils.js';
import { createRequire } from 'module';
import {
  CONNECTION_FAILED_MESSAGE,
  DISCONNECTED_RMQ_MESSAGE,
  NO_MESSAGE_HANDLER,
  RMQ_SEPARATOR,
  RMQ_WILDCARD_ALL,
  RMQ_WILDCARD_SINGLE,
  RQM_DEFAULT_IS_GLOBAL_PREFETCH_COUNT,
  RQM_DEFAULT_NOACK,
  RQM_DEFAULT_NO_ASSERT,
  RQM_DEFAULT_PREFETCH_COUNT,
  RQM_DEFAULT_QUEUE,
  RQM_DEFAULT_QUEUE_OPTIONS,
  RQM_DEFAULT_URL,
  RQM_NO_EVENT_HANDLER,
  RQM_NO_MESSAGE_HANDLER,
} from '../constants.js';
import { RmqContext } from '../ctx-host/index.js';
import { Transport } from '../enums/index.js';
import { RmqEvents, RmqEventsMap, RmqStatus } from '../events/rmq.events.js';
import { RmqUrl } from '../external/rmq-url.interface.js';
import {
  MessageHandler,
  RmqOptions,
  TransportId,
} from '../interfaces/index.js';
import {
  IncomingRequest,
  OutgoingResponse,
  ReadPacket,
} from '../interfaces/packet.interface.js';
import { RmqRecordSerializer } from '../serializers/rmq-record.serializer.js';
import { Server } from './server.js';

// To enable type safety for RMQ. This cant be uncommented by default
// because it would require the user to install the amqplib package even if they dont use RabbitMQ
// Otherwise, TypeScript would fail to compile the code.
//
// type AmqpConnectionManager =
//   import('amqp-connection-manager').AmqpConnectionManager;
// type ChannelWrapper = import('amqp-connection-manager').ChannelWrapper;
// type Message = import('amqplib').Message;
// type Channel = import('amqplib').Channel | import('amqplib').ConfirmChannel;

type AmqpConnectionManager = any;
type ChannelWrapper = any;
type Message = any;
type Channel = any;

const INFINITE_CONNECTION_ATTEMPTS = -1;

/**
 * @publicApi
 */
export class ServerRMQ extends Server<RmqEvents, RmqStatus> {
  public transportId: TransportId = Transport.RMQ;

  protected server: AmqpConnectionManager | null = null;
  protected channel: ChannelWrapper | null = null;
  protected connectionAttempts = 0;
  protected readonly urls: string[] | RmqUrl[];
  protected readonly queue: string;
  protected readonly noAck: boolean;
  protected readonly queueOptions: any;
  protected readonly wildcardHandlers = new Map<string, MessageHandler>();
  protected pendingEventListeners: Array<{
    event: keyof RmqEvents;
    callback: RmqEvents[keyof RmqEvents];
  }> = [];

  constructor(protected readonly options: Required<RmqOptions>['options']) {
    super();
    this.urls = this.getOptionsProp(this.options, 'urls') || [RQM_DEFAULT_URL];
    this.queue =
      this.getOptionsProp(this.options, 'queue') || RQM_DEFAULT_QUEUE;
    this.noAck = this.getOptionsProp(this.options, 'noAck', RQM_DEFAULT_NOACK);
    this.queueOptions =
      this.getOptionsProp(this.options, 'queueOptions') ||
      RQM_DEFAULT_QUEUE_OPTIONS;

    this.loadPackageSynchronously('amqplib', ServerRMQ.name, () =>
      createRequire(import.meta.url)('amqplib'),
    );

    this.initializeSerializer(options);
    this.initializeDeserializer(options);
  }

  public async listen(
    callback: (err?: unknown, ...optionalParams: unknown[]) => void,
  ): Promise<void> {
    try {
      await this.start(callback);
    } catch (err) {
      callback(err);
    }
  }

  public async close(): Promise<void> {
    this.channel && (await this.channel.close());
    this.server && (await this.server.close());
    this.pendingEventListeners = [];
  }

  public async start(
    callback?: (err?: unknown, ...optionalParams: unknown[]) => void,
  ) {
    this.server = await this.createClient();
    this.server!.once(RmqEventsMap.CONNECT, () => {
      if (this.channel) {
        return;
      }
      this._status$.next(RmqStatus.CONNECTED);
      this.channel = this.server!.createChannel({
        json: false,
        setup: (channel: Channel) => this.setupChannel(channel, callback!),
      });
    });

    const maxConnectionAttempts = this.getOptionsProp(
      this.options,
      'maxConnectionAttempts',
      INFINITE_CONNECTION_ATTEMPTS,
    );

    this.registerConnectListener();
    this.registerDisconnectListener();
    this.pendingEventListeners.forEach(({ event, callback }) =>
      this.server!.on(event, callback),
    );
    this.pendingEventListeners = [];

    const connectFailedEvent = 'connectFailed';
    this.server!.once(
      connectFailedEvent,
      async (error: Record<string, unknown>) => {
        this._status$.next(RmqStatus.DISCONNECTED);

        this.logger.error(CONNECTION_FAILED_MESSAGE);
        if (error?.err) {
          this.logger.error(error.err);
        }
        const isReconnecting = !!this.channel;
        if (
          maxConnectionAttempts === INFINITE_CONNECTION_ATTEMPTS ||
          isReconnecting
        ) {
          return;
        }
        if (++this.connectionAttempts === maxConnectionAttempts) {
          await this.close();
          callback?.(error.err ?? new Error(CONNECTION_FAILED_MESSAGE));
        }
      },
    );
  }

  public async createClient<T = any>(): Promise<T> {
    const rmqPackage = await this.loadPackage(
      'amqp-connection-manager',
      ServerRMQ.name,
      () => import('amqp-connection-manager'),
    );
    const socketOptions = this.getOptionsProp(this.options, 'socketOptions');
    return rmqPackage.connect(this.urls, {
      connectionOptions: socketOptions?.connectionOptions,
      heartbeatIntervalInSeconds: socketOptions?.heartbeatIntervalInSeconds,
      reconnectTimeInSeconds: socketOptions?.reconnectTimeInSeconds,
    }) as T;
  }

  private registerConnectListener() {
    this.server!.on(RmqEventsMap.CONNECT, (err: any) => {
      this._status$.next(RmqStatus.CONNECTED);
    });
  }

  private registerDisconnectListener() {
    this.server!.on(RmqEventsMap.DISCONNECT, (err: any) => {
      this._status$.next(RmqStatus.DISCONNECTED);
      this.logger.error(DISCONNECTED_RMQ_MESSAGE);
      this.logger.error(err);
    });
  }

  public async setupChannel(channel: Channel, callback: Function) {
    const noAssert =
      this.getOptionsProp(this.options, 'noAssert') ??
      this.queueOptions.noAssert ??
      RQM_DEFAULT_NO_ASSERT;

    let createdQueue: string;

    if (this.queue === RQM_DEFAULT_QUEUE || !noAssert) {
      const { queue } = await channel.assertQueue(
        this.queue,
        this.queueOptions,
      );
      createdQueue = queue;
    } else {
      createdQueue = this.queue;
    }

    const isGlobalPrefetchCount = this.getOptionsProp(
      this.options,
      'isGlobalPrefetchCount',
      RQM_DEFAULT_IS_GLOBAL_PREFETCH_COUNT,
    );
    const prefetchCount = this.getOptionsProp(
      this.options,
      'prefetchCount',
      RQM_DEFAULT_PREFETCH_COUNT,
    );

    if (this.options.exchange || this.options.wildcards) {
      // Use queue name as exchange name if exchange is not provided and "wildcards" is set to true
      const exchange = this.getOptionsProp(
        this.options,
        'exchange',
        this.options.queue,
      );
      const exchangeType = this.getOptionsProp(
        this.options,
        'exchangeType',
        'topic',
      );
      await channel.assertExchange(exchange, exchangeType, {
        durable: true,
        arguments: this.getOptionsProp(this.options, 'exchangeArguments', {}),
      });

      if (this.options.routingKey || this.options.exchangeType === 'fanout') {
        await channel.bindQueue(
          createdQueue,
          exchange,
          this.options.exchangeType === 'fanout' ? '' : this.options.routingKey,
        );
      }

      if (this.options.wildcards) {
        const routingKeys = Array.from(this.getHandlers().keys());
        await Promise.all(
          routingKeys.map(routingKey =>
            channel.bindQueue(createdQueue, exchange, routingKey),
          ),
        );

        // When "wildcards" is set to true,  we need to initialize wildcard handlers
        // otherwise we would not be able to associate the incoming messages with the handlers
        this.initializeWildcardHandlersIfExist();
      }
    }

    await channel.prefetch(prefetchCount, isGlobalPrefetchCount);
    channel.consume(
      createdQueue,
      (msg: Record<string, any> | null) => this.handleMessage(msg!, channel),
      {
        noAck: this.noAck,
        consumerTag: this.getOptionsProp(
          this.options,
          'consumerTag',
          undefined,
        ),
      },
    );
    callback();
  }

  public async handleMessage(
    message: Record<string, any>,
    channel: any,
  ): Promise<void> {
    if (isNil(message)) {
      return;
    }
    const { content, properties } = message;
    const rawMessage = this.parseMessageContent(content);
    const packet = await this.deserializer.deserialize(rawMessage, properties);
    const pattern = isString(packet.pattern)
      ? packet.pattern
      : JSON.stringify(packet.pattern);

    const rmqContext = new RmqContext([message, channel, pattern]);
    if (isUndefined((packet as IncomingRequest).id)) {
      return this.handleEvent(pattern, packet, rmqContext);
    }
    const handler = this.getHandlerByPattern(pattern);

    if (!handler) {
      if (!this.noAck) {
        this.logger.warn(RQM_NO_MESSAGE_HANDLER`${pattern}`);
        this.channel!.nack(rmqContext.getMessage() as Message, false, false);
      }
      const status = 'error';
      const noHandlerPacket = {
        id: (packet as IncomingRequest).id,
        err: NO_MESSAGE_HANDLER,
        status,
      };
      return this.sendMessage(
        noHandlerPacket,
        properties.replyTo,
        properties.correlationId,
        rmqContext,
      );
    }
    return this.onProcessingStartHook(
      this.transportId,
      rmqContext,
      async () => {
        const response$ = this.transformToObservable(
          await handler(packet.data, rmqContext),
        );

        const publish = <T>(data: T) =>
          this.sendMessage(
            data,
            properties.replyTo,
            properties.correlationId,
            rmqContext,
          );

        response$ && this.send(response$, publish);
      },
    );
  }

  public async handleEvent(
    pattern: string,
    packet: ReadPacket,
    context: RmqContext,
  ): Promise<any> {
    const handler = this.getHandlerByPattern(pattern);
    if (!handler && !this.noAck) {
      this.channel!.nack(context.getMessage() as Message, false, false);
      return this.logger.warn(RQM_NO_EVENT_HANDLER`${pattern}`);
    }
    return super.handleEvent(pattern, packet, context);
  }

  public sendMessage<T = any>(
    message: T,
    replyTo: any,
    correlationId: string,
    context: RmqContext,
  ): void {
    const outgoingResponse = this.serializer.serialize(
      message as unknown as OutgoingResponse,
    );
    const options = outgoingResponse.options;
    delete outgoingResponse.options;

    const buffer = Buffer.from(JSON.stringify(outgoingResponse));
    const sendOptions = { correlationId, ...options };

    this.onProcessingEndHook?.(this.transportId, context);
    this.channel!.sendToQueue(replyTo, buffer, sendOptions);
  }

  public unwrap<T>(): T {
    if (!this.server) {
      throw new Error(
        'Not initialized. Please call the "listen"/"startAllMicroservices" method before accessing the server.',
      );
    }
    return this.server as T;
  }

  public on<
    EventKey extends keyof RmqEvents = keyof RmqEvents,
    EventCallback extends RmqEvents[EventKey] = RmqEvents[EventKey],
  >(event: EventKey, callback: EventCallback) {
    if (this.server) {
      this.server.addListener(event, callback);
    } else {
      this.pendingEventListeners.push({ event, callback });
    }
  }

  public getHandlerByPattern(pattern: string): MessageHandler | null {
    if (!this.options.wildcards) {
      return super.getHandlerByPattern(pattern);
    }

    // Search for non-wildcard handler first
    const handler = super.getHandlerByPattern(pattern);
    if (handler) {
      return handler;
    }

    // Search for wildcard handler
    if (this.wildcardHandlers.size === 0) {
      return null;
    }
    for (const [wildcardPattern, handler] of this.wildcardHandlers) {
      if (this.matchRmqPattern(wildcardPattern, pattern)) {
        return handler;
      }
    }
    return null;
  }

  protected initializeSerializer(options: RmqOptions['options']) {
    this.serializer = options?.serializer ?? new RmqRecordSerializer();
  }

  private parseMessageContent(content: Buffer) {
    try {
      return JSON.parse(content.toString());
    } catch {
      return content.toString();
    }
  }

  private initializeWildcardHandlersIfExist() {
    if (this.wildcardHandlers.size !== 0) {
      return;
    }
    const handlers = this.getHandlers();

    handlers.forEach((handler, pattern) => {
      if (
        pattern.includes(RMQ_WILDCARD_ALL) ||
        pattern.includes(RMQ_WILDCARD_SINGLE)
      ) {
        this.wildcardHandlers.set(pattern, handler);
      }
    });
  }

  private matchRmqPattern(pattern: string, routingKey: string): boolean {
    if (!routingKey) {
      return pattern === RMQ_WILDCARD_ALL;
    }

    const patternSegments = pattern.split(RMQ_SEPARATOR);
    const routingKeySegments = routingKey.split(RMQ_SEPARATOR);

    const patternSegmentsLength = patternSegments.length;
    const routingKeySegmentsLength = routingKeySegments.length;
    const lastIndex = patternSegmentsLength - 1;

    for (const [i, currentPattern] of patternSegments.entries()) {
      const currentRoutingKey = routingKeySegments[i];

      if (!currentRoutingKey && !currentPattern) {
        continue;
      }
      if (!currentRoutingKey && currentPattern !== RMQ_WILDCARD_ALL) {
        return false;
      }
      if (currentPattern === RMQ_WILDCARD_ALL) {
        return i === lastIndex;
      }
      if (
        currentPattern !== RMQ_WILDCARD_SINGLE &&
        currentPattern !== currentRoutingKey
      ) {
        return false;
      }
    }
    return patternSegmentsLength === routingKeySegmentsLength;
  }
}