bitvid/scripts/agent/load-test.mjs

import WebSocket from 'ws';
import { finalizeEvent, getPublicKey, generateSecretKey } from 'nostr-tools/pure';
import fs from 'fs';
import path from 'path';
import { spawn } from 'child_process';
import { fileURLToPath } from 'url';
import { buildVideoPostEvent, buildViewEvent } from '../../js/nostrEventSchemas.js';

// --- Configuration ---
let RELAY_URL = process.env.RELAY_URL;
const CLIENTS = parseInt(process.env.CLIENTS || '1000', 10);
const DURATION_SEC = parseInt(process.env.DURATION_SEC || '600', 10);
const RATE_EPS = parseFloat(process.env.RATE_EPS || '10');
const MIX_VIDEO_RATIO = parseFloat(process.env.MIX || '0.5'); // 0.0 to 1.0
const DRY_RUN = process.env.DRY_RUN === '1';
const VERBOSE = process.env.VERBOSE === '1';
const FORCE_UNSAFE = process.env.FORCE_UNSAFE === '1';

// --- Safety Checks ---
function isSafeRelayUrl(url) {
  if (!url) return false;
  try {
    const parsed = new URL(url);
    const hostname = parsed.hostname;

    if (hostname === 'localhost' || hostname === '127.0.0.1' || hostname === '::1') return true;

    // Private IP ranges
    // 10.0.0.0 - 10.255.255.255
    // 172.16.0.0 - 172.31.255.255
    // 192.168.0.0 - 192.168.255.255
    const parts = hostname.split('.').map(Number);
    if (parts.length === 4) {
      if (parts[0] === 10) return true;
      if (parts[0] === 192 && parts[1] === 168) return true;
      if (parts[0] === 172 && parts[1] >= 16 && parts[1] <= 31) return true;
    }

    return false;
  } catch (e) {
    return false;
  }
}

// --- Metrics ---
const metrics = {
  startTime: Date.now(),
  sent: 0,
  accepted: 0,
  rejected: 0,
  errors: 0,
  latencies: [], // ms
  connectedClients: 0,
  resources: [], // { time, rss, heapTotal, heapUsed, external, cpuUser, cpuSystem }
};

// --- Helpers ---
const delay = (ms) => new Promise((resolve) => setTimeout(resolve, ms));

function generateKeypair() {
  const secretKey = generateSecretKey();
  const sk = Buffer.from(secretKey).toString('hex');
  const pk = getPublicKey(secretKey);
  return { sk, pk };
}

// Event builders are imported from js/nostrEventSchemas.js

// --- Client Class ---
class LoadClient {
  constructor(id, relayUrl) {
    this.id = id;
    this.relayUrl = relayUrl;
    this.ws = null;
    this.keys = generateKeypair();
    this.isOpen = false;
    this.pendingPublishes = new Map(); // eventId -> startTime
  }

  connect() {
    return new Promise((resolve, reject) => {
      if (DRY_RUN) {
        this.isOpen = true;
        metrics.connectedClients++;
        resolve();
        return;
      }

      this.ws = new WebSocket(this.relayUrl);

      this.ws.on('open', () => {
        this.isOpen = true;
        metrics.connectedClients++;
        resolve();
      });

      this.ws.on('error', (err) => {
        if (VERBOSE) console.error(`Client ${this.id} error:`, err.message);
        metrics.errors++;
        if (!this.isOpen) reject(err);
      });

      this.ws.on('close', () => {
        if (this.isOpen) metrics.connectedClients--;
        this.isOpen = false;
      });

      this.ws.on('message', (data) => {
        try {
          const msg = JSON.parse(data);
          // ["OK", eventId, true, "message"]
          if (Array.isArray(msg) && msg[0] === 'OK') {
            const eventId = msg[1];
            const accepted = msg[2];
            const message = msg[3];
            this.handleOk(eventId, accepted, message);
          }
        } catch (err) {
          // ignore malformed
        }
      });
    });
  }

  handleOk(eventId, accepted, message) {
    const startTime = this.pendingPublishes.get(eventId);
    if (startTime) {
      const latency = Date.now() - startTime;
      metrics.latencies.push(latency);
      this.pendingPublishes.delete(eventId);

      if (accepted) {
        metrics.accepted++;
      } else {
        metrics.rejected++;
        if (VERBOSE) console.warn(`Event ${eventId} rejected: ${message}`);
      }
    }
  }

