MK
摩柯社区 - 一个极简的技术知识社区
AI 面试

Rust 原子存储操作的性能优化

2021-12-067.7k 阅读

Rust 原子存储操作基础

在 Rust 中，原子存储操作允许在多线程环境下对数据进行安全的共享访问。原子类型（Atomic*）位于 std::sync::atomic 模块中，提供了对基本数据类型的原子操作。例如，AtomicI32 用于对 32 位整数进行原子操作。

use std::sync::atomic::{AtomicI32, Ordering};

let num = AtomicI32::new(0);
num.store(10, Ordering::SeqCst);
let result = num.load(Ordering::SeqCst);
println!("Loaded value: {}", result);

在上述代码中，首先创建了一个 AtomicI32 类型的变量 num 并初始化为 0。然后使用 store 方法将值设置为 10，load 方法用于读取该值。Ordering 参数决定了内存访问的顺序，这里使用的 SeqCst（顺序一致性）是最严格的顺序。

内存顺序（Memory Ordering）

内存顺序是原子操作性能优化的关键因素之一。Rust 提供了多种内存顺序选项，每种选项在性能和同步保证上各有不同。

顺序一致性（SeqCst）

SeqCst 提供了最强的同步保证，所有线程以相同的顺序观察到所有 SeqCst 操作。然而，这种严格的保证也带来了较高的性能开销。

use std::sync::atomic::{AtomicBool, Ordering};

let flag = AtomicBool::new(false);

// 线程 1
std::thread::spawn(move || {
    flag.store(true, Ordering::SeqCst);
});

// 线程 2
std::thread::spawn(move || {
    while!flag.load(Ordering::SeqCst) {
        std::thread::yield_now();
    }
    println!("Flag is true");
});

在这个例子中，两个线程通过 AtomicBool 类型的 flag 进行同步。由于使用了 SeqCst 顺序，线程 2 能确保在线程 1 设置 flag 为 true 后，才会退出循环并打印信息。

释放 - 获取（Release - Acquire）

Release 和 Acquire 顺序用于在性能和同步保证之间取得平衡。Release 顺序确保在该操作之前的所有写操作对其他线程可见，而 Acquire 顺序确保在该操作之后的所有读操作能看到之前线程的写操作。

use std::sync::atomic::{AtomicUsize, Ordering};

let data = AtomicUsize::new(0);
let ready = AtomicBool::new(false);

// 线程 1
std::thread::spawn(move || {
    data.store(42, Ordering::Release);
    ready.store(true, Ordering::Release);
});

// 线程 2
std::thread::spawn(move || {
    while!ready.load(Ordering::Acquire) {
        std::thread::yield_now();
    }
    let value = data.load(Ordering::Acquire);
    println!("Loaded value: {}", value);
});

在线程 1 中，先将 data 设置为 42，然后将 ready 设置为 true，都使用 Release 顺序。在线程 2 中，先通过 Acquire 顺序读取 ready，确保看到线程 1 的写操作，然后再以 Acquire 顺序读取 data。这样能保证线程 2 读取到正确的 data 值。

松弛顺序（Relaxed）

Relaxed 顺序提供了最小的同步保证，只保证操作的原子性，不保证内存访问顺序。这在一些不需要严格同步的场景下可以提供较好的性能。

use std::sync::atomic::{AtomicUsize, Ordering};

let counter = AtomicUsize::new(0);

// 多个线程并发操作
let handles: Vec<_> = (0..10).map(|_| {
    std::thread::spawn(move || {
        for _ in 0..1000 {
            counter.fetch_add(1, Ordering::Relaxed);
        }
    })
}).collect();

for handle in handles {
    handle.join().unwrap();
}

let final_value = counter.load(Ordering::Relaxed);
println!("Final value: {}", final_value);

在这个例子中，多个线程并发地对 counter 进行 fetch_add 操作，使用 Relaxed 顺序。虽然这种顺序不能保证每个线程看到的操作顺序一致，但能保证 fetch_add 操作的原子性。

原子引用计数（Atomic Reference Counting）

std::sync::Arc（原子引用计数指针）常用于多线程环境下共享数据。Arc 内部使用原子操作来管理引用计数。

use std::sync::Arc;

let data = Arc::new(42);
let data_clone = data.clone();

println!("Data in original: {}", *data);
println!("Data in clone: {}", *data_clone);

在上述代码中，Arc::new 创建了一个指向整数 42 的 Arc 实例，clone 方法增加了引用计数。Arc 的实现确保了引用计数的更新是线程安全的，通过原子操作来完成。

性能优化策略

减少不必要的原子操作

在多线程环境中，并非所有数据访问都需要原子操作。如果某个数据只在单个线程内部使用，将其声明为普通变量而非原子类型，可以避免原子操作带来的性能开销。

// 仅在单个线程内使用的数据
let mut local_var = 0;
for _ in 0..1000 {
    local_var += 1;
}

