twilight_lavalink/

client.rs

//! Client to manage nodes and players.

use crate::{
    model::VoiceUpdate,
    node::{IncomingEvents, Node, NodeConfig, NodeError},
    player::{Player, PlayerManager},
};
use dashmap::DashMap;
use std::{
    error::Error,
    fmt::{Display, Formatter, Result as FmtResult},
    net::SocketAddr,
    sync::Arc,
};
use twilight_model::{
    gateway::{ShardId, event::Event, payload::incoming::VoiceServerUpdate},
    id::{
        Id,
        marker::{ChannelMarker, GuildMarker, UserMarker},
    },
};

/// An error that can occur while interacting with the client.
#[derive(Debug)]
pub struct ClientError {
    kind: ClientErrorType,
    source: Option<Box<dyn Error + Send + Sync>>,
}

impl ClientError {
    /// Immutable reference to the type of error that occurred.
    pub const fn kind(&self) -> &ClientErrorType {
        &self.kind
    }

    /// Consume the error, returning the source error if there is any.
    pub fn into_source(self) -> Option<Box<dyn Error + Send + Sync>> {
        self.source
    }

    /// Consume the error, returning the owned error type and the source error.
    #[must_use = "consuming the error into its parts has no effect if left unused"]
    pub fn into_parts(self) -> (ClientErrorType, Option<Box<dyn Error + Send + Sync>>) {
        (self.kind, self.source)
    }
}

impl Display for ClientError {
    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
        match &self.kind {
            ClientErrorType::NodesUnconfigured => f.write_str("no node has been configured"),
            ClientErrorType::SendingVoiceUpdate => {
                f.write_str("couldn't send voice update to node")
            }
        }
    }
}

impl Error for ClientError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        self.source
            .as_ref()
            .map(|source| &**source as &(dyn Error + 'static))
    }
}

/// Type of [`ClientError`] that occurred.
#[derive(Debug)]
#[non_exhaustive]
pub enum ClientErrorType {
    /// A node isn't configured, so the operation isn't possible to fulfill.
    NodesUnconfigured,
    /// Sending a voice update event to the node failed because the node's
    /// connection was shutdown.
    SendingVoiceUpdate,
}

#[derive(Debug)]
struct DiscordSession {
    channel_id: Option<Id<ChannelMarker>>,
    id: Box<str>,
}

/// The lavalink client that manages nodes, players, and processes events from
/// Discord to tie it all together.
///
/// **Note**: You must call the [`process`] method with every Voice State Update
/// and Voice Server Update event you receive from Discord. It will
/// automatically forward these events to Lavalink. See its documentation for
/// more information.
///
/// You can retrieve players using the [`player`] method. Players contain
/// information about the active playing information of a guild and allows you to send events to the
/// connected node, such as [`Play`] events.
///
/// # Using a Lavalink client in multiple tasks
///
/// To use a Lavalink client instance in multiple tasks, consider wrapping it in
/// an [`std::sync::Arc`] or [`std::rc::Rc`].
///
/// [`Play`]: crate::model::outgoing::Play
/// [`player`]: Self::player
/// [`process`]: Self::process
#[derive(Debug)]
pub struct Lavalink {
    nodes: DashMap<SocketAddr, Arc<Node>>,
    players: PlayerManager,
    shard_count: u32,
    user_id: Id<UserMarker>,
    server_updates: DashMap<Id<GuildMarker>, VoiceServerUpdate>,
    discord_sessions: DashMap<Id<GuildMarker>, DiscordSession>,
}

impl Lavalink {
    /// Create a new Lavalink client instance.
    ///
    /// The user ID and number of shards provided may not be modified during
    /// runtime, and the client must be re-created. These parameters are
    /// automatically passed to new nodes created via [`add`].
    ///
    /// [`add`]: Self::add
    pub fn new(user_id: Id<UserMarker>, shard_count: u32) -> Self {
        Self {
            nodes: DashMap::new(),
            players: PlayerManager::new(),
            shard_count,
            user_id,
            server_updates: DashMap::new(),
            discord_sessions: DashMap::new(),
        }
    }

