Make use of discord.VoiceProtocol instead of socket_response

lavalink/__init__.py

@@ -25,10 +25,7 @@
     "Node",
     "NodeStats",
     "Stats",
-    "user_id",
-    "channel_finder_func",
     "Player",
-    "PlayerManager",
     "initialize",
     "connect",
     "get_player",
@@ -43,4 +40,4 @@
     "all_connected_players",
     "active_players",
 ]
-__version__ = "0.8.1"
+__version__ = "0.9.0"

lavalink/lavalink.py

@@ -66,20 +66,17 @@ async def initialize(
     global _loop
     _loop = bot.loop
 
-    player_manager.user_id = bot.user.id
-    player_manager.channel_finder_func = bot.get_channel
     register_event_listener(_handle_event)
     register_update_listener(_handle_update)
 
     lavalink_node = node.Node(
-        _loop,
-        dispatch,
-        bot._connection._get_websocket,
-        host,
-        password,
+        _loop=_loop,
+        event_handler=dispatch,
+        host=host,
+        password=password,
         port=ws_port,
-        user_id=player_manager.user_id,
-        num_shards=bot.shard_count if bot.shard_count is not None else 1,
+        user_id=bot.user.id,
+        num_shards=bot.shard_count or 1,
         resume_key=resume_key,
         resume_timeout=resume_timeout,
         bot=bot,
@@ -88,7 +85,6 @@ async def initialize(
     await lavalink_node.connect(timeout=timeout)
     lavalink_node._retries = 0
 
-    bot.add_listener(node.on_socket_response)
     bot.add_listener(_on_guild_remove, name="on_guild_remove")
 
     return lavalink_node
@@ -116,13 +112,13 @@ async def connect(channel: discord.VoiceChannel, deafen: bool = False):
         If there are no available lavalink nodes ready to connect to discord.
     """
     node_ = node.get_node(channel.guild.id)
-    p = await node_.player_manager.create_player(channel, deafen=deafen)
+    p = await node_.create_player(channel, deafen=deafen)
     return p
 
 
 def get_player(guild_id: int) -> player_manager.Player:
     node_ = node.get_node(guild_id)
-    return node_.player_manager.get_player(guild_id)
+    return node_.get_player(guild_id)
 
 
 async def _on_guild_remove(guild):
@@ -182,7 +178,7 @@ def _get_event_args(data: enums.LavalinkEvents, raw_data: dict):
 
     try:
         node_ = node.get_node(guild_id, ignore_ready_status=True)
-        player = node_.player_manager.get_player(guild_id)
+        player = node_.get_player(guild_id)
     except (IndexError, KeyError):
         if data != enums.LavalinkEvents.TRACK_END:
             log.debug(
@@ -345,7 +341,6 @@ async def close(bot):
     """
     unregister_event_listener(_handle_event)
     unregister_update_listener(_handle_update)
-    bot.remove_listener(node.on_socket_response)
     bot.remove_listener(_on_guild_remove, name="on_guild_remove")
     await node.disconnect()
 
@@ -355,13 +350,13 @@ async def close(bot):
 
 def all_players() -> Tuple[player_manager.Player]:
     nodes = node._nodes
-    ret = tuple(p for n in nodes for p in n.player_manager.players)
+    ret = tuple(p for n in nodes for p in n.players)
     return ret
 
 
 def all_connected_players() -> Tuple[player_manager.Player]:
     nodes = node._nodes
-    ret = tuple(p for n in nodes for p in n.player_manager.players if p.connected)
+    ret = tuple(p for n in nodes for p in n.players if p.connected)
     return ret
 
 

lavalink/node.py

@@ -1,26 +1,28 @@
 from __future__ import annotations
+
 import asyncio
 import contextlib
 import secrets
 import string
+import typing
 from collections import namedtuple
-from typing import Awaitable, List, Optional, cast
+from typing import Awaitable, KeysView, List, Optional, ValuesView, cast
 
 import aiohttp
-import typing
 from discord.backoff import ExponentialBackoff
 from discord.ext.commands import Bot
 
-from . import ws_discord_log, ws_ll_log
-from .enums import *
-from .player_manager import PlayerManager
+from . import log, ws_ll_log, ws_rll_log
+from .enums import LavalinkEvents, LavalinkIncomingOp, LavalinkOutgoingOp, NodeState, PlayerState
+from .player_manager import Player
 from .rest_api import Track
+from .utils import VoiceChannel
 
-__all__ = ["Stats", "Node", "NodeStats", "get_node", "get_nodes_stats", "join_voice"]
+__all__ = ["Stats", "Node", "NodeStats", "get_node", "get_nodes_stats"]
 
 _nodes: List[Node] = []
 
-PlayerState = namedtuple("PlayerState", "position time")
+PositionTime = namedtuple("PositionTime", "position time")
 MemoryInfo = namedtuple("MemoryInfo", "reservable used free allocated")
 CPUInfo = namedtuple("CPUInfo", "cores systemLoad lavalinkLoad")
 
@@ -103,7 +105,6 @@ def __init__(
         self,
         _loop: asyncio.BaseEventLoop,
         event_handler: typing.Callable,
-        voice_ws_func: typing.Callable,
         host: str,
         password: str,
         port: int,
@@ -122,8 +123,6 @@ def __init__(
             The event loop of the bot.
         event_handler
             Function to dispatch events to.
-        voice_ws_func : typing.Callable
-            Function that takes one argument, guild ID, and returns a websocket.
         host : str
             Lavalink player host.
         password : str
@@ -146,7 +145,6 @@ def __init__(
         self.loop = _loop
         self.bot = bot
         self.event_handler = event_handler
-        self.get_voice_ws = voice_ws_func
         self.host = host
         self.port = port
         self.password = password
@@ -165,13 +163,12 @@ def __init__(
         self.session = aiohttp.ClientSession()
 
         self._queue: List = []
+        self._players_dict = {}
 
         self.state = NodeState.CONNECTING
         self._state_handlers: List = []
         self._retries = 0
 
-        self.player_manager = PlayerManager(self)
-
         self.stats = None
 
         if self not in _nodes:
@@ -183,6 +180,8 @@ def __init__(
             aiohttp.WSMsgType.CLOSED,
         )
 
+        self.register_state_handler(self.node_state_handler)
+
     def __repr__(self):
         return (
             "<Node: "
@@ -197,6 +196,14 @@ def __repr__(self):
     def headers(self) -> dict:
         return self._get_connect_headers()
 
+    @property
+    def players(self) -> ValuesView[Player]:
+        return self._players_dict.values()
+
+    @property
+    def guild_ids(self) -> KeysView[int]:
+        return self._players_dict.keys()
+
     def _gen_key(self):
         if self._resume_key is None:
             return _Key()
@@ -367,7 +374,7 @@ async def _handle_op(self, op: LavalinkIncomingOp, data):
                 self.event_handler(op, event, data)
         elif op == LavalinkIncomingOp.PLAYER_UPDATE:
             state = data.get("state", {})
-            state = PlayerState(position=state.get("position", 0), time=state.get("time", 0))
+            state = PositionTime(position=state.get("position", 0), time=state.get("time", 0))
             self.event_handler(op, state, data)
         elif op == LavalinkIncomingOp.STATS:
             stats = Stats(
@@ -413,7 +420,7 @@ async def _reconnect(self):
                 self._retries = 0
 
     def dispatch_reconnect(self):
-        for guild_id in self.player_manager.guild_ids:
+        for guild_id in self.guild_ids:
             self.event_handler(
                 LavalinkIncomingOp.EVENT,
                 LavalinkEvents.WEBSOCKET_CLOSED,
@@ -449,12 +456,91 @@ def register_state_handler(self, func):
     def unregister_state_handler(self, func):
         self._state_handlers.remove(func)
 
-    async def join_voice_channel(self, guild_id, channel_id, deafen: bool = False):
+    async def create_player(self, channel: VoiceChannel, deafen: bool = False) -> Player:
+        """
+        Connects to a discord voice channel.
+
+        This function is safe to repeatedly call as it will return an existing
+        player if there is one.
+
+        Parameters
+        ----------
+        channel
+
+        Returns
+        -------
+        Player
+            The created Player object.
         """
-        Alternative way to join a voice channel if node is known.
+        if self._already_in_guild(channel):
+            p = self.get_player(channel.guild.id)
+            await p.move_to(channel, deafen=deafen)
+        else:
+            p = await channel.connect(cls=Player)
+            if deafen:
+                await p.guild.change_voice_state(channel=p.channel, self_deaf=True)
+            self._players_dict[channel.guild.id] = p
+            await self.refresh_player_state(p)
+        return p
+
+    def _already_in_guild(self, channel: VoiceChannel) -> bool:
+        return channel.guild.id in self._players_dict
+
+    def get_player(self, guild_id: int) -> Player:
+        """
+        Gets a Player object from a guild ID.
+
+        Parameters
+        ----------
+        guild_id : int
+            Discord guild ID.
+
+        Returns
+        -------
+        Player
+
+        Raises
+        ------
+        KeyError
+            If that guild does not have a Player, e.g. is not connected to any
+            voice channel.
         """
-        voice_ws = self.get_voice_ws(guild_id)
-        await voice_ws.voice_state(guild_id, channel_id, self_deaf=deafen)
+        if guild_id in self._players_dict:
+            return self._players_dict[guild_id]
+        raise KeyError("No such player for that guild.")
+
+    async def node_state_handler(self, next_state: NodeState, old_state: NodeState):
+        ws_rll_log.debug("Received node state update: %s -> %s", old_state.name, next_state.name)
+        if next_state == NodeState.READY:
+            await self.update_player_states(PlayerState.READY)
+        elif next_state == NodeState.DISCONNECTING:
+            await self.update_player_states(PlayerState.DISCONNECTING)
+        elif next_state in (NodeState.CONNECTING, NodeState.RECONNECTING):
+            await self.update_player_states(PlayerState.NODE_BUSY)
+
+    async def update_player_states(self, state: PlayerState):
+        for p in self.players:
+            await p.update_state(state)
+
+    async def refresh_player_state(self, player: Player):
+        if self.ready:
+            await player.update_state(PlayerState.READY)
+        elif self.state == NodeState.DISCONNECTING:
+            await player.update_state(PlayerState.DISCONNECTING)
+        else:
+            await player.update_state(PlayerState.NODE_BUSY)
+
+    def remove_player(self, player: Player):
+        if player.state != PlayerState.DISCONNECTING:
+            log.error(
+                "Attempting to remove a player (%r) from player list with state: %s",
+                player,
+                player.state.name,
+            )
+            return
+        guild_id = player.channel.guild.id
+        if guild_id in self._players_dict:
+            del self._players_dict[guild_id]
 
     async def disconnect(self):
         """
@@ -468,7 +554,10 @@ async def disconnect(self):
         if self._resuming_configured:
             await self.send(dict(op="configureResuming", key=None))
         self._resuming_configured = False
-        await self.player_manager.disconnect()
+
+        for p in tuple(self.players):
+            await p.disconnect(force=True)
+        log.debug("Disconnected all players.")
 
         if self._ws is not None and not self._ws.closed:
             await self._ws.close()
@@ -561,7 +650,7 @@ async def seek(self, guild_id: int, position: int):
         )
 
 
-def get_node(guild_id: int, ignore_ready_status: bool = False) -> Node:
+def get_node(guild_id: int = None, ignore_ready_status: bool = False) -> Node:
     """
     Gets a node based on a guild ID, useful for noding separation. If the
     guild ID does not already have a node association, the least used
@@ -580,7 +669,7 @@ def get_node(guild_id: int, ignore_ready_status: bool = False) -> Node:
     least_used = None
 
     for node in _nodes:
-        guild_ids = node.player_manager.guild_ids
+        guild_ids = node.guild_ids
 
         if ignore_ready_status is False and not node.ready:
             continue
@@ -601,41 +690,6 @@ def get_nodes_stats():
     return [node.stats for node in _nodes]
 
 
-async def join_voice(guild_id: int, channel_id: int, deafen: bool = False):
-    """
-    Joins a voice channel by ID's.
-
-    Parameters
-    ----------
-    guild_id : int
-    channel_id : int
-    """
-    node = get_node(guild_id)
-    await node.join_voice_channel(guild_id, channel_id, deafen)
-
-
 async def disconnect():
     for node in _nodes.copy():
         await node.disconnect()
-
-
-async def on_socket_response(data):
-    raw_event = data.get("t")
-    try:
-        event = DiscordVoiceSocketResponses(raw_event)
-    except ValueError:
-        return
-
-    guild_id = data["d"]["guild_id"]
-
-    try:
-        node = get_node(guild_id, ignore_ready_status=True)
-    except IndexError:
-        ws_discord_log.info(
-            f"Received unhandled Discord WS voice response for guild: %d, %s", int(guild_id), data
-        )
-    else:
-        ws_ll_log.debug(
-            f"Received Discord WS voice response for guild: %d, %s", int(guild_id), data
-        )
-        await node.player_manager.on_socket_response(data)