在这个简单的例子中，local_var 仅在单个线程内使用，使用普通的可变变量就足够了，无需使用原子类型。

批量操作

将多个原子操作合并为一个批量操作可以减少原子操作的次数，从而提高性能。例如，在更新多个相关的原子变量时，可以将这些更新操作放在一个临界区（使用 Mutex 或 RwLock）内，然后一次性更新所有变量。

use std::sync::{Arc, Mutex};

let shared_data = Arc::new(Mutex::new((0, 0)));
let shared_data_clone = shared_data.clone();

std::thread::spawn(move || {
    let mut data = shared_data_clone.lock().unwrap();
    data.0 += 1;
    data.1 += 2;
});

let result = shared_data.lock().unwrap();
println!("Values: {}, {}", result.0, result.1);

在这个例子中，使用 Mutex 来保护一个包含两个整数的元组。线程通过获取锁，一次性更新两个值，避免了多次原子操作。

选择合适的内存顺序

如前文所述，不同的内存顺序在性能和同步保证上各有不同。在性能敏感且同步要求不高的场景下，应优先选择 Relaxed 或 Release - Acquire 顺序。而在需要严格同步的场景下，才使用 SeqCst 顺序。

// 性能敏感且同步要求不高的场景
let counter = AtomicUsize::new(0);
for _ in 0..10000 {
    counter.fetch_add(1, Ordering::Relaxed);
}

在这个场景下，使用 Relaxed 顺序可以提高性能，因为对操作顺序的要求并不严格。

利用无锁数据结构

Rust 社区提供了一些无锁数据结构，如 crossbeam 库中的无锁队列和栈。这些数据结构使用原子操作来实现线程安全，同时避免了锁带来的性能开销。

use crossbeam::queue::MsQueue;

let queue = MsQueue::new();

// 生产者线程
std::thread::spawn(move || {
    for i in 0..100 {
        queue.push(i).unwrap();
    }
});

// 消费者线程
std::thread::spawn(move || {
    while let Some(value) = queue.pop() {
        println!("Popped value: {}", value);
    }
});

在这个例子中，MsQueue 是一个无锁队列，生产者和消费者线程可以并发地进行入队和出队操作，无需使用锁，从而提高了性能。

原子操作与缓存一致性

现代 CPU 通常具有多级缓存，以提高数据访问速度。原子操作在多线程环境下需要处理缓存一致性问题，这也会影响性能。

缓存行（Cache Line）

缓存行是 CPU 缓存与内存之间交换数据的最小单位，通常大小为 64 字节。当一个原子操作发生时，可能会导致整个缓存行被刷新，影响其他数据的访问性能。

use std::sync::atomic::{AtomicUsize, Ordering};

// 两个原子变量可能位于同一缓存行
let var1 = AtomicUsize::new(0);
let var2 = AtomicUsize::new(0);

// 线程 1
std::thread::spawn(move || {
    for _ in 0..1000 {
        var1.fetch_add(1, Ordering::Relaxed);
    }
});

// 线程 2
std::thread::spawn(move || {
    for _ in 0..1000 {
        var2.fetch_add(1, Ordering::Relaxed);
    }
});

在这个例子中，如果 var1 和 var2 位于同一缓存行，线程 1 对 var1 的操作可能会导致缓存行被刷新，影响线程 2 对 var2 的操作性能。

缓存行填充（Cache Line Padding）

为了避免缓存行争用，可以使用缓存行填充技术。通过在原子变量之间填充足够的空间，确保它们位于不同的缓存行。

use std::sync::atomic::{AtomicUsize, Ordering};

// 缓存行填充结构体
struct CacheLinePadding {
    _padding: [u8; 64 - std::mem::size_of::<AtomicUsize>()],
    value: AtomicUsize,
}

let var1 = CacheLinePadding {
    _padding: [0; 64 - std::mem::size_of::<AtomicUsize>()],
    value: AtomicUsize::new(0),
};

let var2 = CacheLinePadding {
    _padding: [0; 64 - std::mem::size_of::<AtomicUsize>()],
    value: AtomicUsize::new(0),
};

// 线程 1
std::thread::spawn(move || {
    for _ in 0..1000 {
        var1.value.fetch_add(1, Ordering::Relaxed);
    }
});

// 线程 2
std::thread::spawn(move || {
    for _ in 0..1000 {
        var2.value.fetch_add(1, Ordering::Relaxed);
    }
});

在这个例子中，通过 CacheLinePadding 结构体，确保 var1 和 var2 的 AtomicUsize 成员位于不同的缓存行，减少了缓存行争用带来的性能损失。

与其他语言原子操作的对比

与 C++ 相比，Rust 的原子操作具有更安全的类型系统。在 C++ 中，原子操作需要手动管理内存和类型转换，容易出错。而 Rust 的原子类型是强类型的，并且通过所有权系统来确保内存安全。

// C++ 原子操作示例
#include <atomic>
#include <thread>
#include <iostream>

