本文摘自《深入RUST标准库》,已经全网发售,恳请支持
Once 类型分析
Once是对全局变量的初始化必须在多个线程中(例如,库)竞争执行且只需要执行一次时的需求的方案。
C的pthread库实现了pthread_once来实现这个特性。RUST实现了自己的方案。Once的call_once方法使得可以用闭包的形式初始化全局变量,闭包内的代码不必考虑竞争,由Once确保线程安全且只初始化只被执行一次。
代码如下:
type Masked = ();
pub struct Once {
//用一个变量即实现状态,又实现了等待队列的头节点
//最后两位是Once的状态,前面是* const Waiter的裸指针地址
//Waiter是4字节对齐,因此地址最后两位为0,这个设计技巧不值得倡导,
//但是为了效率在这里不得不如此
state_and_queue: AtomicPtr<Masked>,
//state_and_queue中包含了一个等待的头节点的裸指针
_marker: marker::PhantomData<*const Waiter>,
}
//不能自动生成这两个trait
unsafe impl Sync for Once {}
unsafe impl Send for Once {}
impl UnwindSafe for Once {}
impl RefUnwindSafe for Once {}
pub struct OnceState {
//闭包执行期间出现panic的标识
poisoned: bool,
//给初始化闭包使用,用来标识是否中毒,或者已经顺利完成
set_state_on_drop_to: Cell<*mut Masked>,
}
//所有的静态变量可以使用ONCE_INIT进行赋值
pub const ONCE_INIT: Once = Once::new();
//闭包没有执行
const INCOMPLETE: usize = 0x0;
//闭包执行时线程panic
const POISONED: usize = 0x1;
//闭包正在执行
const RUNNING: usize = 0x2;
//初始化完成
const COMPLETE: usize = 0x3;
// 用来取出最后两位,用来做INCOMPLETE/POISONED/RUNNING/COMPLETE
const STATE_MASK: usize = 0x3;
//用来作为等待的线程队列节点,这些线程都对once做了闭包初始化的调用
#[repr(align(4))] //确保指针的地址的后2位无意义,可以用来作为状态,这是一个不值得推倡的技巧
struct Waiter {
//标识自身
thread: Cell<Option<Thread>>,
signaled: AtomicBool,
//next的节点
next: *const Waiter,
}
// 等待的队列.
struct WaiterQueue<'a> {
//Once的state_and_queue的引用
state_and_queue: &'a AtomicPtr<Masked>,
//初始化闭包的返回结果
set_state_on_drop_to: *mut Masked,
}
impl Once {
//一般直接使用ONCE_INIT
pub const fn new() -> Once {
Once {
//初始赋值
state_and_queue: AtomicPtr::new(ptr::invalid_mut(INCOMPLETE)),
_marker: marker::PhantomData,
}
}
//例如ONCE_INIT.call_once(|| {}), 在函数体内,可以对全局变量执行初始化,不必考虑
//线程间安全问题。
pub fn call_once<F>(&self, f: F)
where
F: FnOnce(),
{
// 是否已经完成初始化
if self.is_completed() {
return;
}
//以下将FnOnce()转换为了FnMut(state)
let mut f = Some(f);
//需要处理panic情况
self.call_inner(false, &mut |_| f.take().unwrap()());
}
//不理会panic的初始化
pub fn call_once_force<F>(&self, f: F)
where
F: FnOnce(&OnceState),
{
if self.is_completed() {
return;
}
let mut f = Some(f);
self.call_inner(true, &mut |p| f.take().unwrap()(p));
}
pub fn is_completed(&self) -> bool {
self.state_and_queue.load(Ordering::Acquire).addr() == COMPLETE
}
fn call_inner(&self, ignore_poisoning: bool, init: &mut dyn FnMut(&OnceState)) {
let mut state_and_queue = self.state_and_queue.load(Ordering::Acquire);
loop {
//判断当前状态
match state_and_queue.addr() {
//初始化成功且没有等待线程
COMPLETE => break,
//没有等待线程且已经POISONED,且不能忽视panic做初始化
POISONED if !ignore_poisoning => {
// Panic to propagate the poison.
panic!("Once instance has previously been poisoned");
}
//没有等待线程,没有初始化,或者已经panic但可以初始化
POISONED | INCOMPLETE => {
// 将状态转为RUNNING
let exchange_result = self.state_and_queue.compare_exchange(
state_and_queue,
ptr::invalid_mut(RUNNING),
Ordering::Acquire,
Ordering::Acquire,
);
//判断是否出现竞争
if let Err(old) = exchange_result {
//有竞争者,再次做循环
state_and_queue = old;
continue;
}
// 本线程获得初始化权利, 后面这段代码不会有竞争出现
// 创建其他线程等待的队列
let mut waiter_queue = WaiterQueue {
//设置Once的state_and_que为队列头部
state_and_queue: &self.state_and_queue,
//默认是POISONED
set_state_on_drop_to: ptr::invalid_mut(POISONED),
};
// 设置初始化状态
let init_state = OnceState {
poisoned: state_and_queue.addr() == POISONED,
//默认为COMPLETE
set_state_on_drop_to: Cell::new(ptr::invalid_mut(COMPLETE)),
};
//调用初始化函数
init(&init_state);
//对等待队列中的状态进行更新, 如果初始化闭包不关心init_state(call_once), 则默认为COMPLETE
waiter_queue.set_state_on_drop_to = init_state.set_state_on_drop_to.get();
//waiter_queue被释放,调用drop
break;
}
_ => {
// RUNNING,进入阻塞状态
assert!(state_and_queue.addr() & STATE_MASK == RUNNING);
wait(&self.state_and_queue, state_and_queue);
//阻塞被唤醒,重新获取新的状态并再次做判断循环
state_and_queue = self.state_and_queue.load(Ordering::Acquire);
}
}
}
}
}
//进入等待队列
fn wait(state_and_queue: &AtomicPtr<Masked>, mut current_state: *mut Masked) {
loop {
//在不是初次循环的情况下,初始化结束后的竞争赋值
//可能导致current_state被更新。
if current_state.addr() & STATE_MASK != RUNNING {
return;
}
// 针对本线程创建一个等待队列的节点
let node = Waiter {
thread: Cell::new(Some(thread::current())),
signaled: AtomicBool::new(false),
//将地址清零后两位后,得到Waiter节点的地址
//如果是头节点,此处的偏移为0,next即是0
next: current_state.with_addr(current_state.addr() & !STATE_MASK) as *const Waiter,
};
//本身作为下一个节点时的地址
let me = &node as *const Waiter as *const Masked as *mut Masked;
//更换当前的state_and_queue,将新创建的node作为队列头
let exchange_result = state_and_queue.compare_exchange(
current_state,
//node的地址与状态做或操作,一个变量即是队列头,又是状态
me.with_addr(me.addr() | RUNNING),
Ordering::Release,
Ordering::Relaxed,
);
//判断是否成功
if let Err(old) = exchange_result {
//不成功,更新current_state,再次循环
//此时
current_state = old;
//此处node会被drop掉
continue;
}
//下面的代码面对一个非常复杂的竞争冲突分析
//作为一个课题留给读者
while !node.signaled.load(Ordering::Acquire) {
//如果没有发信号,则阻塞
thread::park();
//阻塞结束后,进入循环再次判断是否信号已经被接收
}
break;
}
}
impl Drop for WaiterQueue<'_> {
fn drop(&mut self) {
//更新Once的state_and_queue的值,并获取老值,
//这个做法对规避竞争有很大的意义
//即我的孩子我领走。
let state_and_queue =
self.state_and_queue.swap(self.set_state_on_drop_to, Ordering::AcqRel);
// 老值应该只可能是RUNNING状态
assert_eq!(state_and_queue.addr() & STATE_MASK, RUNNING);
unsafe {
//获取等待的队列头地址
let mut queue =
state_and_queue.with_addr(state_and_queue.addr() & !STATE_MASK) as *const Waiter;
while !queue.is_null() {
//保存下一个节点信息
let next = (*queue).next;
//发送信号,唤醒等待的线程
let thread = (*queue).thread.take().unwrap();
(*queue).signaled.store(true, Ordering::Release);
queue = next;
thread.unpark();
}
}
}
}
//留给初始化闭包函数使用。
impl OnceState {
pub fn is_poisoned(&self) -> bool {
self.poisoned
}
pub(crate) fn poison(&self) {
self.set_state_on_drop_to.set(ptr::invalid_mut(POISONED));
}
}
以上展现了RUST为了高效能的一些剑走偏锋的设计技巧。从这些技巧来看,RUST和C绝对是亲爷俩。