trifid/tfclient/src/daemon.rs

use std::sync::mpsc;
use std::sync::mpsc::Sender;
use std::thread;
use log::{error, info};

use crate::apiworker::{apiworker_main, APIWorkerMessage};
use crate::config::load_config;
use crate::main;
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`
    check_server_url(&server);

    info!("Loading config...");
    let config = match load_config(&name) {
        Ok(cfg) => cfg,
        Err(e) => {
            error!("Error loading configuration: {}", e);
            std::process::exit(1);
        }
    };

    info!("Creating transmitter");

    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...");



    let config_api = config.clone();
    let transmitter_api = transmitter.clone();
    let name_api = name.clone();
    let api_thread = thread::spawn(move || {
        apiworker_main(config_api, name_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);
    });

    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");

    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);
        }
    }
    info!("API thread exited");

    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");

    info!("All threads exited");
}

#[derive(Clone)]
pub struct ThreadMessageSender {
    pub socket_thread: Sender<SocketWorkerMessage>,
    pub api_thread: Sender<APIWorkerMessage>,
    pub nebula_thread: Sender<NebulaWorkerMessage>
}