文盤Rust -- tokio綁定cpu實踐

tokio 是 rust 生態(tài)中流行的異步運行框架。在實際生產中我們如果希望 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 環(huán)境，然后派生多個 tokio 并發(fā)，每個并發(fā)執(zhí)行一個無限循環(huán)做overflowing_add。overflowing_add函數(shù)返回一個加法的元組以及一個表示是否會發(fā)生算術溢出的布爾值。如果會發(fā)生溢出，那么將返回包裝好的值。然后取元祖的第一個元素打印。

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

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

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

我們把代碼修改一下：

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親和?？梢钥吹截撦d被綁定到了指定的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，將派生并發(fā)平均綁定在指定的core上。代碼運行的監(jiān)控如下圖。

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