pin a task to a cpu

setaffinity.cpp

#include <iostream>
#include <string>
#include <vector>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <sstream>
#include <sys/resource.h>
#include <sched.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <signal.h>

namespace fs = std::filesystem;

pid_t child_pid;

// Store original IRQ affinities for restoration
struct IrqAffinity {
    std::string path;
    std::string original_value;
};
std::vector<IrqAffinity> original_irq_affinities;

std::vector<int> parse_cpu_list(const std::string& cpu_list) {
    std::vector<int> cpus;
    std::stringstream ss(cpu_list);
    std::string cpu;

    while (std::getline(ss, cpu, ',')) {
        try {
            int cpu_num = std::stoi(cpu);
            cpus.push_back(cpu_num);
        } catch (const std::exception& e) {
            throw std::runtime_error("Invalid CPU number: " + cpu);
        }
    }

    return cpus;
}

int get_cpu_count() {
    std::ifstream cpuinfo("/proc/cpuinfo");
    std::string line;
    int count = 0;

    while (std::getline(cpuinfo, line)) {
        if (line.find("processor") == 0) {
            count++;
        }
    }
    return count;
}

void save_and_set_irq_affinity(unsigned long irq_mask) {
    for (const auto& entry : fs::directory_iterator("/proc/irq")) {
        if (fs::is_directory(entry)) {
            std::string affinity_path = entry.path().string() + "/smp_affinity";
            if (fs::exists(affinity_path)) {
                try {
                    std::ifstream current_affinity(affinity_path);
                    if (current_affinity.is_open()) {
                        std::string original_value;
                        std::getline(current_affinity, original_value);
                        
                        uint64_t ov;
                        std::stringstream ss, ss2;
                        ss << std::hex << original_value;
                        ss >> ov;
                        ov &= irq_mask;

                        ss2 << std::hex << ov;
                        std::string hex_mask = ss2.str();
                        original_irq_affinities.push_back({affinity_path, original_value});

                        // Set new affinity
                        std::ofstream affinity_file(affinity_path);
                        if (affinity_file.is_open()) {
                            affinity_file << hex_mask;
                        }
                    }

                } catch (const std::exception& e) {
                    std::cerr << "Warning: Could not set IRQ affinity for '" << affinity_path <<"': " << e.what() << std::endl;
                }
            }
        }
    }
}

void restore_irq_affinities() {
    for (const auto& irq : original_irq_affinities) {
        try {
            std::ofstream affinity_file(irq.path);
            if (affinity_file.is_open()) {
                affinity_file << irq.original_value;
            }
        } catch (const std::exception& e) {
            std::cerr << "Warning: Could not restore IRQ affinity for " << irq.path 
                << ": " << e.what() << std::endl;
        }
    }
}

void signal_handler(int signum) {
    if(child_pid > 0) {
        kill(child_pid, SIGINT);
    }
    restore_irq_affinities();
}

int main(int argc, char* argv[]) {
    if (argc < 3) {
        std::cerr << "Usage: " << argv[0] << " <cpu_list> <command> [args...]\n";
        std::cerr << "Example: " << argv[0] << " 0,1,2 mycommand arg1 arg2\n";
        return 1;
    }

    // Check if running as root
    if (geteuid() != 0) {
        std::cerr << "This program must be run as root\n";
        return 1;
    }

    // Parse CPU list
    std::vector<int> cpus;
    try {
        cpus = parse_cpu_list(argv[1]);
    } catch (const std::exception& e) {
        std::cerr << "Error parsing CPU list: " << e.what() << "\n";
        return 1;
    }

    if (cpus.empty()) {
        std::cerr << "No CPUs specified\n";
        return 1;
    }

    int cpu_count = get_cpu_count();
    unsigned long process_mask = 0;

    for (int cpu : cpus) {
        if (cpu >= cpu_count) {
            std::cerr << "Invalid CPU number " << cpu << ". System has " << cpu_count << " CPUs\n";
            return 1;
        }
        process_mask |= (1UL << cpu);
    }

    // Calculate masks
    unsigned long all_cpus = (1UL << cpu_count) - 1;
    unsigned long irq_mask = all_cpus ^ process_mask;

    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);
    signal(SIGHUP, signal_handler);

    save_and_set_irq_affinity(irq_mask);

    child_pid = fork();

    if (child_pid == -1) {
        std::cerr << "Failed to fork process\n";
        restore_irq_affinities();
        return 1;
    }

    if (child_pid == 0) {
        cpu_set_t cpu_set;
        CPU_ZERO(&cpu_set);
        for (uint64_t cpu : cpus) {
            CPU_SET(cpu, &cpu_set);
            std::cerr << "Using CPU: " << cpu << std::endl;
        }

        /* struct sched_param param;
        param.sched_priority = 99;  // Maximum RT priority
    
        if (sched_setscheduler(0, SCHED_RR, &param) != 0) {
            perror("Warning: Could not set SCHED_FIFO scheduling policy: ");
        } */
        if (setpriority(PRIO_PROCESS, 0, -19) != 0) {
            std::cerr << "Warning: Could not set priority\n";
        }

        if (sched_setaffinity(0, sizeof(cpu_set_t), &cpu_set) != 0) {
            std::cerr << "Failed to set CPU affinity\n";
            exit(1);
        }

        std::vector<char*> args;
        for (int i = 2; i < argc; i++) {
            args.push_back(argv[i]);
        }
        args.push_back(nullptr);

        execvp(args[0], args.data());

        std::cerr << "Failed to execute command: " << args[0] << std::endl;
        restore_irq_affinities();
        exit(1);
    } else {
        int status;
        waitpid(child_pid, &status, 0);

        restore_irq_affinities();

        if (WIFEXITED(status)) {
            return WEXITSTATUS(status);
        } else if (WIFSIGNALED(status)) {
            return 128 + WTERMSIG(status);
        }
        return 1;
    }
}