k6-script.js.j2

import http from 'k6/http';
import { group, check, sleep } from 'k6';
import exec from 'k6/execution';
import encoding from 'k6/encoding';
import crypto from 'k6/crypto';
import { Gauge } from 'k6/metrics';

// --- Env / config ---
const payloadServerUrl = __ENV.EXPB_PAYLOAD_SERVER_URL;
const payloadsWarmup = parseInt(__ENV.EXPB_PAYLOADS_WARMUP || '0', 10);
const payloadsDelay = parseFloat(__ENV.EXPB_PAYLOADS_DELAY || '0');
const payloadsWarmupDelay = parseFloat(__ENV.EXPB_PAYLOADS_WARMUP_DELAY || '0');
const warmupWait = parseInt(__ENV.EXPB_WARMUP_WAIT || '0', 10);
const abortOnEOF = (__ENV.EXPB_ABORT_ON_EOF || '0') === '1';
const addCorrelationHeader = (__ENV.EXPB_ADD_CID || '1') === '1';
const engineEndpoint = __ENV.EXPB_ENGINE_ENDPOINT;
const perPayloadMetrics = (__ENV.EXPB_PER_PAYLOAD_METRICS || '0') === '1';
const discardResponses = (__ENV.EXPB_DISCARD_RESPONSES || '1') === '1';
const enableLogging = (__ENV.EXPB_ENABLE_LOGGING || '0') === '1';
const perPayloadMetricsLogs = (__ENV.EXPB_PER_PAYLOAD_METRICS_LOGS || '0') === '1';
const verbosePostLogs = (__ENV.EXPB_VERBOSE_POST_LOGS || '0') === '1';

// Load k6 options JSON
const config = JSON.parse(open(__ENV.EXPB_CONFIG_FILE_PATH));
export const options = config["options"];

// --- Metrics ---
const slowestPayloads = new Gauge('expb_slowest_payloads');
const slowestProcessing = new Gauge('expb_slowest_processing');
const clientProcessing = new Gauge('expb_client_processing');

// --- JWT ---
function hex2ArrayBuffer(hex) {
  const buf = new ArrayBuffer(hex.length / 2);
  const view = new Uint8Array(buf);
  for (let i = 0; i < hex.length; i += 2) view[i / 2] = parseInt(hex.slice(i, i + 2), 16);
  return buf;
}
const jwtsecretBytes = hex2ArrayBuffer(open(__ENV.EXPB_JWTSECRET_FILE_PATH).trim());

// Per-VU JWT cache (refresh ~once/minute)
// Pre-compute static JWT parts to avoid repeated JSON.stringify
const jwtHeaderB64 = encoding.b64encode(JSON.stringify({ typ: 'JWT', alg: 'HS256' }), 'rawurl');
let jwtCache = { token: '', exp: 0 };
function getJwtToken() {
  const now = Math.floor(Date.now() / 1000);
  if (jwtCache.token && now < (jwtCache.exp - 2)) return jwtCache.token;

  const iat = now, exp = iat + 60;
  const payload = encoding.b64encode(JSON.stringify({ iat, exp }), 'rawurl');
  const h = crypto.createHMAC('sha256', jwtsecretBytes);
  h.update(jwtHeaderB64 + '.' + payload);
  const sig = h.digest('base64rawurl');
  jwtCache = { token: `${jwtHeaderB64}.${payload}.${sig}`, exp };
  return jwtCache.token;
}

function cid() {
  return `${__ENV.testid || 'expb'}-vu${exec.vu.idInTest}-it${exec.scenario.iterationInTest}`;
}

function logPost(kind, idx, warmup, correlationId = '') {
  if (!enableLogging || !verbosePostLogs) return;
  const ts = new Date().toISOString();
  const warmupTag = warmup ? ' warmup=1' : '';
  const cidTag = correlationId ? ` cid=${correlationId}` : '';
  console.log(`[${ts}] POST ${kind} idx=${idx}${warmupTag}${cidTag}`);
}

function emitPerPayloadMetricsRow(idx, gasUsed, processingMs) {
  if (!enableLogging || !perPayloadMetricsLogs) return;
  const gas = gasUsed !== undefined && gasUsed !== '' ? gasUsed : 'n/a';
  const processing = Number.isFinite(processingMs) ? processingMs.toFixed(2) : 'n/a';
  console.log(
    `EXPB_PER_PAYLOAD_METRIC idx=${idx} gas_used=${gas} processing_ms=${processing}`
  );
}

// --- Fetch next pair from payload server ---
// Line format: {metadata_json}\t{raw_NP}\t{raw_FCU}
// Metadata is pre-extracted by the executor, no regex needed here.

function fetchNextPair() {
  const resp = http.get(payloadServerUrl + '/next', {
    tags: { name: 'payload_fetch' },
    responseType: 'text',
  });

  if (resp.status === 404) {
    return null;
  }

  if (resp.status !== 200) {
    console.error(`[payload-server] unexpected status: ${resp.status}`);
    return null;
  }

  // Parse merged line: {metadata}\t{NP}\t{FCU}
  const body = resp.body;
  const firstTab = body.indexOf('\t');
  if (firstTab < 0) {
    console.error('[payload-server] malformed line: no tab separator');
    return null;
  }
  const secondTab = body.indexOf('\t', firstTab + 1);
  if (secondTab < 0) {
    console.error('[payload-server] malformed line: missing second tab separator');
    return null;
  }

  const meta = JSON.parse(body.substring(0, firstTab));
  const rawPayload = body.substring(firstTab + 1, secondTab);
  const rawFcu = body.substring(secondTab + 1);

  return {
    rawPayload: rawPayload,
    rawFcu: rawFcu,
    idx: meta.idx,
    method: meta.method || 'engine_newPayload',
    jrpcId: meta.jrpc_id,
    gasUsed: meta.gas_used,
    fcuMethod: meta.fcu_method || 'engine_forkchoiceUpdated',
    prevClientProcessingMs: meta.prev_client_processing_ms,
  };
}

