#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, ¤tAffinity, &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, ¤tAffinity, 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, ¤tAffinity, &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上虚拟机上测试可以完美运行。供大家参考。