Fix shutdown timeouts (#695)

Co-authored-by: AAGaming <aagaming@riseup.net>

2 files changed, 35 insertions, 33 deletions

diff --git a/backend/decky_loader/plugin/plugin.py b/backend/decky_loader/plugin/plugin.py
index a9a9ce29..e22eca3f 100644
--- a/backend/decky_loader/plugin/plugin.py
+++ b/backend/decky_loader/plugin/plugin.py
@@ -1,8 +1,10 @@
-from asyncio import CancelledError, Task, create_task, sleep
+from asyncio import CancelledError, Task, create_task, sleep, get_event_loop, wait
 from json import dumps, load, loads
 from logging import getLogger
 from os import path
 from multiprocessing import Process
+from time import time
+from traceback import format_exc
 
 from .sandboxed_plugin import SandboxedPlugin
 from .messages import MethodCallRequest, SocketMessageType
@@ -42,8 +44,7 @@ class PluginWrapper:
 
         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._socket = LocalSocket()
         self._listener_task: Task[Any]
         self._method_call_requests: Dict[str, MethodCallRequest] = {}
 
@@ -65,7 +66,7 @@ class PluginWrapper:
         return self.name
     
     async def _response_listener(self):
-        while True:
+        while self._socket.active:
             try:
                 line = await self._socket.read_single_line()
                 if line != None:
@@ -115,29 +116,40 @@ class PluginWrapper:
         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 self.proc:
-            self.proc.join()
-        await self.kill_if_still_running()
-        if hasattr(self, "_listener_task"):
-            self._listener_task.cancel()
+        try:
+            start_time = time()
+            if self.passive:
+                return
+
+            _, pending = await wait([
+                create_task(self._socket.write_single_line(dumps({ "stop": True, "uninstall": uninstall }, ensure_ascii=False)))
+            ], timeout=1)
+
+            if hasattr(self, "_listener_task"):
+                self._listener_task.cancel()
+            
+            await self.kill_if_still_running()
+
+            for pending_task in pending:
+                pending_task.cancel()
+
+            self.log.info(f"Plugin {self.name} has been stopped in {time() - start_time:.1f}s")
+        except Exception as e:
+            self.log.error(f"Error during shutdown for plugin {self.name}: {str(e)}\n{format_exc()}")
 
     async def kill_if_still_running(self):
-        time = 0
+        start_time = time()
+        sigtermed = False
         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!")
+            elapsed_time = time() - start_time
+            if elapsed_time >= 5 and not sigtermed:
+                sigtermed = True
+                self.log.warn(f"Plugin {self.name} still alive 5 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!")
+            elif elapsed_time >= 10:
+                self.log.warn(f"Plugin {self.name} still alive 10 seconds after stop request! Sending SIGKILL!")
                 self.terminate(True)
+            await sleep(0.1)
 
     def terminate(self, kill: bool = False):
         if self.proc and self.proc.is_alive():
diff --git a/backend/decky_loader/plugin/sandboxed_plugin.py b/backend/decky_loader/plugin/sandboxed_plugin.py
index 93691a44..23575900 100644
--- a/backend/decky_loader/plugin/sandboxed_plugin.py
+++ b/backend/decky_loader/plugin/sandboxed_plugin.py
@@ -1,6 +1,5 @@
 import sys
 from os import path, environ
-from signal import SIG_IGN, SIGINT, SIGTERM, getsignal, signal
 from importlib.util import module_from_spec, spec_from_file_location
 from json import dumps, loads
 from logging import getLogger
@@ -19,8 +18,6 @@ from typing import List, TypeVar, Any
 
 DataType = TypeVar("DataType")
 
-original_term_handler = getsignal(SIGTERM)
-
 class SandboxedPlugin:
     def __init__(self,
                  name: str,
@@ -48,11 +45,6 @@ class SandboxedPlugin:
         self._socket = socket
 
         try:
-            # Ignore signals meant for parent Process
-            # TODO SURELY there's a better way to do this.
-            signal(SIGINT, SIG_IGN)
-            signal(SIGTERM, SIG_IGN)
-
             setproctitle(f"{self.name} ({self.file})")
             setthreadtitle(self.name)
 
@@ -120,7 +112,7 @@ class SandboxedPlugin:
                     get_event_loop().create_task(self.Plugin._main())
                 else:
                     get_event_loop().create_task(self.Plugin._main(self.Plugin))
-            get_event_loop().create_task(socket.setup_server())
+            get_event_loop().create_task(socket.setup_server(self.on_new_message))
         except:
             self.log.error("Failed to start " + self.name + "!\n" + format_exc())
             sys.exit(0)
@@ -167,8 +159,6 @@ class SandboxedPlugin:
         data = loads(message)
 
         if "stop" in data:
-            # Incase the loader needs to terminate our process soon
-            signal(SIGTERM, original_term_handler)
             self.log.info(f"Calling Loader unload function for {self.name}.")
             await self._unload()