  async publish() {
    if (!this.isOpen && !DRY_RUN) return;

    // Decide event type
    const isVideo = Math.random() < MIX_VIDEO_RATIO;
    const dTag = `load-${this.id}-${Date.now()}-${Math.random().toString(36).slice(2)}`;

    let eventTemplate;
    const now = Math.floor(Date.now() / 1000);
    if (isVideo) {
      eventTemplate = buildVideoPostEvent({
        pubkey: this.keys.pk,
        created_at: now,
        dTagValue: dTag,
        content: {
          version: 3,
          title: `Load Test Video ${dTag}`,
          videoRootId: dTag,
          url: `https://example.com/videos/${dTag}.mp4`,
          magnet: `magnet:?xt=urn:btih:${dTag}&dn=video`,
          mode: 'dev',
          description: 'Generated by load-test-agent'
        }
      });
    } else {
      eventTemplate = buildViewEvent({
        pubkey: this.keys.pk,
        created_at: now,
        dedupeTag: dTag,
        pointerTag: ['s', `nostr:${dTag}`],
        content: 'View event'
      });
    }

    // Sign
    const skBytes = Uint8Array.from(Buffer.from(this.keys.sk, 'hex'));
    try {
      const signedEvent = finalizeEvent(eventTemplate, skBytes);
      if (DRY_RUN) {
        metrics.sent++;
        metrics.accepted++; // Simulate acceptance
        metrics.latencies.push(Math.random() * 50); // Simulate latency
        return;
      }

      // Send
      const msg = JSON.stringify(['EVENT', signedEvent]);
      this.pendingPublishes.set(signedEvent.id, Date.now());

      try {
        this.ws.send(msg);
        metrics.sent++;
      } catch (err) {
        metrics.errors++;
        this.pendingPublishes.delete(signedEvent.id);
      }
    } catch (err) {
      console.error('Failed to finalize event:', err);
      console.error('Event template:', JSON.stringify(eventTemplate, null, 2));
      metrics.errors++;
    }

  }

  close() {
    if (this.ws) {
      this.ws.close();
    }
  }
}

