Add initial load tests (#3331)

.eslintignore

@@ -18,3 +18,4 @@ globalSetup.js
 globalTeardown.js
 jest-css-modules-transform-config.js
 cmd/profilecli/**
+tools/k6/**

.gitignore

@@ -42,3 +42,4 @@ yarn-error.log
 /.docker-image-id-pyroscope
 
 /.tmp
+tools/k6/.env

tools/k6/.env.template

@@ -0,0 +1,21 @@
+## Copy this template to .env and replace the values with your own.
+
+## reads.js
+## (REQUIRED) The base URL of the application under test.
+K6_BASE_URL="http://localhost:4040"
+
+## The read token to authenticate with the server. Only used if the server
+## requires authentication.
+K6_READ_TOKEN=""
+
+## The tenant ID to use for the test, if in running against a multi-tenant
+## environment.
+K6_TENANT_ID=""
+
+## (DEFAULT "fire-dev-001/ingester") The service name to use in the read
+## queries.
+K6_QUERY_SERVICE_NAME=""
+
+## (DEFAULT "1h") Comma-separated list of durations to use in the read queries.
+## These durations will be used in "last X" time ranges.
+K6_QUERY_DURATIONS=""

tools/k6/README.md

@@ -0,0 +1,38 @@
+# Load testing for Pyroscope
+
+This directory contains load test scripts and helpers. Load test scripts are stored in the `tests` directory, whilst helper functions are kept in `lib`.
+
+## Quick start
+
+### TL;DR
+
+To run a test against a local Pyroscope, run
+
+```bash
+./run.sh testname
+```
+
+To run a test from k6 cloud, run
+
+```bash
+./run.sh -c testname
+```
+
+Keep in mind when running from k6 cloud, you will need to have a properly configured `.env` file.
+
+### Configuration options
+
+Each load test can be configured with specific environment variables. Descriptions of these variables and their defaults can be found in the `.env.template` file. To modify a test's behavior, copy `.env.template` to `.env`, then modify the variables how you see fit. Once done, `run.sh` will use the new values to configure the test's behavior.
+
+## Tests
+
+### `read.js`
+
+This will run load tests targeting Pyroscope's read API. It issues the following queries in "last 1 hour" and "last 24 hour" time windows:
+
+- `SelectMergeProfile`
+- `/render`
+- `SelectMergeStacktraces`
+- `LabelNames`
+- `Series`
+- `/render-diff`

tools/k6/lib/env.js

@@ -0,0 +1,5 @@
+import { fail } from 'k6';
+
+export const READ_TOKEN = __ENV.K6_READ_TOKEN;
+export const TENANT_ID = __ENV.K6_TENANT_ID;
+export const BASE_URL = __ENV.K6_BASE_URL || fail('K6_BASE_URL environment variable missing');

tools/k6/lib/request.js

@@ -0,0 +1,92 @@
+import { check } from 'k6';
+import http from 'k6/http';
+import encoding from 'k6/encoding';
+
+import { URL } from 'https://jslib.k6.io/url/1.0.0/index.js';
+
+import { READ_TOKEN, TENANT_ID, BASE_URL } from './env.js';
+
+export function doSelectMergeProfileRequest(body, headers) {
+  const res = http.post(`${BASE_URL}/querier.v1.QuerierService/SelectMergeProfile`, JSON.stringify(body), {
+    headers: withHeaders(headers),
+  });
+
+  check(res, {
+    'status is 200': (r) => r.status === 200,
+  });
+}
+
+export function doRenderRequest(body, headers) {
+  const params = new URL(`${BASE_URL}/pyroscope/render`);
+  for (const [k, v] of Object.entries(body)) {
+    params.searchParams.set(k, v);
+  }
+
+  const res = http.get(params.toString(), {
+    headers: withHeaders(headers),
+    tags: { name: '/render' },
+  });
+
+  check(res, {
+    'status is 200': (r) => r.status === 200,
+  });
+}
+
+export function doSelectMergeStacktracesRequest(body, headers) {
+  const res = http.post(`${BASE_URL}/querier.v1.QuerierService/SelectMergeStacktraces`, JSON.stringify(body), {
+    headers: withHeaders(headers),
+  });
+
+  check(res, {
+    'status is 200': (r) => r.status === 200,
+  });
+}
+
+export function doLabelNamesRequest(body, headers) {
+  const res = http.post(`${BASE_URL}/querier.v1.QuerierService/LabelNames`, JSON.stringify(body), {
+    headers: withHeaders(headers),
+  });
+
+  check(res, {
+    'status is 200': (r) => r.status === 200,
+  });
+}
+
+export function doSeriesRequest(body, headers) {
+  const res = http.post(`${BASE_URL}/querier.v1.QuerierService/Series`, JSON.stringify(body), {
+    headers: withHeaders(headers),
+  });
+
+  check(res, {
+    'status is 200': (r) => r.status === 200,
+  });
+}
+
+export function doRenderDiffRequest(body, headers) {
+  const params = new URL(`${BASE_URL}/pyroscope/render-diff`);
+  for (const [k, v] of Object.entries(body)) {
+    params.searchParams.set(k, v);
+  }
+
+  const res = http.get(params.toString(), {
+    headers: withHeaders(headers),
+    tags: { name: '/render' },
+  });
+
+  check(res, {
+    'status is 200': (r) => r.status === 200,
+  });
+}
+
+function withHeaders(headers) {
+  const baseHeaders = {
+    'User-Agent': 'k6-load-test',
+    'Content-Type': 'application/json',
+    'Authorization': `Basic ${encoding.b64encode(`${TENANT_ID}:${READ_TOKEN}`)}`
+  };
+
+  for (const [k, v] of Object.entries(headers || {})) {
+    baseHeaders[k] = v;
+  }
+  return baseHeaders;
+}

tools/k6/lib/time.js

@@ -0,0 +1,15 @@
+export function newRelativeTimeRange(scalar, unit) {
+  const end = Date.now();
+  switch (unit) {
+    case 's':
+      return { start: end - scalar * 1000, end };
+    case 'm':
+      return { start: end - scalar * 60 * 1000, end };
+    case 'h':
+      return { start: end - scalar * 60 * 60 * 1000, end };
+    case 'd':
+      return { start: end - scalar * 24 * 60 * 60 * 1000, end };
+    default:
+      throw new Error(`Invalid unit: ${unit}`);
+  }
+}

tools/k6/run.sh

@@ -0,0 +1,67 @@
+#! /bin/bash
+
+set -euo pipefail
+
+usage() {
+  echo "Usage: $0 [-c] file"
+  echo "    file        File name of the test to run"
+  echo "      -c        Run the test in the cloud"
+  echo "      -h        Print help message and list all available tests"
+}
+
+tests() {
+  echo "Available tests:"
+  for file in tools/k6/tests/*.js; do
+    echo "  $(basename "${file}")"
+  done
+}
+
+IS_CLOUD=
+while getopts "hc" opt; do
+  case "${opt}" in
+    c)
+      IS_CLOUD=1
+      ;;
+    h)
+      usage
+      echo # Add a newline.
+      tests
+      exit 0
+      ;;
+    *)
+      usage
+      exit 1
+      ;;
+  esac
+done
+shift $((OPTIND-1))
+
+if [ "$#" -lt 1 ]; then
+  usage
+  exit 1
+fi
+
+TEST=$1
+if [ -z "${TEST}" ]; then
+  usage
+  exit 1
+fi
+
+DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" ; pwd -P)
+source "${DIR}"/.env
+
+if [ -n "${IS_CLOUD}" ]; then
+  k6 cloud tools/k6/tests/"${TEST}" \
+    -e "K6_BASE_URL=${K6_BASE_URL}" \
+    -e "K6_READ_TOKEN=${K6_READ_TOKEN}" \
+    -e "K6_TENANT_ID=${K6_TENANT_ID}" \
+    -e "K6_QUERY_SERVICE_NAME=${K6_QUERY_SERVICE_NAME}" \
+    -e "K6_QUERY_DURATIONS=${K6_QUERY_DURATIONS}"
+else
+  K6_BASE_URL="${K6_BASE_URL}" \
+  K6_READ_TOKEN="${K6_READ_TOKEN}" \
+  K6_TENANT_ID="${K6_TENANT_ID}" \
+  K6_QUERY_SERVICE_NAME="${K6_QUERY_SERVICE_NAME}" \
+  K6_QUERY_DURATIONS="${K6_QUERY_DURATIONS}" \
+  k6 run tools/k6/tests/"${TEST}"
+fi