From 8a913201195ddd04bd356493b7b905736b54f76c Mon Sep 17 00:00:00 2001
From: Andrei Baciu <8437201+NiceAesth@users.noreply.github.com>
Date: Mon, 19 Feb 2024 20:59:17 +0200
Subject: [PATCH] feat: disconnect upon multiple consecutive errors (#136)

---
 src/classes/pomice.py | 12 ++++++++++++
 src/cogs/music.py     | 22 +++++++++++++++++++---
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/src/classes/pomice.py b/src/classes/pomice.py
index ba28417..b463a15 100644
--- a/src/classes/pomice.py
+++ b/src/classes/pomice.py
@@ -15,6 +15,8 @@
 from models.guild_settings import GuildSettings
 from ui.embeds.music import MusicTrackEmbed
 
+FLAKY_EXCEPTION_CUTOFF = 2
+
 
 class Player(pomice.Player):
     """Pomice guild player."""
@@ -26,6 +28,7 @@ class Player(pomice.Player):
         "guild_settings",
         "dj",
         "auto_play",
+        "consecutive_exception_count",
         "loop_mode",
         "pause_votes",
         "resume_votes",
@@ -43,6 +46,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.guild_settings: GuildSettings | None = None
         self.dj: Member | None = None
         self.auto_play: bool = False
+        self.consecutive_exception_count: int = 0
 
         self.pause_votes = set()
         self.resume_votes = set()
@@ -67,6 +71,14 @@ async def do_next(self) -> None:
         self.shuffle_votes.clear()
         self.stop_votes.clear()
 
+        if self.consecutive_exception_count >= FLAKY_EXCEPTION_CUTOFF:
+            if self.context:
+                await self.context.send(
+                    "⚠️ | Sorry! I've been having trouble playing music. I will disconnect for now.",
+                )
+            await self.teardown()
+            return
+
         if self.controller:
             with suppress(HTTPException):
                 await self.controller.delete()
diff --git a/src/cogs/music.py b/src/cogs/music.py
index c89308c..0ea0a37 100644
--- a/src/cogs/music.py
+++ b/src/cogs/music.py
@@ -71,25 +71,41 @@ def required(self, ctx: commands.Context) -> int:
         return required
 
     @MetadataGroupCog.listener()
-    async def on_pomice_track_end(self, player: Player, track: pomice.Track, _) -> None:
+    async def on_pomice_track_start(
+        self,
+        player: Player,
+        track: pomice.Track | None,
+        _,
+    ) -> None:
+        player.consecutive_exception_count = 0
+
+    @MetadataGroupCog.listener()
+    async def on_pomice_track_end(
+        self,
+        player: Player,
+        track: pomice.Track | None,
+        _,
+    ) -> None:
         await player.do_next()
 
     @MetadataGroupCog.listener()
     async def on_pomice_track_stuck(
         self,
         player: Player,
-        track: pomice.Track,
+        track: pomice.Track | None,
         _,
     ) -> None:
+        player.consecutive_exception_count += 1
         await player.do_next()
 
     @MetadataGroupCog.listener()
     async def on_pomice_track_exception(
         self,
         player: Player,
-        track: pomice.Track,
+        track: pomice.Track | None,
         _,
     ) -> None:
+        player.consecutive_exception_count += 1
         await player.do_next()
 
     async def ensure_voice(self, ctx: commands.Context) -> None: