求助：将网络数据处理流程迁移至Tokio多线程的问题

我有一个处理网络接入数据的流程，这些数据帧约128字节，单帧处理耗时约1.8微秒。虽然单帧处理速度看似足够快，但因待处理帧的数量庞大，需要评估如何进行多线程优化。

以下是我通过随机字节/休眠方法模拟该流程的代码：

• linear函数展示了无线程的单线程数据处理方式
• tokio函数是我尝试多线程化的初步实现，但相比Rust或Tokio文档中的示例复杂得多，是在解决各类编译错误后拼凑的版本

Cargo.toml 文件内容 #

[package]
name = "codereview"
version = "0.1.0"
edition = "2021"

[dependencies]
tokio = {version = "1.38.0", features = ["full"] }
rand = "0.8.5"

main.rs 文件内容 #

use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::thread;
use std::time::Duration;
use tokio::{task, time::{Instant}};
use Vec;
use tokio::sync::mpsc;
use rand::Rng;

fn fill_vec(count: u8) -> Vec<u8> {
    let mut rng = rand::thread_rng();
    (0..count).map(|_| rng.gen()).collect()
}

// simulate the actual process, which takes about 2us.
fn process_data(data: &Vec<u8>) -> bool {
    thread::sleep(Duration::from_micros(2));
    return true;
}

fn do_linear_process(data: &Vec<Vec<u8>>, max_iterations:u32) {
    let mut iteration = 0;
    let iteration_start = Instant::now();
    let mut elements = 0;
    loop {
        iteration += 1;
        if iteration > max_iterations {
            break;
        }
        // simply classify the frame
        let data = data[iteration as usize % 100].clone();
        let val = process_data(&data);
        if val {
            elements += 1;
        }
    }
    let iteration_elapsed = iteration_start.elapsed();
    let time_elements = iteration_elapsed.as_micros() / elements;
    println!("Iterated {} elements in {:?}us at {} us/element", elements, iteration_elapsed.as_micros(), time_elements);
}

fn do_tokio_mpsc(data: &Vec<Vec<u8>>, max_iterations:u32) {

    println!("---- Repeating with tokio mpsc threading ----");
    // now we'll repeat the last part with some threading action...
    let mut iteration = 0;
    let iteration_start = Instant::now();
    let elements = Arc::new(AtomicU64::new(0));
    let (data_tx, mut data_rx) = mpsc::channel::<bool>(2000);;
    let (frame_tx, mut frame_rx) = mpsc::channel::<Vec<u8>>(2000);

    // fully decode the frame
    tokio::spawn(async move {
        while let Some(frame) = frame_rx.recv().await {
            if let was_successful = process_data(&frame) {
                data_tx.send(was_successful).await;
            }
        }
    });

    let elements_clone = elements.clone();
    let mut process = || async move {
        while let Some(_) = data_rx.recv().await {
            elements_clone.fetch_add(1, Ordering::SeqCst);
        }
        println!("Finished receiving {} elements.", elements_clone.load(Ordering::SeqCst));
    };

    let handle = task::spawn(process());

    loop {
        iteration += 1;
        if iteration > max_iterations {
            drop(frame_tx);
            break;
        }
        let data = data[iteration as usize% 100].clone();
        frame_tx.send(data);
    }

    tokio::spawn(async move {
        match handle.await {
            Ok(_) => println!("Task completed successfully."),
            Err(e) => println!("Task failed: {:?}", e),
        }
    });

    let iteration_elapsed = iteration_start.elapsed();
    if elements.load(Ordering::SeqCst) != 0 {
        let time_elements = iteration_elapsed.as_micros() / elements.load(Ordering::SeqCst) as u128;
        println!("Iterated {} elements in {:?}us at {} us/element", elements.load(Ordering::SeqCst), iteration_elapsed.as_micros(), time_elements);
    } else {
        println!("Elements is 0.  Time to complete was: {}us", iteration_elapsed.as_micros());
    }
}

#[tokio::main]
async fn main() {

    // create random data to simulate actual data
    let mut frame_data:Vec<Vec<u8>> = Vec::new();
    for _ in 0..100 {
        frame_data.insert(0,fill_vec(128));
    }

    // maximum iterations for the demo...
    let max_iterations = 1_000_000;

    // process the data in a single, linear thread
    do_linear_process(&frame_data, max_iterations);
    do_tokio_mpsc(&frame_data, max_iterations);

}

当前遇到的问题 #

