trifid/tfclient/src/daemon.rs

130 lines
3.8 KiB
Rust
Raw Normal View History

2023-03-22 18:34:06 +00:00
use std::sync::mpsc;
use std::sync::mpsc::Sender;
use std::thread;
2023-03-22 18:36:36 +00:00
use log::{error, info};
2023-03-22 18:34:06 +00:00
use crate::apiworker::{apiworker_main, APIWorkerMessage};
use crate::config::load_config;
use crate::main;
2023-03-22 18:34:06 +00:00
use crate::nebulaworker::{nebulaworker_main, NebulaWorkerMessage};
use crate::socketworker::{socketworker_main, SocketWorkerMessage};
use crate::util::check_server_url;
pub fn daemon_main(name: String, server: String) {
// Validate the `server`
2023-03-22 18:34:06 +00:00
check_server_url(&server);
info!("Loading config...");
let config = match load_config(&name) {
Ok(cfg) => cfg,
Err(e) => {
2023-03-22 18:34:06 +00:00
error!("Error loading configuration: {}", e);
std::process::exit(1);
}
};
2023-03-22 18:34:06 +00:00
info!("Creating transmitter");
2023-03-22 18:34:06 +00:00
let (tx_api, rx_api) = mpsc::channel::<APIWorkerMessage>();
let (tx_socket, rx_socket) = mpsc::channel::<SocketWorkerMessage>();
let (tx_nebula, rx_nebula) = mpsc::channel::<NebulaWorkerMessage>();
let transmitter = ThreadMessageSender {
socket_thread: tx_socket,
api_thread: tx_api,
nebula_thread: tx_nebula
};
let mainthread_transmitter = transmitter.clone();
info!("Setting signal trap...");
match ctrlc::set_handler(move || {
info!("Ctrl-C detected. Stopping threads...");
match mainthread_transmitter.nebula_thread.send(NebulaWorkerMessage::Shutdown) {
Ok(_) => (),
Err(e) => {
error!("Error sending shutdown message to nebula worker thread: {}", e);
}
}
match mainthread_transmitter.api_thread.send(APIWorkerMessage::Shutdown) {
Ok(_) => (),
Err(e) => {
error!("Error sending shutdown message to api worker thread: {}", e);
}
}
match mainthread_transmitter.socket_thread.send(SocketWorkerMessage::Shutdown) {
Ok(_) => (),
Err(e) => {
error!("Error sending shutdown message to socket worker thread: {}", e);
}
}
}) {
Ok(_) => (),
Err(e) => {
error!("Unable to set sigtrap: {}", e);
std::process::exit(1);
}
}
info!("Starting API thread...");
2023-03-22 18:34:06 +00:00
let config_api = config.clone();
let transmitter_api = transmitter.clone();
let api_thread = thread::spawn(move || {
apiworker_main(config_api, transmitter_api, rx_api);
});
info!("Starting Nebula thread...");
let config_nebula = config.clone();
let transmitter_nebula = transmitter.clone();
let nebula_thread = thread::spawn(move || {
nebulaworker_main(config_nebula, transmitter_nebula, rx_nebula);
});
info!("Starting socket worker thread...");
let socket_thread = thread::spawn(move || {
socketworker_main(config, name.clone(), transmitter, rx_socket);
2023-03-22 18:34:06 +00:00
});
info!("Waiting for socket thread to exit...");
match socket_thread.join() {
Ok(_) => (),
Err(_) => {
error!("Error waiting for socket thread to exit.");
std::process::exit(1);
}
}
info!("Socket thread exited");
2023-03-22 18:34:06 +00:00
info!("Waiting for API thread to exit...");
match api_thread.join() {
Ok(_) => (),
Err(_) => {
error!("Error waiting for api thread to exit.");
std::process::exit(1);
}
2023-03-22 18:34:06 +00:00
}
info!("API thread exited");
2023-03-22 18:34:06 +00:00
info!("Waiting for Nebula thread to exit...");
match nebula_thread.join() {
Ok(_) => (),
Err(_) => {
error!("Error waiting for nebula thread to exit.");
std::process::exit(1);
}
}
info!("Nebula thread exited");
2023-03-22 18:34:06 +00:00
info!("All threads exited");
}
#[derive(Clone)]
pub struct ThreadMessageSender {
pub socket_thread: Sender<SocketWorkerMessage>,
pub api_thread: Sender<APIWorkerMessage>,
pub nebula_thread: Sender<NebulaWorkerMessage>
}