tokio::spawn() instead of a thicc select!
This commit is contained in:
parent
5577475b6c
commit
ac5f14a21e
108
src/main.rs
108
src/main.rs
|
@ -2,19 +2,19 @@ use std::fmt::Write;
|
||||||
use std::fs::File;
|
use std::fs::File;
|
||||||
use std::io::Read;
|
use std::io::Read;
|
||||||
use std::net::SocketAddr;
|
use std::net::SocketAddr;
|
||||||
|
use std::sync::Arc;
|
||||||
use std::thread;
|
use std::thread;
|
||||||
use std::{collections::HashMap, env};
|
use std::{collections::HashMap, env};
|
||||||
|
|
||||||
use arrayvec::ArrayString;
|
use arrayvec::ArrayString;
|
||||||
use futures_util::stream::StreamExt;
|
use futures_util::stream::StreamExt;
|
||||||
use irc::client::prelude::Config;
|
use irc::client::prelude::Config;
|
||||||
use irc::client::Client;
|
use irc::client::{Client, ClientStream};
|
||||||
use irc::proto::{ChannelExt, Command, Prefix};
|
use irc::proto::{ChannelExt, Command, Prefix};
|
||||||
use rspotify::Credentials;
|
use rspotify::Credentials;
|
||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
use tokio::select;
|
use tokio::sync::broadcast;
|
||||||
use tokio::sync::mpsc::{channel, Receiver, Sender};
|
use tokio::sync::mpsc::{unbounded_channel};
|
||||||
use tracing_log::LogTracer;
|
|
||||||
use tracing_subscriber::EnvFilter;
|
use tracing_subscriber::EnvFilter;
|
||||||
|
|
||||||
use crate::bots::{leek, misc, sed, title};
|
use crate::bots::{leek, misc, sed, title};
|
||||||
|
@ -55,12 +55,12 @@ async fn terminate_signal() {
|
||||||
|
|
||||||
pub struct AppState {
|
pub struct AppState {
|
||||||
prefix: String,
|
prefix: String,
|
||||||
client: Client,
|
client: Arc<Client>,
|
||||||
|
stream: ClientStream,
|
||||||
last_msgs: HashMap<String, String>,
|
last_msgs: HashMap<String, String>,
|
||||||
last_eval: HashMap<String, f64>,
|
last_eval: HashMap<String, f64>,
|
||||||
titlebot: title::Titlebot,
|
titlebot: title::Titlebot,
|
||||||
db: ExecutorConnection,
|
db: ExecutorConnection,
|
||||||
git_channel: String,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Deserialize)]
|
#[derive(Deserialize)]
|
||||||
|
@ -80,9 +80,8 @@ struct ClientConf {
|
||||||
git_channel: String,
|
git_channel: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::main(flavor = "current_thread")]
|
#[tokio::main]
|
||||||
async fn main() -> anyhow::Result<()> {
|
async fn main() -> anyhow::Result<()> {
|
||||||
LogTracer::init()?;
|
|
||||||
tracing_subscriber::fmt()
|
tracing_subscriber::fmt()
|
||||||
.with_env_filter(EnvFilter::from_env("UBERBOT_LOG"))
|
.with_env_filter(EnvFilter::from_env("UBERBOT_LOG"))
|
||||||
.init();
|
.init();
|
||||||
|
@ -112,7 +111,11 @@ async fn main() -> anyhow::Result<()> {
|
||||||
|
|
||||||
let uber_ver = concat!("Überbot ", env!("CARGO_PKG_VERSION"));
|
let uber_ver = concat!("Überbot ", env!("CARGO_PKG_VERSION"));
|
||||||
let irc_config = Config {
|
let irc_config = Config {
|
||||||
nickname: client_config.nickname,
|
nickname: Some(
|
||||||
|
client_config
|
||||||
|
.nickname
|
||||||
|
.unwrap_or_else(|| client_config.username.clone()),
|
||||||
|
),
|
||||||
username: Some(client_config.username.clone()),
|
username: Some(client_config.username.clone()),
|
||||||
realname: Some(client_config.username),
|
realname: Some(client_config.username),
|
||||||
server: Some(client_config.host),
|
server: Some(client_config.host),
|
||||||
|
@ -124,66 +127,76 @@ async fn main() -> anyhow::Result<()> {
|
||||||
version: Some(uber_ver.into()),
|
version: Some(uber_ver.into()),
|
||||||
..Config::default()
|
..Config::default()
|
||||||
};
|
};
|
||||||
let client = Client::from_config(irc_config).await?;
|
let mut client = Client::from_config(irc_config).await?;
|
||||||
|
let stream = client.stream()?;
|
||||||
client.identify()?;
|
client.identify()?;
|
||||||
|
let client = Arc::new(client);
|
||||||
|
|
||||||
|
let (ctx, _) = broadcast::channel(1);
|
||||||
|
let (etx, mut erx) = unbounded_channel();
|
||||||
|
|
||||||
|
let web_task = tokio::spawn(web_service::run(
|
||||||
|
db_conn.clone(),
|
||||||
|
client.clone(),
|
||||||
|
client_config.git_channel,
|
||||||
|
http_listen,
|
||||||
|
ctx.subscribe()
|
||||||
|
));
|
||||||
|
|
||||||
let state = AppState {
|
let state = AppState {
|
||||||
prefix: client_config.prefix,
|
prefix: client_config.prefix,
|
||||||
client,
|
client: client.clone(),
|
||||||
|
stream,
|
||||||
last_msgs: HashMap::new(),
|
last_msgs: HashMap::new(),
|
||||||
last_eval: HashMap::new(),
|
last_eval: HashMap::new(),
|
||||||
titlebot: title::Titlebot::create(spotify_creds).await?,
|
titlebot: title::Titlebot::create(spotify_creds).await?,
|
||||||
db: db_conn,
|
db: db_conn,
|
||||||
git_channel: client_config.git_channel,
|
|
||||||
};
|
};
|
||||||
|
let message_loop_task = tokio::spawn(async move {
|
||||||
|
if let Err(e) = message_loop(state).await {
|
||||||
|
let _ = etx.send(e);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
let (git_tx, git_recv) = channel(512);
|
tokio::select! {
|
||||||
|
_ = terminate_signal() => {
|
||||||
if let Err(e) = executor(state, git_tx, git_recv, http_listen).await {
|
tracing::info!("Received shutdown signal, sending QUIT message");
|
||||||
tracing::error!("Error in message loop: {}", e);
|
client.send_quit("überbot shutting down")?;
|
||||||
}
|
}
|
||||||
|
e = erx.recv() => {
|
||||||
if let Err(e) = exec_thread.join() {
|
if let Some(e) = e {
|
||||||
tracing::error!("Error while shutting down the database: {:?}", e);
|
tracing::error!("An error has occurred, shutting down: {}", e);
|
||||||
}
|
} else {
|
||||||
tracing::info!("Shutting down");
|
tracing::error!("Error channel has been dropped due to an unknown error, shutting down");
|
||||||
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn executor(
|
|
||||||
mut state: AppState,
|
|
||||||
git_tx: Sender<String>,
|
|
||||||
mut git_recv: Receiver<String>,
|
|
||||||
http_listen: SocketAddr,
|
|
||||||
) -> anyhow::Result<()> {
|
|
||||||
let web_db = state.db.clone();
|
|
||||||
select! {
|
|
||||||
r = web_service::run(web_db, git_tx, http_listen) => r?,
|
|
||||||
r = message_loop(&mut state) => r?,
|
|
||||||
r = git_recv.recv() => {
|
|
||||||
if let Some(message) = r {
|
|
||||||
state.client.send_privmsg(&state.git_channel, &message)?;
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
_ = terminate_signal() => {
|
|
||||||
tracing::info!("Sending QUIT message");
|
|
||||||
state.client.send_quit("überbot shutting down")?;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
tracing::info!("Closing services...");
|
||||||
|
let _ = ctx.send(());
|
||||||
|
web_task
|
||||||
|
.await
|
||||||
|
.unwrap_or_else(|e| tracing::warn!("Couldn't join the web service: {:?}", e));
|
||||||
|
message_loop_task
|
||||||
|
.await
|
||||||
|
.unwrap_or_else(|e| tracing::warn!("Couldn't join the web service: {:?}", e));
|
||||||
|
exec_thread
|
||||||
|
.join()
|
||||||
|
.unwrap_or_else(|e| tracing::warn!("Couldn't join the database: {:?}", e));
|
||||||
|
tracing::info!("Shutdown complete!");
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn message_loop(state: &mut AppState) -> anyhow::Result<()> {
|
async fn message_loop(mut state: AppState) -> anyhow::Result<()> {
|
||||||
let mut stream = state.client.stream()?;
|
while let Some(message) = state.stream.next().await.transpose()? {
|
||||||
while let Some(message) = stream.next().await.transpose()? {
|
|
||||||
if let Command::PRIVMSG(ref origin, content) = message.command {
|
if let Command::PRIVMSG(ref origin, content) = message.command {
|
||||||
if origin.is_channel_name() {
|
if origin.is_channel_name() {
|
||||||
if let Some(author) = message.prefix.as_ref().and_then(|p| match p {
|
if let Some(author) = message.prefix.as_ref().and_then(|p| match p {
|
||||||
Prefix::Nickname(name, _, _) => Some(&name[..]),
|
Prefix::Nickname(name, _, _) => Some(&name[..]),
|
||||||
_ => None,
|
_ => None,
|
||||||
}) {
|
}) {
|
||||||
if let Err(e) = handle_privmsg(state, author, origin, content).await {
|
if let Err(e) = handle_privmsg(&mut state, author, origin, content).await {
|
||||||
state
|
state
|
||||||
.client
|
.client
|
||||||
.send_privmsg(origin, &format!("Error: {}", e))?;
|
.send_privmsg(origin, &format!("Error: {}", e))?;
|
||||||
|
@ -194,6 +207,7 @@ async fn message_loop(state: &mut AppState) -> anyhow::Result<()> {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
tracing::info!("Message loop finished");
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
Loading…
Reference in a new issue