CustomStartupRunner/executionmanager.cpp

#include "executionmanager.h"
#include <iostream>
#include <format>
#include <syncstream>
#include <filesystem>

extern "C" {
#include <fcntl.h>
#include <unistd.h>
#include <sys/wait.h>
#include <pwd.h>
#include <grp.h>
#include <stdlib.h>
}

ExecutionManager::ExecutionManager(int threads) :
    log_file{nullptr}
{
    this->threads = threads;
    if(this->threads <= 1) {
        this->threads = 2; // Does not work with threads = 1;
    }
}

void ExecutionManager::setLogFile(const Log &log_file)
{
    this->log_file.copyFromOther(log_file);
}

void ExecutionManager::setLogFileMaxSize(long long size)
{
    log_file_max_size = size;
}

void ExecutionManager::setLogDirectoy(std::string log_dir)
{
    this->log_dir = log_dir;
    if(!std::filesystem::exists(log_dir)) {
        std::filesystem::create_directories(log_dir);
        log_file << "Created log dir: " << log_dir << std::endl;
    }
}

void ExecutionManager::setJobs(std::vector<std::shared_ptr<AppInfo> > &job_list)
{
    this->job_list = std::move(job_list);
}


//This function assumes that job_list_mutex is locked
void ExecutionManager::printActiveJobs()
{
    for(auto &i: job_list){
        log_file << "Main Thread: " + i->getModuleName() + " mode: " + i->getStringRunMode() + " is running.";
    }
}

void ExecutionManager::run()
{
    finished = false;
    seperateScheduledTasks();
    log_file << "Main Thread: Job list: " << job_list.size() <<
        " Scheduled list: " << scheduled_list.size() << std::endl;
    last_poll_time = std::chrono::system_clock::now();
    while(true) {
        // Clean up completed worker threads
        for(auto it = thread_list.begin(); it != thread_list.end(); ) {
            if((*it)->done.load()) {
                if((*it)->thread.joinable())
                    (*it)->thread.join();
                it = thread_list.erase(it);
            } else {
                ++it;
            }
        }

        pollScheduledTasks();

        // Dispatch jobs from job_list to worker threads
        while(active_thread_count.load() < threads) {
            std::shared_ptr<AppInfo> ptr;
            int val = 0;
            size_t remaining_size = 0;
            {
                std::lock_guard<std::mutex> lock_get(this->job_list_mutex);
                if(job_list.empty()) {
                    break;
                }
                ptr = job_list.front();
                job_list.erase(job_list.begin());
                remaining_size = job_list.size();
                val = index_count++;
            }

            active_thread_count.fetch_add(1);
            auto worker = std::make_shared<WorkerThread>();
            worker->thread = std::thread([this, val, ptr, remaining_size, worker](){
                struct ThreadGuard {
                    std::atomic<int>& count;
                    std::atomic<bool>& done;
                    ~ThreadGuard() {
                        count.fetch_sub(1);
                        done.store(true);
                    }
                } guard{active_thread_count, worker->done};

                std::osyncstream os(log_file.rdbuf());
                std::string module_name = ptr->getModuleName();
                os << std::vformat("Thread[{}] running: ", std::make_format_args(val))
                          << module_name << " current list size: " << remaining_size << std::endl;
                int old = state_change.load();
                while(!ptr->isParentDependencyCleared(os))
                {
                    state_change.wait(old);
                    old = state_change.load();
                    os << std::vformat("Thread[{}] waiting for dependency: ", std::make_format_args(val))
                              << module_name << " current state count: " << old << std::endl;
                }

                // Create variables required for forking
                std::string dir = ptr->getWorkingDir();
                std::string program_name = ptr->getProgramName();
                std::string program_log_name_temp = module_name + ".log";
                std::string program_log_name = "";
                std::string run_as = ptr->getRunAs();
                struct passwd *pw = nullptr;
                if(run_as.size() != 0)
                    pw = getpwnam(run_as.c_str());
                for(auto &ch: program_log_name_temp) {
                    if(ch != ' ') {
                        program_log_name += ch;
                    }
                    else {
                        program_log_name += '_';
                    }
                }
                std::string program_log_file = std::filesystem::path(this->log_dir) / program_log_name;
                bool create_log = ptr->getCreateLog();
                std::deque<std::string> pargs = ptr->getArgs();
                int total_args = pargs.size() + 2;
                std::vector<char*> args(total_args);
                args[0] = program_name.data();
                for(int i=1,index=0; index < pargs.size(); i++,index++) {
                    args[i] = pargs[index].data();
                }
                args[total_args - 1] = (char *)NULL;
                bool log_file_exists = std::filesystem::exists(program_log_file);
                long long size = 0;
                if(log_file_exists)
                    size = std::filesystem::file_size(program_log_file);
                // Do fork and run
                auto pid = fork();
                FILE *popen_file = nullptr;
                std::unordered_map<std::string,std::string> env = ptr->getEnv();
                switch(pid) {
                    case 0: {
                        if(!dir.empty()) {
                            int err = chdir(dir.c_str());
                        }
                        if(pw) {
                            initgroups(pw->pw_name, pw->pw_gid);
                            setgid(pw->pw_gid);
                            setuid(pw->pw_uid);
                        }
                        for(auto &[key,value]: env) {
                            setenv(key.data(), value.data(), 1);
                        }
                        if(create_log) {
                            int output_fd = -99;
                            if(!log_file_exists)
                                output_fd = open(program_log_file.c_str(), O_WRONLY | O_CREAT | O_TRUNC, S_IRUSR | S_IWUSR);
                            else {
                                if(size < log_file_max_size)
                                {
                                    output_fd = open(program_log_file.c_str(), O_WRONLY | O_APPEND, S_IRUSR | S_IWUSR);
                                }
                                else {
                                    output_fd = open(program_log_file.c_str(), O_WRONLY | O_TRUNC | O_CREAT, S_IRUSR | S_IWUSR);
                                }
                            }
                            if (output_fd < 0) {
                                _exit(3);
                            }
                            if (dup2(output_fd, STDOUT_FILENO) < 0) {
                                _exit(3);
                            }
                            if (dup2(output_fd, STDERR_FILENO) < 0) {
                                _exit(3);
                            }
                            close(output_fd);
                        }
                        execv(program_name.c_str(), args.data());
                        _exit(1);
                    }
                    case -1: {
                        os << "Fork failed for module: " << module_name << std::endl;
                        ptr->setFinished(); // Or a dedicated failed state
                        state_change.fetch_add(1);
                        state_change.notify_all();
                        break;
                    }
                    default: {
                        // After fork in parent do
                        ptr->setRunning();
                        // If it was a scheduled task recalculate its next execution time
                        // and send it back to the scheduled job list
                        if(ptr->getScheduled()) {
                            ptr->addToPreviousExecutionList(pid);
                            ptr->updateNextTimeToRun();
                            if(ptr->getScheduled()) {
                                std::lock_guard<std::mutex> lock_get(this->job_list_mutex);
                                scheduled_list.push_back(ptr);
                            }
                        }
                        state_change.fetch_add(1);
                        state_change.notify_all();
                        break;
                    }
                }

                // If script wait here
                AppInfo::RunMode mode = ptr->getRunMode();
                int status;
                if(mode == AppInfo::RunMode::script) {
                    os << "Run mode is script so will wait for child to finish: " << module_name << std::endl;
                    int wait = waitpid(pid, &status, WUNTRACED);
                    ptr->setFinished();
                    os << "Parent wait finished: " << wait << " child status: " << status << std::endl;
                }

                // Set finished
                state_change.fetch_add(1);
                state_change.notify_all();
                os << std::vformat("Thread[{}] finished: ", std::make_format_args(val))
                          << module_name << std::endl;
            });
            thread_list.push_back(worker);
        }

        // Check if all work is done (no active jobs, no pending jobs, no scheduled tasks)
        {
            std::lock_guard<std::mutex> lock_get(this->job_list_mutex);
            if(job_list.empty() && scheduled_list.empty() && active_thread_count.load() == 0) {
                log_file << "Main Thread: All jobs finished." << std::endl;
                break;
            }
        }
        std::this_thread::sleep_for(std::chrono::milliseconds(50));
    }

    for(auto &worker: thread_list) {
        if(worker->thread.joinable())
            worker->thread.join();
    }
    thread_list.clear();
    finished = true;
}