// --- Main Execution ---
async function run() {
  let relayProcess = null;

  // Auto-start logic
  if (!RELAY_URL && !DRY_RUN) {
    const scriptDir = path.dirname(fileURLToPath(import.meta.url));
    const relayScript = path.resolve(scriptDir, 'simple-relay.mjs');

    if (fs.existsSync(relayScript)) {
      console.log('No RELAY_URL provided. Starting local relay...');
      RELAY_URL = 'ws://localhost:8888';

      relayProcess = spawn('node', [relayScript], {
        stdio: VERBOSE ? 'inherit' : 'ignore',
        detached: false
      });

      relayProcess.unref(); // Allow independent life, but we kill it on exit
      await delay(2000); // Wait for boot
    } else {
      console.error('Error: RELAY_URL not set and local relay script not found.');
      process.exit(1);
    }
  }

  if (!DRY_RUN && !isSafeRelayUrl(RELAY_URL) && !FORCE_UNSAFE) {
    console.error(`
SAFETY BLOCK: The target relay (${RELAY_URL}) appears to be public.
Load testing against public relays is strictly forbidden without explicit overrides.
To force this run, set FORCE_UNSAFE=1 (use with extreme caution).
    `);
    if (relayProcess) relayProcess.kill();
    process.exit(1);
  }

  console.log(`
Load Test Configuration:
  Relay: ${RELAY_URL || '(Dry Run)'}
  Clients: ${CLIENTS}
  Duration: ${DURATION_SEC}s
  Rate: ${RATE_EPS} events/sec
  Mix: ${MIX_VIDEO_RATIO * 100}% Video Posts
  Dry Run: ${DRY_RUN}
`);

  const clients = [];

  // 1. Init Clients
  console.log(`Initializing ${CLIENTS} clients...`);
  const connectPromises = [];
  for (let i = 0; i < CLIENTS; i++) {
    const client = new LoadClient(i, RELAY_URL);
    clients.push(client);
    connectPromises.push(client.connect().catch(e => e)); // Catch connection errors individually
  }

  await Promise.all(connectPromises);
  console.log(`Connected ${metrics.connectedClients}/${CLIENTS} clients.`);

  if (metrics.connectedClients === 0 && !DRY_RUN) {
    console.error("No clients connected. Aborting.");
    if (relayProcess) relayProcess.kill();
    process.exit(1);
  }

  // 2. Load Loop
  console.log('Starting load generation...');
  const intervalMs = 1000 / RATE_EPS;
  const startTime = Date.now();
  const endTime = startTime + (DURATION_SEC * 1000);

  const tick = async () => {
    if (Date.now() >= endTime) {
      clearInterval(timer);
      finish();
      return;
    }

    // Pick a random connected client
    const availableClients = clients.filter(c => c.isOpen || DRY_RUN);
    if (availableClients.length > 0) {
      const client = availableClients[Math.floor(Math.random() * availableClients.length)];
      client.publish();
    }
  };

  const timer = setInterval(tick, intervalMs);

  // Progress Logger
  const progressTimer = setInterval(() => {
    const elapsed = (Date.now() - startTime) / 1000;
    const sentRate = metrics.sent / elapsed;
    const mem = process.memoryUsage();
    const cpu = process.cpuUsage();

    metrics.resources.push({
      time: elapsed,
      rss: mem.rss,
      heapTotal: mem.heapTotal,
      heapUsed: mem.heapUsed,
      external: mem.external,
      cpuUser: cpu.user,
      cpuSystem: cpu.system
    });

    console.log(`[${elapsed.toFixed(1)}s] Sent: ${metrics.sent} (${sentRate.toFixed(1)}/s) | Accepted: ${metrics.accepted} | Rejected: ${metrics.rejected} | Errors: ${metrics.errors} | RSS: ${(mem.rss / 1024 / 1024).toFixed(1)}MB`);
  }, 5000);

  // 3. Finish & Report
  async function finish() {
    clearInterval(progressTimer);
    console.log('Load generation finished. Waiting for pending confirmations...');
    await delay(2000); // Wait for trailing OKs

    clients.forEach(c => c.close());
    if (relayProcess) {
        console.log('Stopping local relay...');
        relayProcess.kill();
    }

    const elapsed = (Date.now() - metrics.startTime) / 1000;

    // Calculate Latency Stats
    metrics.latencies.sort((a, b) => a - b);
    const p50 = metrics.latencies[Math.floor(metrics.latencies.length * 0.50)] || 0;
    const p95 = metrics.latencies[Math.floor(metrics.latencies.length * 0.95)] || 0;
    const p99 = metrics.latencies[Math.floor(metrics.latencies.length * 0.99)] || 0;
    const avg = metrics.latencies.reduce((a, b) => a + b, 0) / (metrics.latencies.length || 1);

    const report = {
      config: {
        relay: RELAY_URL,
        clients: CLIENTS,
        duration: DURATION_SEC,
        rate: RATE_EPS,
        mix: MIX_VIDEO_RATIO,
        dryRun: DRY_RUN
      },
      summary: {
        totalSent: metrics.sent,
        totalAccepted: metrics.accepted,
        totalRejected: metrics.rejected,
        totalErrors: metrics.errors,
        throughput: metrics.accepted / elapsed,
        latency: {
          avg: parseFloat(avg.toFixed(2)),
          p50,
          p95,
          p99
        },
        resources: metrics.resources
      },
      timestamp: new Date().toISOString()
    };

    console.log('\n--- Final Report ---');
    console.log(JSON.stringify(report, null, 2));

    // Ensure artifacts directory exists
    const artifactsDir = path.resolve('artifacts');
    if (!fs.existsSync(artifactsDir)) {
      fs.mkdirSync(artifactsDir, { recursive: true });
    }

    const filename = `load-report-${new Date().toISOString().split('T')[0].replace(/-/g, '')}.json`;
    const filepath = path.join(artifactsDir, filename);
    fs.writeFileSync(filepath, JSON.stringify(report, null, 2));
    console.log(`Report saved to ${filepath}`);

    process.exit(0);
  }
}

run().catch(err => {
  console.error('Fatal error:', err);
  process.exit(1);
});