Shade/packages/shade-sdk/tests/webrtc-throughput.test.ts

/**
 * V3.11 acceptance criterion (loopback flavour): a multi-lane payload
 * over the in-process WebRTC transport completes faster than the same
 * payload over HTTP-loopback.
 *
 * The MemoryRtcFactory short-circuits the network entirely, so this is
 * effectively comparing "in-process pipe" vs "HTTP-loopback round-trip"
 * — P2P should still win because every chunk goes through the OS TCP
 * stack on the HTTP side. This stand-in test validates the wiring; the
 * "real" same-LAN comparison runs in `webrtc-native.test.ts` when
 * `globalThis.RTCPeerConnection` exists.
 */
import { describe, test, expect, beforeAll, afterAll } from 'bun:test';
import {
  createShade,
  type Shade,
  type TransferHandle,
  type TransferResult,
} from '../src/index.js';
import {
  createPrekeyServer,
  MemoryPrekeyStore,
  PrekeyServerEvents,
} from '@shade/server';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import { sha256Once } from '@shade/streams';
import { MemoryRtcFactory } from '@shade/transport-webrtc';

const crypto = new SubtleCryptoProvider();

interface Rig {
  alice: Shade;
  bob: Shade;
  prekeyStop: () => void;
  aliceServerStop: () => void;
  bobServerStop: () => void;
}

async function startPrekeyServer(): Promise<{ url: string; stop: () => void }> {
  const events = new PrekeyServerEvents();
  const server = createPrekeyServer({
    crypto,
    store: new MemoryPrekeyStore(),
    disableRateLimit: true,
    events,
  });
  const port = 24500 + Math.floor(Math.random() * 500);
  const handle = Bun.serve({ port, fetch: server.fetch });
  return { url: `http://localhost:${port}`, stop: () => handle.stop() };
}

async function setupRig(opts: { withWebRTC: boolean }): Promise<Rig> {
  const prekey = await startPrekeyServer();
  const alice = await createShade({ prekeyServer: prekey.url, address: 'alice' });
  const bob = await createShade({ prekeyServer: prekey.url, address: 'bob' });

  const baseUrls = new Map<string, string>();
  const resolveBaseUrl = async (addr: string): Promise<string> => {
    const url = baseUrls.get(addr);
    if (url === undefined) throw new Error(`unknown peer ${addr}`);
    return url;
  };
  alice.configureTransfers({ resolveBaseUrl });
  bob.configureTransfers({ resolveBaseUrl });

  if (opts.withWebRTC) {
    const factory = new MemoryRtcFactory();
    alice.configureWebRTC({ factory, connectTimeoutMs: 10_000 });
    bob.configureWebRTC({ factory, connectTimeoutMs: 10_000 });
  }

  const bobApp = await bob.transferRoute();
  const bobPort = 25000 + Math.floor(Math.random() * 500);
  const bobServer = Bun.serve({ port: bobPort, fetch: bobApp.fetch });
  baseUrls.set('bob', `http://localhost:${bobPort}`);

  const aliceApp = await alice.transferRoute();
  const alicePort = 25500 + Math.floor(Math.random() * 500);
  const aliceServer = Bun.serve({ port: alicePort, fetch: aliceApp.fetch });
  baseUrls.set('alice', `http://localhost:${alicePort}`);

  return {
    alice,
    bob,
    prekeyStop: prekey.stop,
    aliceServerStop: () => aliceServer.stop(),
    bobServerStop: () => bobServer.stop(),
  };
}

async function teardownRig(rig: Rig): Promise<void> {
  await rig.alice.shutdown();
  await rig.bob.shutdown();
  rig.bobServerStop();
  rig.aliceServerStop();
  rig.prekeyStop();
  MemoryRtcFactory.reset();
}

function hex(b: Uint8Array): string {
  return Array.from(b, (x) => x.toString(16).padStart(2, '0')).join('');
}

async function uploadAndAwait(
  rig: Rig,
  input: Uint8Array,
  opts: { lanes: number; chunkSize: number },
): Promise<{ senderResult: TransferResult; received: Uint8Array; elapsed: number }> {
  let resolveRecv!: (h: TransferHandle) => void;
  const recvHandlePromise = new Promise<TransferHandle>((r) => {
    resolveRecv = r;
  });
  const unsubscribe = await rig.bob.onIncomingTransfer(async (incoming) => {
    const h = await incoming.accept({ output: { kind: 'buffer' } });
    resolveRecv(h);
  });

  const t0 = performance.now();
  const handle = await rig.alice.upload({
    to: 'bob',
    input,
    lanes: opts.lanes,
    chunkSize: opts.chunkSize,
    metadata: { name: 'throughput.bin' },
  });
  const recvHandle = await recvHandlePromise;
  const [senderResult, recvResult] = await Promise.all([
    handle.done(),
    recvHandle.done(),
  ]);
  const elapsed = performance.now() - t0;
  unsubscribe();
  const received =
    (recvResult as TransferResult & { bytes?: Uint8Array }).bytes ?? new Uint8Array();
  return { senderResult, received, elapsed };
}

describe('V3.11 throughput — WebRTC loopback vs HTTP loopback', () => {
  let webrtcRig: Rig;
  let httpRig: Rig;

  beforeAll(async () => {
    webrtcRig = await setupRig({ withWebRTC: true });
    httpRig = await setupRig({ withWebRTC: false });
  });

  afterAll(async () => {
    await teardownRig(webrtcRig);
    await teardownRig(httpRig);
  });

  test(
    'integrity match across both transports for 4 MiB / 4 lanes',
    async () => {
      const input = crypto.randomBytes(4 * 1024 * 1024);
      const expectedHash = hex(sha256Once(input));

      const w = await uploadAndAwait(webrtcRig, input, { lanes: 4, chunkSize: 64 * 1024 });
      expect(w.received).toEqual(input);
      expect(w.senderResult.sha256).toBe(expectedHash);

      const h = await uploadAndAwait(httpRig, input, { lanes: 4, chunkSize: 64 * 1024 });
      expect(h.received).toEqual(input);
      expect(h.senderResult.sha256).toBe(expectedHash);

      // Diagnostic logging — not a hard assertion since loopback is
      // dominated by crypto cost rather than transport. We do assert
      // that WebRTC is the primary on the WebRTC rig and that no fallback
      // happened.
      const runtime = webrtcRig.alice.getWebRtcRuntime();
      expect(runtime!.fallback.activeName).toBe('webrtc');
      expect(runtime!.fallback.hasFallenBack).toBe(false);

      // eslint-disable-next-line no-console
      console.log(
        `[throughput] webrtc=${w.elapsed.toFixed(0)}ms  http=${h.elapsed.toFixed(0)}ms  ` +
          `(speedup ×${(h.elapsed / w.elapsed).toFixed(2)})`,
      );
    },
    60_000,
  );
});