tokio 是 rust 生態中流行的異步運行框架。在實際生產中我們如果希望 tokio 應用程序與特定的 cpu core 綁定該怎么處理呢？這次我們來聊聊這個話題。

首先我們先寫一段簡單的多任務程序。

use tokio::runtime;
pub fn main() {
    let rt = runtime::new_multi_thread()
        .enable_all()
        .build()
        .unwrap();


    rt.block_on(async {
        for i in 0..8 {
            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);
                }
            });
        }
    });
}

程序非常簡單，首先構造一個tokio runtime 環境，然后派生多個 tokio 并發，每個并發執行一個無限循環做overflowing_add。overflowing_add函數返回一個加法的元組以及一個表示是否會發生算術溢出的布爾值。如果會發生溢出，那么將返回包裝好的值。然后取元祖的第一個元素打印。

這個程序運行在 Ubuntu 20 OS，4 core cpu。通過nmon的監控如下：

可以看到每個 core 都有負載。

要想把負載綁定在某一 core 上，需要使用core_affinity_rs（https://github.com/Elzair/core_affinity_rs）。core_affinity_rs是一個用于管理CPU親和力的Rust crate。目前支持Linux、Mac OSX和Windows。官方宣稱支持多平臺，本人只做了linux 操作系統的測試。

我們把代碼修改一下：

use tokio::runtime;


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


    let rt = runtime::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..8 {
            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);           
                }
            });
        }
    });
}

在構建多線程runtime時，在on_thread_start 設置cpu親和。可以看到負載被綁定到了指定的core上。

上面的代碼只是把負載綁定到了一個core上，那么要綁定多個核怎么辦呢

我們看看下面的代碼

pub fn main() {
    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 = 2;


    rt.block_on(async {
        for i in 0..8 {
            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);
                    }
            });
        }
    });
}

代碼需要把所有負載綁在 core3和core4上。原理是在派生任務中加入 core_affinity 設置.通過調整idx，將派生并發平均綁定在指定的core上。代碼運行的監控如下圖。

本期關于cpu親和的話題就聊到這兒，下期見