open_source / daemonbase

/* ****************************************************************************
 * Copyright 2019 Open Systems Development BV                                 *
 *                                                                            *
 * Permission is hereby granted, free of charge, to any person obtaining a    *
 * copy of this software and associated documentation files (the "Software"), *
 * to deal in the Software without restriction, including without limitation  *
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,   *
 * and/or sell copies of the Software, and to permit persons to whom the      *
 * Software is furnished to do so, subject to the following conditions:       *
 *                                                                            *
 * The above copyright notice and this permission notice shall be included in *
 * all copies or substantial portions of the Software.                        *
 *                                                                            *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   *
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    *
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER *
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING    *
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER        *
 * DEALINGS IN THE SOFTWARE.                                                  *
 * ****************************************************************************/
#pragma once

#include <cstdlib>
#include <iostream>
#include <sys/types.h>
#include <unistd.h>
#include <cerrno>
#include <cstring>
#include <string>
#include <csignal>
#include <sys/stat.h>
#include <stdexcept>
#include <chrono>
#include <thread>
#include <atomic>
#include <condition_variable>

#include "daemonlog.h"
#include "daemonconfig.h"

namespace osdev::components::daemon
{
class DaemonBase
{
private:
    static DaemonBase *instance;

public:
    /**
     *  Construct a new daemon process.
     *  @note: only one daemon per application is possible.
     *  @param  name:   name of daemon process
     *  @param  cwd:    daemon current working directory, root "/" directory by default.
     *  @param  update_duration:    duration to sleep before waking up the on_update() callback every time, deafult 10 seconds.
     */
    DaemonBase(const std::string &name, 
               const std::string &cwd = "/", 
               const std::chrono::high_resolution_clock::duration &update_duration = std::chrono::seconds(10))
        : m_name(name)
        , m_cwd(cwd)
        , m_update_duration(update_duration)
        , m_is_running(false)
        , m_exit_code(EXIT_SUCCESS)
    {
        if (instance)
        {
            daemonlog::error("Only one daemon instance is possible.");
            std::exit(EXIT_FAILURE);
        }
        instance = this;
    }

    DaemonBase() 
        : m_name("<unknown_daemon>")
        , m_cwd("/")
        , m_update_duration(std::chrono::seconds(10))
        , m_is_running(false)
    {
        if(instance) 
        {
            daemonlog::error("Only one daemon instance is possible.");
            std::exit(EXIT_FAILURE);
        }
        instance = this;
    }

    void run(int argc, char *argv[])
    {
        if (m_is_running.load())
        {
            daemonlog::error("Daemon '" + m_name + "' is already running.");
            return;
        }

        // Get config file path from cmd args passed by ExecStart=/usr/bin/my_daemon --config /etc/my_daemon/my_daemon.conf
        // since we need it for a reload./
        for (std::int32_t i = 0; i < argc; i++)
        {
            if (!std::strcmp(argv[i], "--config"))
            {
                if (i + 1 < argc)
                {
                    m_config_file = argv[i + 1];
                }
                else
                {
                    daemonlog::error("Missing config file. Did ytou forget to specify a config file in your .serve file's ExecStart ?");
                }
                break;
            }
        }

        // daemonize this program by forking the parent process.
        daemonize();

        // Mark as running (better to have it before on_start() as a user may call stop() inside on_start()).
        m_is_running = true;
        on_start(daemonconfig::from_file(m_config_file));
        while (m_is_running.load())
        {
            on_update();

            // On long sleeps, if we want to exit we need a condition_variable to wake up the thread from sleep to carry on exiting.
            std::unique_lock<std::mutex> lock(m_mutex);
            m_update_cv.wait_for(lock, m_update_duration, [this]()
            {
                return !m_is_running.load();
            });
        }
        on_stop();
    }

    void stop(std::int32_t code = EXIT_SUCCESS)
    {
        m_exit_code = code;
        m_is_running.store(false);
        m_update_cv.notify_all();
    }

    virtual ~DaemonBase()
    {
        daemonlog::shutdown();
        // Terminate the child process when the daemon completes (loop stopped)
        // @note that calling std::exit() inside the run function will not call DTor,
        std::exit(m_exit_code); 
    }

    // Getters & Setters
    void set_update_duration(const std::chrono::high_resolution_clock::duration &duration) noexcept
    { m_update_duration = duration; }

    const std::chrono::high_resolution_clock::duration& get_update_duration() const noexcept
    { return m_update_duration; }

    void set_name(const std::string &daemon_name) noexcept 
    { m_name = daemon_name; }
    
    const std::string& get_name() const noexcept
    { return m_name; }

    void set_cwd(const std::string &current_working_dir) noexcept
    { 
        // Change the current working directory to a directory guaranteed to exist, provided by the user.
        if (chdir(current_working_dir.c_str()) < 0)
        {
            daemonlog::error("Could not change current working directory to'" 
                                      + current_working_dir + "': " 
                                      + std::string(std::strerror(errno)));
            return;
        }
        m_cwd = current_working_dir;
    }

