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

void ExecutionManager::run()
{
    finished = false;
    initial_list_size = job_list.size();
    while(true) {
        for(int i=active_thread_count.load(); i < threads ;i++) {
            int val = 0;
            {
                std::lock_guard<std::mutex> lock_get(this->job_list_mutex);
                val = index_count++;
                if(job_list.empty()) {
                    log_file << "Main Thread: Job list is empty, waiting for finish." << std::endl;
                    break;
                }
                if(val > (initial_list_size - 1)) {
                    log_file << "Main Thread: All jobs queued." << std::endl;
                    break;
                }
            }
            active_thread_count.fetch_add(1);
            std::shared_ptr<std::thread> thread = std::make_shared<std::thread>([this, val](){
                std::shared_ptr<AppInfo> ptr;
                std::osyncstream os(log_file.rdbuf());
                size_t list_size = 0;
                {
                    std::lock_guard<std::mutex> lock_get(this->job_list_mutex);
                    if(job_list.size() == 0) {
                        os << std::vformat("Thread[{}] called while the job list was empty: ", std::make_format_args(val))
                                  << std::endl;
                        return;
                    }
                    ptr = this->job_list[0];
                    job_list.erase(job_list.begin());
                    list_size = job_list.size();
                }
                std::string module_name = ptr->getModuleName();
                os << std::vformat("Thread[{}] running: ", std::make_format_args(val))
                          << module_name << " current list size: " << list_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();
                        state_change.fetch_add(1);
                        state_change.notify_all();
                        break;
                    }
                }

                // If script daemon 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;
                active_thread_count.fetch_sub(1);
            });
            thread_list.push_back(thread);
        }
        {
            std::lock_guard<std::mutex> lock_get(this->job_list_mutex);
            if(job_list.empty()) {
                log_file << "Main Thread: Job list is empty, waiting for finish." << std::endl;
                break;
            }
            if((index_count) > (initial_list_size - 1)) {
                log_file << "Main Thread: All jobs queued." << std::endl;
                break;
            }
        }
        if (active_thread_count.load() >= threads) {
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    }
    for(auto i: thread_list) {
        if(i->joinable())
            i->join();
    }
    finished = true;
}