Linux编程之《进程/线程绑定CPU》

Intro
-----

通常我们在编写服务器代码时，可以通过将当前进程绑定到固定的CPU核心或者线程绑定到固定的CPU核心来提高系统调度程序的效率来提高程序执行的效率，下面将完整代码贴上。

/************************************************
 * 该例程讲解了进程、线程绑定到固定的cpu核心上运行
 * 来提高程序运行效率
************************************************/
#include <unistd.h>
#ifndef __USE_GNU
#define __USE_GNU   // 为了使用SET_SET()等宏定义，但unistd.h里面好像已经定义了
#endif
#include <sched.h>
#include <pthread.h>
#include <stdio.h>
#include <vector>

unsigned int systemCPUNum()
{
    // _SC_NPROCESSORS_CONF的值为CPU个数，基于0开始编号
    return sysconf(_SC_NPROCESSORS_CONF);
}

bool currentProcessAffinity(std::vector<unsigned int>& runningCPUVector)
{
    cpu_set_t cpuSet;

    // 清空一个CPU集合
    CPU_ZERO(&cpuSet);

    // 得到指定进程ID绑定到哪个CPU
    int ret = sched_getaffinity(0, // 0代表当前进程
                               sizeof(cpuSet),
                               &cpuSet);
    if (ret < 0)
    {
        return false;
    }

    unsigned int cpuNum = systemCPUNum();
    runningCPUVector.clear();
    for (unsigned int i = 0; i < cpuNum; ++i)
    {
        // 检查一个CPU号是否在一个集合中
        if (CPU_ISSET(i, &cpuSet))
        {
            runningCPUVector.push_back(i);
        }
    }

    return true;
}

bool setCurrentProcessAffinity(const std::vector<unsigned int>& needBindCPUVector)
{
    cpu_set_t cpuSet;

    // 清空一个CPU集合
    CPU_ZERO(&cpuSet);
    
    for (auto& iter : needBindCPUVector)
    {
        CPU_SET(iter, &cpuSet);
    }

    // 将指定进程ID绑定到CPU
    int ret = sched_setaffinity(0, // 0代表当前进程
                               sizeof(cpuSet),
                               &cpuSet);
    if (ret < 0)
    {
        return false;
    }

    return true;
}

bool currentThreadAffinity(std::vector<unsigned int>& runningCPUVector)
{
    cpu_set_t cpuSet;

    // 清空一个CPU集合
    CPU_ZERO(&cpuSet);

    // 得到指定线程ID绑定到哪个CPU
    int ret = pthread_getaffinity_np(pthread_self(),
                                     sizeof(cpuSet),
                                     &cpuSet);
    if (ret < 0)
    {
        return false;
    }

    unsigned int cpuNum = systemCPUNum();
    runningCPUVector.clear();
    for (unsigned int i = 0; i < cpuNum; ++i)
    {
        // 检查一个CPU号是否在一个集合中
        if (CPU_ISSET(i, &cpuSet))
        {
            runningCPUVector.push_back(i);
        }
    }

    return true;
}

bool setCurrentThreadAffinity(const std::vector<unsigned int>& needBindCPUVector)
{
    cpu_set_t cpuSet;

    // 清空一个CPU集合
    CPU_ZERO(&cpuSet);
    
    for (auto& iter : needBindCPUVector)
    {
        CPU_SET(iter, &cpuSet);
    }

    // 将指定线程ID绑定到CPU
    int ret = pthread_setaffinity_np(pthread_self(),
                                     sizeof(cpuSet),
                                     &cpuSet);
    if (ret < 0)
    {
        return false;
    }

    return true;
}

int main()
{
    printf("*****Process bind CPU sample*****
");
    unsigned int cpuNum = systemCPUNum();
    printf("Current system has %u CPU(s)
", cpuNum);

    std::vector<unsigned int> runningCPUVector;
    if (!currentProcessAffinity(runningCPUVector))
    {
        printf("Get current process was bound witch CPU failed
");
        return 1;
    }

    for (auto& iter : runningCPUVector)
    {
        printf("Current process is running at %u CPU
", iter);
    }

    std::vector<unsigned int> needBindCPUVector {0, 2};
    if (!setCurrentProcessAffinity(needBindCPUVector))
    {
        printf("Current process bind CPU failed
");
        return 1;
    }

    printf("Current process bind CPU success
");

    runningCPUVector.clear();
    if (!currentProcessAffinity(runningCPUVector))
    {
        printf("Get current process was bound witch CPU failed
");
        return 1;
    }

    for (auto& iter : runningCPUVector)
    {
        printf("Current process is running at %u CPU
", iter);
    }

    printf("
*****Thread bind CPU sample*****
");
    runningCPUVector.clear();
    if (!currentThreadAffinity(runningCPUVector))
    {
        printf("Get current thread was bound witch CPU failed
");
        return 1;
    }

    for (auto& iter : runningCPUVector)
    {
        printf("Thread %lu is running at %u CPU
", pthread_self(), iter);
    }

    needBindCPUVector.clear();
    needBindCPUVector.push_back(1);
    if (!setCurrentThreadAffinity(needBindCPUVector))
    {
        printf("Current thread bind CPU failed
");
        return 1;
    }

    printf("Thread %lu bind CPU success
", pthread_self());

    runningCPUVector.clear();
    if (!currentThreadAffinity(runningCPUVector))
    {
        printf("Get current thread was bound witch CPU failed
");
        return 1;
    }

    for (auto& iter : runningCPUVector)
    {
        printf("Thread %lu is running at %u CPU
", pthread_self(), iter);
    }

    return 0;
}

程序执行的输出结果：
*****Process bind CPU sample*****
Current system has 4 CPU(s)
Current process is running at 0 CPU
Current process is running at 1 CPU
Current process is running at 2 CPU
Current process is running at 3 CPU
Current process bind CPU success
Current process is running at 0 CPU
Current process is running at 2 CPU

*****Thread bind CPU sample*****
Thread 139871129114432 is running at 0 CPU
Thread 139871129114432 is running at 2 CPU
Thread 139871129114432 bind CPU success
Thread 139871129114432 is running at 1 CPU

该例子的github地址：https://github.com/chxuan/samples/blob/master/BindCPU/BindCPU.cpp