artillery/lib/platform/local/artillery-worker-local.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 path from 'node:path';
import { Worker } from 'node:worker_threads';
import { EventEmitter } from 'eventemitter3';
import awaitOnEE from '../../util/await-on-ee.ts';
import STATES from '../worker-states.ts';
import type { WorkerState } from '../worker-states.ts';
import type {
  PrepareWorkerOptions,
  WorkerCommand,
  WorkerEnvelope
} from './protocol.ts';
import type { StashDetails } from '../../stash.ts';

// Emitted output runs from dist/ where the sibling module is worker.js;
// when this file is loaded directly from source (tests importing .ts),
// the sibling is worker.ts (worker threads strip types natively).
const workerModulePath = path.join(
  import.meta.dirname,
  import.meta.filename.endsWith('.ts') ? 'worker.ts' : 'worker.js'
);

const returnWorkerEnv = (needsSourcemap: unknown) => {
  const env = { ...process.env };

  if (needsSourcemap) {
    env.NODE_OPTIONS = process.env.NODE_OPTIONS
      ? `${process.env.NODE_OPTIONS} --enable-source-maps`
      : '--enable-source-maps';
  }

  return env;
};

class ArtilleryWorker {
  declare opts: unknown;
  declare events: EventEmitter;
  declare workerEvents: EventEmitter;
  declare worker: Worker;
  declare workerId: number;
  declare state: WorkerState;

  constructor(opts?: unknown) {
    this.opts = opts;
    this.events = new EventEmitter(); // events for consumers of this object
    this.workerEvents = new EventEmitter(); // turn events delivered via 'message' events into their own messages
  }

  async init(_opts?: unknown) {
    this.state = STATES.initializing;

    const workerEnv = returnWorkerEnv(global.artillery.hasTypescriptProcessor);

    this.worker = new Worker(workerModulePath, {
      env: workerEnv
    });
    this.workerId = this.worker.threadId;
    this.worker.on('error', this.onError.bind(this));
    // TODO:
    this.worker.on('exit', (exitCode: number) => {
      this.events.emit('exit', exitCode);
    });

    this.worker.on('messageerror', (_err: Error) => {});

    // TODO: Expose performance metrics via getHeapSnapshot() and performance object.

    await awaitOnEE(this.worker, 'online', 10);

    // Relay messages onto the real event emitter:
    this.worker.on('message', (message: WorkerEnvelope) => {
      switch (message.event) {
        case 'log':
          this.events.emit('log', message);
          this.workerEvents.emit('log', message);
          break;
        case 'workerError':
          this.events.emit('workerError', message);
          this.workerEvents.emit('workerError', message);
          break;
        case 'phaseStarted':
          this.events.emit('phaseStarted', message);
          this.workerEvents.emit('phaseStarted', message);
          break;
        case 'phaseCompleted':
          this.events.emit('phaseCompleted', message);
          this.workerEvents.emit('phaseCompleted', message);
          break;
        case 'stats':
          this.events.emit('stats', message);
          this.workerEvents.emit('stats', message);
          break;
        case 'done':
          this.events.emit('done', message);
          this.workerEvents.emit('done', message);
          break;
        case 'running':
          this.events.emit('running', message);
          this.workerEvents.emit('running', message);
          break;
        case 'readyWaiting':
          this.events.emit('readyWaiting', message);
          this.workerEvents.emit('readyWaiting', message);
          break;
        case 'setSuggestedExitCode':
          this.events.emit('setSuggestedExitCode', message);
          break;
        default:
          global.artillery.log(
            `Unknown message from worker ${message}`,
            'error'
          );
      }
    });

    this.state = STATES.online;
  }

  async prepare(opts: {
    script: Record<string, any>;
    payload: unknown;
    options: Record<string, any>;
    stashDetails?: StashDetails | null;
  }) {
    this.state = STATES.preparing;

    const { script, payload, options, stashDetails } = opts;
    let scriptForWorker: Record<string, any> = script;

    if (script.__transpiledTypeScriptPath && script.__originalScriptPath) {
      scriptForWorker = {
        __transpiledTypeScriptPath: script.__transpiledTypeScriptPath,
        __originalScriptPath: script.__originalScriptPath,
        __phases: script.config?.phases
      };
    }

    const prepareOpts: PrepareWorkerOptions = {
      script: scriptForWorker,
      payload,
      options,
      testRunId: global.artillery.testRunId,
      stashDetails
    };
    const command: WorkerCommand = {
      command: 'prepare',
      opts: prepareOpts
    };
    this.worker.postMessage(command);

    await awaitOnEE(this.workerEvents, 'readyWaiting', 50);
    this.state = STATES.readyWaiting;
  }

  async run(opts: string) {
    const command: WorkerCommand = {
      command: 'run',
      opts: JSON.parse(opts)
    };
    this.worker.postMessage(command);

    await awaitOnEE(this.workerEvents, 'running', 50);
    this.state = STATES.running;
  }

  async stop() {
    const command: WorkerCommand = { command: 'stop' };
    this.worker.postMessage(command);
  }

  onError(err: Error) {
    // TODO: set state, clean up
    this.events.emit('error', err);
    console.log('worker error, id:', this.workerId, err);
  }
}

export { ArtilleryWorker, STATES };