use std::sync::Arc;
use futures::channel::mpsc::{UnboundedReceiver as Receiver, UnboundedSender as Sender};
use futures::StreamExt;
use tokio::sync::Mutex;
use tracing::{debug, instrument, warn};
use super::{ShardManager, ShardManagerMessage};
use crate::client::bridge::gateway::ShardId;
#[derive(Debug)]
pub struct ShardManagerMonitor {
pub manager: Arc<Mutex<ShardManager>>,
pub rx: Receiver<ShardManagerMessage>,
pub shutdown: Sender<ShardId>,
}
#[derive(Debug)]
pub enum ShardManagerError {
InvalidToken,
InvalidGatewayIntents,
DisallowedGatewayIntents,
}
type Result<T> = std::result::Result<T, ShardManagerError>;
impl ShardManagerMonitor {
#[instrument(skip(self))]
pub async fn run(&mut self) -> Result<()> {
debug!("Starting shard manager worker");
while let Some(value) = self.rx.next().await {
match value {
ShardManagerMessage::Restart(shard_id) => {
self.manager.lock().await.restart(shard_id).await;
drop(self.shutdown.unbounded_send(shard_id));
},
ShardManagerMessage::ShardUpdate {
id,
latency,
stage,
} => {
let manager = self.manager.lock().await;
let mut runners = manager.runners.lock().await;
if let Some(runner) = runners.get_mut(&id) {
runner.latency = latency;
runner.stage = stage;
}
},
ShardManagerMessage::Shutdown(shard_id, code) => {
self.manager.lock().await.shutdown(shard_id, code).await;
drop(self.shutdown.unbounded_send(shard_id));
},
ShardManagerMessage::ShutdownAll => {
self.manager.lock().await.shutdown_all().await;
break;
},
ShardManagerMessage::ShutdownInitiated => break,
ShardManagerMessage::ShutdownFinished(shard_id) => {
if let Err(why) = self.shutdown.unbounded_send(shard_id) {
warn!(
"[ShardMonitor] Could not forward Shutdown signal to ShardManager for shard {}: {:#?}",
shard_id,
why
);
}
},
ShardManagerMessage::ShardInvalidAuthentication => {
self.manager.lock().await.shutdown_all().await;
return Err(ShardManagerError::InvalidToken);
},
ShardManagerMessage::ShardInvalidGatewayIntents => {
self.manager.lock().await.shutdown_all().await;
return Err(ShardManagerError::InvalidGatewayIntents);
},
ShardManagerMessage::ShardDisallowedGatewayIntents => {
self.manager.lock().await.shutdown_all().await;
return Err(ShardManagerError::DisallowedGatewayIntents);
},
}
}
Ok(())
}
}