pub mod bridge;
mod context;
#[cfg(feature = "gateway")]
mod dispatch;
mod error;
#[cfg(feature = "gateway")]
mod event_handler;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context as FutContext, Poll};
use futures::future::BoxFuture;
use tokio::sync::{Mutex, RwLock};
use tracing::{debug, error, info, instrument};
use typemap_rev::{TypeMap, TypeMapKey};
#[cfg(feature = "gateway")]
use self::bridge::gateway::{
ShardManager,
ShardManagerError,
ShardManagerMonitor,
ShardManagerOptions,
};
#[cfg(feature = "voice")]
use self::bridge::voice::VoiceGatewayManager;
pub use self::context::Context;
pub use self::error::Error as ClientError;
#[cfg(feature = "gateway")]
pub use self::event_handler::{EventHandler, RawEventHandler};
#[cfg(feature = "gateway")]
use super::gateway::GatewayError;
#[cfg(feature = "cache")]
pub use crate::cache::Cache;
#[cfg(feature = "cache")]
use crate::cache::Settings as CacheSettings;
#[cfg(feature = "framework")]
use crate::framework::Framework;
use crate::http::Http;
use crate::internal::prelude::*;
#[cfg(feature = "gateway")]
use crate::model::gateway::GatewayIntents;
use crate::model::id::ApplicationId;
pub use crate::CacheAndHttp;
#[cfg(feature = "gateway")]
#[must_use = "Builders do nothing unless they are awaited"]
pub struct ClientBuilder {
data: Option<TypeMap>,
http: Option<Http>,
fut: Option<BoxFuture<'static, Result<Client>>>,
intents: GatewayIntents,
#[cfg(feature = "cache")]
cache_settings: Option<CacheSettings>,
#[cfg(feature = "framework")]
framework: Option<Arc<dyn Framework + Send + Sync + 'static>>,
#[cfg(feature = "voice")]
voice_manager: Option<Arc<dyn VoiceGatewayManager + Send + Sync + 'static>>,
event_handler: Option<Arc<dyn EventHandler>>,
raw_event_handler: Option<Arc<dyn RawEventHandler>>,
}
#[cfg(feature = "gateway")]
impl ClientBuilder {
fn _new(http: Http, intents: GatewayIntents) -> Self {
Self {
data: Some(TypeMap::new()),
http: Some(http),
fut: None,
intents,
#[cfg(feature = "cache")]
cache_settings: Some(CacheSettings::new()),
#[cfg(feature = "framework")]
framework: None,
#[cfg(feature = "voice")]
voice_manager: None,
event_handler: None,
raw_event_handler: None,
}
}
pub fn new(token: impl AsRef<str>, intents: GatewayIntents) -> Self {
Self::_new(Http::new(token.as_ref()), intents)
}
pub fn new_with_http(http: Http, intents: GatewayIntents) -> Self {
Self::_new(http, intents)
}
pub fn token(mut self, token: impl AsRef<str>) -> Self {
self.http = Some(Http::new(token.as_ref()));
self
}
pub fn get_token(&self) -> Option<&str> {
self.http.as_ref().map(|http| http.token.as_str())
}
pub fn application_id(self, application_id: u64) -> Self {
if let Some(http) = &self.http {
http.set_application_id(application_id);
}
self
}
pub fn get_application_id(&self) -> Option<ApplicationId> {
self.http.as_ref().and_then(|h| h.application_id().map(ApplicationId))
}
pub fn type_map(mut self, type_map: TypeMap) -> Self {
self.data = Some(type_map);
self
}
pub fn get_type_map(&self) -> Option<&TypeMap> {
self.data.as_ref()
}
pub fn type_map_insert<T: TypeMapKey>(mut self, value: T::Value) -> Self {
self.data.get_or_insert_with(TypeMap::new).insert::<T>(value);
self
}
#[cfg(feature = "cache")]
pub fn cache_settings<F>(mut self, f: F) -> Self
where
F: FnOnce(&mut CacheSettings) -> &mut CacheSettings,
{
if let Some(ref mut settings) = self.cache_settings {
f(settings);
}
self
}
#[cfg(feature = "cache")]
pub fn get_cache_settings(&self) -> Option<&CacheSettings> {
self.cache_settings.as_ref()
}
#[cfg(feature = "framework")]
pub fn framework<F>(mut self, framework: F) -> Self
where
F: Framework + Send + Sync + 'static,
{
self.framework = Some(Arc::new(framework));
self
}
#[cfg(feature = "framework")]
pub fn framework_arc<T: Framework + Send + Sync + 'static>(
mut self,
framework: Arc<T>,
) -> Self {
self.framework = Some(framework as Arc<dyn Framework + Send + Sync + 'static>);
self
}
#[cfg(feature = "framework")]
pub fn get_framework(&self) -> Option<Arc<dyn Framework + Send + Sync>> {
self.framework.clone()
}
#[cfg(feature = "voice")]
pub fn voice_manager<V>(mut self, voice_manager: V) -> Self
where
V: VoiceGatewayManager + Send + Sync + 'static,
{
self.voice_manager = Some(Arc::new(voice_manager));
self
}
#[cfg(feature = "voice")]
pub fn voice_manager_arc(
mut self,
voice_manager: Arc<dyn VoiceGatewayManager + Send + Sync + 'static>,
) -> Self {
self.voice_manager = Some(voice_manager);
self
}
#[cfg(feature = "voice")]
pub fn get_voice_manager(&self) -> Option<Arc<dyn VoiceGatewayManager + Send + Sync>> {
self.voice_manager.clone()
}
pub fn intents(mut self, intents: GatewayIntents) -> Self {
self.intents = intents;
self
}
pub fn get_intents(&self) -> GatewayIntents {
self.intents
}
pub fn event_handler<H: EventHandler + 'static>(mut self, event_handler: H) -> Self {
self.event_handler = Some(Arc::new(event_handler));
self
}
pub fn event_handler_arc<H: EventHandler + 'static>(
mut self,
event_handler_arc: Arc<H>,
) -> Self {
self.event_handler = Some(event_handler_arc);
self
}
pub fn get_event_handler(&self) -> Option<Arc<dyn EventHandler>> {
self.event_handler.clone()
}
pub fn raw_event_handler<H: RawEventHandler + 'static>(mut self, raw_event_handler: H) -> Self {
self.raw_event_handler = Some(Arc::new(raw_event_handler));
self
}
pub fn get_raw_event_handler(&self) -> Option<Arc<dyn RawEventHandler>> {
self.raw_event_handler.clone()
}
}
#[cfg(feature = "gateway")]
impl Future for ClientBuilder {
type Output = Result<Client>;
#[allow(clippy::unwrap_used)] #[instrument(skip(self))]
fn poll(mut self: Pin<&mut Self>, ctx: &mut FutContext<'_>) -> Poll<Self::Output> {
if self.fut.is_none() {
let data = Arc::new(RwLock::new(self.data.take().unwrap()));
#[cfg(feature = "framework")]
let framework = self.framework.take()
.expect("The `framework`-feature is enabled (it's on by default), but no framework was provided.\n\
If you don't want to use the command framework, disable default features and specify all features you want to use.");
let event_handler = self.event_handler.take();
let raw_event_handler = self.raw_event_handler.take();
let intents = self.intents;
let mut http = self.http.take().unwrap();
if let Some(event_handler) = event_handler.clone() {
http.ratelimiter.set_ratelimit_callback(Box::new(move |info| {
let event_handler = event_handler.clone();
tokio::spawn(async move { event_handler.ratelimit(info).await });
}));
}
let http = Arc::new(http);
#[cfg(feature = "voice")]
let voice_manager = self.voice_manager.take();
let cache_and_http = Arc::new(CacheAndHttp {
#[cfg(feature = "cache")]
cache: Arc::new(Cache::new_with_settings(self.cache_settings.take().unwrap())),
http: Arc::clone(&http),
});
self.fut = Some(Box::pin(async move {
let ws_url = Arc::new(Mutex::new(match http.get_gateway().await {
Ok(response) => response.url,
Err(err) => {
tracing::warn!("HTTP request to get gateway URL failed: {}", err);
"wss://gateway.discord.gg".to_string()
},
}));
let (shard_manager, shard_manager_worker) = {
ShardManager::new(ShardManagerOptions {
data: &data,
event_handler: &event_handler,
raw_event_handler: &raw_event_handler,
#[cfg(feature = "framework")]
framework: &framework,
shard_index: 0,
shard_init: 0,
shard_total: 0,
#[cfg(feature = "voice")]
voice_manager: &voice_manager,
ws_url: &ws_url,
cache_and_http: &cache_and_http,
intents,
})
.await
};
Ok(Client {
data,
shard_manager,
shard_manager_worker,
#[cfg(feature = "voice")]
voice_manager,
ws_url,
cache_and_http,
})
}));
}
self.fut.as_mut().unwrap().as_mut().poll(ctx)
}
}
#[cfg(feature = "gateway")]
pub struct Client {
pub data: Arc<RwLock<TypeMap>>,
pub shard_manager: Arc<Mutex<ShardManager>>,
shard_manager_worker: ShardManagerMonitor,
#[cfg(feature = "voice")]
pub voice_manager: Option<Arc<dyn VoiceGatewayManager + Send + Sync + 'static>>,
pub ws_url: Arc<Mutex<String>>,
pub cache_and_http: Arc<CacheAndHttp>,
}
impl Client {
pub fn builder(token: impl AsRef<str>, intents: GatewayIntents) -> ClientBuilder {
ClientBuilder::new(token, intents)
}
#[instrument(skip(self))]
pub async fn start(&mut self) -> Result<()> {
self.start_connection([0, 0, 1]).await
}
#[instrument(skip(self))]
pub async fn start_autosharded(&mut self) -> Result<()> {
let (x, y) = {
let res = self.cache_and_http.http.get_bot_gateway().await?;
(res.shards - 1, res.shards)
};
self.start_connection([0, x, y]).await
}
#[instrument(skip(self))]
pub async fn start_shard(&mut self, shard: u64, shards: u64) -> Result<()> {
self.start_connection([shard, shard, shards]).await
}
#[instrument(skip(self))]
pub async fn start_shards(&mut self, total_shards: u64) -> Result<()> {
self.start_connection([0, total_shards - 1, total_shards]).await
}
#[instrument(skip(self))]
pub async fn start_shard_range(&mut self, range: [u64; 2], total_shards: u64) -> Result<()> {
self.start_connection([range[0], range[1], total_shards]).await
}
#[instrument(skip(self))]
async fn start_connection(&mut self, shard_data: [u64; 3]) -> Result<()> {
#[cfg(feature = "voice")]
if let Some(voice_manager) = &self.voice_manager {
let user = self.cache_and_http.http.get_current_user().await?;
voice_manager.initialise(shard_data[2], user.id).await;
}
{
let mut manager = self.shard_manager.lock().await;
let init = shard_data[1] - shard_data[0] + 1;
manager.set_shards(shard_data[0], init, shard_data[2]).await;
debug!("Initializing shard info: {} - {}/{}", shard_data[0], init, shard_data[2],);
if let Err(why) = manager.initialize() {
error!("Failed to boot a shard: {:?}", why);
info!("Shutting down all shards");
manager.shutdown_all().await;
return Err(Error::Client(ClientError::ShardBootFailure));
}
}
if let Err(why) = self.shard_manager_worker.run().await {
let err = match why {
ShardManagerError::DisallowedGatewayIntents => {
GatewayError::DisallowedGatewayIntents
},
ShardManagerError::InvalidGatewayIntents => GatewayError::InvalidGatewayIntents,
ShardManagerError::InvalidToken => GatewayError::InvalidAuthentication,
};
return Err(Error::Gateway(err));
}
Ok(())
}
}