use super::{
builder::ShardBuilder,
command::Command,
config::Config,
emitter::Emitter,
event::Events,
json,
processor::{ConnectingErrorType, Latency, Session, ShardProcessor},
raw_message::Message,
stage::Stage,
};
use crate::Intents;
use serde::{Deserialize, Serialize};
use std::{
borrow::Cow,
error::Error,
fmt::{Display, Formatter, Result as FmtResult},
sync::{atomic::Ordering, Arc, Mutex},
};
use tokio::{
sync::{watch::Receiver as WatchReceiver, OnceCell},
task::JoinHandle,
};
use tokio_tungstenite::tungstenite::protocol::{
frame::coding::CloseCode, CloseFrame as TungsteniteCloseFrame,
};
#[derive(Debug)]
pub struct CommandError {
kind: CommandErrorType,
source: Option<Box<dyn Error + Send + Sync>>,
}
impl CommandError {
#[must_use = "retrieving the type has no effect if left unused"]
pub const fn kind(&self) -> &CommandErrorType {
&self.kind
}
#[must_use = "consuming the error and retrieving the source has no effect if left unused"]
pub fn into_source(self) -> Option<Box<dyn Error + Send + Sync>> {
self.source
}
#[must_use = "consuming the error into its parts has no effect if left unused"]
pub fn into_parts(self) -> (CommandErrorType, Option<Box<dyn Error + Send + Sync>>) {
(self.kind, None)
}
}
impl CommandError {
pub(crate) fn from_send(error: SendError) -> Self {
let (kind, source) = error.into_parts();
let new_kind = match kind {
SendErrorType::ExecutorShutDown => CommandErrorType::ExecutorShutDown,
SendErrorType::HeartbeaterNotStarted => CommandErrorType::HeartbeaterNotStarted,
SendErrorType::Sending => CommandErrorType::Sending,
SendErrorType::SessionInactive => CommandErrorType::SessionInactive,
};
Self {
kind: new_kind,
source,
}
}
}
impl Display for CommandError {
fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
match &self.kind {
CommandErrorType::ExecutorShutDown => f.write_str("runtime executor has shut down"),
CommandErrorType::HeartbeaterNotStarted => {
f.write_str("heartbeater task hasn't been started yet")
}
CommandErrorType::Sending => {
f.write_str("sending the message over the websocket failed")
}
CommandErrorType::Serializing => f.write_str("serializing the value as json failed"),
CommandErrorType::SessionInactive => Display::fmt(&SessionInactiveError, f),
}
}
}
impl Error for CommandError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
self.source
.as_ref()
.map(|source| &**source as &(dyn Error + 'static))
}
}
#[derive(Debug)]
#[non_exhaustive]
pub enum CommandErrorType {
ExecutorShutDown,
HeartbeaterNotStarted,
Sending,
Serializing,
SessionInactive,
}
#[derive(Debug)]
#[non_exhaustive]
pub struct SessionInactiveError;
impl Display for SessionInactiveError {
fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
f.write_str("the shard session is inactive and was not started")
}
}
impl Error for SessionInactiveError {}
#[derive(Debug)]
pub struct SendError {
kind: SendErrorType,
source: Option<Box<dyn Error + Send + Sync>>,
}
impl SendError {
#[must_use = "retrieving the type has no effect if left unused"]
pub const fn kind(&self) -> &SendErrorType {
&self.kind
}
#[must_use = "consuming the error and retrieving the source has no effect if left unused"]
pub fn into_source(self) -> Option<Box<dyn Error + Send + Sync>> {
self.source
}
#[must_use = "consuming the error into its parts has no effect if left unused"]
pub fn into_parts(self) -> (SendErrorType, Option<Box<dyn Error + Send + Sync>>) {
(self.kind, None)
}
}
impl Display for SendError {
fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
match &self.kind {
SendErrorType::ExecutorShutDown { .. } => f.write_str("runtime executor has shut down"),
SendErrorType::HeartbeaterNotStarted { .. } => {
f.write_str("heartbeater task hasn't been started yet")
}
SendErrorType::Sending { .. } => {
f.write_str("sending the message over the websocket failed")
}
SendErrorType::SessionInactive { .. } => f.write_str("shard hasn't been started"),
}
}
}
impl Error for SendError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
self.source
.as_ref()
.map(|source| &**source as &(dyn Error + 'static))
}
}
#[derive(Debug)]
#[non_exhaustive]
pub enum SendErrorType {
ExecutorShutDown,
HeartbeaterNotStarted,
Sending,
SessionInactive,
}
#[derive(Debug)]
pub struct ShardStartError {
kind: ShardStartErrorType,
source: Option<Box<dyn Error + Send + Sync>>,
}
impl ShardStartError {
#[must_use = "retrieving the type has no effect if left unused"]
pub const fn kind(&self) -> &ShardStartErrorType {
&self.kind
}
#[must_use = "consuming the error and retrieving the source has no effect if left unused"]
pub fn into_source(self) -> Option<Box<dyn Error + Send + Sync>> {
self.source
}
#[must_use = "consuming the error into its parts has no effect if left unused"]
pub fn into_parts(self) -> (ShardStartErrorType, Option<Box<dyn Error + Send + Sync>>) {
(self.kind, None)
}
}
impl Display for ShardStartError {
fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
match &self.kind {
ShardStartErrorType::AlreadyStarted => {
f.write_str("shard has already been previously started")
}
ShardStartErrorType::Establishing => f.write_str("establishing the connection failed"),
ShardStartErrorType::ParsingGatewayUrl { url } => {
f.write_str("the gateway url `")?;
f.write_str(url)?;
f.write_str("` is invalid")
}
ShardStartErrorType::RetrievingGatewayUrl => {
f.write_str("retrieving the gateway URL via HTTP failed")
}
}
}
}
impl Error for ShardStartError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
self.source
.as_ref()
.map(|source| &**source as &(dyn Error + 'static))
}
}
#[derive(Debug)]
#[non_exhaustive]
pub enum ShardStartErrorType {
AlreadyStarted,
Establishing,
ParsingGatewayUrl {
url: String,
},
RetrievingGatewayUrl,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Information {
id: u64,
latency: Latency,
session_id: Option<Box<str>>,
seq: u64,
stage: Stage,
}
impl Information {
pub const fn id(&self) -> u64 {
self.id
}
pub const fn latency(&self) -> &Latency {
&self.latency
}
pub fn session_id(&self) -> Option<&str> {
self.session_id.as_deref()
}
pub const fn seq(&self) -> u64 {
self.seq
}
pub const fn stage(&self) -> Stage {
self.stage
}
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ResumeSession {
pub session_id: String,
pub sequence: u64,
}
#[derive(Debug)]
pub struct Shard {
config: Arc<Config>,
emitter: Mutex<Option<Emitter>>,
processor_handle: OnceCell<JoinHandle<()>>,
session: OnceCell<WatchReceiver<Arc<Session>>>,
}
impl Shard {
pub fn new(token: impl Into<String>, intents: Intents) -> (Self, Events) {
Self::builder(token, intents).build()
}
pub(crate) fn new_with_config(config: Config) -> (Self, Events) {
let config = Arc::new(config);
let event_types = config.event_types();
let (emitter, rx) = Emitter::new(event_types);
let this = Self {
config,
emitter: Mutex::new(Some(emitter)),
processor_handle: OnceCell::new(),
session: OnceCell::new(),
};
(this, Events::new(event_types, rx))
}
pub fn builder(token: impl Into<String>, intents: Intents) -> ShardBuilder {
ShardBuilder::new(token, intents)
}
pub fn config(&self) -> &Config {
&self.config
}
pub async fn start(&self) -> Result<(), ShardStartError> {
let url = if let Some(u) = self.config.gateway_url.clone() {
u.into_string()
} else {
self.config
.http_client()
.gateway()
.authed()
.exec()
.await
.map_err(|source| ShardStartError {
source: Some(Box::new(source)),
kind: ShardStartErrorType::RetrievingGatewayUrl,
})?
.model()
.await
.map_err(|source| ShardStartError {
source: Some(Box::new(source)),
kind: ShardStartErrorType::RetrievingGatewayUrl,
})?
.url
};
let emitter = self
.emitter
.lock()
.expect("emitter poisoned")
.take()
.ok_or(ShardStartError {
kind: ShardStartErrorType::AlreadyStarted,
source: None,
})?;
let config = Arc::clone(&self.config);
let (processor, wrx) =
ShardProcessor::new(config, url, emitter)
.await
.map_err(|source| {
let (kind, source) = source.into_parts();
let new_kind = match kind {
ConnectingErrorType::Establishing => ShardStartErrorType::Establishing,
ConnectingErrorType::ParsingUrl { url } => {
ShardStartErrorType::ParsingGatewayUrl { url }
}
};
ShardStartError {
source,
kind: new_kind,
}
})?;
let handle = tokio::spawn(async move {
processor.run().await;
#[cfg(feature = "tracing")]
tracing::debug!("shard processor future ended");
});
let _res = self.processor_handle.set(handle);
let _session = self.session.set(wrx);
Ok(())
}
pub fn info(&self) -> Result<Information, SessionInactiveError> {
let session = self.session()?;
Ok(Information {
id: self.config().shard()[0],
latency: session.heartbeats.latency(),
session_id: session.id(),
seq: session.seq(),
stage: session.stage(),
})
}
pub async fn command(&self, value: &impl Command) -> Result<(), CommandError> {
let json = json::to_vec(value).map_err(|source| CommandError {
source: Some(Box::new(source)),
kind: CommandErrorType::Serializing,
})?;
self.send(Message::Binary(json))
.await
.map_err(CommandError::from_send)
}
pub async fn send(&self, message: Message) -> Result<(), SendError> {
let session = self.session().map_err(|source| SendError {
source: Some(Box::new(source)),
kind: SendErrorType::SessionInactive,
})?;
match session.tx.send(message.into_tungstenite()) {
Ok(()) => {
if let Some(limiter) = session.ratelimit.get() {
limiter.acquire_one().await.map_err(|source| SendError {
kind: SendErrorType::ExecutorShutDown,
source: Some(Box::new(source)),
})
} else {
Err(SendError {
kind: SendErrorType::HeartbeaterNotStarted,
source: None,
})
}
}
Err(source) => Err(SendError {
source: Some(Box::new(source)),
kind: SendErrorType::Sending,
}),
}
}
pub fn shutdown(&self) {
if let Some(processor_handle) = self.processor_handle.get() {
processor_handle.abort();
}
if let Ok(session) = self.session() {
let _res = session.close(Some(TungsteniteCloseFrame {
code: CloseCode::Normal,
reason: "".into(),
}));
session.stop_heartbeater();
}
}
pub fn shutdown_resumable(&self) -> (u64, Option<ResumeSession>) {
if let Some(processor_handle) = self.processor_handle.get() {
processor_handle.abort();
}
let shard_id = self.config().shard()[0];
let session = match self.session() {
Ok(session) => session,
Err(_) => return (shard_id, None),
};
let _res = session.close(Some(TungsteniteCloseFrame {
code: CloseCode::Restart,
reason: Cow::from("Closing in a resumable way"),
}));
let session_id = session.id();
let sequence = session.seq.load(Ordering::Relaxed);
session.stop_heartbeater();
let data = session_id.map(|id| ResumeSession {
session_id: id.into_string(),
sequence,
});
(shard_id, data)
}
fn session(&self) -> Result<Arc<Session>, SessionInactiveError> {
let session = self.session.get().ok_or(SessionInactiveError)?;
Ok(Arc::clone(&session.borrow()))
}
}
#[cfg(test)]
mod tests {
use super::{
CommandError, CommandErrorType, Information, ResumeSession, SendError, SendErrorType,
SessionInactiveError, Shard, ShardStartError, ShardStartErrorType,
};
use static_assertions::{assert_fields, assert_impl_all};
use std::{error::Error, fmt::Debug};
assert_impl_all!(CommandErrorType: Debug, Send, Sync);
assert_impl_all!(CommandError: Error, Send, Sync);
assert_impl_all!(Information: Clone, Debug, Send, Sync);
assert_impl_all!(ResumeSession: Clone, Debug, Send, Sync);
assert_impl_all!(SendErrorType: Debug, Send, Sync);
assert_impl_all!(SendError: Error, Send, Sync);
assert_impl_all!(SessionInactiveError: Error, Send, Sync);
assert_fields!(ShardStartErrorType::ParsingGatewayUrl: url);
assert_impl_all!(ShardStartErrorType: Debug, Send, Sync);
assert_impl_all!(ShardStartError: Error, Send, Sync);
assert_impl_all!(Shard: Debug, Send, Sync);
}