owen
b8eb40bc64
Replace the drained asyncio.Queue with an explicit, always-inspectable model: GuildPlayer.pending (waiting tracks) + current (now playing). - queue command shows Now playing plus numbered Up next list - play reports queue position when adding behind an active track - stop/leave use clear_pending(); skip relies on the TRACK_END event to advance exactly one track (no double-skip) - queue alias q, nowplaying alias np
521 lines
19 KiB
Python
521 lines
19 KiB
Python
"""
|
|
YTAudio - YouTube audio for Red-DiscordBot via yt-dlp -> mp3 -> Lavalink (local source).
|
|
|
|
This cog OWNS the shared Red-Lavalink client connection:
|
|
* On startup it calls ``lavalink.initialize()`` (once the bot is ready),
|
|
connecting the in-process Lavalink websocket client to the standalone
|
|
Lavalink node running on localhost:2333 (managed by systemd).
|
|
* On unload it calls ``lavalink.close()``.
|
|
|
|
Because the Red-Lavalink client is a process-global singleton, exactly ONE cog
|
|
should own initialize()/close(). Other cogs (e.g. ttstoy) simply call
|
|
``lavalink.get_player(guild_id)`` and share this connection.
|
|
|
|
Playback path:
|
|
yt-dlp downloads bestaudio -> ffmpeg transcodes to mp3 -> the mp3 file is
|
|
loaded through Lavalink's ``local`` source and played. The mp3 is deleted
|
|
when the track finishes / is skipped / on stop.
|
|
"""
|
|
|
|
import asyncio
|
|
import functools
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import time as _time
|
|
from typing import Dict, List, Optional
|
|
|
|
import discord
|
|
import lavalink
|
|
import yt_dlp
|
|
from redbot.core import Config, commands
|
|
from redbot.core.bot import Red
|
|
from redbot.core.data_manager import cog_data_path
|
|
|
|
log = logging.getLogger("red.scrapyard.ytaudio")
|
|
|
|
# --- Lavalink node connection (must match the standalone application.yml) ---
|
|
LL_HOST = "localhost"
|
|
LL_PORT = 2333
|
|
LL_PASSWORD = "youshallnotpass"
|
|
LL_SECURED = False
|
|
|
|
MP3_BITRATE = "192"
|
|
|
|
YTDLP_BASE_OPTS = {
|
|
"format": "bestaudio/best",
|
|
"quiet": True,
|
|
"no_warnings": True,
|
|
"noplaylist": True,
|
|
"ignoreerrors": False,
|
|
"default_search": "ytsearch",
|
|
"source_address": "0.0.0.0",
|
|
"postprocessors": [
|
|
{
|
|
"key": "FFmpegExtractAudio",
|
|
"preferredcodec": "mp3",
|
|
"preferredquality": MP3_BITRATE,
|
|
}
|
|
],
|
|
}
|
|
|
|
|
|
class QueuedTrack:
|
|
"""A track queued for a guild, resolved lazily at play time."""
|
|
|
|
__slots__ = ("query", "requester", "channel", "title", "mp3_path", "status_msg")
|
|
|
|
def __init__(self, query: str, requester: discord.abc.User, channel):
|
|
self.query = query
|
|
self.requester = requester
|
|
self.channel = channel
|
|
self.title: Optional[str] = None
|
|
self.mp3_path: Optional[str] = None
|
|
self.status_msg = None
|
|
|
|
|
|
class GuildPlayer:
|
|
"""Per-guild queue + player-loop state. One per guild.
|
|
|
|
The queue model is explicit and always inspectable:
|
|
* ``pending`` is the list of tracks waiting to play (index 0 is next).
|
|
* ``current`` is the track playing right now (or None).
|
|
Commands read these directly, so nowplaying/queue always reflect reality.
|
|
"""
|
|
|
|
def __init__(self, cog: "YTAudio", guild_id: int):
|
|
self.cog = cog
|
|
self.guild_id = guild_id
|
|
self.pending: List[QueuedTrack] = []
|
|
self.current: Optional[QueuedTrack] = None
|
|
self.task: Optional[asyncio.Task] = None
|
|
self.next_event = asyncio.Event() # fired when the current track ends
|
|
self.wakeup = asyncio.Event() # fired when a track is added
|
|
self.volume = 100
|
|
|
|
def start(self):
|
|
if self.task is None or self.task.done():
|
|
self.task = asyncio.create_task(self._player_loop())
|
|
|
|
def enqueue(self, track: "QueuedTrack"):
|
|
self.pending.append(track)
|
|
self.wakeup.set()
|
|
|
|
def clear_pending(self):
|
|
cleared = self.pending
|
|
self.pending = []
|
|
for t in cleared:
|
|
self.cog._cleanup_file(t)
|
|
|
|
async def _player_loop(self):
|
|
while True:
|
|
if not self.pending:
|
|
self.wakeup.clear()
|
|
if not self.pending:
|
|
try:
|
|
await self.wakeup.wait()
|
|
except asyncio.CancelledError:
|
|
return
|
|
continue
|
|
|
|
track = self.pending.pop(0)
|
|
self.current = track
|
|
self.next_event.clear()
|
|
try:
|
|
await self.cog._play_track(self.guild_id, track, self.next_event)
|
|
await self.next_event.wait()
|
|
except asyncio.CancelledError:
|
|
self.cog._cleanup_file(track)
|
|
self.current = None
|
|
return
|
|
except Exception:
|
|
log.exception("Error playing track in guild %s", self.guild_id)
|
|
finally:
|
|
self.cog._cleanup_file(track)
|
|
self.current = None
|
|
|
|
|
|
class YTAudio(commands.Cog):
|
|
"""YouTube audio via yt-dlp -> mp3 -> Lavalink."""
|
|
|
|
def __init__(self, bot: Red):
|
|
self.bot = bot
|
|
self.config = Config.get_conf(self, identifier=0x59544155_44494F, force_registration=True)
|
|
self.config.register_guild(volume=100)
|
|
self.players: Dict[int, GuildPlayer] = {}
|
|
self._ll_ready = False
|
|
self._connect_task: Optional[asyncio.Task] = None
|
|
self.mp3_dir = cog_data_path(self) / "mp3cache"
|
|
self.mp3_dir.mkdir(parents=True, exist_ok=True)
|
|
self._ytdlp_opts = dict(YTDLP_BASE_OPTS)
|
|
self._ytdlp_opts["outtmpl"] = str(self.mp3_dir / "%(id)s.%(ext)s")
|
|
|
|
# ------------------------------------------------------------------
|
|
# Lavalink client ownership
|
|
# ------------------------------------------------------------------
|
|
def start_up_task(self):
|
|
"""Called from setup() after add_cog."""
|
|
self._connect_task = asyncio.create_task(self._connect_lavalink())
|
|
|
|
async def _connect_lavalink(self):
|
|
await self.bot.wait_until_red_ready()
|
|
# If a node is already connected (e.g. built-in Audio still loaded),
|
|
# don't double-initialize, just share it.
|
|
if lavalink.get_all_nodes():
|
|
self._ll_ready = True
|
|
log.info("YTAudio: a Lavalink node is already connected; sharing it.")
|
|
return
|
|
max_retries = 5
|
|
for attempt in range(1, max_retries + 1):
|
|
try:
|
|
await lavalink.initialize(
|
|
bot=self.bot,
|
|
host=LL_HOST,
|
|
password=LL_PASSWORD,
|
|
port=LL_PORT,
|
|
timeout=30,
|
|
secured=LL_SECURED,
|
|
)
|
|
self._ll_ready = True
|
|
lavalink.register_event_listener(self._track_event_handler)
|
|
log.info("YTAudio: connected Lavalink client to %s:%s", LL_HOST, LL_PORT)
|
|
return
|
|
except lavalink.AbortingNodeConnection:
|
|
await self._safe_close()
|
|
log.warning("YTAudio: Lavalink connection aborted.")
|
|
return
|
|
except asyncio.TimeoutError:
|
|
await self._safe_close()
|
|
log.warning(
|
|
"YTAudio: Lavalink connect timed out (attempt %s/%s), retrying...",
|
|
attempt,
|
|
max_retries,
|
|
)
|
|
await asyncio.sleep(2)
|
|
except Exception:
|
|
log.exception("YTAudio: unexpected error connecting to Lavalink.")
|
|
await self._safe_close()
|
|
return
|
|
log.critical("YTAudio: failed to connect to Lavalink node after %s attempts.", max_retries)
|
|
|
|
async def _safe_close(self):
|
|
try:
|
|
await lavalink.close(self.bot)
|
|
except Exception:
|
|
pass
|
|
|
|
# ------------------------------------------------------------------
|
|
# Track end handling
|
|
# ------------------------------------------------------------------
|
|
async def _track_event_handler(self, player, event, extra):
|
|
if event in (
|
|
lavalink.LavalinkEvents.TRACK_END,
|
|
lavalink.LavalinkEvents.TRACK_EXCEPTION,
|
|
lavalink.LavalinkEvents.TRACK_STUCK,
|
|
):
|
|
gp = self.players.get(player.guild.id)
|
|
if gp is not None:
|
|
gp.next_event.set()
|
|
|
|
# ------------------------------------------------------------------
|
|
# yt-dlp resolution (blocking -> executor) with live progress
|
|
# ------------------------------------------------------------------
|
|
def _make_progress_hook(self, loop, track: QueuedTrack):
|
|
# Runs inside the executor thread. Schedules throttled message edits
|
|
# back onto the event loop. Discord rate-limits edits, so we edit at
|
|
# most once per second and only when the whole-number percent changes.
|
|
state = {"last_edit": 0.0, "last_pct": -1}
|
|
|
|
def hook(d):
|
|
if track.status_msg is None:
|
|
return
|
|
status = d.get("status")
|
|
if status == "downloading":
|
|
total = d.get("total_bytes") or d.get("total_bytes_estimate")
|
|
downloaded = d.get("downloaded_bytes") or 0
|
|
if not total:
|
|
return
|
|
pct = int(downloaded * 100 / total)
|
|
now = _time.monotonic()
|
|
if pct == state["last_pct"] or (now - state["last_edit"]) < 1.0:
|
|
if pct < 100:
|
|
return
|
|
state["last_pct"] = pct
|
|
state["last_edit"] = now
|
|
content = "Downloading: {} ({}%)".format(track.query, pct)
|
|
loop.call_soon_threadsafe(
|
|
lambda c=content: asyncio.ensure_future(self._safe_edit(track, c))
|
|
)
|
|
elif status == "finished":
|
|
content = "Converting to mp3: {}".format(track.query)
|
|
loop.call_soon_threadsafe(
|
|
lambda c=content: asyncio.ensure_future(self._safe_edit(track, c))
|
|
)
|
|
|
|
return hook
|
|
|
|
async def _safe_edit(self, track: QueuedTrack, content: str):
|
|
if track.status_msg is None:
|
|
return
|
|
try:
|
|
await track.status_msg.edit(content=content)
|
|
except Exception:
|
|
pass
|
|
|
|
def _blocking_extract(self, query: str, progress_hook=None) -> dict:
|
|
opts = dict(self._ytdlp_opts)
|
|
if progress_hook is not None:
|
|
opts["progress_hooks"] = [progress_hook]
|
|
with yt_dlp.YoutubeDL(opts) as ydl:
|
|
info = ydl.extract_info(query, download=True)
|
|
if "entries" in info: # search/playlist -> first entry
|
|
info = info["entries"][0]
|
|
return info
|
|
|
|
async def _resolve_to_mp3(self, track: QueuedTrack):
|
|
loop = asyncio.get_running_loop()
|
|
hook = self._make_progress_hook(loop, track)
|
|
info = await loop.run_in_executor(
|
|
None, functools.partial(self._blocking_extract, track.query, hook)
|
|
)
|
|
track.title = info.get("title", track.query)
|
|
vid = info.get("id")
|
|
mp3_path = self.mp3_dir / f"{vid}.mp3"
|
|
track.mp3_path = str(mp3_path)
|
|
return track
|
|
|
|
def _cleanup_file(self, track: Optional[QueuedTrack]):
|
|
if track and track.mp3_path and os.path.isfile(track.mp3_path):
|
|
try:
|
|
os.remove(track.mp3_path)
|
|
except OSError:
|
|
log.warning("Could not delete mp3: %s", track.mp3_path)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Playback
|
|
# ------------------------------------------------------------------
|
|
async def _play_track(self, guild_id: int, track: QueuedTrack, done_event: asyncio.Event):
|
|
status_msg = None
|
|
if track.channel is not None:
|
|
try:
|
|
status_msg = await track.channel.send(
|
|
"Downloading: {} (0%)".format(track.query)
|
|
)
|
|
except Exception:
|
|
status_msg = None
|
|
track.status_msg = status_msg
|
|
await self._resolve_to_mp3(track)
|
|
player = lavalink.get_player(guild_id)
|
|
gp = self.players.get(guild_id)
|
|
if gp:
|
|
await player.set_volume(gp.volume)
|
|
loaded = await player.load_tracks(track.mp3_path)
|
|
if loaded.has_error or loaded.load_type != lavalink.enums.LoadType.TRACK_LOADED:
|
|
log.error("YTAudio: failed to load mp3 into Lavalink: %s", track.mp3_path)
|
|
if status_msg is not None:
|
|
try:
|
|
await status_msg.edit(content="Failed to load: {}".format(track.query))
|
|
except Exception:
|
|
pass
|
|
done_event.set()
|
|
return
|
|
player.add(requester=track.requester, track=loaded.tracks[0])
|
|
if not player.is_playing:
|
|
await player.play()
|
|
if status_msg is not None:
|
|
try:
|
|
await status_msg.edit(content="Now playing: {}".format(track.title or track.query))
|
|
except Exception:
|
|
pass
|
|
|
|
def _get_gp(self, guild_id: int) -> GuildPlayer:
|
|
gp = self.players.get(guild_id)
|
|
if gp is None:
|
|
gp = GuildPlayer(self, guild_id)
|
|
self.players[guild_id] = gp
|
|
gp.start()
|
|
return gp
|
|
|
|
async def _ensure_connected(self, ctx: commands.Context) -> bool:
|
|
if not self._ll_ready and not lavalink.get_all_nodes():
|
|
await ctx.send("Lavalink is not connected yet. Try again in a moment.")
|
|
return False
|
|
if ctx.author.voice is None or ctx.author.voice.channel is None:
|
|
await ctx.send("You need to be in a voice channel.")
|
|
return False
|
|
target = ctx.author.voice.channel
|
|
me_voice = ctx.guild.me.voice
|
|
if me_voice is None or me_voice.channel is None:
|
|
await lavalink.connect(target)
|
|
elif me_voice.channel.id != target.id:
|
|
# Already connected to a different channel: move instead of reconnecting.
|
|
try:
|
|
player = lavalink.get_player(ctx.guild.id)
|
|
await player.move_to(target)
|
|
except Exception:
|
|
await lavalink.connect(target)
|
|
return True
|
|
|
|
# ------------------------------------------------------------------
|
|
# Commands
|
|
# ------------------------------------------------------------------
|
|
@commands.command()
|
|
@commands.guild_only()
|
|
async def play(self, ctx: commands.Context, *, query: str):
|
|
"""Play a YouTube URL or search query."""
|
|
if not await self._ensure_connected(ctx):
|
|
return
|
|
gp = self._get_gp(ctx.guild.id)
|
|
gp.volume = await self.config.guild(ctx.guild).volume()
|
|
track = QueuedTrack(query, ctx.author, ctx.channel)
|
|
already_active = gp.current is not None or bool(gp.pending)
|
|
gp.enqueue(track)
|
|
if already_active:
|
|
position = len(gp.pending)
|
|
await ctx.send("Added to queue (position {}): {}".format(position, query))
|
|
|
|
@commands.command()
|
|
@commands.guild_only()
|
|
async def skip(self, ctx: commands.Context):
|
|
"""Skip the current track."""
|
|
gp = self.players.get(ctx.guild.id)
|
|
if gp is None or gp.current is None:
|
|
return await ctx.send("Nothing is playing.")
|
|
try:
|
|
player = lavalink.get_player(ctx.guild.id)
|
|
except Exception:
|
|
return await ctx.send("Nothing is playing.")
|
|
# Stopping fires TRACK_END, which the event handler uses to advance
|
|
# the queue. Do not set next_event here or it would skip two tracks.
|
|
await player.stop()
|
|
await ctx.send("Skipped.")
|
|
|
|
@commands.command()
|
|
@commands.guild_only()
|
|
async def stop(self, ctx: commands.Context):
|
|
"""Stop playback and clear the queue."""
|
|
gp = self.players.get(ctx.guild.id)
|
|
if gp:
|
|
gp.clear_pending()
|
|
gp.next_event.set()
|
|
try:
|
|
player = lavalink.get_player(ctx.guild.id)
|
|
await player.stop()
|
|
except Exception:
|
|
pass
|
|
await ctx.send("Stopped and cleared the queue.")
|
|
|
|
@commands.command()
|
|
@commands.guild_only()
|
|
async def pause(self, ctx: commands.Context):
|
|
"""Pause playback."""
|
|
try:
|
|
player = lavalink.get_player(ctx.guild.id)
|
|
except Exception:
|
|
return await ctx.send("Nothing is playing.")
|
|
await player.pause(True)
|
|
await ctx.send("Paused.")
|
|
|
|
@commands.command()
|
|
@commands.guild_only()
|
|
async def resume(self, ctx: commands.Context):
|
|
"""Resume playback."""
|
|
try:
|
|
player = lavalink.get_player(ctx.guild.id)
|
|
except Exception:
|
|
return await ctx.send("Nothing is playing.")
|
|
await player.pause(False)
|
|
await ctx.send("Resumed.")
|
|
|
|
@commands.command(name="queue", aliases=["q"])
|
|
@commands.guild_only()
|
|
async def queue_cmd(self, ctx: commands.Context):
|
|
"""Show the current track and the queue."""
|
|
gp = self.players.get(ctx.guild.id)
|
|
if not gp or (gp.current is None and not gp.pending):
|
|
return await ctx.send("The queue is empty.")
|
|
lines = []
|
|
if gp.current:
|
|
lines.append("Now playing: {}".format(gp.current.title or gp.current.query))
|
|
if gp.pending:
|
|
lines.append("Up next:")
|
|
for i, t in enumerate(gp.pending, 1):
|
|
lines.append("{}. {}".format(i, t.title or t.query))
|
|
else:
|
|
lines.append("Queue is empty.")
|
|
# Keep within Discord message limits.
|
|
out = "\n".join(lines)
|
|
if len(out) > 1900:
|
|
out = out[:1900] + "\n..."
|
|
await ctx.send(out)
|
|
|
|
@commands.command(name="nowplaying", aliases=["np"])
|
|
@commands.guild_only()
|
|
async def nowplaying(self, ctx: commands.Context):
|
|
"""Show the current track."""
|
|
gp = self.players.get(ctx.guild.id)
|
|
if not gp or gp.current is None:
|
|
return await ctx.send("Nothing is playing.")
|
|
await ctx.send("Now playing: {}".format(gp.current.title or gp.current.query))
|
|
|
|
@commands.command()
|
|
@commands.guild_only()
|
|
async def volume(self, ctx: commands.Context, level: int):
|
|
"""Set volume (0-150)."""
|
|
level = max(0, min(150, level))
|
|
await self.config.guild(ctx.guild).volume.set(level)
|
|
gp = self.players.get(ctx.guild.id)
|
|
if gp:
|
|
gp.volume = level
|
|
try:
|
|
player = lavalink.get_player(ctx.guild.id)
|
|
await player.set_volume(level)
|
|
except Exception:
|
|
pass
|
|
await ctx.send("Volume set to {}.".format(level))
|
|
|
|
@commands.command(name="summon", aliases=["ytjoin"])
|
|
@commands.guild_only()
|
|
async def join(self, ctx: commands.Context):
|
|
"""Summon the bot to your voice channel."""
|
|
if await self._ensure_connected(ctx):
|
|
await ctx.send("Joined your voice channel.")
|
|
|
|
@commands.command(name="ytleave", aliases=["ytdc"])
|
|
@commands.guild_only()
|
|
async def leave(self, ctx: commands.Context):
|
|
"""Leave the voice channel and clear state."""
|
|
gp = self.players.get(ctx.guild.id)
|
|
if gp:
|
|
gp.clear_pending()
|
|
gp.next_event.set()
|
|
try:
|
|
player = lavalink.get_player(ctx.guild.id)
|
|
await player.stop()
|
|
await player.disconnect()
|
|
except Exception:
|
|
pass
|
|
await ctx.send("Left the channel.")
|
|
|
|
# ------------------------------------------------------------------
|
|
# Cleanup
|
|
# ------------------------------------------------------------------
|
|
async def cog_unload(self):
|
|
if self._connect_task:
|
|
self._connect_task.cancel()
|
|
for gp in self.players.values():
|
|
if gp.task:
|
|
gp.task.cancel()
|
|
try:
|
|
lavalink.unregister_event_listener(self._track_event_handler)
|
|
except Exception:
|
|
pass
|
|
# This cog owns the connection: close it on unload.
|
|
await self._safe_close()
|
|
# Sweep leftover mp3s.
|
|
try:
|
|
shutil.rmtree(self.mp3_dir, ignore_errors=True)
|
|
self.mp3_dir.mkdir(parents=True, exist_ok=True)
|
|
except Exception:
|
|
pass
|