implement a shutdown routine instead of just waiting for all plugins to stop on their own

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}