RUST的Once源代码分析

本文摘自《深入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绝对是亲爷俩。

评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

任成珺

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值