diff --git a/tests/test_link.py b/tests/test_link.py new file mode 100644 index 0000000..e842a40 --- /dev/null +++ b/tests/test_link.py @@ -0,0 +1,201 @@ +"""Delivery order of calls over the websocket link. + +``TestServer`` drives a bare WebsocketLinkServer with a Python websocket +client, the other tests go through link.js in the browser. +""" + +import json +import threading +import time + +import pytest + + +def _wait(cond, timeout=10.0): + t0 = time.time() + while time.time() - t0 < timeout: + if cond(): + return True + time.sleep(0.01) + return False + + +class Recorder: + def __init__(self, work=None): + self.started = [] + self.done = [] + self.work = work + self.running = 0 + self.overlap = False + self._lock = threading.Lock() + + def ev(self, i, *rest): + with self._lock: + self.running += 1 + self.overlap |= self.running > 1 + self.started.append(i) + try: + if self.work: + self.work(i) + finally: + with self._lock: + self.running -= 1 + self.done.append(i) + + def answer(self): + return 42 + + +@pytest.fixture +def server(): + from websockets.sync.client import connect + + from webgpu.link.websocket import WebsocketLinkServer + + srv = WebsocketLinkServer() + srv.wait_for_server_running() + ws = connect(f"ws://127.0.0.1:{srv.port}?token={srv.auth_token}") + srv.wait_for_connection() + srv.client = ws + yield srv + ws.close() + srv.stop() + + +def _send_calls(ws, id_, n, prop=None, **flags): + for i in range(n): + msg = {"type": "call", "id": id_, "args": [i]} | flags + if prop: + msg["prop"] = prop + ws.send(json.dumps(msg)) + + +class TestServer: + def test_proxy_calls_keep_wire_order(self, server): + rec = Recorder() + handle = server.create_proxy(rec.ev, True) + _send_calls(server.client, handle["id"], 500) + assert _wait(lambda: len(rec.done) == 500) + assert rec.done == list(range(500)) + + def test_ordered_calls_run_serially_in_order(self, server): + rec = Recorder(lambda i: time.sleep(0.001)) + server.expose("rec", rec) + _send_calls(server.client, "rec", 300, prop="ev", ordered=True) + assert _wait(lambda: len(rec.done) == 300) + assert rec.started == list(range(300)) + assert not rec.overlap + + def test_ordered_methods(self, server): + rec = Recorder(lambda i: time.sleep(0.001)) + server.expose("rec", rec) + server.ordered_methods.add("ev") + _send_calls(server.client, "rec", 200, prop="ev") + assert _wait(lambda: len(rec.done) == 200) + assert rec.started == list(range(200)) + assert not rec.overlap + + def test_unordered_calls_still_concurrent(self, server): + rec = Recorder(lambda i: time.sleep(0.2)) + server.expose("rec", rec) + t0 = time.time() + _send_calls(server.client, "rec", 4, prop="ev") + assert _wait(lambda: len(rec.done) == 4) + assert time.time() - t0 < 0.6 + assert rec.overlap + + def test_slow_ordered_handler_releases_lane(self, server): + server.ordered_timeout = 0.2 + rec = Recorder(lambda i: time.sleep(1.5) if i == 0 else None) + server.expose("rec", rec) + _send_calls(server.client, "rec", 5, prop="ev", ordered=True) + assert _wait(lambda: len(rec.done) == 4, timeout=1.0) + assert rec.started == list(range(5)) + assert rec.done == [1, 2, 3, 4] + assert _wait(lambda: len(rec.done) == 5) + + def test_coalesce_keeps_newest(self, server): + rec = Recorder(lambda i: time.sleep(0.05)) + server.expose("rec", rec) + _send_calls(server.client, "rec", 40, prop="ev", ordered=True, coalesce="move") + assert _wait(lambda: rec.done and rec.done[-1] == 39) + time.sleep(0.1) + assert len(rec.done) < 40 + assert rec.done == sorted(rec.done) + + def test_disconnect_and_stop(self, server): + rec = Recorder(lambda i: time.sleep(0.05)) + server.expose("rec", rec) + _send_calls(server.client, "rec", 5, prop="ev", ordered=True) + server.client.close() + # already received events still run, in order + assert _wait(lambda: len(rec.done) == 5) + assert rec.done == list(range(5)) + assert _wait(lambda: server._connection is None) + lane = server._ordered._thread + server.stop() + lane.join(timeout=2) + assert not lane.is_alive() + server._ordered.push("{}") + assert not server._ordered._queue + + +@pytest.fixture +def js_rec(webgpu_env): + pl = webgpu_env.platform + rec = Recorder() + pl.js.window.linkTestRec = rec + pl.js.window.linkTestSink = pl.create_proxy(rec.ev, True) + pl.js.eval("0") # sets are fire-and-forget, a round trip flushes them + yield rec + webgpu_env.page.evaluate("delete window.linkTestRec; delete window.linkTestSink") + + +class TestBrowser: + def test_proxy_calls_keep_js_call_order(self, webgpu_env, js_rec): + # args of different depth take different numbers of microtasks to serialize + webgpu_env.page.evaluate("""() => { + for (let i = 0; i < 400; i++) + if (i % 2) window.linkTestSink(i, [[[[i]]]], {a: [[i]]}); + else window.linkTestSink(i); + }""") + assert _wait(lambda: len(js_rec.done) == 400) + assert js_rec.done == list(range(400)) + + def test_ordered_method_calls(self, webgpu_env, js_rec): + js_rec.work = lambda i: time.sleep(0.001) + webgpu_env.page.evaluate("""() => { + for (let i = 0; i < 300; i++) + window.linkTestRec.callMethodIgnoreResult('ev', [i], {ordered: true}); + }""") + assert _wait(lambda: len(js_rec.done) == 300) + assert js_rec.started == list(range(300)) + assert not js_rec.overlap + + def test_nested_round_trips_in_ordered_handler(self, webgpu_env, js_rec): + pl = webgpu_env.platform + pl.js.eval("window.linkTestReenter = async () => await window.linkTestRec.callMethod('answer', [])") + results = [] + + def work(i): + results.append((pl.js.eval(f"{i} + 1"), pl.js.linkTestReenter())) + + js_rec.work = work + webgpu_env.page.evaluate("""() => { + for (let i = 0; i < 20; i++) + window.linkTestRec.callMethodIgnoreResult('ev', [i], {ordered: true}); + }""") + assert _wait(lambda: len(js_rec.done) == 20) + assert results == [(i + 1, 42) for i in range(20)] + assert js_rec.done == list(range(20)) + + def test_serialized_event_has_timestamp(self, webgpu_env, js_rec): + events = [] + js_rec.work = None + webgpu_env.platform.js.window.linkTestEvent = webgpu_env.platform.create_proxy( + lambda ev: events.append(ev), True + ) + webgpu_env.platform.js.eval("0") + webgpu_env.page.evaluate("() => window.linkTestEvent(new KeyboardEvent('keydown', {key: 'a'}))") + assert _wait(lambda: events) + assert events[0]["key"] == "a" and events[0]["timeStamp"] > 0 diff --git a/webgpu/link/base.py b/webgpu/link/base.py index 8e8c091..e5bbd9e 100644 --- a/webgpu/link/base.py +++ b/webgpu/link/base.py @@ -652,6 +652,8 @@ def wrapper(*args): self._callback_queue.put((func, args)), self._callback_loop ) + # only enqueues, so the reader may call it inline (keeps wire order) + wrapper._link_enqueue = True id_ = id(wrapper) self._objects[id_] = wrapper return { diff --git a/webgpu/link/link.js b/webgpu/link/link.js index 2aba551..713657d 100644 --- a/webgpu/link/link.js +++ b/webgpu/link/link.js @@ -25,6 +25,7 @@ function serializeEvent(event) { 'deltaMode', 'movementX', 'movementY', + 'timeStamp', ]; let obj = Object.fromEntries(keys.map((k) => [k, event[k]])); if (event.x !== undefined) { @@ -116,8 +117,8 @@ function createProxy(link, id, parent_id, ignore_return = false) { target.handleEvent = async (eventType, event) => { return link.callIgnoreResult(id, [eventType, event], parent_id); }; - target.callMethodIgnoreResult = (prop, args) => { - link.callMethodIgnoreResult(id, prop, args); + target.callMethodIgnoreResult = (prop, args, options) => { + link.callMethodIgnoreResult(id, prop, args, options); }; if (ignore_return) { target.callFunction = (self, arg) => { @@ -183,6 +184,7 @@ class CrossLink { this.objects = {}; this.connection = connection; + this._sendTail = Promise.resolve(); this.connection.onMessage((data) => this.onMessage(data)); this.connected = new Promise((resolve) => { this.connection.onOpen(() => { @@ -244,13 +246,27 @@ class CrossLink { }); } - async callIgnoreResult(id, args = [], parent_id = undefined) { - return await this.connection.send({ - type: 'call', - id, - parent_id, - args: await this._dumpData(args), - }); + // Fire-and-forget calls leave in call order, even if their args take a + // different number of microtasks to serialize. + _sendCall(data, args, options) { + const dumped = this._dumpData(args); + const sent = this._sendTail + .then(() => dumped) + .then((a) => { + data.args = a; + if (options && options.ordered) data.ordered = true; + if (options && options.coalesce) data.coalesce = options.coalesce; + return this.connection.send(data); + }); + this._sendTail = sent.catch((e) => console.error('send failed', e)); + return sent; + } + + // options: { ordered: true } runs the Python handler in wire order with + // other ordered calls, { coalesce: key } lets a newer queued call with the + // same key replace this one + async callIgnoreResult(id, args = [], parent_id = undefined, options = undefined) { + return await this._sendCall({ type: 'call', id, parent_id }, args, options); } async call(id, args = [], parent_id = undefined, prop = undefined) { @@ -272,13 +288,8 @@ class CrossLink { }); } - async callMethodIgnoreResult(id, prop, args = []) { - return await this.connection.send({ - type: 'call', - id, - args: await this._dumpData(args), - prop, - }); + async callMethodIgnoreResult(id, prop, args = [], options = undefined) { + return await this._sendCall({ type: 'call', id, prop }, args, options); } expose(name, obj) { diff --git a/webgpu/link/websocket.py b/webgpu/link/websocket.py index 3856055..0ab3a5f 100644 --- a/webgpu/link/websocket.py +++ b/webgpu/link/websocket.py @@ -1,9 +1,10 @@ import asyncio +import collections import json import os import secrets import threading -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import ThreadPoolExecutor, wait from urllib.parse import parse_qs, urlparse os.environ.setdefault("WEBSOCKETS_MAX_LINE_LENGTH", "1048576") @@ -49,6 +50,61 @@ def _connect(self): raise NotImplementedError +class _OrderedLane: + """FIFO for fire-and-forget calls flagged ``ordered`` (UI events). + + Handlers start in wire order and run one at a time. One still running after + ``timeout`` seconds keeps running, but no longer blocks later events + (``None``: strictly serial). A queued message with the same ``coalesce`` key + as the newest queued one replaces it. + """ + + def __init__(self, run, timeout=0.5): + self._run = run + self.timeout = timeout + self._queue = collections.deque() + self._cond = threading.Condition() + self._executor = ThreadPoolExecutor(max_workers=8, thread_name_prefix="link-ordered") + self._thread = None + self._stopped = False + + def push(self, message, coalesce=None): + with self._cond: + if self._stopped: + return + q = self._queue + if coalesce is not None and q and q[-1][1] == coalesce: + q[-1] = (message, coalesce) + else: + q.append((message, coalesce)) + if self._thread is None: + self._thread = threading.Thread( + target=self._work, daemon=True, name="link-ordered-lane" + ) + self._thread.start() + self._cond.notify() + + def _work(self): + while True: + with self._cond: + while not self._queue and not self._stopped: + self._cond.wait() + if self._stopped: + return + message, _ = self._queue.popleft() + try: + wait([self._executor.submit(self._run, message)], self.timeout) + except RuntimeError: + return # executor shut down + + def stop(self): + with self._cond: + self._stopped = True + self._queue.clear() + self._cond.notify() + self._executor.shutdown(wait=False, cancel_futures=True) + + class WebsocketLinkServer(WebsocketLinkBase): _stop: asyncio.Future _port: int = None @@ -58,6 +114,9 @@ def __init__(self): self._port = 8700 self._auth_token = secrets.token_urlsafe(32) self._executor = ThreadPoolExecutor(max_workers=8) + self._ordered = _OrderedLane(self._on_message) + # fire-and-forget calls to these methods are ordered without the flag + self.ordered_methods = set() self._chunk_buffers = {} self._stop = None super().__init__() @@ -70,6 +129,15 @@ def auth_token(self): def port(self): return self._port + @property + def ordered_timeout(self): + """Seconds an ordered handler may block later ordered events.""" + return self._ordered.timeout + + @ordered_timeout.setter + def ordered_timeout(self, value): + self._ordered.timeout = value + def _check_auth(self, connection, request): """Reject WebSocket connections that don't carry a valid token.""" params = parse_qs(urlparse(request.path).query) @@ -79,18 +147,22 @@ def _check_auth(self, connection, request): return None @staticmethod - def _message_type(message): - """Return the top-level message type, parsing only the JSON header - (not buffer payloads). Returns None on malformed input.""" + def _message_header(message): + """Return the top-level JSON header, without parsing buffer payloads. + Returns {} on malformed input.""" try: if isinstance(message, (memoryview, bytes)): prefix_size = 4 + int.from_bytes(message[:4], byteorder="little") header = json.loads(bytes(message[4:prefix_size]).decode("utf-8")) else: header = json.loads(message) - return header.get("type") if isinstance(header, dict) else None + return header if isinstance(header, dict) else {} except Exception: - return None + return {} + + @classmethod + def _message_type(cls, message): + return cls._message_header(message).get("type") def _is_response(self, message): return self._message_type(message) == "response" @@ -114,11 +186,32 @@ def _reassemble_chunk(self, message): return None def _dispatch(self, message): - if self._is_response(message): + header = self._message_header(message) + kind = header.get("type") + if kind == "response" or kind == "call" and self._is_enqueue_call(header): self._on_message(message) + elif ( + kind == "call" + and header.get("request_id") is None + and (header.get("ordered") or header.get("prop") in self.ordered_methods) + ): + # only fire-and-forget: JS never waits on the lane, so no deadlock + coalesce = header.get("coalesce") + if coalesce is not None: + coalesce = (header.get("id"), header.get("prop"), coalesce) + self._ordered.push(message, coalesce) else: self._executor.submit(self._on_message, message) + def _is_enqueue_call(self, header): + if header.get("prop") is not None: + return False + try: + obj = self._objects.get(header.get("id")) + except TypeError: + return False + return getattr(obj, "_link_enqueue", False) + async def _websocket_handler(self, websocket, path=""): if self._connection is not None: await websocket.close(4000, "Another session is already active") @@ -184,6 +277,7 @@ async def start_websocket(): def stop(self): self._executor.shutdown(wait=False) + self._ordered.stop() try: if self._stop is not None and not self._stop.done(): self._send_loop.call_soon_threadsafe(self._stop.set_result, None)