GitHub/decky-loader
1
0
Fork 0
mirror of https://github.com/SteamDeckHomebrew/decky-loader.git synced 2026-08-03 13:27:01 +03:00
Code Issues Packages Projects Releases Wiki Activity

fix(wsrouter): ensure no message reply is lost (#906)

Browse Source
This commit is contained in:
Lukas Senionis
2026-07-13 22:02:13 -04:00
committed by GitHub
parent 69475d9a0f
commit 91abff7811
4 changed files with 144 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/decky_loader/plugin/messages.py
+10 -1
View File
@@ -19,18 +19,27 @@ class MethodCallResponse:
self.success = success
self.result = result
class PluginStopped(Exception):
pass
class MethodCallRequest:
def __init__(self) -> None:
self.id = str(uuid4())
self.event = Event()
self.response: MethodCallResponse
self.response: MethodCallResponse | PluginStopped
def set_result(self, dc: SocketResponseDict):
self.response = MethodCallResponse(dc["success"], dc["res"])
self.event.set()
def cancel(self):
self.response = PluginStopped("Plugin has been stopped")
self.event.set()
async def wait_for_result(self):
await self.event.wait()
if isinstance(self.response, PluginStopped):
raise self.response
if not self.response.success:
raise Exception(self.response.result)
return self.response.result
backend/decky_loader/plugin/plugin.py
+6
View File
@@ -94,6 +94,12 @@ class PluginWrapper:
except CancelledError:
self.log.info(f"Stopping response listener for {self.name}")
await self._socket.close_socket_connection()
# Cancel the request so that WSRouter can perform cleanup
for request in self._method_call_requests.values():
request.cancel()
self._method_call_requests = {}
raise
except:
pass
backend/decky_loader/wsrouter.py
+64 -31
View File
@@ -1,25 +1,28 @@
from logging import getLogger
from asyncio import AbstractEventLoop
from aiohttp import WSCloseCode, WSMsgType, WSMessage
from aiohttp.web import Application, WebSocketResponse, Request, Response, get
from enum import IntEnum
from typing import Callable, Coroutine, Dict, Any, cast
from traceback import format_exc
from .helpers import get_csrf_token
from .plugin.messages import PluginStopped
class MessageType(IntEnum):
ERROR = -1
# Call-reply, Frontend -> Backend -> Frontend
# Call-reply, Frontend (CALL) -> Backend (REPLY|DISCARD) -> Frontend (RECEIVED_RESPONSE) -> Backend
CALL = 0
REPLY = 1
DISCARD = 2
RECEIVED_RESPONSE = 3
# Call-reply sync on connection lost,
# Frontend (FULL_SYNC) -> Backend (REPLY|DISCARD...) -> Frontend (RECEIVED_RESPONSE...) -> Backend
FULL_SYNC = 4
# Pub/Sub, Backend -> Frontend
EVENT = 3
EVENT = 5
# WSMessage with slightly better typings
class WSMessageExtra(WSMessage):
@@ -35,9 +38,8 @@ class WSRouter:
def __init__(self, loop: AbstractEventLoop, server_instance: Application) -> None:
self.loop = loop
self.ws: WebSocketResponse | None = None
self.instance_id = 0
self.routes: Dict[str, Route] = {}
# self.subscriptions: Dict[str, Callable[[Any]]] = {}
self.pending_responses: Dict[int, Dict[str, Any] | None] = {}
self.logger = getLogger("WSRouter")
server_instance.add_routes([
@@ -45,9 +47,16 @@ class WSRouter:
])
async def write(self, data: Dict[str, Any]):
# Cache all of the pending responses in case the connection is lost.
# Cleanup will happen during full-sync or when frontend confirms it has received the message
can_drop_message = True
if "id" in data and data["id"] in self.pending_responses:
self.pending_responses[data["id"]] = data
can_drop_message = False
if self.ws != None:
await self.ws.send_json(data)
else:
elif can_drop_message:
self.logger.warning("Dropping message as there is no connected socket: %s", data)
def add_route(self, name: str, route: Route):
@@ -57,26 +66,16 @@ class WSRouter:
del self.routes[name]
async def _call_route(self, route: str, args: ..., call_id: int):
instance_id = self.instance_id
error = None
try:
res = await self.routes[route](*args)
message = {"type": MessageType.REPLY.value, "id": call_id, "result": res}
except PluginStopped as err:
message = {"type": MessageType.DISCARD.value, "id": call_id}
except Exception as err:
error = {"name":err.__class__.__name__, "message":str(err), "traceback":format_exc()}
res = None
message = {"type": MessageType.ERROR.value, "id": call_id, "error": error}
if instance_id != self.instance_id:
try:
self.logger.warning("Ignoring %s reply from stale instance %d with args %s and response %s", route, instance_id, args, res)
except:
self.logger.warning("Ignoring %s reply from stale instance %d (failed to log event data)", route, instance_id)
finally:
return
if error:
await self.write({"type": MessageType.ERROR.value, "id": call_id, "error": error})
else:
await self.write({"type": MessageType.REPLY.value, "id": call_id, "result": res})
await self.write(message)
async def handle(self, request: Request):
# Auth is a query param as JS WebSocket doesn't support headers
@@ -85,7 +84,6 @@ class WSRouter:
self.logger.debug('Websocket connection starting')
ws = WebSocketResponse()
await ws.prepare(request)
self.instance_id += 1
self.logger.debug('Websocket connection ready')
if self.ws != None:
@@ -109,13 +107,11 @@ class WSRouter:
data = msg.json()
match data["type"]:
case MessageType.CALL.value:
# do stuff with the message
if data["route"] in self.routes:
self.logger.debug(f'Started PY call {data["route"]} ID {data["id"]}')
self.loop.create_task(self._call_route(data["route"], data["args"], data["id"]))
else:
error = {"error":f'Route {data["route"]} does not exist.', "name": "RouteNotFoundError", "traceback": None}
self.loop.create_task(self.write({"type": MessageType.ERROR.value, "id": data["id"], "error": error}))
self.handle_call_message(data)
case MessageType.RECEIVED_RESPONSE.value:
self.handle_received_response_message(data["id"])
case MessageType.FULL_SYNC.value:
self.handle_full_sync_message(data)
case _:
self.logger.error("Unknown message type", data)
finally:
@@ -128,6 +124,43 @@ class WSRouter:
self.logger.debug('Websocket connection closed')
return ws
def handle_call_message(self, data: Dict[str, Any]):
call_id = data["id"]
if call_id in self.pending_responses:
reply = self.pending_responses[call_id]
if reply:
self.logger.debug(f'Found pending PY call matching ID {call_id}. Sending reply...')
self.loop.create_task(self.write(reply))
else:
self.logger.debug(f'Found pending PY call matching ID {call_id}. Waiting for reply from plugin...')
return
# Prepare a call entry for caching the reply once we have it, this will also
# help us identify (above) if we are already handling the call message or not.
self.pending_responses[call_id] = None
if data["route"] in self.routes:
self.logger.debug(f'Started PY call {data["route"]} ID {call_id}')
self.loop.create_task(self._call_route(data["route"], data["args"], call_id))
else:
error = {"error":f'Route {data["route"]} does not exist.', "name": "RouteNotFoundError", "traceback": None}
self.loop.create_task(self.write({"type": MessageType.ERROR.value, "id": call_id, "error": error}))
def handle_received_response_message(self, call_id: int):
self.logger.debug(f'Removing pending response with ID {call_id}')
self.pending_responses.pop(call_id, None)
def handle_full_sync_message(self, data: Dict[str, Any]):
messages = data["messages"]
outdated_response_ids = set(self.pending_responses.keys())
for message in messages:
outdated_response_ids.discard(message["id"])
self.handle_call_message(message)
for outdated_id in outdated_response_ids:
self.handle_received_response_message(outdated_id)
async def emit(self, event: str, *args: Any):
self.logger.debug(f'Firing frontend event {event} with args {args}')
await self.write({ "type": MessageType.EVENT.value, "event": event, "args": args })
frontend/src/wsrouter.ts
+64 -22
View File
@@ -8,11 +8,16 @@ declare global {
enum MessageType {
ERROR = -1,
// Call-reply, Frontend -> Backend -> Frontend
// Call-reply, Frontend (CALL) -> Backend (REPLY|DISCARD) -> Frontend (RECEIVED_RESPONSE) -> Backend
CALL = 0,
REPLY = 1,
DISCARD = 2,
RECEIVED_RESPONSE = 3,
// Call-reply sync on connection lost,
// Frontend (FULL_SYNC) -> Backend (REPLY|DISCARD...) -> Frontend (RECEIVED_RESPONSE...) -> Backend
FULL_SYNC = 4,
// Pub/Sub, Backend -> Frontend
EVENT = 3,
EVENT = 5,
}
interface CallMessage {
@@ -28,6 +33,21 @@ interface ReplyMessage {
id: number;
}
interface DiscardMessage {
type: MessageType.DISCARD;
id: number;
}
interface ReceivedResponseMessage {
type: MessageType.RECEIVED_RESPONSE;
id: number;
}
interface FullSyncMessage {
type: MessageType.FULL_SYNC;
messages: CallMessage[];
}
interface ErrorMessage {
type: MessageType.ERROR;
error: { name: string; error: string; traceback: string | null };
@@ -58,20 +78,21 @@ interface EventMessage {
args: any;
}
type Message = CallMessage | ReplyMessage | ErrorMessage | EventMessage;
type MessageToBackend = CallMessage | ReceivedResponseMessage | FullSyncMessage;
type MessageFromBackend = ReplyMessage | ErrorMessage | DiscardMessage | EventMessage;
// Helper to resolve a promise from the outside
interface PromiseResolver<T> {
resolve: (res: T) => void;
reject: (error: PyError) => void;
promise: Promise<T>;
message: CallMessage;
}
export class WSRouter extends Logger {
runningCalls: Map<number, PromiseResolver<any>> = new Map();
eventListeners: Map<string, Set<(...args: any) => any>> = new Map();
ws?: WebSocket;
connectPromise?: Promise<void>;
// Used to map results and errors to calls
reqId: number = 0;
constructor() {
@@ -79,23 +100,26 @@ export class WSRouter extends Logger {
}
connect() {
return (this.connectPromise = new Promise<void>((resolve) => {
return new Promise<void>((resolve) => {
// Auth is a query param as JS WebSocket doesn't support headers
this.ws = new WebSocket(`ws://127.0.0.1:1337/ws?auth=${deckyAuthToken}`);
this.ws.addEventListener('open', () => {
this.debug('WS Connected');
resolve();
delete this.connectPromise;
// Synchronize frontend and backend by forwarding all the running calls again
// and letting backend reply to them accordingly.
const messages = Array.from(this.runningCalls.values()).map((resolver) => resolver.message);
this.write({ type: MessageType.FULL_SYNC, messages });
});
this.ws.addEventListener('message', this.onMessage.bind(this));
this.ws.addEventListener('close', this.onError.bind(this));
// this.ws.addEventListener('error', this.onError.bind(this));
}));
});
}
createPromiseResolver<T>(): PromiseResolver<T> {
let resolver: Partial<PromiseResolver<T>> = {};
createPromiseResolver<T>(message: CallMessage): PromiseResolver<T> {
let resolver: Partial<PromiseResolver<T>> = { message };
const promise = new Promise<T>((resolve, reject) => {
resolver.resolve = resolve;
resolver.reject = reject;
@@ -104,11 +128,6 @@ export class WSRouter extends Logger {
return resolver as PromiseResolver<T>;
}
async write(data: Message) {
if (this.connectPromise) await this.connectPromise;
this.ws?.send(JSON.stringify(data));
}
addEventListener(event: string, listener: (...args: any) => any) {
if (!this.eventListeners.has(event)) {
this.eventListeners.set(event, new Set([listener]));
@@ -130,7 +149,7 @@ export class WSRouter extends Logger {
async onMessage(msg: MessageEvent) {
try {
const data = JSON.parse(msg.data) as Message;
const data = JSON.parse(msg.data) as MessageFromBackend;
switch (data.type) {
case MessageType.REPLY:
if (this.runningCalls.has(data.id)) {
@@ -138,6 +157,8 @@ export class WSRouter extends Logger {
this.runningCalls.delete(data.id);
this.debug(`[${data.id}] Resolved PY call with value`, data.result);
}
this.write({ type: MessageType.RECEIVED_RESPONSE, id: data.id });
break;
case MessageType.ERROR:
@@ -147,6 +168,21 @@ export class WSRouter extends Logger {
this.runningCalls.delete(data.id);
this.debug(`[${data.id}] Rejected PY call with error`, data.error);
}
this.write({ type: MessageType.RECEIVED_RESPONSE, id: data.id });
break;
case MessageType.DISCARD:
if (this.runningCalls.has(data.id)) {
// We are explicitly not resolving the promise here as this message
// is only received (at the moment) when plugin is being unloaded and
// the promise should be garbage-collected, unless the plugin is doing
// very bad things!
this.runningCalls.delete(data.id);
this.debug(`[${data.id}] Discarding PY call`);
}
this.write({ type: MessageType.RECEIVED_RESPONSE, id: data.id });
break;
case MessageType.EVENT:
@@ -177,15 +213,12 @@ export class WSRouter extends Logger {
// this.call<[number, number], string>('methodName', 1, 2);
call<Args extends any[] = [], Return = void>(route: string, ...args: Args): Promise<Return> {
const resolver = this.createPromiseResolver<Return>();
const id = ++this.reqId;
const resolver = this.createPromiseResolver<Return>({ type: MessageType.CALL, route, args, id });
this.runningCalls.set(id, resolver);
this.debug(`[${id}] Calling PY method ${route} with args`, args);
this.write({ type: MessageType.CALL, route, args, id });
this.write(resolver.message);
return resolver.promise;
}
@@ -196,8 +229,17 @@ export class WSRouter extends Logger {
async onError(error: any) {
this.error('WS DISCONNECTED', error);
// TODO queue up lost messages and send them once we connect again
await sleep(5000);
await this.connect();
}
private write(data: MessageToBackend) {
// In case the message is discarded here:
// - CallMessage will be resent during FULL_SYNC
// - Backend cleanup via missed ReceivedReplyMessage will be done during FULL_SYNC
// - FullSyncMessage will be resent during FULL_SYNC
if (this.ws?.readyState === WebSocket.OPEN) {
this.ws?.send(JSON.stringify(data));
}
}
}