    const std::string& get_cwd() const noexcept { return m_cwd; }

    pid_t get_pid() const noexcept { return m_pid; }
    pid_t get_sid() const noexcept { return m_sid; }

protected: // Callbacks
    /**
     *  @brief  Called once on daemon starts
     *  @scenarios:
     *      - when system starts
     *      - when you run `$ systemctl start your_daemon` manually
     *  @param  cfg: Installed daemon config file
     *  Initialize your code here...
     */
    virtual void on_start(const daemonconfig &cfg) = 0;

    /**
     *  @brief  Called every DURATION which was set by set_update_duration(DURATION).
     *  Update your code here...
     */
    virtual void on_update() = 0;

    /**
     *  @brief  Called once before daemon is about to exit.
     *  @scenarios:
     *      - when you call stop(ewxit_code)
     *      - when you run `$ systemctl stop your_daemon` manually
     *      - when the system kills your daemon for some reason
     *  Cleanup your code here...
     */
    virtual void on_stop() = 0;

    /**
     *  @brief  Called once when daemon's config or service files are updated.
     *  @scenarios:
     *      - when you run `$systemctl daemon-reload` after you have changed your .conf or .service files
     *        (after reinstalling your daemon with `$ sudo make install` for example)
     *  Reinitialize your code here...
     */
    virtual void on_reload(const daemonconfig &cfg) = 0;

private:
    static void signal_handler(std::int32_t sig)
    {
        daemonlog::info("Signal " + std::to_string(sig) + " received.");

        switch(sig)
        {
            // daemon.service handler : ExecStop=/bin/kill -s SIGTERM $MAINPID
            // When daemon is stopped, system sends SIGTERM first. 
            // If daemon didn't respond during 90 seconds, it will send a SIGKILL signal
            case SIGTERM:
            case SIGKILL:
            {
                instance->stop();
                break;
            }
            // daemon.service handler : ExecReload=/bin/kill -S SIGHUB $MAINPID
            // When a daemon is reloaded due updates in .service or .conf, system sends SIGHUP signal.
            case SIGHUP:
            {
                instance->on_reload(daemonconfig::from_file(instance->m_config_file));
                break;
            }
            default:
            {
                break;
            }
        }
    }

    /**
     *  @brief  Daemonize this program
     *  @note: It is also possible to use glibc function daemon()
     *  at this point, but it is useful to customize your daemon.
     *  Like for example handle signals, set working directory...
     */
    void daemonize()
    {
        // Fork off the parent process (https://linux.die.net/man/3/fork)
        m_pid = fork();
        // Success: The parent process continues with a PID > 0
        if (m_pid > 0)
        {
            std::exit(EXIT_SUCCESS);
        }
        else if (m_pid < 0)
        {
            // An error occurred. A process ID lower than 0 indicates a failure in either process
            std::exit(EXIT_FAILURE);
        }
        // The parent process has now terminated, and the forked child process will continue
        // (the pid of the child process was 0)

        // Since the child process is a daemon, the unask needs to be set so files and logs can be written
        umask(0);

        // Initialize syslog for this daemon, here it's a good place to do so.
        daemonlog::init(m_name);

        // On success: The child process becomes session leader. Generate a session ID for the child process.
        m_sid = setsid();
        if (m_sid < 0)
        {
            daemonlog::error("Could not set SID to child process: " + std::string(std::strerror(errno)));
            std::exit(EXIT_FAILURE);
        }

        // Ignore the Child terminated or stopped signal.
        std::signal(SIGCHLD, SIG_IGN);

        // Set signal handlers to detect daemon interrupt, restart..
        std::signal(SIGHUP, signal_handler);

        // When a sudo systemctl stop my_daemon is ran, by default, a SIGTERM is sent, 
        // followed by 90 seconds of waiting followed by a SIGKILL
        std::signal(SIGTERM, signal_handler);
        std::signal(SIGKILL, signal_handler);

        // Change the current working directory to a directory guaranteed to exist, procided by the user
        if (chdir(m_cwd.c_str()) < 0)
        {
            daemonlog::error("Could not change current working directory to `" + m_cwd + "': " + std::string(std::strerror(errno)));
            std::exit(EXIT_FAILURE);
        }

        // A daemon cannot use the terminal, so close standard file descriptors for security reasons
        close(STDIN_FILENO);
        close(STDOUT_FILENO);
        close(STDERR_FILENO);
    }
    
    // Member variables
    pid_t   m_pid;
    pid_t   m_sid;
    std::string m_config_file;
    std::string m_name;
    std::string m_cwd;
    std::chrono::high_resolution_clock::duration m_update_duration;
    std::atomic<bool> m_is_running;
    std::condition_variable m_update_cv;
    std::mutex m_mutex;
    std::int32_t m_exit_code;

};

DaemonBase* DaemonBase::instance = nullptr;

}   /* End namespace osdev::components::daemon */