• 因未在send调用后使用await，数据未实际发送至线程（存在编译器警告）
• 尝试调用await时，提示当前环境非异步（尽管是从async main中调用）
• 实现简单的处理计数共享状态是否真的需要如此复杂？

核心问题分析 #

1. send未加await导致数据未发送：Tokio的MPSC通道send是异步方法，必须await才能完成发送逻辑，否则只是返回一个未执行的Future，数据根本没进入通道。
2. 同步循环中无法调用await：do_tokio_mpsc是普通同步函数，里面的循环是同步的，不能直接在循环里用await，需要把整个函数改成异步的。
3. 状态管理过度复杂：不需要两层MPSC通道加原子计数，直接用Tokio的任务池批量处理，结合JoinHandle等待所有任务完成，用Arc<AtomicU64>即可完成简单计数。

修复后的代码实现 #

use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::thread;
use std::time::Duration;
use tokio::{task, time::Instant};
use rand::Rng;

fn fill_vec(count: u8) -> Vec<u8> {
    let mut rng = rand::thread_rng();
    (0..count).map(|_| rng.gen()).collect()
}

// 注意：实际处理是阻塞操作，用spawn_blocking放到线程池执行，避免阻塞Tokio异步调度器
fn process_data(data: &Vec<u8>) -> bool {
    thread::sleep(Duration::from_micros(2));
    true
}

fn do_linear_process(data: &Vec<Vec<u8>>, max_iterations: u32) {
    let mut iteration = 0;
    let iteration_start = Instant::now();
    let mut elements = 0;
    loop {
        iteration += 1;
        if iteration > max_iterations {
            break;
        }
        let data = data[iteration as usize % 100].clone();
        let val = process_data(&data);
        if val {
            elements += 1;
        }
    }
    let iteration_elapsed = iteration_start.elapsed();
    let time_elements = iteration_elapsed.as_micros() / elements as u128;
    println!(
        "Iterated {} elements in {}us at {} us/element",
        elements,
        iteration_elapsed.as_micros(),
        time_elements
    );
}

async fn do_tokio_optimized(data: &Vec<Vec<u8>>, max_iterations: u32) {
    println!("---- Repeating with optimized tokio threading ----");
    let iteration_start = Instant::now();
    let elements = Arc::new(AtomicU64::new(0));
    let mut handles = Vec::new();

    // 批量生成任务，利用Tokio多线程任务池
    for iteration in 1..=max_iterations {
        let data_clone = data[iteration as usize % 100].clone();
        let elements_clone = elements.clone();
        // 阻塞型任务用spawn_blocking，避免占用异步调度线程
        let handle = task::spawn_blocking(move || {
            let success = process_data(&data_clone);
            if success {
                elements_clone.fetch_add(1, Ordering::Relaxed);
            }
        });
        handles.push(handle);
    }

    // 等待所有任务完成
    for handle in handles {
        handle.await.expect("Task failed");
    }

    let iteration_elapsed = iteration_start.elapsed();
    let element_count = elements.load(Ordering::Relaxed);
    let time_elements = iteration_elapsed.as_micros() / element_count as u128;
    println!(
        "Iterated {} elements in {}us at {} us/element",
        element_count,
        iteration_elapsed.as_micros(),
        time_elements
    );
}

#[tokio::main]
async fn main() {
    let mut frame_data: Vec<Vec<u8>> = Vec::new();
    for _ in 0..100 {
        frame_data.push(fill_vec(128));
    }

    let max_iterations = 1_000_000;

    do_linear_process(&frame_data, max_iterations);
    do_tokio_optimized(&frame_data, max_iterations).await;
}

关键优化点说明 #

1. 异步函数改造：将do_tokio_mpsc改为async fn，允许内部使用await等待任务完成。
2. 简化任务调度：直接用task::spawn_blocking处理阻塞型的process_data，批量生成任务并收集JoinHandle，避免复杂的通道逻辑。
3. 状态管理简化：仅用Arc<AtomicU64>完成计数，每个任务完成后更新计数，最后统一等待所有任务结束再统计结果。
4. 性能优化：移除无用的MPSC通道，减少额外开销；用spawn_blocking处理阻塞任务，避免影响Tokio异步调度器的效率。

额外建议 #

• 如果实际process_data是CPU密集型而非阻塞型，推荐使用rayon库的并行迭代器，比Tokio更适合这类场景。
• 减少不必要的内存拷贝：可以用Arc<Vec<u8>>共享数据，替代频繁的clone操作。

内容的提问来源于stack exchange，提问作者PilotGuy