Rust Rayon并行计算API的使用

Rayon简介

Rayon是Rust语言中的一个轻量级并行计算库。它的设计初衷是为了帮助开发者轻松地将顺序计算转换为并行计算，同时确保数据争用不会发生。Rayon的API非常直观，主要通过并行迭代器实现高效的数据处理，适合多核CPU的现代计算环境。

并行计算的基本概念

并行计算是一种同时使用多个计算资源解决计算问题的技术。在当今多核处理器广泛普及的背景下，并行计算尤为重要。它能够显著提高程序的执行效率，通过分配任务到不同的处理器核心来加速计算过程。Rayon正是利用这种技术，使得Rust程序能够在不牺牲安全性和准确性的前提下，充分发挥硬件性能。

Rayon工作窃取算法

Rayon使用的核心技术是工作窃取算法。这种算法通过动态地分配任务，使得线程可以从其他线程的任务队列中“窃取”任务来执行，从而实现负载均衡。工作窃取算法的主要优势在于其高度的灵活性和效率，能够在保证任务调度均衡的同时，最大程度地利用系统资源。

工作窃取的实现

工作窃取的实现依赖于一个全局的任务队列，线程在完成自身任务后，会从队列中获取新的任务。如果队列为空，线程将尝试从其他线程的任务队列中窃取任务。这种机制确保了即使在任务不均匀分配的情况下，所有线程都能保持忙碌状态。

工作窃取的优势

这种算法的优势在于：1. 提高了资源利用率；2. 减少了线程在等待新任务时的空闲时间；3. 在复杂的任务调度场景下，能够显著提升整体执行效率。

Rayon并行迭代器的使用

Rayon的使用主要体现在其并行迭代器上。通过将标准的迭代器转换为并行迭代器，开发者可以在不改动大量代码的前提下，轻松实现并行计算。

并行迭代器的基本操作

并行迭代器的基本操作包括：

use rayon::prelude::*;

fn sum_of_squares(input: &[i32]) -> i32 {
    input.par_iter() // <-- 只需要更改这里
         .map(|&i| i * i)
         .sum()
}

上述代码中，par_iter替换了iter，从而将顺序迭代转换为并行迭代。

并行迭代器的灵活性

并行迭代器允许用户通过join和scope函数自行创建并行任务，这为开发者提供了更大的灵活性。在复杂场景下，用户可以创建自定义线程池以取代Rayon的默认线程池，进一步优化性能。

示例1：文件夹大小计算

在计算文件夹大小的任务中，Rayon的并行迭代器能够显著提高性能。通过使用原子计数器和并行迭代器，避免数据竞态问题，同时实现高效的文件遍历。

模拟大量任务

use rayon::prelude::*;
use std::fs;
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Instant;

