From 016955c8baf66819dc931de504f5695532e83589 Mon Sep 17 00:00:00 2001
From: DevilXD <4180725+DevilXD@users.noreply.github.com>
Date: Tue, 16 Jul 2024 18:04:41 +0200
Subject: [PATCH] Critical tasks now cause the app to exit, instead of letting
 it continue with broken functionality

---
 twitch.py    | 11 +++++------
 utils.py     | 42 ++++++++++++++++++++++++++++++------------
 websocket.py |  2 +-
 3 files changed, 36 insertions(+), 19 deletions(-)

diff --git a/twitch.py b/twitch.py
index d87576c..f29fd3c 100644
--- a/twitch.py
+++ b/twitch.py
@@ -849,7 +849,7 @@ class Twitch:
         with suppress(asyncio.TimeoutError):
             await asyncio.wait_for(self._watching_restart.wait(), timeout=delay)
 
-    @task_wrapper
+    @task_wrapper(critical=True)
     async def _watch_loop(self) -> NoReturn:
         interval: float = WATCH_INTERVAL.total_seconds()
         while True:
@@ -898,7 +898,7 @@ class Twitch:
                         logger.log(CALL, "No active drop could be determined")
             await self._watch_sleep(interval)
 
-    @task_wrapper
+    @task_wrapper(critical=True)
     async def _maintenance_task(self) -> None:
         claim_period = timedelta(minutes=30)
         max_period = timedelta(hours=1)
@@ -911,10 +911,9 @@ class Twitch:
                 break
             next_trigger = min(now + claim_period, next_period)
             trigger_cleanup = False
-            while self._mnt_triggers and (switch_trigger := self._mnt_triggers[0]) <= next_trigger:
+            while self._mnt_triggers and self._mnt_triggers[0] <= next_trigger:
+                next_trigger = self._mnt_triggers.popleft()
                 trigger_cleanup = True
-                self._mnt_triggers.popleft()
-                next_trigger = switch_trigger
             if next_trigger == next_period:
                 trigger_type: str = "Reload"
             elif trigger_cleanup:
@@ -943,7 +942,7 @@ class Twitch:
                     await watching_channel.claim_bonus()
                 except Exception:
                     pass  # we intentionally silently skip anything else
-        # this triggers this task restart every (up to) 60 minutes
+        # this triggers a restart of this task every (up to) 60 minutes
         logger.log(CALL, "Maintenance task requests a reload")
         self.change_state(State.INVENTORY_FETCH)
 
diff --git a/utils.py b/utils.py
index bb297b5..a95fa37 100644
--- a/utils.py
+++ b/utils.py
@@ -131,18 +131,36 @@ def deduplicate(iterable: abc.Iterable[_T]) -> list[_T]:
 
 
 def task_wrapper(
-    afunc: abc.Callable[_P, abc.Coroutine[Any, Any, _T]]
-) -> abc.Callable[_P, abc.Coroutine[Any, Any, _T]]:
-    @wraps(afunc)
-    async def wrapper(*args: _P.args, **kwargs: _P.kwargs):
-        try:
-            await afunc(*args, **kwargs)
-        except (ExitRequest, ReloadRequest):
-            pass
-        except Exception:
-            logger.exception(f"Exception in {afunc.__name__} task")
-            raise  # raise up to the wrapping task
-    return wrapper
+    afunc: abc.Callable[_P, abc.Coroutine[Any, Any, _T]] | None = None, *, critical: bool = False
+):
+    def decorator(
+        afunc: abc.Callable[_P, abc.Coroutine[Any, Any, _T]]
+    ) -> abc.Callable[_P, abc.Coroutine[Any, Any, _T]]:
+        @wraps(afunc)
+        async def wrapper(*args: _P.args, **kwargs: _P.kwargs):
+            try:
+                await afunc(*args, **kwargs)
+            except (ExitRequest, ReloadRequest):
+                pass
+            except Exception:
+                logger.exception(f"Exception in {afunc.__name__} task")
+                if critical:
+                    # critical task's death should trigger a termination.
+                    # there isn't an easy and sure way to obtain the Twitch instance here,
+                    # but we can improvise finding it
+                    from twitch import Twitch  # cyclic import
+                    probe = args and args[0] or None  # extract from 'self' arg
+                    if isinstance(probe, Twitch):
+                        probe.close()
+                    elif probe is not None:
+                        probe = getattr(probe, "_twitch", None)  # extract from '_twitch' attr
+                        if isinstance(probe, Twitch):
+                            probe.close()
+                raise  # raise up to the wrapping task
+        return wrapper
+    if afunc is None:
+        return decorator
+    return decorator(afunc)
 
 
 def invalidate_cache(instance, *attrnames):
diff --git a/websocket.py b/websocket.py
index f384143..598e568 100644
--- a/websocket.py
+++ b/websocket.py
@@ -141,7 +141,7 @@ class Websocket:
                 )
                 break
 
-    @task_wrapper
+    @task_wrapper(critical=True)
     async def _handle(self):
         # ensure we're logged in before connecting
         self.set_status(_("gui", "websocket", "initializing"))