std::atomic<int> num(0);

void increment() {
    for (int i = 0; i < 1000; ++i) {
        num.fetch_add(1, std::memory_order_relaxed);
    }
}

int main() {
    std::thread t1(increment);
    std::thread t2(increment);

    t1.join();
    t2.join();

    std::cout << "Final value: " << num.load(std::memory_order_relaxed) << std::endl;
    return 0;
}

// Rust 原子操作示例
use std::sync::atomic::{AtomicI32, Ordering};
use std::thread;

let num = AtomicI32::new(0);

let handle1 = thread::spawn(move || {
    for _ in 0..1000 {
        num.fetch_add(1, Ordering::Relaxed);
    }
});

let handle2 = thread::spawn(move || {
    for _ in 0..1000 {
        num.fetch_add(1, Ordering::Relaxed);
    }
});

handle1.join().unwrap();
handle2.join().unwrap();

let result = num.load(Ordering::Relaxed);
println!("Final value: {}", result);

在上述代码中，虽然功能相似，但 Rust 通过类型系统和所有权机制提供了更安全的编程模型。

与 Java 相比，Rust 的原子操作更加底层和灵活。Java 的原子操作主要通过 java.util.concurrent.atomic 包提供，并且与 Java 的内存模型紧密结合。而 Rust 允许开发者根据具体需求更细粒度地控制内存顺序。

// Java 原子操作示例
import java.util.concurrent.atomic.AtomicInteger;

public class AtomicExample {
    private static AtomicInteger num = new AtomicInteger(0);

    public static void increment() {
        for (int i = 0; i < 1000; ++i) {
            num.getAndIncrement();
        }
    }

    public static void main(String[] args) throws InterruptedException {
        Thread t1 = new Thread(() -> increment());
        Thread t2 = new Thread(() -> increment());

        t1.start();
        t2.start();

        t1.join();
        t2.join();

        System.out.println("Final value: " + num.get());
    }
}

// Rust 原子操作示例
use std::sync::atomic::{AtomicI32, Ordering};
use std::thread;

let num = AtomicI32::new(0);

let handle1 = thread::spawn(move || {
    for _ in 0..1000 {
        num.fetch_add(1, Ordering::Relaxed);
    }
});

let handle2 = thread::spawn(move || {
    for _ in 0..1000 {
        num.fetch_add(1, Ordering::Relaxed);
    }
});

handle1.join().unwrap();
handle2.join().unwrap();

let result = num.load(Ordering::Relaxed);
println!("Final value: {}", result);

在这个对比中，Rust 给予开发者更多的控制权，能根据实际场景选择合适的内存顺序，而 Java 的原子操作在一定程度上封装了这些细节。

实际应用场景

多线程计数器

在服务器性能监控等场景中，需要使用多线程计数器来统计请求数量等指标。原子操作可以确保计数器在多线程环境下的安全更新。

use std::sync::atomic::{AtomicUsize, Ordering};
use std::thread;

let counter = AtomicUsize::new(0);

// 模拟多个请求处理线程
let handles: Vec<_> = (0..10).map(|_| {
    thread::spawn(move || {
        for _ in 0..1000 {
            counter.fetch_add(1, Ordering::Relaxed);
        }
    })
}).collect();

for handle in handles {
    handle.join().unwrap();
}

let total_requests = counter.load(Ordering::Relaxed);
println!("Total requests: {}", total_requests);

在这个例子中，多个线程并发地更新 counter，通过原子操作保证了数据的一致性。

分布式系统中的状态同步

在分布式系统中，不同节点需要同步状态信息。原子操作可以用于确保状态更新的原子性和一致性。

// 假设这里是分布式系统中的一个节点
use std::sync::atomic::{AtomicBool, Ordering};

let is_ready = AtomicBool::new(false);

// 节点启动时的初始化操作
std::thread::spawn(move || {
    // 初始化完成后设置为 true
    std::thread::sleep(std::time::Duration::from_secs(2));
    is_ready.store(true, Ordering::Release);
});

// 其他节点等待该节点准备好
std::thread::spawn(move || {
    while!is_ready.load(Ordering::Acquire) {
        std::thread::yield_now();
    }
    println!("Node is ready, proceeding with distributed operation.");
});

在这个简单的模拟中，一个节点通过原子变量 is_ready 向其他节点同步其准备状态，其他节点通过原子操作等待该状态的更新。

总结

Rust 的原子存储操作为多线程编程提供了安全和高效的基础。通过合理选择内存顺序、减少不必要的原子操作、利用批量操作和无锁数据结构等优化策略，可以显著提高原子操作的性能。同时，了解原子操作与缓存一致性的关系，以及与其他语言原子操作的对比，有助于开发者在实际应用中更好地运用 Rust 的原子操作，构建高效、稳定的多线程应用程序。在实际应用场景中，如多线程计数器和分布式系统状态同步，原子操作发挥着关键作用，确保数据的一致性和系统的稳定性。