fn main() {
    let begin = Instant::now();
    let (size, count) = dir_size(r#"C:Mll"#);
    let size = (size as f64) / (1 << 30) as f64;
    println!("数量:{count} 大小:{size:.2}G 耗时:{:.2?}", begin.elapsed())
}

pub fn dir_size<P: AsRef>(path: P) -> (u64, u64) {
    let size = AtomicU64::new(0);
    let count = AtomicU64::new(0);
    fs::read_dir(path)
        .unwrap()
        .par_bridge()
        .filter_map(Result::ok)
        .for_each(|entry| {
            let meta = entry.metadata().unwrap();
            if meta.is_symlink() {
                return;
            }

            count.fetch_add(1, Ordering::SeqCst);
            if meta.is_dir() {
                let (s, c) = dir_size(entry.path());
                size.fetch_add(s, Ordering::SeqCst);
                count.fetch_add(c, Ordering::SeqCst);
            } else {
                size.fetch_add(meta.len(), Ordering::SeqCst);
            }
        });
    (size.into_inner(), count.into_inner())
}

结果分析

在大量文件的情况下，Rayon的并行处理能力可以显著减少计算时间。这种提升在需要遍历和处理大量文件的应用中尤为重要。

示例2：素数计算

计算素数是一个计算密集型任务，通过Rayon并行迭代器，可以有效提升计算速度。

暴力素数算法

use rayon::prelude::*;
use std::time::Instant;

fn main() {
    let begin = Instant::now();

    let n:usize = 1_000_000;
    let res = (2..n)
        // .par_bridge()
        .filter(|&i| (2..i).find(|f| i % f == 0).is_none())
        .count();

    println!("范围:2-{n} 数量:{res} 耗时: {:.2?}", begin.elapsed());
}

并行处理前后的对比

通过对比可以发现，并行处理在大规模素数计算中极大减少了计算时间，但同时也大幅增加了CPU的使用率。

示例3：文件内容检索

在需要处理大量I/O任务的场景中，Rayon也能展现出其强大的并行处理能力。

模拟大量IO任务

use rayon::iter::*;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::sync::mpsc;
use std::sync::mpsc::Sender;
use std::time::Instant;
use std::{fs, io, thread};

const FLAG: &str = "FLAG-123456";

fn main() {
    let begin = Instant::now();

    let (tx, rx) = mpsc::channel();
    let res = thread::spawn(|| rx.into_iter().collect());
    file_content_search(".", FLAG.as_bytes(), tx);
    let res: Vec = res.join().unwrap();

    println!("找到:{} 耗时:{:.2?}", res.len(), begin.elapsed())
}

fn file_content_search<P: AsRef>(path: P, search: &[u8], tx: Sender) {
    fs::read_dir(path)
        .unwrap()
        .par_bridge()
        .filter_map(Result::ok)
        .for_each(|entry| {
            if entry.metadata().unwrap().is_dir() {
                return file_content_search(entry.path(), search, tx.clone());
            }
            let path = entry.path();
            if let Ok(true) = bytes_search(&path, search) {
                tx.send(path).unwrap()
            }
        });
}

fn bytes_search<P: AsRef>(path: P, search: &[u8]) -> io::Result {
    let mut file = fs::File::open(path)?;
    let mut buf = [0; 1 << 10];
    let mut offset = 0;

    loop {
        let n = file.read(&mut buf)?;
        if n == 0 {
            break Ok(false);
        }

        unsafe {
            for item in buf.get_unchecked(0..n).iter() {
                if offset == search.len() {
                    return Ok(true);
                }
                if item == search.get_unchecked(offset) {
                    offset += 1
                } else {
                    offset = 0
                }
            }
        }
    }
}

性能提升与风险

并行处理使得文件内容检索任务的时间显著缩短，但也可能导致高I/O负载，需根据实际需求进行调整。

Rayon线程池

Rayon提供了默认的全局线程池，但在某些应用场景下，自定义线程池可以提供更优秀的性能调节能力。

自定义线程池的使用

在需要更高效的资源分配时，可以考虑创建自定义线程池。

use rayon::ThreadPoolBuilder;

fn main() {
    let pool = ThreadPoolBuilder::new().num_threads(4).build().unwrap();
    pool.install(|| {
        // 在这里执行并行任务
    });
}

线程池的配置选项

Rayon允许用户通过ThreadPoolBuilder配置线程池的线程数量、栈大小等参数，以便更好地适应不同的计算需求。

Rayon的应用场景

Rayon适合于多种需要高效并行计算的场景，例如：

数据处理与分析
图像处理
科学计算
大规模文件处理

在这些场景中，Rayon能够显著缩短任务的执行时间，提高程序的整体效率。

小结

Rayon是一个功能强大且易于使用的并行计算库。通过简单的API和高效的任务调度机制，它使Rust程序能够在保证安全性的同时，最大限度地发挥硬件性能。在计算密集型和I/O密集型任务中，Rayon都能表现出色。

FAQ

问：Rayon如何避免数据争用？
答： Rayon通过其API保证无数据争用的发生，只要代码能够通过编译，通常就能确保与顺序执行相同的结果。
问：如何使用Rayon创建并行任务？
答：可以通过将顺序迭代器转换为并行迭代器，或者使用join、scope函数来创建并行任务。
问：Rayon适合哪些类型的应用？
答： Rayon适合需要高效并行计算的应用场景，如数据处理、图像处理、科学计算等。
问：是否可以在Rayon中使用自定义线程池？
答：是的，Rayon允许用户通过ThreadPoolBuilder创建自定义线程池，以满足特定的计算需求。
问：Rayon的并行迭代器有什么优势？
答：并行迭代器使得数据处理更加高效，能够自动适应任务大小和系统资源，提供最佳性能。