diff --git a/system/athena/athenad.py b/system/athena/athenad.py index 6bccaae63e..870d271624 100755 --- a/system/athena/athenad.py +++ b/system/athena/athenad.py @@ -4,6 +4,7 @@ from __future__ import annotations import base64 import hashlib import io +import itertools import json import os import queue @@ -57,6 +58,9 @@ WS_FRAME_SIZE = 4096 DEVICE_STATE_UPDATE_INTERVAL = 1.0 # in seconds DEFAULT_UPLOAD_PRIORITY = 99 # higher number = lower priority +SEND_PRIORITY_HIGH = 0 +SEND_PRIORITY_LOW = 1 + # https://bytesolutions.com/dscp-tos-cos-precedence-conversion-chart, # https://en.wikipedia.org/wiki/Differentiated_services UPLOAD_TOS = 0x20 # CS1, low priority background traffic @@ -126,14 +130,18 @@ class UploadItem: dispatcher["echo"] = lambda s: s recv_queue: Queue[str] = queue.Queue() -send_queue: Queue[str] = queue.Queue() +send_queue: Queue[tuple[int, int, str]] = queue.PriorityQueue() upload_queue: Queue[UploadItem] = queue.PriorityQueue() -low_priority_send_queue: Queue[str] = queue.Queue() log_recv_queue: Queue[str] = queue.Queue() cancelled_uploads: set[str] = set() cur_upload_items: dict[int, UploadItem | None] = {} +send_seq = itertools.count() +def send_queue_push(data: str, priority: int) -> None: + assert priority is not None, "send queue priority must be specified" + send_queue.put_nowait((priority, next(send_seq), data)) # tie-break with a monotonic counter + def strip_zst_extension(fn: str) -> str: if fn.endswith('.zst'): @@ -208,7 +216,7 @@ def jsonrpc_handler(end_event: threading.Event) -> None: if "method" in data: cloudlog.event("athena.jsonrpc_handler.call_method", data=data) response = JSONRPCResponseManager.handle(data, dispatcher) - send_queue.put_nowait(response.json) + send_queue_push(response.json, SEND_PRIORITY_HIGH) elif "id" in data and ("result" in data or "error" in data): log_recv_queue.put_nowait(data) else: @@ -217,7 +225,7 @@ def jsonrpc_handler(end_event: threading.Event) -> None: pass except Exception as e: cloudlog.exception("athena jsonrpc handler failed") - send_queue.put_nowait(json.dumps({"error": str(e)})) + send_queue_push(json.dumps({"error": str(e)}), SEND_PRIORITY_HIGH) def retry_upload(tid: int, end_event: threading.Event, increase_count: bool = True) -> None: @@ -596,7 +604,6 @@ def startStream(sdp: str) -> dict: except requests.ConnectionError as e: raise Exception("webrtc is not running. turn on comma body ignition.") from e - @dispatcher.add_method def takeSnapshot() -> str | dict[str, str] | None: from openpilot.system.camerad.snapshot import jpeg_write, snapshot @@ -666,7 +673,7 @@ def log_handler(end_event: threading.Event) -> None: "jsonrpc": "2.0", "id": log_entry } - low_priority_send_queue.put_nowait(json.dumps(jsonrpc)) + send_queue_push(json.dumps(jsonrpc), SEND_PRIORITY_LOW) curr_log = log_entry except OSError: pass # file could be deleted by log rotation @@ -717,7 +724,7 @@ def stat_handler(end_event: threading.Event) -> None: "jsonrpc": "2.0", "id": stat_filenames[0] } - low_priority_send_queue.put_nowait(json.dumps(jsonrpc)) + send_queue_push(json.dumps(jsonrpc), SEND_PRIORITY_LOW) os.remove(stat_path) last_scan = curr_scan except Exception: @@ -799,10 +806,7 @@ def ws_recv(ws: WebSocket, end_event: threading.Event) -> None: def ws_send(ws: WebSocket, end_event: threading.Event) -> None: while not end_event.is_set(): try: - try: - data = send_queue.get_nowait() - except queue.Empty: - data = low_priority_send_queue.get(timeout=1) + _, _, data = send_queue.get(timeout=1) for i in range(0, len(data), WS_FRAME_SIZE): frame = data[i:i+WS_FRAME_SIZE] last = i + WS_FRAME_SIZE >= len(data) diff --git a/system/athena/tests/test_athenad.py b/system/athena/tests/test_athenad.py index b0c20a26a9..33b0113b32 100644 --- a/system/athena/tests/test_athenad.py +++ b/system/athena/tests/test_athenad.py @@ -422,11 +422,11 @@ class TestAthenadMethods: try: # with params athenad.recv_queue.put_nowait(json.dumps({"method": "echo", "params": ["hello"], "jsonrpc": "2.0", "id": 0})) - resp = athenad.send_queue.get(timeout=3) + _, _, resp = athenad.send_queue.get(timeout=3) assert json.loads(resp) == {'result': 'hello', 'id': 0, 'jsonrpc': '2.0'} # without params athenad.recv_queue.put_nowait(json.dumps({"method": "getNetworkType", "jsonrpc": "2.0", "id": 0})) - resp = athenad.send_queue.get(timeout=3) + _, _, resp = athenad.send_queue.get(timeout=3) assert json.loads(resp) == {'result': 1, 'id': 0, 'jsonrpc': '2.0'} # log forwarding athenad.recv_queue.put_nowait(json.dumps({'result': {'success': 1}, 'id': 0, 'jsonrpc': '2.0'}))