Zenn
🐙

Rustのonce_cell::OnceCellとstd::cell::OnceCellの違い

2023/06/07に公開1件
Rust
thread
tech

概要

  • stdとonce_cellではThread Safeに対する表現方法が異なる
  • いくつかのメソッドがstdではNightly
  • stdではLazyCellがNightlyのみ

std::cell::OnceCell

  • get_or_try_init がnightlyのみに存在する
  • !Syncマークが実装されているため、Thread Safeではない。Thread Safeにしたい場合はstd::sync::OnceLockを使う
  • OnceLockのSetはブロックされる
  • OnceLockはStableだが、内部で使用している一部の関数はnightlyとなっている

once_cell::sync::OnceCell

Lazy関係

なぜstd::sync::OnceLockのsetはブロックされるのに対し、once_cell::sync::OnceCellのsetはブロックされない?

once_cell::sync::OnceCellのset実装

    pub fn set(&self, value: T) -> Result<(), T> {
        // SAFETY: Safe because we cannot have overlapping mutable borrows
        let slot = unsafe { &*self.inner.get() };
        if slot.is_some() {
            return Err(value);
        }

        // SAFETY: This is the only place where we set the slot, no races
        // due to reentrancy/concurrency are possible, and we've
        // checked that slot is currently `None`, so this write
        // maintains the `inner`'s invariant.
        let slot = unsafe { &mut *self.inner.get() };
        *slot = Some(value);
        Ok(())
    }

このメソッドの実装を見ると、まず最初にis_someメソッドを使ってセルが既に値を保持しているかどうかを確認します。もし値が既にある場合は、エラーを返して早期に終了します。

そうでない場合、Cellがまだ空の場合は、新しい値をセットします。このとき、一度だけsetが呼び出され、初期化が一度だけ行われることが保証される状態になっています。

つまり、このメソッドはアトミックな操作を行うため、データ競合を引き起こすことなく複数のスレッドから安全に呼び出すことが可能です。ブロックやスピンロックを使用せずにこれを達成しているため、ブロックすることなく処理を進行させることができます。

std::sync::OnceLockのset実装

    #[inline]
    #[stable(feature = "once_cell", since = "1.70.0")]
    pub fn set(&self, value: T) -> Result<(), T> {
        let mut value = Some(value);
        self.get_or_init(|| value.take().unwrap());
        match value {
            None => Ok(()),
            Some(value) => Err(value),
        }
    }

このなかでget_or_initが呼ばれ、get_or_try_initが呼ばれます。その内部で、initializeが実行されています。

  #[cold]
    fn initialize<F, E>(&self, f: F) -> Result<(), E>
    where
        F: FnOnce() -> Result<T, E>,
    {
        let mut res: Result<(), E> = Ok(());
        let slot = &self.value;

        // Ignore poisoning from other threads
        // If another thread panics, then we'll be able to run our closure
        self.once.call_once_force(|p| {
            match f() {
                Ok(value) => {
                    unsafe { (&mut *slot.get()).write(value) };
                }
                Err(e) => {
                    res = Err(e);

                    // Treat the underlying `Once` as poisoned since we
                    // failed to initialize our value. Calls
                    p.poison();
                }
            }
        });
        res
    }

std::lazy::OnceCellのsetメソッドがブロックする可能性があるのは、その内部でget_or_initメソッドを呼び出しているからです。

get_or_initメソッドは、セルが未初期化の場合にのみ初期化関数を呼び出します。しかし、複数のスレッドから同時にget_or_initが呼び出された場合、どのスレッドがセルを初期化するべきかを決定するためにブロッキングが発生します。つまり、このメソッドはセルの初期化が完了するまで他のすべてのスレッドをブロックします。

initializeではcall_once_forceでクロージャを実行されています。call_once_forceでは、AtomicUsizeで状態を管理しています。AtomicUsizeは同期プリミティブなのでスレッド間で操作が競合することはありません。これで原始的なロックを実現しているというわけです。

コラム #[cold] ってなに?

Rustの#[cold]アトリビュートは、その関数があまり頻繁に呼び出されないことをコンパイラにヒントとして与えるものです。このヒントは、コンパイラが生成するコードの最適化に影響を与えます。

具体的には、#[cold]が付けられた関数は、呼び出し回数が少ないと予想されるため、その呼び出しに関連するコードをプログラムの「冷たい」部分（つまり、あまり頻繁に実行されない部分）に配置することがあります。このようにすることで、頻繁に実行されるコードパス（「ホット」パス）のパフォーマンスを最適化するための資源を確保できます。

このようなアトリビュートは、コンパイラが自動的に判断することが難しい場合に、プログラマが手動でパフォーマンスのヒントを提供するための手段です。

逆に、Rustには直接#[hot]アトリビュートが提供されていません。プログラマが特定の関数やコードパスが頻繁に使用される（"ホット"な）と示すための組み込みの手段はありません。

それにもかかわらず、Rustコンパイラはプロファイルガイド付き最適化（PGO）をサポートしており、これを使用すると、プログラムの実行プロファイルに基づいて最適化を行うことができます。これにより、コンパイラはどの部分が"ホット"であるかを自動的に判断することができます。

Discussion