who_need_help/load/k6/public-and-websocket.js

import { check, sleep } from "k6";
import http from "k6/http";
import { Counter } from "k6/metrics";
import { WebSocket } from "k6/websockets";

function required(name) {
  const value = __ENV[name];

  if (value === undefined || value === "") {
    throw new Error(`${name} is required`);
  }

  return value;
}

function positiveInteger(name) {
  const value = Number(required(name));

  if (!Number.isInteger(value) || value <= 0) {
    throw new Error(`${name} must be a positive integer`);
  }

  return value;
}

function nonNegativeNumber(name) {
  const value = Number(required(name));

  if (!Number.isFinite(value) || value < 0) {
    throw new Error(`${name} must be a non-negative number`);
  }

  return value;
}

const baseUrl = required("BASE_URL").replace(/\/+$/, "");

if (!/^https?:\/\/[^/?#]+$/.test(baseUrl)) {
  throw new Error("BASE_URL must be an HTTP(S) origin without a path");
}

const websocketUrl =
  baseUrl.replace(/^http/, "ws") + "/live/websocket?vsn=2.0.0";
const websocketHoldMs = positiveInteger("WS_HOLD_MS");
const websocketConnectTimeoutMs = positiveInteger("WS_CONNECT_TIMEOUT_MS");
const httpThinkSeconds = nonNegativeNumber("HTTP_THINK_SECONDS");

const websocketOpened = new Counter("wnh_websocket_opened");
const websocketErrors = new Counter("wnh_websocket_errors");
const heartbeatReplies = new Counter("wnh_websocket_heartbeat_replies");

export const options = {
  discardResponseBodies: true,
  scenarios: {
    public_http: {
      executor: "constant-vus",
      exec: "publicHttp",
      vus: positiveInteger("HTTP_VUS"),
      duration: required("DURATION"),
    },
    phoenix_websocket: {
      executor: "constant-vus",
      exec: "phoenixWebsocket",
      vus: positiveInteger("WS_VUS"),
      duration: required("DURATION"),
    },
  },
};

export function publicHttp() {
  const responses = http.batch([
    {
      method: "GET",
      url: `${baseUrl}/`,
      tags: { endpoint: "home" },
    },
    {
      method: "GET",
      url: `${baseUrl}/safety`,
      tags: { endpoint: "safety" },
    },
    {
      method: "GET",
      url: `${baseUrl}/healthz/ready`,
      tags: { endpoint: "readiness" },
    },
  ]);

  check(responses[0], { "home returned 200": (response) => response.status === 200 });
  check(responses[1], { "safety returned 200": (response) => response.status === 200 });
  check(responses[2], {
    "readiness returned 200": (response) => response.status === 200,
  });

  sleep(httpThinkSeconds);
}

export function phoenixWebsocket() {
  const socket = new WebSocket(websocketUrl, [], {
    tags: { endpoint: "phoenix_live_socket" },
  });
  const ref = `${__VU}-${__ITER}`;
  let opened = false;
  let closeTimer;
  const connectTimer = setTimeout(() => {
    if (!opened) {
      websocketErrors.add(1);
      socket.close();
    }
  }, websocketConnectTimeoutMs);

  socket.addEventListener("open", () => {
    opened = true;
    clearTimeout(connectTimer);
    websocketOpened.add(1);
    socket.send(JSON.stringify([null, ref, "phoenix", "heartbeat", {}]));
    closeTimer = setTimeout(() => socket.close(), websocketHoldMs);
  });

  socket.addEventListener("message", (event) => {
    let frame;

    try {
      frame = JSON.parse(event.data);
    } catch (_error) {
      return;
    }

    if (
      Array.isArray(frame) &&
      frame[1] === ref &&
      frame[2] === "phoenix" &&
      frame[3] === "phx_reply" &&
      frame[4] &&
      frame[4].status === "ok"
    ) {
      heartbeatReplies.add(1);
    }
  });

  socket.addEventListener("error", () => {
    websocketErrors.add(1);
    clearTimeout(connectTimer);
    socket.close();
  });

  socket.addEventListener("close", () => {
    clearTimeout(connectTimer);

    if (closeTimer !== undefined) {
      clearTimeout(closeTimer);
    }
  });
}