Simulate latency and intermittent connection drops on WebSockets.

WebSocketLatency.js

// ==UserScript==
// @name         WebSocket latency tester
// @namespace    https://jespertheend.com/
// @version      0.0.2
// @description  Simulate latency and intermittent connection drops on WebSockets.
// @author       Jesper van den Ende
// @match        *://*/*
// @grant        none
// @run-at       document-start
// ==/UserScript==

// Usage:
// Press W + T to toggle the active network connection
// Press W + +/- to increase/decrease latency

// TODO:
// - notifications when pressing keys
// - Simulating an intermittent network connection which only drops for brief moments.

(function () {
	"use strict";

	let networkDisabled = false;
	let simulatedLatency = 0;
	/** @type {Set<NewWebSocket>} */
	const createdSockets = new Set();

	const MODIFIER_KEY = "KeyW";
	const INCREASE_LATENCY_KEY = "Equal";
	const DECREASE_LATENCY_KEY = "Minus";
	const TOGGLE_NETWORK_KEY = "KeyT";

	function scheduleDrainEventBuffers() {
		for (const socket of createdSockets) {
			socket.scheduleDrainEventBuffer();
		}
	}

	let lastModifierPress = -Infinity;
	document.addEventListener("keydown", e => {
		if (e.code == MODIFIER_KEY) {
			lastModifierPress = performance.now();
		} else {
			if (performance.now() - lastModifierPress > 3_000) return;

			if (e.code == INCREASE_LATENCY_KEY) {
				simulatedLatency += 10;
				console.log("increase", simulatedLatency);
				scheduleDrainEventBuffers();
			} else if (e.code == DECREASE_LATENCY_KEY) {
				simulatedLatency -= 10;
				simulatedLatency = Math.max(0, simulatedLatency);
				console.log("decrease", simulatedLatency);
				scheduleDrainEventBuffers();
			} else if (e.code == TOGGLE_NETWORK_KEY) {
				console.log("toggle");
				networkDisabled = !networkDisabled;
				scheduleDrainEventBuffers();
			}

			// Consume the modifier key
			lastModifierPress = -Infinity;
		}
	});

	/**
	 * @typedef BufferedWebSocketEventSend
	 * @property {"send"} type
	 * @property {Parameters<WebSocket["send"]>} args
	 */
	/**
	 * @typedef BufferedWebSocketEventClose
	 * @property {"close"} type
	 * @property {Parameters<WebSocket["close"]>} args
	 */
	/**
	 * @typedef BufferedWebSocketEventTargetEvent
	 * @property {"targetEvent"} type
	 * @property {Event} event
	 */
	/**
	 * @typedef {BufferedWebSocketEventSend | BufferedWebSocketEventClose| BufferedWebSocketEventTargetEvent} BufferedWebSocketEventData
	 */
	/**
	 * @typedef BufferedWebSocketEvent
	 * @property {number} time When the event was fired.
	 * @property {BufferedWebSocketEventData} eventData
	 */

	const OriginalWebSocket = globalThis.WebSocket;

	class NewWebSocket extends EventTarget {
		#original;
		/** @type {BufferedWebSocketEvent[]} */
		#eventBuffer = [];

		#drainEventBufferTimeout = 0;

		/** @type {WebSocket["onopen"]} */
		onopen = null;
		/** @type {WebSocket["onclose"]} */
		onclose = null;
		/** @type {WebSocket["onerror"]} */
		onerror = null;
		/** @type {WebSocket["onmessage"]} */
		onmessage = null;

		/**
		 * @param  {ConstructorParameters<typeof WebSocket>} args
		 */
		constructor(...args) {
			super();
			this.#original = new OriginalWebSocket(...args);
			this.#monitorEvents("open");
			this.#monitorEvents("close");
			this.#monitorEvents("error");
			this.#monitorEvents("message");

			createdSockets.add(this);

			this.addEventListener("open", e => {
				if (this.onopen) this.onopen.bind(this)(e);
			});
			this.addEventListener("close", e => {
				if (this.onclose) this.onclose.bind(this)(e);
			});
			this.addEventListener("error", e => {
				if (this.onerror) this.onerror.bind(this)(e);
			});
			this.addEventListener("message", e => {
				if (this.onmessage) this.onmessage.bind(this)(e);
			});
		}

		static get CONNECTING() {
			return OriginalWebSocket.CONNECTING;
		}

		static get OPEN() {
			return OriginalWebSocket.OPEN;
		}

		static get CLOSING() {
			return OriginalWebSocket.CLOSING;
		}

		static get CLOSED() {
			return OriginalWebSocket.CLOSED;
		}

		get binaryType() {
			return this.#original.binaryType;
		}

		set binaryType(type) {
			this.#original.binaryType = type;
		}

		get bufferedAmount() {
			return this.#original.bufferedAmount;
		}

		get extensions() {
			return this.#original.extensions;
		}

		get protocol() {
			return this.#original.protocol;
		}

		get url() {
			return this.#original.url;
		}

		get readyState() {
			return this.#original.readyState;
		}

		/**
		 * @param {BufferedWebSocketEventData} eventData
		 */
		#addEvent(eventData) {
			this.#eventBuffer.push({
				time: performance.now(),
				eventData,
			});
		}

		/**
		 * @param {Parameters<WebSocket["send"]>} args
		 */
		send(...args) {
			this.#addEvent({
				type: "send",
				args,
			});
			this.#drainEventBuffer();
		}

		/**
		 * @param {Parameters<WebSocket["close"]>} args
		 */
		close(...args) {
			this.#addEvent({
				type: "close",
				args,
			});
			this.#drainEventBuffer();
		}

		/**
		 * @param {string} type
		 */
		#monitorEvents(type) {
			this.#original.addEventListener(type, event => {
				this.#addEvent({
					type: "targetEvent",
					event,
				});
				this.#drainEventBuffer();
			});
		}

		#drainEventBuffer() {
			if (networkDisabled) return;

			while (true) {
				const event = this.#getFirstBufferEvent();
				if (!event || !event.shouldBeHandled) break;
				this.#eventBuffer.shift();
				this.#handleBufferEvent(event.firstEvent.eventData);
			}
			this.scheduleDrainEventBuffer();
		}

		/**
		 * Returns the first event and whether it should be handled. Returns null when no event exists.
		 */
		#getFirstBufferEvent() {
			const firstEvent = this.#eventBuffer[0];
			if (!firstEvent) return null;
			const now = performance.now() - simulatedLatency;
			const delay = firstEvent.time - now;
			return {
				delay,
				shouldBeHandled: delay <= 0,
				firstEvent,
			};
		}

		/**
		 * Schedules a call to #drainEventBuffer so that it fires right when the first event needs to be handled.
		 */
		scheduleDrainEventBuffer() {
			// Don't need to schedule anything when the network is disabled.
			if (networkDisabled) return;

			if (this.#drainEventBufferTimeout) {
				globalThis.clearInterval(this.#drainEventBufferTimeout);
				this.#drainEventBufferTimeout = 0;
			}
			const event = this.#getFirstBufferEvent();
			if (!event) {
				// Event buffer is empty, we don't need to schedule anything
				return;
			}
			if (event.shouldBeHandled) {
				// The event should already have been fired, so we drain right away
				this.#drainEventBuffer();
			} else {
				this.#drainEventBufferTimeout = globalThis.setTimeout(() => {
					this.#drainEventBuffer();
				}, event.delay);
			}
		}

		/**
		 * @param {BufferedWebSocketEventData} event
		 */
		#handleBufferEvent(event) {
			if (event.type == "send") {
				this.#original.send(...event.args);
			} else if (event.type == "close") {
				this.#original.close(...event.args);
			} else if (event.type == "targetEvent") {
				if (event.event.type == "open") {
					this.dispatchEvent(new Event("open"));
				} else if (event.event.type == "error") {
					this.dispatchEvent(new Event("error"));
				} else if (event.event.type == "close" && event.event instanceof CloseEvent) {
					this.dispatchEvent(
						new CloseEvent("close", {
							code: event.event.code,
							wasClean: event.event.wasClean,
							reason: event.event.reason,
						})
					);
				} else if (event.event.type == "message" && event.event instanceof MessageEvent) {
					this.dispatchEvent(
						new MessageEvent("message", {
							data: event.event.data,
							origin: event.event.origin,
						})
					);
				}
			}
		}
	}

	globalThis.WebSocket = NewWebSocket;
})();