    /// Process an event into the Lavalink client.
    ///
    /// **Note**: calling this method in your event loop is required. See the
    /// [crate documentation] for an example.
    ///
    /// This requires the `VoiceServerUpdate` and `VoiceStateUpdate` events that
    /// you receive from Discord over the gateway to send voice updates to
    /// nodes. For simplicity in some applications' event loops, any event can
    /// be provided, but they will just be ignored.
    ///
    /// The Ready event can optionally be provided to do some cleaning of
    /// stalled voice states that never received their voice server update half
    /// or vice versa. It is recommended that you process Ready events.
    ///
    /// # Errors
    ///
    /// Returns a [`ClientErrorType::NodesUnconfigured`] error type if no nodes
    /// have been added to the client when attempting to retrieve a guild's
    /// player.
    ///
    /// [crate documentation]: crate#examples
    pub async fn process(&self, event: &Event) -> Result<(), ClientError> {
        tracing::trace!("processing event: {event:?}");

        let guild_id = match event {
            Event::Ready(e) => {
                let shard_id = e.shard.map_or(0, ShardId::number);

                self.clear_shard_states(shard_id);

                return Ok(());
            }
            Event::VoiceServerUpdate(e) => {
                self.server_updates.insert(e.guild_id, e.clone());
                e.guild_id
            }
            Event::VoiceStateUpdate(e) => {
                if e.user_id != self.user_id {
                    tracing::trace!("got voice state update from another user");

                    return Ok(());
                }

                if let Some(guild_id) = e.guild_id {
                    // Update player if it exists and update the connected channel ID.
                    if let Some(player) = self.players.get(&guild_id) {
                        player.set_channel_id(e.channel_id);
                    }

                    if e.channel_id.is_none() {
                        self.discord_sessions.remove(&guild_id);
                        self.server_updates.remove(&guild_id);
                    } else {
                        let session = DiscordSession {
                            channel_id: e.channel_id,
                            id: e.session_id.clone().into_boxed_str(),
                        };
                        self.discord_sessions.insert(guild_id, session);
                    }
                    guild_id
                } else {
                    tracing::trace!("event has no guild ID: {e:?}");
                    return Ok(());
                }
            }
            _ => return Ok(()),
        };

        tracing::debug!("got voice server/state update for {guild_id:?}: {event:?}");

        let update = {
            let server = self.server_updates.get(&guild_id);
            let session = self.discord_sessions.get(&guild_id);
            match (server, session) {
                (Some(server), Some(session)) => {
                    let server = server.value();
                    let DiscordSession {
                        channel_id,
                        id: session_id,
                    } = session.value();
                    tracing::debug!(
                        "got both halves for {guild_id}: {server:?}; Session ID: {session_id:?}; Channel ID: {channel_id:?}",
                    );
                    VoiceUpdate::new(guild_id, session_id.as_ref(), *channel_id, server.clone())
                }
                (Some(server), None) => {
                    tracing::debug!(
                        "guild {guild_id} is now waiting for other half; got: {:?}",
                        server.value()
                    );
                    return Ok(());
                }
                (None, Some(session)) => {
                    tracing::debug!(
                        "guild {guild_id} is now waiting for other half; got session ID: {:?}",
                        session.value()
                    );
                    return Ok(());
                }
                _ => return Ok(()),
            }
        };

        tracing::debug!("getting player for guild {guild_id}");

        let player = self.player(guild_id).await?;

        tracing::debug!("sending voice update for guild {guild_id}: {update:?}");

        player.send(update).map_err(|source| ClientError {
            kind: ClientErrorType::SendingVoiceUpdate,
            source: Some(Box::new(source)),
        })?;

        tracing::debug!("sent voice update for guild {guild_id}");

        Ok(())
    }

    /// Add a new node to be managed by the Lavalink client.
    ///
    /// If a node already exists with the provided address, then it will be
    /// replaced.
    ///
    /// Pass `None` for `session_id` to create a fresh session. Pass
    /// `Some(id)` to attempt resuming an existing Lavalink session. Resume
    /// success is indicated by the `resumed` field on the
    /// [`Ready`] event received via the returned [`IncomingEvents`] stream.
    ///
    /// If a `session_id` is provided and the session cannot be resumed,
    /// the connection will fail with a [`Connecting`] error. The caller
    /// can then retry with `session_id` set to `None` for a fresh session.
    ///
    /// [`Ready`]: crate::model::incoming::Ready
    /// [`Connecting`]: crate::node::NodeErrorType::Connecting
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use twilight_lavalink::Lavalink;
    /// # use twilight_model::id::Id;
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// # let lavalink = Lavalink::new(Id::new(1), 1);
    /// # let address = "127.0.0.1:2333".parse()?;
    /// # let auth = "youshallnotpass";
    /// // Fresh session (no resume)
    /// let (node, events) = lavalink.add(address, auth, None).await?;
    ///
    /// // Resume an existing session
    /// let old_session_id = Some("existing-session-id".to_string());
    /// let (node, events) = lavalink.add(address, auth, old_session_id).await?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// # Errors
    ///
    /// See the errors section of [`Node::connect`].
    pub async fn add(
        &self,
        address: SocketAddr,
        authorization: impl Into<String>,
        session_id: Option<String>,
    ) -> Result<(Arc<Node>, IncomingEvents), NodeError> {
        let config = NodeConfig {
            address,
            authorization: authorization.into(),
            user_id: self.user_id,
            enable_tls: cfg!(feature = "tls"),
            session_id,
        };

        let (node, rx) = Node::connect(config, self.players.clone()).await?;
        let node = Arc::new(node);
        self.nodes.insert(address, Arc::clone(&node));

        Ok((node, rx))
    }

