From edca2bf11af027b26fbbf44ffd10e1a99107ac24 Mon Sep 17 00:00:00 2001
From: dmunozv04 <39565245+dmunozv04@users.noreply.github.com>
Date: Mon, 6 Jan 2025 18:23:46 +0100
Subject: [PATCH] Improve exit logic

---
 src/iSponsorBlockTV/main.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/iSponsorBlockTV/main.py b/src/iSponsorBlockTV/main.py
index 09a8a66..c1e1699 100644
--- a/src/iSponsorBlockTV/main.py
+++ b/src/iSponsorBlockTV/main.py
@@ -133,15 +133,23 @@ async def skip(self, time_to, position, uuids):
 
     async def cancel(self):
         self.cancelled = True
-        try:
+        await self.lounge_controller.disconnect()
+        if self.task:
             self.task.cancel()
-        except Exception:
-            pass
+        if self.lounge_controller.subscribe_task_watchdog:
+            self.lounge_controller.subscribe_task_watchdog.cancel()
+        if self.lounge_controller.subscribe_task:
+            self.lounge_controller.subscribe_task.cancel()
+        await asyncio.gather(
+            self.task,
+            self.lounge_controller.subscribe_task_watchdog,
+            self.lounge_controller.subscribe_task,
+            return_exceptions=True
+        )
 
 
 async def finish(devices, web_session, tcp_connector):
-    for device in devices:
-        await device.cancel()
+    await asyncio.gather(*(device.cancel() for device in devices), return_exceptions=True)
     await web_session.close()
     await tcp_connector.close()
 
@@ -169,9 +177,10 @@ def main(config, debug):
         loop.run_forever()
     except KeyboardInterrupt:
         print("Cancelling tasks and exiting...")
+        loop.run_until_complete(finish(devices, web_session, tcp_connector))
         for task in tasks:
             task.cancel()
         loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
-        loop.run_until_complete(finish(devices, web_session, tcp_connector))
     finally:
         loop.close()
+        print("Exited")