设置和获取系统中进程的CPU亲和性代码

原创 于 2025-03-29 15:46:36 发布 · 229 阅读 · 0 ·
CC 4.0 BY-SA版权
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
文章标签：

#windows 

#include <windows.h>
#include <tlhelp32.h>
#include <iostream>
#include <fstream>
#include <string>
#include <iomanip>
#include <sstream>

// 获取当前时间（精确到秒）
std::string GetFormattedCurrentTime() {
    SYSTEMTIME st;
    GetLocalTime(&st);

    std::ostringstream oss;
    oss << std::setfill('0')
        << std::setw(4) << st.wYear << "-"
        << std::setw(2) << st.wMonth << "-"
        << std::setw(2) << st.wDay << " "
        << std::setw(2) << st.wHour << ":"
        << std::setw(2) << st.wMinute << ":"
        << std::setw(2) << st.wSecond;

    return oss.str();
}

// 将GROUP_AFFINITY转换为CPU核心名称表示
std::string GroupAffinityToReadable(const GROUP_AFFINITY& ga, int maxCores = 64) {
    std::stringstream ss;
    bool first = true;

    for (int i = 0; i < maxCores; ++i) {
        if (ga.Mask & (1ULL << i)) {
            if (!first) {
                ss << ", ";
            }
            ss << "CPU-" << i;
            first = false;
        }
    }

    if (ss.str().empty()) {
        ss << "None";
    }

    return ss.str();
}

// 将DWORD_PTR转换为CPU核心名称表示
std::string AffinityMaskToReadable(DWORD_PTR affinityMask, int maxCores = 64) {
    std::stringstream ss;
    bool first = true;

    for (int i = 0; i < maxCores; ++i) {
        if (affinityMask & (1ULL << i)) {
            if (!first) {
                ss << ", ";
            }
            ss << "CPU-" << i;
            first = false;
        }
    }

    if (ss.str().empty()) {
        ss << "None";
    }

    return ss.str();
}

// 将WCHAR数组转换为std::string
std::string WideCharToMultiByteString(const WCHAR* wideStr) {
    // Calculate the length of the string excluding the null terminator
    int len = WideCharToMultiByte(CP_ACP, 0, wideStr, -1, NULL, 0, NULL, NULL);
    if (len == 0) {
        return "";
    }

    // Allocate buffer and convert
    std::string result(len - 1, '\0');  // Exclude the null terminator
    WideCharToMultiByte(CP_ACP, 0, wideStr, -1, &result[0], len, NULL, NULL);
    return result;
}

// 设置进程的CPU亲和性并记录日志
bool SetProcessAffinityAndLog(DWORD processId, DWORD_PTR targetAffinity, const std::string& processName) {
    HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_SET_INFORMATION, FALSE, processId);
    if (hProcess == NULL) {
        std::cerr << "OpenProcess failed for PID: " << processId << " (" << GetLastError() << ")" << std::endl;

        // 记录日志
        std::ofstream logFile("process_affinity_set_log.txt", std::ios_base::app); // 追加模式
        if (logFile.is_open()) {
            logFile << "--------------------------------------------------\n";
            logFile << GetFormattedCurrentTime() << " - PID: " << processId << ", Name: "
                << processName
                << ", Failed to get process handle (" << GetLastError() << ")\n";
            logFile.close();
        }
        return false;
    }

    DWORD_PTR currentAffinity;
    DWORD_PTR systemAffinity;

    if (!GetProcessAffinityMask(hProcess, &currentAffinity, &systemAffinity)) {
        std::cerr << "GetProcessAffinityMask failed for PID: " << processId << " (" << GetLastError() << ")" << std::endl;

        // 记录日志
        std::ofstream logFile("process_affinity_set_log.txt", std::ios_base::app); // 追加模式
        if (logFile.is_open()) {
            logFile << "--------------------------------------------------\n";
            logFile << GetFormattedCurrentTime() << " - PID: " << processId << ", Name: "
                << processName
                << ", Failed to get process affinity mask (" << GetLastError() << ")\n";
            logFile.close();
        }

        CloseHandle(hProcess);
        return false;
    }

    bool needsAdjustment = (currentAffinity & targetAffinity) != currentAffinity;
    bool success = false;
    DWORD lastError = 0;

    if (needsAdjustment) {
        success = SetProcessAffinityMask(hProcess, targetAffinity);
        lastError = GetLastError();
    }

    // 记录日志
    std::ofstream logFile("process_affinity_set_log.txt", std::ios_base::app); // 追加模式
    if (logFile.is_open()) {
        logFile << "--------------------------------------------------\n";
        logFile << GetFormattedCurrentTime() << " - PID: " << processId << ", Name: "
            << processName
            << ", Original Affinity: " << AffinityMaskToReadable(currentAffinity);

        if (needsAdjustment) {
            if (success) {
                logFile << ", New Affinity: " << AffinityMaskToReadable(targetAffinity);
            }
            else {
                logFile << ", New Affinity: Adjust Failed (" << lastError << ")";
            }
        }
        else {
            logFile << ", No Adjustment Needed";
        }

        logFile << "\n";
        logFile.close();
    }

    CloseHandle(hProcess);
    return success;
}

