daggy/utils/daggyd/daggyd.cpp

#include <iostream>
#include <fstream>
#include <atomic>

#include <sys/stat.h>
#include <signal.h>

#include <argparse.hpp>

#include <daggy/Server.hpp>

// Add executors here
#ifdef DAGGY_ENABLE_SLURM

#include <daggy/executors/task/SlurmTaskExecutor.hpp>

#else
#include <daggy/executors/task/ForkingTaskExecutor.hpp>
#endif

// Add loggers here
#include <daggy/loggers/dag_run/OStreamLogger.hpp>

/*
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/types.h>
#include <syslog.h>
*/

static std::atomic<bool> running{true};

void signalHandler(int signal) {
    switch (signal) {
        case SIGHUP:
            break;
        case SIGINT:
        case SIGTERM:
            running = false;
            break;
    }
}

void daemonize() {
    pid_t pid;

    struct sigaction newSigAction;
    sigset_t newSigSet;

    /* Check if parent process id is set */
    if (getppid() == 1) { return; }

    /* Set signal mask - signals we want to block */
    sigemptyset(&newSigSet);
    sigaddset(&newSigSet, SIGCHLD);  /* ignore child - i.e. we don't need to wait for it */
    sigaddset(&newSigSet, SIGTSTP);  /* ignore Tty stop signals */
    sigaddset(&newSigSet, SIGTTOU);  /* ignore Tty background writes */
    sigaddset(&newSigSet, SIGTTIN);  /* ignore Tty background reads */
    sigprocmask(SIG_BLOCK, &newSigSet, NULL);   /* Block the above specified signals */

    /* Set up a signal handler */
    newSigAction.sa_handler = signalHandler;
    sigemptyset(&newSigAction.sa_mask);
    newSigAction.sa_flags = 0;

    /* Signals to handle */
    sigaction(SIGHUP, &newSigAction, NULL);     /* catch hangup signal */
    sigaction(SIGTERM, &newSigAction, NULL);    /* catch term signal */
    sigaction(SIGINT, &newSigAction, NULL);     /* catch interrupt signal */

    // Fork once
    pid = fork();
    if (pid < 0) { exit(EXIT_FAILURE); }
    if (pid > 0) { exit(EXIT_SUCCESS); }

    /* On success: The child process becomes session leader */
    if (setsid() < 0) {
        std::cerr << "Unable to setsid" << std::endl;
        exit(EXIT_FAILURE);
    }

    /* Catch, ignore and handle signals */
    signal(SIGCHLD, SIG_IGN);
    signal(SIGHUP, SIG_IGN);

    /* Fork off for the second time*/
    pid = fork();
    if (pid < 0)
        exit(EXIT_FAILURE);
    if (pid > 0)
        exit(EXIT_SUCCESS);

    umask(0);

    /* Change the working directory to the root directory */
    /* or another appropriated directory */
    auto rc = chdir("/");
    (void)rc;

    /* Close all open file descriptors */
    for (auto x = sysconf(_SC_OPEN_MAX); x >= 0; x--) { close(x); }
}

int main(int argc, char **argv) {
    argparse::ArgumentParser args("Daggy");

    args.add_argument("-v", "--verbose")
            .default_value(false)
            .implicit_value(true);
    args.add_argument("-d", "--daemon")
            .default_value(false)
            .implicit_value(true);
    args.add_argument("--ip")
            .help("IP address to listen to")
            .default_value(std::string{"127.0.0.1"});
    args.add_argument("--log-file")
            .help("File to log to.")
            .default_value(std::string{"daggyd.log"});
    args.add_argument("--port")
            .help("Port to listen to")
            .default_value(2503)
            .action([](const std::string &value) { return std::stoi(value); });
    args.add_argument("--dag-threads")
            .help("Number of DAGs to run concurrently")
            .default_value(10UL)
            .action([](const std::string &value) { return std::stoull(value); });
    args.add_argument("--web-threads")
            .help("Number of web requests to support concurrently")
            .default_value(30UL)
            .action([](const std::string &value) { return std::stoull(value); });
    args.add_argument("--executor-threads")
            .help("Number of tasks to run concurrently")
            .default_value(30UL)
            .action([](const std::string &value) { return std::stoull(value); });

    try {
        args.parse_args(argc, argv);
    } catch (std::exception &e) {
        std::cout << "Error: " << e.what() << std::endl;
        std::cout << args;
        exit(1);
    }

    bool verbose = args.get<bool>("--verbose");
    bool asDaemon = args.get<bool>("--daemon");
    std::string logFileName = args.get<std::string>("--log-file");
    std::string listenIP = args.get<std::string>("--ip");
    uint16_t listenPort = args.get<int>("--port");
    size_t executorThreads = args.get<size_t>("--executor-threads");
    size_t webThreads = args.get<size_t>("--web-threads");
    size_t dagThreads = args.get<size_t>("--dag-threads");

    if (logFileName == "-") {
        if (asDaemon) {
            std::cout << "Unable to daemonize if logging to stdout" << std::endl;
            exit(1);
        }
    } else {
        fs::path logFn{logFileName};
        if (!logFn.is_absolute()) {
            logFileName = (fs::current_path() / logFileName).string();
        }
    }

    if (verbose) {
        std::cout << "Server running at http://" << listenIP << ':' << listenPort << std::endl
                  << "Max DAG Processing: " << dagThreads << std::endl
                  << "Max Task Execution: " << executorThreads << std::endl
                  << "Max Web Clients: " << webThreads << std::endl
                  << "Logging to: " << logFileName << std::endl
                  << std::endl << "Ctrl-C to exit" << std::endl;
    }

    if (asDaemon) {
        daemonize();
    }

    std::ofstream logFH;
    std::unique_ptr<daggy::loggers::dag_run::DAGRunLogger> logger;
    if (logFileName == "-") {
        logger = std::make_unique<daggy::loggers::dag_run::OStreamLogger>(std::cout);
    } else {
        logFH.open(logFileName, std::ios::app);
        logger = std::make_unique<daggy::loggers::dag_run::OStreamLogger>(logFH);
    }

#ifdef DAGGY_ENABLE_SLURM
    daggy::executors::task::SlurmTaskExecutor executor;
#else
    daggy::executors::task::ForkingTaskExecutor executor(executorThreads);
#endif
    Pistache::Address listenSpec(listenIP, listenPort);

    daggy::Server server(listenSpec, *logger, executor, dagThreads);
    server.init(webThreads);
    server.start();


    running = true;
    while (running) {
        std::this_thread::sleep_for(std::chrono::seconds(30));
    }
    server.shutdown();
}