From f7f820691b8ea3e50a499190c4995b1b5caff556 Mon Sep 17 00:00:00 2001 From: Marcin Macierzanka Date: Sun, 13 Sep 2026 16:29:31 +0200 Subject: [PATCH] feat: add voice channel cleanup and signal handling during application shutdown --- tests/test_shutdown_voice_cleanup.py | 79 ++++++++++++++++++++++++++++ zyrarbot/application.py | 67 ++++++++++++++++++++--- zyrarbot/ui/bot/client.py | 74 +++++++++++++++++++++++++- 3 files changed, 211 insertions(+), 9 deletions(-) create mode 100644 tests/test_shutdown_voice_cleanup.py diff --git a/tests/test_shutdown_voice_cleanup.py b/tests/test_shutdown_voice_cleanup.py new file mode 100644 index 0000000..166f9ef --- /dev/null +++ b/tests/test_shutdown_voice_cleanup.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import unittest +from unittest.mock import AsyncMock, MagicMock, patch + +import discord + +from zyrarbot.settings import Settings +from zyrarbot.ui.bot.client import MusicBot + + +class ShutdownVoiceCleanupTests(unittest.IsolatedAsyncioTestCase): + def setUp(self): + self.settings = Settings(discord_token="fake_token") + self.tracks = MagicMock() + self.uploads = MagicMock() + self.config = MagicMock() + self.history = MagicMock() + + self.bot = MusicBot( + self.settings, + self.tracks, + self.uploads, + self.config, + self.history, + ) + self.player = MagicMock() + self.bot.player = self.player + + async def test_disconnect_all_voice_disconnects_vcs_and_cleans_guild_state(self): + mock_vc = MagicMock() + mock_vc.guild.id = 12345 + mock_vc.channel.name = "Test Channel" + mock_vc.is_connected.return_value = True + mock_vc.disconnect = AsyncMock() + + mock_guild = MagicMock(spec=discord.Guild) + mock_guild.id = 12345 + mock_guild.name = "Test Guild" + mock_guild.me.voice.channel = MagicMock() + mock_guild.change_voice_state = AsyncMock() + + with patch.object(type(self.bot), "voice_clients", [mock_vc]), patch.object( + type(self.bot), "guilds", [mock_guild] + ): + await self.bot.disconnect_all_voice() + + self.player.save_queue.assert_called_once_with(12345) + self.player.clear.assert_called_once_with(12345) + mock_vc.disconnect.assert_awaited_once_with(force=False) + mock_guild.change_voice_state.assert_awaited_once_with(channel=None) + + async def test_close_calls_disconnect_all_voice_before_super_close(self): + disconnect_mock = AsyncMock() + self.bot.disconnect_all_voice = disconnect_mock + + with patch("discord.ext.commands.Bot.close", new_callable=AsyncMock) as super_close_mock: + await self.bot.close() + + disconnect_mock.assert_awaited_once() + super_close_mock.assert_awaited_once() + + # Calling close again should be a no-op due to _is_closing guard + await self.bot.close() + self.assertEqual(disconnect_mock.await_count, 1) + + async def test_on_ready_cleans_dangling_voice_connection(self): + mock_guild = MagicMock(spec=discord.Guild) + mock_guild.id = 12345 + mock_guild.name = "Test Guild" + mock_guild.me.voice.channel.name = "Ghost Channel" + mock_guild.voice_client = MagicMock(spec=discord.VoiceClient) + mock_guild.voice_client.disconnect = AsyncMock() + mock_guild.change_voice_state = AsyncMock() + + with patch.object(type(self.bot), "guilds", [mock_guild]): + await self.bot.on_ready() + + mock_guild.voice_client.disconnect.assert_awaited_once_with(force=True) diff --git a/zyrarbot/application.py b/zyrarbot/application.py index f98f28e..5a87ebd 100644 --- a/zyrarbot/application.py +++ b/zyrarbot/application.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import logging import os +import signal from collections.abc import AsyncIterator from contextlib import asynccontextmanager @@ -89,6 +90,8 @@ class Application: finally: self.player.on_state_change = None await broadcaster.stop() + if not self.bot.is_closed(): + await self.bot.close() self.web.router.lifespan_context = lifespan @@ -141,19 +144,69 @@ class Application: server = uvicorn.Server(config) self._server = server - async with self.bot: - server_task = asyncio.create_task(server.serve()) + loop = asyncio.get_running_loop() + + def _handle_shutdown_signal() -> None: + logger.info("Odebrano sygnał zakończenia, zamykanie serwera i bota...") + server.should_exit = True + if not self.bot.is_closed(): + asyncio.create_task(self.bot.close()) + + for sig in (signal.SIGINT, signal.SIGTERM): try: + loop.add_signal_handler(sig, _handle_shutdown_signal) + except (NotImplementedError, RuntimeError): + pass + + async with self.bot: + server_task = asyncio.create_task(server.serve(), name="web-server") + + async def run_bot() -> None: if not self.settings.discord_token: raise RuntimeError("Brak DISCORD_TOKEN w środowisku") await self.bot.start(self.settings.discord_token) - except Exception as error: # noqa: BLE001 - panel ma przetrwać awarię bota - logger.error("Nie udało się wystartować bota Discord: %s", error) - logger.info("Serwer webowy działa dalej...") - await self._wait_for_shutdown(server) + + bot_task = asyncio.create_task(run_bot(), name="discord-bot") + + try: + done, pending = await asyncio.wait( + [server_task, bot_task], + return_when=asyncio.FIRST_COMPLETED, + ) + + if server_task in done: + # Serwer webowy zakończył działanie (np. przez sygnał wyjścia lub restart) + logger.info("Serwer webowy zakończył pracę, zamykanie bota Discord...") + if not self.bot.is_closed(): + await self.bot.close() + try: + await bot_task + except (asyncio.CancelledError, Exception) as error: + logger.debug("Zakończenie zadania bota: %s", error) + + elif bot_task in done: + # Bot Discord zakończył działanie (błąd startu lub wyjątek) + try: + bot_task.result() + except Exception as error: # noqa: BLE001 + logger.error("Bot Discord zakończył pracę z błędem: %s", error) + logger.info("Serwer webowy działa dalej...") + await self._wait_for_shutdown(server) + + except asyncio.CancelledError: + pass finally: server.should_exit = True - await server_task + if not self.bot.is_closed(): + await self.bot.close() + if not server_task.done(): + await server_task + if not bot_task.done(): + bot_task.cancel() + try: + await bot_task + except (asyncio.CancelledError, Exception): + pass @staticmethod async def _wait_for_shutdown(server: uvicorn.Server) -> None: diff --git a/zyrarbot/ui/bot/client.py b/zyrarbot/ui/bot/client.py index 8a160d5..fc9ede7 100644 --- a/zyrarbot/ui/bot/client.py +++ b/zyrarbot/ui/bot/client.py @@ -1,12 +1,12 @@ from __future__ import annotations +import asyncio import logging +from collections.abc import Callable import discord from discord.ext import commands -from collections.abc import Callable - from zyrarbot.audio.track_factory import TrackFactory from zyrarbot.music.music_player import MusicPlayer from zyrarbot.settings import Settings @@ -51,6 +51,7 @@ class MusicBot(commands.Bot): self.history = history self.player: MusicPlayer | None = None self._restart_callback = restart_callback + self._is_closing = False def schedule_restart(self) -> None: """Żąda restartu bota za pomocą zarejestrowanego wywołania zwrotnego.""" @@ -96,11 +97,80 @@ class MusicBot(commands.Bot): await self.tree.sync() logger.info("Zsynchronizowano komendy globalnie") + async def disconnect_all_voice(self) -> None: + """Rozłącza bota ze wszystkich kanałów głosowych na wszystkich serwerach.""" + if self.voice_clients: + logger.info( + "Odłączanie bota ze wszystkich kanałów głosowych (%d)...", + len(self.voice_clients), + ) + + async def _disconnect_vc(vc: discord.VoiceClient) -> None: + try: + guild_id = vc.guild.id + if self.player: + try: + self.player.save_queue(guild_id) + self.player.clear(guild_id) + except Exception as err: + logger.warning("Błąd podczas czyszczenia odtwarzacza dla serwera %s: %s", guild_id, err) + if vc.is_connected(): + await asyncio.wait_for(vc.disconnect(force=False), timeout=2.0) + except Exception as error: + logger.warning("Błąd podczas łagodnego odłączania z kanału %s, wymuszanie: %s", getattr(vc, "channel", None), error) + try: + await vc.disconnect(force=True) + except Exception: + pass + + if self.voice_clients: + await asyncio.gather(*[_disconnect_vc(vc) for vc in list(self.voice_clients)], return_exceptions=True) + + async def _clear_guild_voice(guild: discord.Guild) -> None: + try: + if guild.me and guild.me.voice and guild.me.voice.channel: + await guild.change_voice_state(channel=None) + except Exception as error: + logger.warning("Błąd podczas czyszczenia stanu głosowego na serwerze %s: %s", guild.name, error) + + # Dla każdego serwera upewnij się przez gateway, że stan głosowy bota to None + if self.guilds: + await asyncio.gather(*[_clear_guild_voice(guild) for guild in self.guilds], return_exceptions=True) + + async def close(self) -> None: + """Zamykanie bota: najpierw odłącz się od wszystkich kanałów głosowych.""" + if getattr(self, "_is_closing", False) or self.is_closed(): + return + self._is_closing = True + try: + await self.disconnect_all_voice() + except Exception as error: + logger.error("Błąd podczas odłączania bota przed zamknięciem: %s", error) + await super().close() + async def on_ready(self) -> None: if self.guilds: first_guild_id = self.guilds[0].id self.uploads.migrate_legacy(first_guild_id) self.config.migrate_legacy(first_guild_id) self.history.migrate_legacy(first_guild_id) + + # Jeśli po uruchomieniu bot z jakiegoś powodu znajduje się w kanale głosowym (np. po awaryjnym restarcie), + # odłącz go, aby stan był czysty i zgodny z API. + for guild in self.guilds: + if guild.me and guild.me.voice and guild.me.voice.channel: + logger.warning( + "Wykryto wiszące połączenie głosowe po starcie na serwerze %s (%s). Odłączanie...", + guild.name, + guild.me.voice.channel.name, + ) + try: + if guild.voice_client: + await guild.voice_client.disconnect(force=True) + else: + await guild.change_voice_state(channel=None) + except Exception as error: + logger.error("Błąd podczas czyszczenia wiszącego kanału głosowego: %s", error) + if self.user: logger.info("Zalogowano jako %s (%s)", self.user.name, self.user.id)