diff options
| -rw-r--r-- | backend/decky_loader/loader.py | 5 | ||||
| -rw-r--r-- | backend/decky_loader/localplatform/localsocket.py | 20 | ||||
| -rw-r--r-- | backend/decky_loader/main.py | 56 | ||||
| -rw-r--r-- | backend/decky_loader/plugin/plugin.py | 35 | ||||
| -rw-r--r-- | backend/decky_loader/plugin/sandboxed_plugin.py | 19 | ||||
| -rw-r--r-- | backend/decky_loader/wsrouter.py | 6 | ||||
| -rw-r--r-- | dist/plugin_loader-prerelease.service | 1 | ||||
| -rw-r--r-- | dist/plugin_loader-release.service | 1 |
8 files changed, 121 insertions, 22 deletions
diff --git a/backend/decky_loader/loader.py b/backend/decky_loader/loader.py index e8a073a9..e7abb889 100644 --- a/backend/decky_loader/loader.py +++ b/backend/decky_loader/loader.py @@ -1,5 +1,5 @@ from __future__ import annotations -from asyncio import AbstractEventLoop, Queue, sleep +from asyncio import AbstractEventLoop, Queue, gather, sleep from logging import getLogger from os import listdir, path from pathlib import Path @@ -98,6 +98,9 @@ class Loader: server_instance.ws.add_route("loader/call_plugin_method", self.handle_plugin_method_call) server_instance.ws.add_route("loader/call_legacy_plugin_method", self.handle_plugin_method_call_legacy) + async def shutdown_plugins(self): + await gather(*[self.plugins[plugin_name].stop() for plugin_name in self.plugins]) + async def enable_reload_wait(self): if self.live_reload: await sleep(10) diff --git a/backend/decky_loader/localplatform/localsocket.py b/backend/decky_loader/localplatform/localsocket.py index e0ef196a..96dbbab6 100644 --- a/backend/decky_loader/localplatform/localsocket.py +++ b/backend/decky_loader/localplatform/localsocket.py @@ -21,8 +21,12 @@ class UnixSocket: self.server_writer = None async def setup_server(self): - self.socket = await asyncio.start_unix_server(self._listen_for_method_call, path=self.socket_addr, limit=BUFFER_LIMIT) - + try: + self.socket = await asyncio.start_unix_server(self._listen_for_method_call, path=self.socket_addr, limit=BUFFER_LIMIT) + except asyncio.CancelledError: + await self.close_socket_connection() + raise + async def _open_socket_if_not_exists(self): if not self.reader: retries = 0 @@ -49,6 +53,10 @@ class UnixSocket: self.reader = None + if self.socket: + self.socket.close() + await self.socket.wait_closed() + async def read_single_line(self) -> str|None: reader, _ = await self.get_socket_connection() @@ -121,8 +129,12 @@ class PortSocket (UnixSocket): self.port = random.sample(range(40000, 60000), 1)[0] async def setup_server(self): - self.socket = await asyncio.start_server(self._listen_for_method_call, host=self.host, port=self.port, limit=BUFFER_LIMIT) - + try: + self.socket = await asyncio.start_server(self._listen_for_method_call, host=self.host, port=self.port, limit=BUFFER_LIMIT) + except asyncio.CancelledError: + await self.close_socket_connection() + raise + async def _open_socket_if_not_exists(self): if not self.reader: retries = 0 diff --git a/backend/decky_loader/main.py b/backend/decky_loader/main.py index fe93c11b..a9b78026 100644 --- a/backend/decky_loader/main.py +++ b/backend/decky_loader/main.py @@ -1,6 +1,6 @@ # Change PyInstaller files permissions import sys -from typing import Dict +from typing import Any, Dict from .localplatform.localplatform import (chmod, chown, service_stop, service_start, ON_WINDOWS, ON_LINUX, get_log_level, get_live_reload, get_server_port, get_server_host, get_chown_plugin_path, @@ -8,7 +8,7 @@ from .localplatform.localplatform import (chmod, chown, service_stop, service_st if hasattr(sys, '_MEIPASS'): chmod(sys._MEIPASS, 755) # type: ignore # Full imports -from asyncio import AbstractEventLoop, new_event_loop, set_event_loop, sleep +from asyncio import AbstractEventLoop, CancelledError, Task, all_tasks, current_task, gather, new_event_loop, set_event_loop, sleep from logging import basicConfig, getLogger from os import path from traceback import format_exc @@ -55,6 +55,8 @@ if get_chown_plugin_path() == True: class PluginManager: def __init__(self, loop: AbstractEventLoop) -> None: self.loop = loop + self.reinject: bool = True + self.js_ctx_tab: Tab | None = None self.web_app = Application() self.web_app.middlewares.append(csrf_middleware) self.cors = aiohttp_cors.setup(self.web_app, defaults={ @@ -82,6 +84,7 @@ class PluginManager: self.loop.create_task(self.load_plugins()) self.web_app.on_startup.append(startup) + self.web_app.on_shutdown.append(self.shutdown) self.loop.set_exception_handler(self.exception_handler) self.web_app.add_routes([get("/auth/token", self.get_auth_token)]) @@ -90,6 +93,40 @@ class PluginManager: self.cors.add(route) # pyright: ignore [reportUnknownMemberType] self.web_app.add_routes([static("/static", path.join(path.dirname(__file__), 'static'))]) + async def shutdown(self, _: Application): + try: + logger.info(f"Shutting down...") + await self.plugin_loader.shutdown_plugins() + await self.ws.disconnect() + self.reinject = False + if self.js_ctx_tab: + await self.js_ctx_tab.close_websocket() + self.js_ctx_tab = None + except: + logger.info("Error during shutdown:\n" + format_exc()) + pass + finally: + logger.info("Cancelling tasks...") + tasks = all_tasks() + current = current_task() + async def cancel_task(task: Task[Any]): + logger.debug(f"Cancelling task {task}") + try: + task.cancel() + try: + await task + except CancelledError: + pass + logger.debug(f"Task {task} finished") + except: + logger.warn(f"Failed to cancel task {task}:\n" + format_exc()) + pass + if current: + tasks.remove(current) + await gather(*[cancel_task(task) for task in tasks]) + + logger.info("Shutdown finished.") + def exception_handler(self, loop: AbstractEventLoop, context: Dict[str, str]): if context["message"] == "Unclosed connection": return @@ -107,11 +144,13 @@ class PluginManager: logger.debug("Did not find pluginOrder setting, set it to default") async def loader_reinjector(self): - while True: + while self.reinject: tab = None nf = False dc = False while not tab: + if not self.reinject: + return try: tab = await get_gamepadui_tab() except (client_exceptions.ClientConnectorError, client_exceptions.ServerDisconnectedError): @@ -127,6 +166,7 @@ class PluginManager: if not tab: await sleep(5) await tab.open_websocket() + self.js_ctx_tab = tab await tab.enable() await self.inject_javascript(tab, True) try: @@ -135,16 +175,22 @@ class PluginManager: if not await tab.has_global_var("deckyHasLoaded", False): await self.inject_javascript(tab) elif msg.get("method", None) == "Inspector.detached": + if not self.reinject: + return logger.info("CEF has requested that we detach.") await tab.close_websocket() + self.js_ctx_tab = None break # If this is a forceful disconnect the loop will just stop without any failure message. In this case, injector.py will handle this for us so we don't need to close the socket. # This is because of https://github.com/aio-libs/aiohttp/blob/3ee7091b40a1bc58a8d7846e7878a77640e96996/aiohttp/client_ws.py#L321 logger.info("CEF has disconnected...") # At this point the loop starts again and we connect to the freshly started Steam client once it is ready. except Exception: + if not self.reinject: + return logger.error("Exception while reading page events " + format_exc()) await tab.close_websocket() + self.js_ctx_tab = None pass # while True: # await sleep(5) @@ -157,6 +203,8 @@ class PluginManager: try: # if first: if ON_LINUX and await tab.has_global_var("deckyHasLoaded", False): + await tab.close_websocket() + self.js_ctx_tab = None await restart_webhelper() return # We'll catch the next tab in the main loop await tab.evaluate_js("try{if (window.deckyHasLoaded){setTimeout(() => SteamClient.Browser.RestartJSContext(), 100)}else{window.deckyHasLoaded = true;(async()=>{try{await import('http://localhost:1337/frontend/index.js?v=%s')}catch(e){console.error(e)};})();}}catch(e){console.error(e)}" % (get_loader_version(), ), False, False, False) @@ -165,7 +213,7 @@ class PluginManager: pass def run(self): - return run_app(self.web_app, host=get_server_host(), port=get_server_port(), loop=self.loop, access_log=None) + run_app(self.web_app, host=get_server_host(), port=get_server_port(), loop=self.loop, access_log=None, handle_signals=False, shutdown_timeout=15) def main(): if ON_WINDOWS: diff --git a/backend/decky_loader/plugin/plugin.py b/backend/decky_loader/plugin/plugin.py index 94e411bc..f4b01ad3 100644 --- a/backend/decky_loader/plugin/plugin.py +++ b/backend/decky_loader/plugin/plugin.py @@ -1,4 +1,4 @@ -from asyncio import Task, create_task +from asyncio import CancelledError, Task, create_task, sleep from json import dumps, load, loads from logging import getLogger from os import path @@ -41,6 +41,7 @@ class PluginWrapper: self.log = getLogger("plugin") self.sandboxed_plugin = SandboxedPlugin(self.name, self.passive, self.flags, self.file, self.plugin_directory, self.plugin_path, self.version, self.author, self.api_version) + self.proc: Process | None = None # TODO: Maybe make LocalSocket not require on_new_message to make this cleaner self._socket = LocalSocket(self.sandboxed_plugin.on_new_message) self._listener_task: Task[Any] @@ -73,6 +74,10 @@ class PluginWrapper: create_task(self.emitted_event_callback(res["event"], res["args"])) elif res["type"] == SocketMessageType.RESPONSE.value: self._method_call_requests.pop(res["id"]).set_result(res) + except CancelledError: + self.log.info(f"Stopping response listener for {self.name}") + await self._socket.close_socket_connection() + raise except: pass @@ -104,13 +109,37 @@ class PluginWrapper: def start(self): if self.passive: return self - Process(target=self.sandboxed_plugin.initialize, args=[self._socket]).start() + self.proc = Process(target=self.sandboxed_plugin.initialize, args=[self._socket], daemon=True) + self.proc.start() self._listener_task = create_task(self._response_listener()) return self async def stop(self, uninstall: bool = False): + self.log.info(f"Stopping plugin {self.name}") + if self.passive: + return if hasattr(self, "_socket"): await self._socket.write_single_line(dumps({ "stop": True, "uninstall": uninstall }, ensure_ascii=False)) await self._socket.close_socket_connection() if hasattr(self, "_listener_task"): - self._listener_task.cancel()
\ No newline at end of file + self._listener_task.cancel() + await self.kill_if_still_running() + + async def kill_if_still_running(self): + time = 0 + while self.proc and self.proc.is_alive(): + await sleep(0.1) + time += 1 + if time == 100: + self.log.warn(f"Plugin {self.name} still alive 10 seconds after stop request! Sending SIGTERM!") + self.terminate() + elif time == 200: + self.log.warn(f"Plugin {self.name} still alive 20 seconds after stop request! Sending SIGKILL!") + self.terminate(True) + + def terminate(self, kill: bool = False): + if self.proc and self.proc.is_alive(): + if kill: + self.proc.kill() + else: + self.proc.terminate()
\ No newline at end of file diff --git a/backend/decky_loader/plugin/sandboxed_plugin.py b/backend/decky_loader/plugin/sandboxed_plugin.py index 7e618590..083a6749 100644 --- a/backend/decky_loader/plugin/sandboxed_plugin.py +++ b/backend/decky_loader/plugin/sandboxed_plugin.py @@ -1,5 +1,5 @@ from os import path, environ -from signal import SIGINT, signal +from signal import SIG_IGN, SIGINT, SIGTERM, signal from importlib.util import module_from_spec, spec_from_file_location from json import dumps, loads from logging import getLogger @@ -39,13 +39,14 @@ class SandboxedPlugin: self.author = author self.api_version = api_version - self.log = getLogger("plugin") + self.log = getLogger("sandboxed_plugin") def initialize(self, socket: LocalSocket): self._socket = socket try: - signal(SIGINT, lambda s, f: exit(0)) + signal(SIGINT, SIG_IGN) + signal(SIGTERM, SIG_IGN) set_event_loop(new_event_loop()) if self.passive: @@ -112,10 +113,10 @@ class SandboxedPlugin: else: get_event_loop().create_task(self.Plugin._main(self.Plugin)) get_event_loop().create_task(socket.setup_server()) - get_event_loop().run_forever() except: self.log.error("Failed to start " + self.name + "!\n" + format_exc()) exit(0) + get_event_loop().run_forever() async def _unload(self): try: @@ -130,7 +131,7 @@ class SandboxedPlugin: self.log.info("Could not find \"_unload\" in " + self.name + "'s main.py" + "\n") except: self.log.error("Failed to unload " + self.name + "!\n" + format_exc()) - exit(0) + pass async def _uninstall(self): try: @@ -145,13 +146,13 @@ class SandboxedPlugin: self.log.info("Could not find \"_uninstall\" in " + self.name + "'s main.py" + "\n") except: self.log.error("Failed to uninstall " + self.name + "!\n" + format_exc()) - exit(0) + pass async def on_new_message(self, message : str) -> str|None: data = loads(message) if "stop" in data: - self.log.info("Calling Loader unload function.") + self.log.info(f"Calling Loader unload function for {self.name}.") await self._unload() if data.get('uninstall'): @@ -160,9 +161,9 @@ class SandboxedPlugin: get_event_loop().stop() while get_event_loop().is_running(): - await sleep(0) + await sleep(0.1) get_event_loop().close() - raise Exception("Closing message listener") + exit(0) d: SocketResponseDict = {"type": SocketMessageType.RESPONSE, "res": None, "success": True, "id": data["id"]} try: diff --git a/backend/decky_loader/wsrouter.py b/backend/decky_loader/wsrouter.py index 96e61daf..b513723f 100644 --- a/backend/decky_loader/wsrouter.py +++ b/backend/decky_loader/wsrouter.py @@ -2,7 +2,7 @@ from logging import getLogger from asyncio import AbstractEventLoop, create_task -from aiohttp import WSMsgType, WSMessage +from aiohttp import WSCloseCode, WSMsgType, WSMessage from aiohttp.web import Application, WebSocketResponse, Request, Response, get from enum import IntEnum @@ -133,3 +133,7 @@ class WSRouter: 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 }) + + async def disconnect(self): + if self.ws: + await self.ws.close(code=WSCloseCode.GOING_AWAY, message=b"Loader is shutting down") diff --git a/dist/plugin_loader-prerelease.service b/dist/plugin_loader-prerelease.service index 61c7c553..089d801a 100644 --- a/dist/plugin_loader-prerelease.service +++ b/dist/plugin_loader-prerelease.service @@ -4,6 +4,7 @@ Description=SteamDeck Plugin Loader Type=simple User=root Restart=always +TimeoutStopSec=45 ExecStart=${HOMEBREW_FOLDER}/services/PluginLoader WorkingDirectory=${HOMEBREW_FOLDER}/services Environment=UNPRIVILEGED_PATH=${HOMEBREW_FOLDER} diff --git a/dist/plugin_loader-release.service b/dist/plugin_loader-release.service index 9e10e5ff..d2e6be37 100644 --- a/dist/plugin_loader-release.service +++ b/dist/plugin_loader-release.service @@ -4,6 +4,7 @@ Description=SteamDeck Plugin Loader Type=simple User=root Restart=always +TimeoutStopSec=45 ExecStart=${HOMEBREW_FOLDER}/services/PluginLoader WorkingDirectory=${HOMEBREW_FOLDER}/services Environment=UNPRIVILEGED_PATH=${HOMEBREW_FOLDER} |