bool ExecutionManager::checkScheduledTaskIsReady(std::chrono::system_clock::time_point next_time)
{
    return (next_time <= std::chrono::system_clock::now());
}

void ExecutionManager::seperateScheduledTasks()
{
    std::lock_guard<std::mutex> lock_get(this->job_list_mutex);
    for(auto task = job_list.begin(); task != job_list.end(); ) {
        // If task is not running on start and task is scheduled.
        if(!((*task)->getRunOnStart()) && ((*task)->getScheduled())) {
            log_file << "Main Thread: Job[" + (*task)->getModuleName() +
                            "] is scheduled for execution." << std::endl;
            scheduled_list.push_back(*task);
            task = job_list.erase(task);
        } else {
            ++task;
        }
    }
}

void ExecutionManager::pollScheduledTasks()
{
    auto current_poll_time = std::chrono::system_clock::now();
    if(std::chrono::duration_cast<std::chrono::seconds>(current_poll_time - last_poll_time).count() < polling_interval)
        return;
    last_poll_time = current_poll_time;
    {
        std::lock_guard<std::mutex> lock_get(this->job_list_mutex);
        for(auto task = scheduled_list.begin(); task != scheduled_list.end(); ) {
            if(checkScheduledTaskIsReady((*task)->getNextTimeToRun())) {
                log_file << "Main Thread: Job[" + (*task)->getModuleName() +
                                "] scheduled time arrived, adding to execution queue." << std::endl;
                job_list.push_back(*task);
                task = scheduled_list.erase(task);
            } else {
                ++task;
            }
        }
    }
}