    /// Remove a node from the list of nodes being managed by the Lavalink
    /// client.
    ///
    /// This does not disconnect the node. Use [`Lavalink::disconnect`] instead.
    /// or drop all [`Node`]s.
    ///
    /// The node is returned if it existed.
    pub fn remove(&self, address: SocketAddr) -> Option<(SocketAddr, Arc<Node>)> {
        self.nodes.remove(&address)
    }

    /// Remove a node from the list of nodes being managed by the Lavalink
    /// client and terminates the connection.
    ///
    /// Use [`Lavalink::remove`] if detaching a node from a Lavalink instance
    /// is required without closing the underlying connection.
    ///
    /// Returns whether the node has been removed and disconnected.
    pub fn disconnect(&self, address: SocketAddr) -> bool {
        self.nodes.remove(&address).is_some()
    }

    /// Determine the "best" node for new players according to available nodes'
    /// penalty scores. Disconnected nodes will not be considered.
    ///
    /// Refer to [`Node::penalty`] for how this is calculated.
    ///
    /// # Errors
    ///
    /// Returns a [`ClientErrorType::NodesUnconfigured`] error type if there are
    /// no connected nodes available in the client.
    ///
    /// [`Node::penalty`]: crate::node::Node::penalty
    pub async fn best(&self) -> Result<Arc<Node>, ClientError> {
        let mut lowest = i32::MAX;
        let mut best = None;

        for node in &self.nodes {
            if node.sender().is_closed() {
                continue;
            }

            let penalty = node.value().penalty().await;

            if penalty < lowest {
                lowest = penalty;
                best.replace(node.clone());
            }
        }

        best.ok_or(ClientError {
            kind: ClientErrorType::NodesUnconfigured,
            source: None,
        })
    }

    /// Retrieve an immutable reference to the player manager.
    pub const fn players(&self) -> &PlayerManager {
        &self.players
    }

    /// Retrieve a player for the guild.
    ///
    /// Creates a player configured to use the best available node if a player
    /// for the guild doesn't already exist. Use [`PlayerManager::get`] to only
    /// retrieve and not create.
    ///
    /// # Errors
    ///
    /// Returns a [`ClientError`] with a [`ClientErrorType::NodesUnconfigured`]
    /// type if no node has been configured via [`add`].
    ///
    /// [`PlayerManager::get`]: crate::player::PlayerManager::get
    /// [`add`]: Self::add
    pub async fn player(&self, guild_id: Id<GuildMarker>) -> Result<Arc<Player>, ClientError> {
        if let Some(player) = self.players().get(&guild_id) {
            return Ok(player);
        }

        let node = self.best().await?;

        Ok(self.players().get_or_insert(guild_id, node))
    }

    /// Clear out the map of guild states/updates for a shard that are waiting
    /// for their other half.
    ///
    /// We can do this by iterating over the map and removing the ones that we
    /// can calculate came from a shard.
    ///
    /// This map should be small or empty, and if it isn't, then it needs to be
    /// cleared out anyway.
    fn clear_shard_states(&self, shard_id: u32) {
        let shard_count = u64::from(self.shard_count);

        self.server_updates
            .retain(|k, _| (k.get() >> 22) % shard_count != u64::from(shard_id));
        self.discord_sessions
            .retain(|k, _| (k.get() >> 22) % shard_count != u64::from(shard_id));
    }
}

#[cfg(test)]
mod tests {
    use super::{ClientError, ClientErrorType, Lavalink};
    use static_assertions::assert_impl_all;
    use std::{error::Error, fmt::Debug};

    assert_impl_all!(ClientErrorType: Debug, Send, Sync);
    assert_impl_all!(ClientError: Error, Send, Sync);
    assert_impl_all!(Lavalink: Debug, Send, Sync);
}