Rust Crossbeam 之 AtomicCell

atomic_cell.rs

AtomicCell<T>的数据结构很简单

#[repr(transparent)]
pub struct AtomicCell<T> {
    value: UnsafeCell<MaybeUninit<T>>,
}

AtomicCell在C++中实际上就是std::atomic, 如果类型T的大小可以被当前硬件平台原子访问, 那就会按照原子访问, 如果太大了的话, 就会退化为锁定后访问(在Crossbeam中, 这个锁就是seq_lock)

这里吐槽一下, Rust的模板特化什么时候才能上... 这段原子类型访问和加锁访问的代码, 由于没有模板特化, 代码比C++啰嗦好多

它的value, 其类型为UnsafeCell<MaybeUninit<T>>, 我们分开说, 首先是UnsafeCell, UnsafeCell是Rust中唯一合法地, 可以从不可变借用中获取可变借用的方法.

像是开了个洞, 对吧? 其实就是C++中的mutable, 可以参考[[Rust UnsafeCell VS C++ mutable]]

这里比较有意思的东西有三个: #[repr(xxx)], UnsafeCell 以及 MaybeUninit, 我们分开讲一讲

repr

在Rust中, 形如#[repr(...)]可以用来指定一个类型的内存布局.
如果没有这个标记时, rust编译器默认使用的是#[repr(Rust)], 这是Rust的默认布局, 与C/C++的内存布局结构不同, Rust会自行对类型内的field进行重排以实现更高效的内存利用.

repr内还可使用其他的一些参数

• transparent 保证单字段结构体/枚举的内存布局与其唯一字段的完全相同.

这个东东可以说是Rust零开销抽象的一个工具.
为了类型安全, 我们可能会对一些类型做一个别名, 实际上就是一个struct, 包装一个字段.
假如这样一个结构体 struct Foo(f64), 可以视为对f64的包装, 但是, 我们知道, 在参数传递的时候, 一些ABI会将比较小的数据放在通用的寄存器上, 过大的放在栈上, 或者浮点数, 会放在浮点寄存器上传参.
对于编译器来说, Foo是一个结构体, 而不是一个浮点数, 那得放在栈上或者通用寄存器上, 当需要使用的时候, 还得把f64从内存上或者通用寄存器上拷贝到浮点寄存器使用. 但是有了transparent, 编译器就会认出它就是一个f64, 从而避免这么一套折腾

• C 保证布局与C的一致

如果需要进行跨语言调用, 那么C的ABI是很合适的.
或者不做FFI, 但是要手动设置结构体的布局, 也可以用repr(C)实现(比如实现一个侵入式数据结构?)

• u8, i32... etc

这个可以用在枚举上, 指定存储枚举的整数类型.
基本相当于C++的 enum XXX: int这样的写法

• packed

不要padding, 相当于C的struct __attribute__((__packed__)) XXX

• align(N)

强制结构体对齐. 一般用来解决伪共享问题以及SIMD对齐比较多

MaybeUninit

在Rust中, 所有变量在绑定之前都是要被初始化好的, 但是某些情况下, 我们可能需要延迟这个初始化行为(比如动态数组, 数组占用的实际空间可以与存放的元素使用的空间大小不一致, 剩下的内存可以用于后续加入的元素使用)

说这个可能有些复杂, 但假如我说C++ placement new的使用场景, 大家是不是就了解了?

不过AtomicCell并不会延迟初始化, 也不会存储非初始化的值. 这里使用MaybeUninit其实是利用了它的一个副作用----禁止"niche"优化

救命.. 我也不知道"niche"怎么翻译, 就这么直接用吧

Rust 对于内存的使用是很细的, 经常会根据一些类型的特性在其内存中没使用的位置扣出几个bit用用.
就用上面的struct Foo举个例子, Rust中Box指向的元素必定有效不为空, 所以Box的值永远不为零.
那么比如Option<Box<Foo>>, 由于Box不为零, 所以编译器就可以用Box位置内存全零来表示Option::None.

或者, 由于Foo大小为8bytes, 因为对齐的需要, 指向Foo的指针的低3位将是永远为0的, 这3个没用到的bit可能会被装有Box<Foo>的枚举或者什么类型用到, 来标记其他状态

像这种扣bit位的优化, 在Rust中被称为"niche optimization".

但由于 MaybeUninit 内部有两种状态, 已初始化和未初始化, 这迫使编译器不能对内部字段的内存使用做任何假设, 也就无从谈起这种扣bit的niche优化了

niche优化会在一些场景下影响到AtomicCell的功能, 我们就用其github issue中的问题举例.

enum Enum {
    NeverConstructed,
    Cell(AtomicCell<NonZeroU128>),
}

128位的整数比较大, 大多数平台都不支持对如此大的内存做原子访问, 所以AtomicCell会退化为使用SeqLock, 并且对于NonZeroU128的访问一般都会被拆分为两次64位访问. 不过有锁保证同步, 倒也没什么.

但问题出在Enum上, 由于NonZeroU128永不为0, 如果没有MaybeUninit, niche优化会将这16字节全为0时视为NeverConstructed. Enum不是AtomicCell, 所以对于Enum的读取并不是原子的(也没有锁保证同步).

假如有这样的修改

let x = NonZeroU128::new(0xFFFFFFFF_FFFFFFFF_00000000_00000000).unwrap();
let y = NonZeroU128::new(0x00000000_00000000_FFFFFFFF_FFFFFFFF).unwrap();
loop {
	cell.store(x);
    cell.store(y);
}

其他线程在读取Enum时, 是可能观测到Enum的内存全为0的, 从而导致Enum被错误识别为NeverConstructed. 这个问题就在于, 由于读取写入的非原子性, 编译器不应当利用这16字节内部情况进行niche优化, 避免这个问题的方法就是使用MaybeUninit, 迫使编译器将Enum的Tag放在另一个位置, 而不是这16字节内存里面.