artillery/lib/platform/aws-ecs/legacy/plugins/artillery-plugin-sqs-reporter/index.ts

Cloud-scale load testing. https://www.artillery.io

/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

import { PutObjectCommand, S3Client } from '@aws-sdk/client-s3';

import { SendMessageCommand, SQSClient } from '@aws-sdk/client-sqs';
import createDebug from 'debug';

const _debug = createDebug('plugin:sqsReporter');

import { randomUUID as uuid } from 'node:crypto';
import { getAQS, sendMessage } from './azure-aqs.ts';

// SQS has 1MB message limit. Use 950KB threshold for safety margin.
const SQS_SIZE_LIMIT = 950 * 1024;

type MessageTag = { key: string; value: string };

class ArtillerySQSPlugin {
  declare script: Record<string, any>;
  declare events: {
    on(event: string, listener: (...args: any[]) => void): unknown;
  };
  declare unsent: number;
  declare tags: MessageTag[];
  declare testId: string | null;
  declare messageAttributes: Record<
    string,
    { DataType: string; StringValue: string }
  >;
  declare sqs: SQSClient | null;
  declare aqs: ReturnType<typeof getAQS> | null;
  declare queueUrl: string | undefined;
  declare s3: S3Client | null;
  declare s3Bucket: string | null;

  constructor(
    script: Record<string, any>,
    events: { on(event: string, listener: (...args: any[]) => void): unknown }
  ) {
    this.script = script;
    this.events = events;

    this.unsent = 0;

    // List of objects: [{key: 'SomeKey', value: 'SomeValue'}, ...]
    this.tags = process.env.SQS_TAGS ? JSON.parse(process.env.SQS_TAGS) : [];
    this.testId = null;
    const messageAttributes: Record<
      string,
      { DataType: string; StringValue: string }
    > = {};

    this.tags.forEach((tag) => {
      if (tag.key === 'testId') {
        this.testId = tag.value;
      }
      messageAttributes[tag.key] = {
        DataType: 'String',
        StringValue: tag.value
      };
    });

    this.messageAttributes = messageAttributes;

    this.sqs = null;
    this.aqs = null;

    if (process.env.SQS_QUEUE_URL) {
      this.sqs = new SQSClient({
        region:
          process.env.SQS_REGION || script.config.plugins['sqs-reporter'].region
      });

      this.queueUrl =
        process.env.SQS_QUEUE_URL ||
        script.config.plugins['sqs-reporter'].queueUrl;
    }

    this.s3 = null;
    this.s3Bucket = process.env.ARTILLERY_S3_BUCKET || null;
    if (this.sqs && this.s3Bucket) {
      this.s3 = new S3Client({
        region:
          process.env.SQS_REGION || script.config.plugins['sqs-reporter'].region
      });
    }

    if (process.env.AZURE_STORAGE_QUEUE_URL) {
      this.aqs = getAQS();
    }

    events.on('stats', (statsOriginal: Record<string, any>) => {
      const serialized =
        global.artillery.__SSMS.serializeMetrics(statsOriginal);
      const body = {
        event: 'workerStats',
        stats: serialized
      };

      this.sendMessage(body);
    });

    //TODO: reconcile some of this code with how lambda does sqs reporting
    events.on('phaseStarted', (phaseContext: Record<string, any>) => {
      const body = {
        event: 'phaseStarted',
        phase: phaseContext
      };

      this.sendMessage(body);
    });

    //TODO: reconcile some of this code with how lambda does sqs reporting
    events.on('phaseCompleted', (phaseContext: Record<string, any>) => {
      const body = {
        event: 'phaseCompleted',
        phase: phaseContext
      };
      this.sendMessage(body);
    });

    events.on('done', (_stats: Record<string, any>) => {
      const body = {
        event: 'done',
        stats: global.artillery.__SSMS.serializeMetrics(_stats)
      };
      this.sendMessage(body);
    });

    global.artillery.globalEvents.on(
      'log',
      (opts: unknown, ...args: unknown[]) => {
        if (process.env.SHIP_LOGS) {
          const body = {
            event: 'artillery.log',
            log: {
              opts,
              args: [...args]
            }
          };

          this.sendMessage(body);
        }
      }
    );
  }

  cleanup(done: (err: Error | null) => void) {
    const interval = setInterval(() => {
      if (this.unsent <= 0) {
        clearInterval(interval);
        done(null);
      }
    }, 200).unref();
  }

  sendMessage(body: Record<string, any>) {
    if (this.sqs) {
      this.sendSQS(body);
    } else {
      this.sendAQS(body);
    }
  }

  async sendSQS(body: Record<string, any>) {
    this.unsent++;

    const payload = JSON.stringify(body);
    const payloadSize = Buffer.byteLength(payload, 'utf8');

    try {
      let messageBody = payload;

      // Upload to S3 if payload exceeds SQS limit
      if (payloadSize > SQS_SIZE_LIMIT && this.s3 && this.s3Bucket) {
        const workerId = this.tags.find((t) => t.key === 'workerId')?.value;
        const messageId = uuid();
        const s3Key = `tests/${this.testId}/overflow/${workerId}/${messageId}.json`;

        await this.s3.send(
          new PutObjectCommand({
            Bucket: this.s3Bucket,
            Key: s3Key,
            Body: payload,
            ContentType: 'application/json'
          })
        );

        messageBody = JSON.stringify({
          event: body.event,
          _overflowRef: s3Key
        });

        _debug(
          'Payload %d bytes exceeded limit, uploaded to S3: %s',
          payloadSize,
          s3Key
        );
      }

      const params = {
        MessageBody: messageBody,
        QueueUrl: this.queueUrl,
        MessageAttributes: this.messageAttributes,
        MessageDeduplicationId: uuid(),
        MessageGroupId: this.testId as string
      };

      // sendSQS is only called when this.sqs is set:
      await (this.sqs as SQSClient).send(new SendMessageCommand(params));
    } catch (err) {
      console.error(err);
    } finally {
      this.unsent--;
    }
  }

  async sendAQS(body: Record<string, any>) {
    this.unsent++;
    // sendAQS is only called when the Azure queue client is configured:
    sendMessage(this.aqs as NonNullable<typeof this.aqs>, body, this.tags)
      .then((_res) => {
        this.unsent--;
      })
      .catch((err) => {
        console.error(err);
        this.unsent--;
      });
  }
}

export default {
  Plugin: ArtillerySQSPlugin,
  LEGACY_METRICS_FORMAT: false
};