// Track previous pair so we can tag client_processing metrics correctly
// (the metric for block N arrives in block N+1's metadata)
let prevPair = null;

function processNextPayload(base_tags = {}, warmup = false) {
  const pair = fetchNextPair();

  if (!pair) {
    if (abortOnEOF) exec.test.abort('No more payloads or fcus found');
    return;
  }

  if (!pair.rawPayload || !pair.rawFcu) {
    if (abortOnEOF) exec.test.abort('Received empty payload or FCU from server');
    return;
  }

  const idx = pair.idx;

  // Track start time to calculate elapsed time
  const startTime = Date.now();

  const correlationId = (!warmup && addCorrelationHeader) ? cid() : '';

  // --- engine_newPayload ---
  group('engine_newPayload', function () {
    // Generate fresh JWT token just before the request
    const jwtToken = getJwtToken();
    const headers = {
      'Authorization': 'Bearer ' + jwtToken,
      'Content-Type': 'application/json',
    };
    if (correlationId) {
      headers['X-Expb-Cid'] = correlationId;
    }

    const tags = {
      jrpc_method: pair.method,
      kind: 'newPayload',
      ...base_tags
    };

    logPost('engine_newPayload', idx, warmup, correlationId);
    const r = http.post(engineEndpoint, pair.rawPayload, {
      headers: headers,
      tags: tags,
      responseType: discardResponses ? 'none' : 'text',
    });
    check(r, { 'status_200': (x) => x.status === 200 }, tags);

    if (!warmup) {
      emitPerPayloadMetricsRow(idx, pair.gasUsed, r.timings.waiting);
    }

    if (!warmup && perPayloadMetrics) {
      const perPayloadMetricsTags = { ...tags };
      if (pair.jrpcId !== undefined) {
        perPayloadMetricsTags.jrpc_id = pair.jrpcId;
      }
      if (pair.gasUsed !== undefined) {
        perPayloadMetricsTags.gas_used = pair.gasUsed;
      }
      slowestPayloads.add(r.timings.duration, perPayloadMetricsTags);
      slowestProcessing.add(r.timings.waiting, perPayloadMetricsTags);
    }
  });

  // --- engine_forkchoiceUpdated ---
  group('engine_forkchoiceUpdated', function () {
    // Generate fresh JWT token just before the request
    const jwtToken = getJwtToken();
    const headers = {
      'Authorization': 'Bearer ' + jwtToken,
      'Content-Type': 'application/json',
    };
    if (correlationId) {
      headers['X-Expb-Cid'] = correlationId;
    }

    const tags = {
      jrpc_method: pair.fcuMethod,
      kind: 'forkchoiceUpdated',
      ...base_tags
    };

    logPost('engine_forkchoiceUpdated', idx, warmup, correlationId);
    const r = http.post(engineEndpoint, pair.rawFcu, {
      headers: headers,
      tags: tags,
      responseType: discardResponses ? 'none' : 'text',
    });
    check(r, { 'status_200': (x) => x.status === 200 }, tags);
  });

  // Emit client-side processing metric for the PREVIOUS block.
  // The metric for block N-1 arrives in block N's metadata from the payload server.
  // Skip 0 values — the gauge may not have been updated yet when scraped.
  // Always tag with jrpc_id/gas_used so Grafana can align by block, not timestamp.
  if (!warmup && pair.prevClientProcessingMs > 0 && prevPair) {
    const cpTags = { kind: 'newPayload', jrpc_method: prevPair.method, ...base_tags };
    if (prevPair.jrpcId !== undefined) cpTags.jrpc_id = prevPair.jrpcId;
    if (prevPair.gasUsed !== undefined) cpTags.gas_used = prevPair.gasUsed;
    clientProcessing.add(pair.prevClientProcessingMs, cpTags);
  }
  prevPair = warmup ? null : pair;

  // delay between payloads - wait only the difference between configured delay and elapsed time
  const delay = warmup ? payloadsWarmupDelay : payloadsDelay;
  if (delay > 0) {
    const elapsedTime = (Date.now() - startTime) / 1000; // Convert to seconds
    const remainingDelay = delay - elapsedTime;
    if (remainingDelay > 0) {
      sleep(remainingDelay);
    }
  }
}

const base_tags = { testid: __ENV.testid || 'expb', tool: 'expb' };

export function setup() {
  // execute warmup payloads (server counter starts at skip, so first N are warmup)
  for (let i = 0; i < payloadsWarmup; i++) {
    processNextPayload(base_tags, true);
  }
}

export default function () {
  if (exec.scenario.iterationInTest === 0 && warmupWait > 0) {
    sleep(warmupWait);
  }

  processNextPayload(base_tags, false);
}