daggy/utils/daggyd/daggyd.cpp

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

#include <argparse.hpp>
#include <atomic>
#include <daggy/Server.hpp>
#include <fstream>
#include <iostream>

// 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 <sys/types.h>
#include <syslog.h>
#include <unistd.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();
}