// 根据指定的核心数计算目标CPU亲和性掩码
DWORD_PTR CalculateTargetAffinity(int coreCount) {
    if (coreCount <= 0 || coreCount > 64) {
        std::cerr << "Invalid core count. Must be between 1 and 64." << std::endl;
        return 0;
    }
    return ((DWORD_PTR)1 << coreCount) - 1;
}

void LogProcessAffinityChanges(DWORD_PTR targetAffinity) {
    HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (snapshot == INVALID_HANDLE_VALUE) {
        std::cerr << "CreateToolhelp32Snapshot failed (" << GetLastError() << ")" << std::endl;
        return;
    }

    PROCESSENTRY32 pe;
    pe.dwSize = sizeof(PROCESSENTRY32);

    DWORD totalProcesses = 0;
    DWORD successfulModifications = 0;
    DWORD failedModifications = 0;
    DWORD noModificationsNeeded = 0;

    if (Process32First(snapshot, &pe)) {
        do {
            totalProcesses++;
            std::string processName = WideCharToMultiByteString(pe.szExeFile);

            // 设置进程的CPU亲和性并记录日志
            bool success = SetProcessAffinityAndLog(pe.th32ProcessID, targetAffinity, processName);
            if (success) {
                successfulModifications++;
            }
            else {
                HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, pe.th32ProcessID);
                if (hProcess != NULL) {
                    DWORD_PTR currentAffinity;
                    if (GetProcessAffinityMask(hProcess, &currentAffinity, NULL)) {
                        if ((currentAffinity & targetAffinity) == currentAffinity) {
                            noModificationsNeeded++;
                        }
                        else {
                            failedModifications++;
                        }
                    }
                    CloseHandle(hProcess);
                }
                else {
                    failedModifications++;
                }
            }
        } while (Process32Next(snapshot, &pe));
    }
    else {
        std::cerr << "Process32First failed (" << GetLastError() << ")" << std::endl;
    }

    CloseHandle(snapshot);

    // 记录统计结果
    std::ofstream logFile("process_affinity_set_log.txt", std::ios_base::app); // 追加模式
    if (logFile.is_open()) {
        logFile << "--------------------------------------------------\n";
        logFile << GetFormattedCurrentTime() << " - Total Processes: " << totalProcesses
            << ", Successfully Modified Affinity for: " << successfulModifications
            << " processes, Failed Modifications: " << failedModifications
            << " processes, No Modifications Needed: " << noModificationsNeeded
            << " processes\n";
        logFile.close();
    }
}

