tokio多任务绑定cpu(绑核)

tokio 是 rust 生态中流行的异步运行时框架。在实际生产中我们如果希望 tokio 应用程序与特定的 cpu core 绑定该怎么处理呢？

首先我们先写一段简单的多任务程序。

use tokio;
use tokio::runtime;
use core_affinity;

fn tokio_sample() {
    let rt = runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap();

    rt.block_on(async {
        for i in 0..4 {
            println!("num {}", i);
            tokio::spawn(async move {
                loop {
                    let mut sum: i32 = 0;
                    for i in 0..100000000 {
                        sum = sum.overflowing_add(i).0;
                    }
                    println!("sum {}", sum);
                }
            });
        }
    });
}

fn main() {
    tokio_sample();   
}

程序非常简单，首先构造一个 tokio runtime 环境，然后派生多个 tokio 并发，每个并发执行一个无限循环做 overflowing_add。overflowing_add 函数返回一个加法的元组以及一个表示是否会发生算术溢出的布尔值。如果会发生溢出，那么将返回包装好的值。然后取元祖的第一个元素打印。

这个程序运行在 Centos8，4 core cpu。通过 pidstat -t 1 -p $(pidof cpu_affinity_tokio) 的监控如下：

可以看到4个 core 负载几乎都是 100%。

要想把负载绑定在某一 core 上，需要使用 core_affinity_rs。core_affinity_rs 是一个用于管理 CPU 亲和力的 Rust crate。目前支持 Linux、Mac OSX 和 Windows。官方宣称支持多平台。

我们把代码修改一下：

fn tokio_affinity_cpu1() {
    let core_ids = core_affinity::get_core_ids().unwrap();
    println!("core num {}", core_ids.len());
    let core_id = core_ids[1];

    let rt = runtime::Builder::new_multi_thread()
        .on_thread_start(move || {
            core_affinity::set_for_current(core_id.clone());
        })
    .enable_all()
        .build()
        .unwrap();

    rt.block_on(async {
        for i in 0..4 {
            println!("num {}", i);
            tokio::spawn(async move {
                loop {
                    let mut sum: i32 = 0;
                    for i in 0..100000000 {
                        sum = sum.overflowing_add(i).0;
                    }
                    println!("sum {}", sum);
                }
            });
        }
    });
}

fn main() {  
   tokio_affinity_cpu1();
}

在构建多线程 runtime 时，在 on_thread_start 设置 cpu 亲和 cpu1。可以看到负载被绑定到了指定的 core 上，4 个线程各占 1/4 cpu。

上面的代码只是把负载绑定到了一个 core 上，那么要绑定多个核怎么办呢？
我们看看下面的代码

fn tokio_affinity_cpu1_cpu2() {
    let core_ids = core_affinity::get_core_ids().unwrap();
    println!("core num {}", core_ids.len());

    let rt = runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap();

    let mut idx = 1;

    rt.block_on(async {
        for i in 0..4 {
            println!("num {}", i);
            let core_id = core_ids[idx];
            if idx.eq(&(core_ids.len() - 1)) {
                idx = 2;
            } else {
                idx += 1;
            }

            tokio::spawn(async move {
                let res = core_affinity::set_for_current(core_id);
                println!("{}", res);
                loop {
                    let mut sum: i32 = 0;
                    for i in 0..100000000 {
                        sum = sum.overflowing_add(i).0;
                    }
                    println!("sum {}", sum);
                }
            });
        }
    });
}


fn main() {
    tokio_affinity_cpu1_cpu2();
}

代码需要把所有负载绑在 core1 和 core2 上。原理是在派生任务中加入 core_affinity 设置。通过调整 idx，将派生并发平均绑定在指定的 core 上。代码运行的监控如下图。

参考：
文盘 Rust – tokio 绑定 cpu 实践
Rust入门秘籍-tokio简介
Rust语言圣经(Rust Course)-tokio
https://tokio.rs/tokio/tutorial