void LogAllProcessAffinities() {
    HANDLE processSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (processSnapshot == INVALID_HANDLE_VALUE) {
        std::cerr << "CreateToolhelp32Snapshot failed (" << GetLastError() << ")" << std::endl;
        return;
    }

    PROCESSENTRY32 pe;
    pe.dwSize = sizeof(PROCESSENTRY32);

    if (Process32First(processSnapshot, &pe)) {
        do {
            std::string processName = WideCharToMultiByteString(pe.szExeFile);
            HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, pe.th32ProcessID);
            if (hProcess != NULL) {
                DWORD_PTR currentAffinity;
                DWORD_PTR systemAffinity;

                if (GetProcessAffinityMask(hProcess, &currentAffinity, &systemAffinity)) {
                    std::ofstream logFile("process_affinity_get_log.txt", std::ios_base::app); // 追加模式
                    if (logFile.is_open()) {
                        logFile << "--------------------------------------------------\n";
                        logFile << GetFormattedCurrentTime() << " - PID: " << pe.th32ProcessID
                            << ", Name: " << processName
                            << ", Process Affinity: " << AffinityMaskToReadable(currentAffinity)
                            << ", System Affinity: " << AffinityMaskToReadable(systemAffinity)
                            << "\n";

                        // 获取并记录线程的CPU亲和性
                        HANDLE threadSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
                        THREADENTRY32 te;
                        te.dwSize = sizeof(THREADENTRY32);

                        if (Thread32First(threadSnapshot, &te)) {
                            do {
                                if (te.th32OwnerProcessID == pe.th32ProcessID) {
                                    HANDLE hThread = OpenThread(THREAD_QUERY_INFORMATION, FALSE, te.th32ThreadID);
                                    if (hThread != NULL) {
                                        GROUP_AFFINITY ga;
                                        if (GetThreadGroupAffinity(hThread, &ga)) {
                                            logFile << "  - TID: " << te.th32ThreadID
                                                << ", Thread Affinity: " << GroupAffinityToReadable(ga)
                                                << "\n";
                                        }
                                        CloseHandle(hThread);
                                    }
                                }
                            } while (Thread32Next(threadSnapshot, &te));
                        }

                        logFile.close();
                    }
                }
                else {
                    std::cerr << "GetProcessAffinityMask failed for PID: " << pe.th32ProcessID << " (" << GetLastError() << ")" << std::endl;
                }

                CloseHandle(hProcess);
            }
            else {
                std::cerr << "OpenProcess failed for PID: " << pe.th32ProcessID << " (" << GetLastError() << ")" << std::endl;
            }
        } while (Process32Next(processSnapshot, &pe));
    }
    else {
        std::cerr << "Process32First failed (" << GetLastError() << ")" << std::endl;
    }

    CloseHandle(processSnapshot);
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        std::cerr << "Usage: " << argv[0] << " <command> [parameter]" << std::endl;
        std::cerr << "Commands: get, set" << std::endl;
        return 1;
    }

    std::string command(argv[1]);
    if (command == "get") {
        LogAllProcessAffinities();
        std::cout << "Affinity settings have been logged to 'process_affinity_get_log.txt'." << std::endl;
    }
    else if (command == "set") {
        if (argc < 3) {
            std::cerr << "Please provide a parameter for setting CPU affinity." << std::endl;
            return 1;
        }

        int affinityParam = std::stoi(argv[2]);

        DWORD_PTR targetAffinity = CalculateTargetAffinity(affinityParam);
        if (targetAffinity == 0) {
            return 1;
        }

        LogProcessAffinityChanges(targetAffinity);
        std::cout << "Affinity settings have been adjusted and logged to 'process_affinity_set_log.txt'." << std::endl;
    }
    else {
        std::cerr << "Unknown command: " << command << std::endl;
        return 1;
    }

    return 0;
}

在windows11上虚拟机上测试可以完美运行。供大家参考。

评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

danxuezx

如果对你有用是我的快乐

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付：¥1
获取中
扫码支付

您的余额不足，请更换扫码支付或充值

打赏作者

实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值