富士通の非公式Advent Calendar 2022 15日目の記事です。記事の内容は個人の見解であり、所属する組織を代表するものではありません。筆者の理解不足により正確でない部分もあるかもしれませんがご了承ください。

はじめに

kernel 6.1 に Rust サポートの実装がマージされました。6.1 にマージされたのは Rust 言語でドライバ(カーネルモジュール)を書くための最小限の機能であり、今のところ、これだけでは Rust を使った本格的なドライバは作れません。Rust for Linux で先行して開発されている Rust サポートの機能が今後もカーネルの master ブランチへマージされていくものと思われます。

今回は (kernel 6.1 ではなく) Rust for Linux の中でも、workqueue の実装を追うことによって、kernel に対する Rust 言語のサポートがどのように実装されているかを調べました。API の使い方というよりも、Rust でドライバを書く人はわざわざ気にする必要のない、 API の内部がどのように実装されているかに注目して調べた記事となります。 Rust for Linux で提供する API は https://rust-for-linux.github.io/docs/kernel/ で公開されています。

workqueue とは

workqueue は Linux kernel における非同期実行の仕組みです。非同期に実行したい処理を work として作成しておき、それをキューに入れると、専用のカーネルスレッド(kworker) がキューからその work を取り出して実行します。

C 言語で書くと以下のような形になります。struct work_struct 型の変数 work を用意しておき、INIT_WORK マクロにより、work を初期化します。初期化の際に、非同期に実行したい関数(ここでは work_handler という名前の自作の関数)を渡します。その work を queue_work() 関数を使ってキューに追加することで、自作の work_handler (例では hello を出力するだけの処理) が非同期に実行されます。queue_work() に渡した system_wq はシステムが用意したキューですが、自作のキューを作成し渡すことも可能です。

#include <linux/module.h>
#include <linux/workqueue.h>
#include <linux/slab.h>

MODULE_LICENSE("GPL");

struct work_struct *work;

static void work_handler(struct work_struct *work) {
        pr_info("hello\n");
}

static int __init my_module_init(void) {
        work = (struct work_struct*)kmalloc(sizeof(struct work_struct), GFP_KERNEL);
        if (!work) {
                return -ENOMEM;
        }
        INIT_WORK(work, work_handler);
        queue_work(system_wq, work);
        return 0;
}

static void __exit my_module_exit(void) {
        if (work) {
                cancel_work_sync(work);
                kfree(work);
        }
}

module_init(my_module_init);
module_exit(my_module_exit);

上記を Rust の API を使って記述すると以下のようになります。

use kernel::prelude::*;
use kernel::workqueue;
use kernel::spawn_work_item;

module! {
    type: RustWorkqueue,
    name: "rust_workqueue_sample",
    license: "GPL",
}

struct RustWorkqueue;

impl kernel::Module for RustWorkqueue {
    fn init(_name: &'static CStr, _module: &'static ThisModule) -> Result<Self> {
        spawn_work_item!(workqueue::system(), || pr_info!("hello\n"))?;
        Ok(RustWorkqueue {})
    }
}

Rust を使った実装でもシスログ(/var/log/messages) に hello というメッセージが出力される結果となります。

前提知識

Linux kernel の workqueue およびカーネルモジュールの基礎知識と、以下のような Rust の基本的な文法については既知のものとします。

構造体(struct)
トレイト(trait)
参照
クロージャ
ライフタイム
クレート
ジェネリクス

確認したソース

Rust for Linux リポジトリの以下の 2022/11/22 時点のコミットのソースを確認対象としました。

https://github.com/Rust-for-Linux/linux/tree/bd123471269354fdd504b65b1f1fe5167cb555fc

今回扱う workqueue のソースはまだカーネルの master ブランチにはマージされておらず、また、今後マージされるかどうかは分かりません。

詳細

Rust 言語で書いた上記のサンプルのコードの内、 workqueue を利用している(workの投入) のは以下の部分です。

spawn_work_item!(workqueue::system(), || pr_info!("hello\n"))?;

Rust を使うと、この1行の中で、 work_struct の初期化と、キューへの投入ができます。

本記事ではこの1行を、以下の3つの要素に分割して、それぞれを詳細に見ていきます。

|| pr_info!("hello\n") → 非同期に実行したい処理
workqueue::system() → カーネルが用意しているキューの取得
spawn_work_item! → キューに対して「非同期に実行したい処理」を投入

`|| pr_info!("hello\n")`

pr_info! は C で書かれた pr_info() マクロに対応する Rust のマクロです。これを呼び出すと hello がシスログに出力されます。その左にある || も含めた、|| pr_info!("hello\n") はクロージャです。別の言い方をすると、引数0で、pr_info!() を呼び出すだけの、無名の関数である、ということです。この関数を、workqueue の機能を利用して非同期に実行したい処理としています。

上の例はクロージャの例としては少し分かりにくいと思いますので、クロージャの別のサンプルを以下に載せておきます。

fn main() {
    let x = 2;
    let func = |y, z| x + y + z;   // =の右側がクロージャ。クロージャを func 変数に保存。
    println!("{}", func(3, 4));    // 9が表示される。
}

func = の右側の部分がクロージャです。|y, z| の部分は、y と z という名前の2つの引数を受け取る、という意味です。その右側に式あるいはブロックを書くことができます(参考)。式を書いた場合は、式を評価した結果がこのクロージャの戻り値となります。最初の例のように || と書けば、それは、定義しようとしているクロージャの引数の数が 0 個であることを意味します。

`workqueue::system()`

workqueue::system() は以下のような定義です。

rust/kernel/workqueue.rs

pub fn system() -> &'static Queue {
    // SAFETY: `system_wq` is a C global, always available.
    unsafe { &*bindings::system_wq.cast() }
}

まず bindings::system_wq に注目します。Rust にはC言語などの別の言語で書かれた関数を呼び出すための機能があります。Rust for Linux では、bindgen を利用して、C の関数呼び出しのインターフェース(FFI)を自動生成し、生成したものを (lib.rs も配置した) rust/bindings ディレクトリ配下に配置しています。system_wq に関しては以下のような実装になっています。

kernel/workqueue.c (Cでの実装)

struct workqueue_struct *system_wq __read_mostly;

rust/bindings/bindings_generated.rs (bindgen によって生成されたインターフェース)

extern "C" {
    pub static mut system_wq: *mut workqueue_struct;
}

この仕組みによって、Rust のコードに bindings::<Cの変数名> や bindings::<Cの関数名> と書くと、C の関数や変数に直接アクセスできます。bindings::system_wq と書いた場合には、C で定義された system_wq 変数に直接アクセスしていることになります。

bindings は APIの内部実装から利用されるもので、API 利用者からは直接使えません。上記の例で言えば、API 利用者からは system() は呼び出せますが、bindings::system_wq は使用できません。

system() 関数は、C 言語における system_wq 変数を型変換だけしてそのまま返却しています。その戻り値は &'static Queue 型です。'static の部分はライフタイムを意味しますので、実質的には、system() は &Queue を返却している、すなわち、「Queue 型の参照」を返している、ということになります。

C言語では struct workqueue_struct * 型であったものを、&Queue 型に変換して返却している、という実装内容から、Queue は、C の struct workqueue_struct に対応する Rust 側の型である、と考えることができます。

'static は、その参照のライフタイムが static であること、つまり、その参照先のデータがプログラム全体に渡って有効であることを意味します。途中で解放される可能性のあるデータを static ライフタイムを持って参照してしまうと、解放済みと知らずにデータにアクセスしてしまう問題が発生する可能性を生んでしまいますが、system_wq は C 言語のグローバルな変数であり、カーネル起動時に初期化される変数であるため、'static なライフタイムでも安全と判断して、 system() の呼び出し元に Queue の参照を返却しています。

system() の実装に話を戻すと、bindings::system_wq は上記の FFI での定義の通り *mut bindings::workqueue_struct 型であり、それを、.cast() により、(戻り値の型に合わせて) Queue の rawポインタ型の *mut Queue に変換しています。さらに、ポインタを参照に変換するために、&* を付けています。rawポインタの解決(deref)は安全な操作ではないこと、および、変換前の型が mutable で static な型であり、多数のスレッドで競合すると問題が発生する可能性があるため、(API実装者が安全であると判断して) unsafe ブロックの中で型の変換を実行しています。

次に Queue 型について見ていきます。

Queue

Queue は上でも触れた通り、C でのstruct workqueue_struct に対応した Rust の型です。以下のように定義されています。

rust/kernel/workqueue.rs

pub struct Queue(Opaque<bindings::workqueue_struct>);

Queue は Opaque<bindings::workqueue_struct> 型の要素を 1つ持つ構造体です。なので、Queue の実体は、Opaque<bindings::workqueue_struct> である、ということになります。 Queue に対して定義されるメソッド(try_spawn() や enqueue()) の実装内容については、後で詳細を説明します。

ここでは先に Queue の定義に使用されている Opaque 型を見てきます。

Opaque

Opaque は rust/kernel/types.rs で定義されている型で、(その使われ方から判断する限りでは) Cの型を Rust における型に対応させるために使用されています。前の例で言うと、C でのstruct workqueue_struct に対応させた Rust の型として Queue を作るために、Opaque を利用しています。

Opaque は以下のような定義です。

rust/kernel/types.rs

pub struct Opaque<T>(MaybeUninit<UnsafeCell<T>>);

Opaque は型パラメータ T を受け取る構造体であり、例えば Opaque<bindings::workqueue_struct> であれば、これ自体が、(C の workqueue_struct に対応する) Rust における型となります。そしてその Opaque は UnsafeCell 型(core ライブラリ)と MaybeUninit 型(core ライブラリ) から成ります。UnsafeCell と MaybeUninit については後で説明します。 Opaque　のメソッドは以下の通り定義されています。

rust/kernel/types.rs

impl<T> Opaque<T> {
    /// Creates a new opaque value.
    pub const fn new(value: T) -> Self {
        Self(MaybeUninit::new(UnsafeCell::new(value)))
    }
        
    /// Creates an uninitialised value.
    pub const fn uninit() -> Self {
        Self(MaybeUninit::uninit())
    }   
    
    /// Returns a raw pointer to the opaque data.
    pub fn get(&self) -> *mut T {
        UnsafeCell::raw_get(self.0.as_ptr())
    }       
}

いずれも UnsafeCell と MaybeUninit のメソッドを呼び出しています。Opaque の特徴は以下の通りです。

Opaque の実体は T 型 (メモリ上は T 型の領域のみ確保)
Opaque の不変参照から、可変なポインタである *mut T を得る手段がある (get)
Opaque の初期化の手段として以下の2つがある
- T型の値を使って初期化(new)
- 値を受け取らずにメモリを未初期化状態として初期化 (uninit)

Opaque が UnsafeCell と MaybeUninit から成る理由としては、筆者の推測も含みますが、以下のような事情からだと思われます。

Rust 側では値を変更する意図はなくても C 側で値を変更される可能性を考慮するため、かつ、Rust から C の関数を呼び出す際にポインタを渡す場合には可変なポインタ(*mut)として渡す必要があるために、UnsafeCell を利用している
Rust では変数を未初期化のまま扱うことはできないが、C ではそれができてしまうので、C に対応して未初期化状態を扱えるようにするために、MaybeUninit を使用している

UnsafeCell (core ライブラリ)

UnsafeCell は core ライブラリが提供する構造体で、内部可変性を扱う型です。Rust では、変数を宣言する時点で、その値あるいは参照を不変なものとして扱うか、可変とするかを、mut の指定の有無で決めます。

fn main() {
    let     val_1 = 1; // 不変
    let mut val_2 = 2; // 可変

    let ref_1 = &val_1;     // 不変 (ref_1 を使用した val_1 の更新は不可)
    let ref_2 = &val_2;     // 不変 (ref_2 を使用した val_2 の更新は不可)
    let ref_3 = &mut val_2; // 可変 (ref_3 を使用した val_2 の更新は可)
}

内部可変性は、変数の宣言時には (mut を指定せずに)その変数を不変なものとして扱うと宣言したにも関わらず、その変数に対して後から可変な操作を可能とする性質のことです。

UnsafeCell が内部可変性を実現するための手段は単純なもので、UnsafeCell 型の変数が不変か可変かに関わらず、その変数に対して、可変な操作が可能な raw ポインタ(*mut T) を得るためのメソッド get() 、raw_get() を定義しているだけです。

pub const fn get(&self) -> *mut T
pub const fn raw_get(this: *const Self) -> *mut T

この raw ポインタのアドレスを解決してそのアドレスに値を書き込むことで可変な操作が可能となります。ポインタの解決は unsafe な操作であるため、UnsafeCell は(コンパイラではなく) その利用者自身が安全性を保証して扱うことを前提に、利用する機能となります。

UnsafeCell のラッパーである、内部可変性を (unsafe ブロック無しに) 安全に扱える Cell 型や RefCell 型というものもあります。

MaybeUninit (core ライブラリ)

MaybeUninit は未初期化状態の変数を扱うための型です。Rust では未初期化変数を読むコードはコンパイル時にエラーとなりますので、未初期化状態のままの変数を扱いたい場合には何かしらの工夫が必要となり、MaybeUninit はその一つの手段です。以下に使用例を載せます。

fn main() {
    let mut x = MaybeUninit::<&i32>::uninit();
    x.write(&3);
    let y = unsafe { x.assume_init() };
}

上記の例では、main 関数内の1行目の let mut x = ...; の宣言の時点での変数 x には MaybeUninit 型の値が入りますので、この意味で Rust から見ると初期化済みの変数に見えており、x のメソッドを呼び出したり他の関数に渡したりということはできますが、その具体的な中身については、未初期化のままです。その次の行で実施しているように、write() メソッドの呼び出し時点で、値を初期化しています。

assume_init() メソッドにより、MaybeUninit の殻を取って、具体的な値を取り出すことができます。この仕組みでは、コンパイラは事前の初期化(write() メソッドの呼び出し)を強制できないため、MaybeUninit の利用者の責任で初期化済みであることを保証し、unsafe な操作によって、assume_init() メソッドを呼び出すことになります。

`spawn_work_item!`

spawn_work_item! は以下のように定義されているマクロです。

rust/kernel/workqueue.rs

macro_rules! spawn_work_item {
    ($queue:expr, $func:expr) => {{
        static CLASS: $crate::sync::LockClassKey = $crate::sync::LockClassKey::new();
        $crate::workqueue::Queue::try_spawn($queue, &CLASS, $func)
    }};
}

マクロの引数として queue および func という名前の2つの式 (expr) を受け取ります。

本筋ではないと思いますので詳細は省略しますが、このマクロの最初に、struct work_struct の初期化に必要なものとして、 C の struct lock_class_key に対応する LockClassKey 型の変数を作成しています。

Queue::try_spawn() が workqueue への処理の投入を実装している部分です。非同期に実行したい処理(func) とそれをどのキュー(queue)で動かしたいかを指定し、また、上記の LockClassKey 型の変数への参照も渡して、Queue::try_spawn() を呼び出しています。

Queue::try_spawn()

Queue::try_spawn は以下のように実装されています。

rust/kernel/workqueue.rs

impl Queue {
...
    pub fn try_spawn<T: 'static + Send + Fn()>(
        &self,
        key: &'static LockClassKey,
        func: T,
    ) -> Result {
        let w = UniqueArc::<ClosureAdapter<T>>::try_new(ClosureAdapter {    //★1 
            // SAFETY: `work` is initialised below.
            work: unsafe { Work::new() },                   //★2
            func,                                           //★3
        })?;
        Work::init(&w, key);                                //★4
        self.enqueue(w.into());                             //★5
        Ok(())
    }

&self で Queue のインスタンスへの参照を受け取ります。関数名の右にある <T: 'static + Send + Fn()> は、 func の型 T が満たすべきトレイト境界が指定されています。それぞれ以下のような性質を持っており、try_spawn() 呼び出し時に指定する T はそのすべてを満す必要があります。

'static は「static ライフタイムを持たない参照」を受け付けないことを意味します。(参考、参考)
Send はスレッド間で安全に移動可能な型であることを意味します(参考)。
Fn() は引数0、戻り値無しのクロージャ(あるいは関数)を意味します。

Send と Fn() の指定は、ClosureAdapter のトレイト境界を満たすために必要です。Fn() を指定しているのは、workqueue に投入したい「非同期に実行したい処理」をクロージャ(あるいは関数)として受け取るためです。ここで受け取ったクロージャは C 言語側の workqueue の実装から呼び出されることになりますが、workqueue の機能(処理の投入元とは別のコンテキストで処理を実行する)から考えて、クロージャがスレッド間を移動して実行されるとみなされ、それでも安全なように　Send を指定しているものと考えられます。また、'static の指定については、C 言語側の workqueue にクロージャを渡している間に、そのクロージャが Rust 側の処理の延長で解放されることが無いようにするためだと考えられます。

try_spawn() 関数内では、UniqueArc::<ClosureAdapter<T>> を try_new() で初期化しています。この中に現れる UniqueArc と ClosureAdapter についてそれぞれ説明しますが、関連する要素も含め、以下の順序で説明します。

Arc
UniqueArc ★
Work と WorkAdapter
ClosureAdapter ★

Arc

Arc (旧名は Ref) は参照カウントの管理およびそれに伴うメモリ管理を行う型です。参照カウンタおよびデータを持つ領域は ArcInner 型のインスタンスとして用意しておき、Arc 自体は ArcInner へのポインタを持っているような構造体です。Arc 型のインスタンスを複製(clone()) すると参照カウントが自動的に1増え、Arc 型のインスタンスを1つ解放 (drop()) すると、参照カウントが自動的に1減ります。参照カウントが 0 になったときに、ArcInner のインスタンスも解放されます。インスタンスに対する drop() 操作は、その利用者が呼び出すのではなく、インスタンスがスコープから外れるタイミングで自動的に呼び出されます。そのため、Arc の利用者から見ると、Arc が不必要になった時点で ArcInner の領域も含め自動的に領域が解放されるように見えます。

実際に実装を見てみると、Arc は以下のように定義されています。

rust/kernel/sync/arc.rs

pub struct Arc<T: ?Sized> {
    ptr: NonNull<ArcInner<T>>,
    _p: PhantomData<ArcInner<T>>,
}

PhantomData についてはここでは詳細を省きますが、PhantomData自体は領域を消費しないため、実質、Arc は NonNull<ArcInner<T>> 型のフィールド ptr を 1つ持つような構造体です。

NonNull は core ライブラリで定義される構造体で、NULL ではないポインタを扱う型です。NonNull<ArcInner<T>> は ArcInner<T> 型の領域へのポインタとなります。

ArcInner の定義は以下の通りです。

rust/kernel/sync/arc.rs

#[repr(C)]
struct ArcInner<T: ?Sized> {
    refcount: Opaque<bindings::refcount_t>,
    data: T,
}

refcount フィールドは、既に現れた Opaque を使って、C における refcount_t 型として定義しています。また、参照カウンタで管理したいデータの型を T で指定し、その T 型の data を持っています。

Arc 型のインスタンスを生成するためのメソッドとして、try_new() が定義されています。実装は以下の通りです。

rust/kernel/sync/arc.rs

impl<T> Arc<T> {
    /// Constructs a new reference counted instance of `T`.
    pub fn try_new(contents: T) -> Result<Self> {
        let layout = Layout::new::<ArcInner<T>>();            //★6
        // SAFETY: The layout size is guaranteed to be non-zero because `ArcInner` contains the
        // reference count.
        let inner = NonNull::new(unsafe { alloc(layout) })    //★7
            .ok_or(ENOMEM)?                                   //★8
            .cast::<ArcInner<T>>();                           //★9

        // INVARIANT: The refcount is initialised to a non-zero value.
        let value = ArcInner {                                //★10
            refcount: Opaque::new(new_refcount()),            //★11
            data: contents,                                   //★12
        };
        // SAFETY: `inner` is writable and properly aligned.
        unsafe { inner.as_ptr().write(value) };               //★13

        // SAFETY: We just created `inner` with a reference count of 1, which is owned by the new
        // `Arc` object.
        Ok(unsafe { Self::from_inner(inner) })                //★14
    }
...
}

★6 の Layout は core ライブラリで定義されたメモリのレイアウトを表現する型で、サイズやアラインメントを表す size() や align() メソッドが定義されています。これらが返却する値は、Layout の初期化関数 new() に指定した型(上のコードでは ArcInner<T>) から決まります。

この Layout 型から、指定の型に対するメモリ獲得に必要な領域サイズなどの情報が得られるため、Layout 型のインスタンスを引数のパラメータとして渡す形でメモリ獲得関数 alloc() を呼び出し、ArcInner<T> 型に対して必要なサイズのメモリを動的に獲得しています(★7)。

alloc 内部では Cのカーネルのインターフェースとして krealloc() が使われており(参考)、カーネルの言葉で言えば、スラブが獲得されることになります。 Rust から見ると安全ではない操作で動的にメモリを獲得していますので、当然ですが、メモリリークが起きないような実装が必要です。これは後に説明する、Clone / Drop トレイトを使って、参照カウントの上げ/下げやメモリ解放をAPI 利用者が意識することなく裏で自動的に行う実装になっています。

alloc が返すのは *mut u8 であるため、NonNull::new(unsafe { alloc(layout) }) は Option<NonNull<u8>> を返します(★7)。これを ok_or() で Result<NonNull<u8>> に変換し、? 演算子により、alloc の成功時には NonNull<u8> を、失敗時には ENOMEM を即時返却するようにしています(★8)。その後に cast() を使って NonNull<u8> を NonNull<ArcInner<T>> 型に変換しています(★9)。

ここまでで、inner 変数の指す先には、ArcInner<T> 型のデータが入る、未初期の領域が作成された状態です。この領域に格納する、ArcInner の中身は ★10 の部分で作っています。

★11 の new_refcount() は以下の通り、 struct refcount_struct (= refcount_t) の refs フィールドを 1 で初期化した値を返すような関数です。これが、Opaque::new() を通じて、 ArcInner の refcount フィールドに設定されます。すなわち、Arc の初期化時点で、その参照カウントは 1 に設定されます。

rust/kernel/sync/arc.rs

pub const fn new_refcount() -> bindings::refcount_struct {
    bindings::refcount_struct {
        refs: bindings::atomic_t { counter: 1 },
    }
}

また、ArcInner の data フィールドは try_new() の引数で受け取った T 型の contents をそのまま与えています(★12)。

このように初期化した ArcInner のインスタンスを、先ほど用意した、未初期化の領域 inner に対して、write で値(value) を書き込んで、初期化しています(★13)。

最後に from_inner という名前のメソッドを利用して、inner を(ptr)フィールドとして持つような新たな Arc のインスタンスを生成して、それを戻り値としてtry_new() メソッドの呼び出し元に返却しています(★14)。

rust/kernel/sync/arc.rs

impl<T: ?Sized> Arc<T> {
...
    unsafe fn from_inner(inner: NonNull<ArcInner<T>>) -> Self {
        // INVARIANT: By the safety requirements, the invariants hold.
        Arc {
            ptr: inner,
            _p: PhantomData,
        }
    }
...
}

以上が Arc の try_new() の実装です。以下、Clone、Drop についても触れます。

Arc 型には Clone トレイトが実装されており、これによって、clone() メソッドの呼び出しにより、Arc 型のインスタンスを複製できます。clone() による複製時に C の refcount_inc() 関数の呼び出しにより、参照カウントを上げています。

rust/kernel/sync/arc.rs

impl<T: ?Sized> Clone for Arc<T> {
    fn clone(&self) -> Self {
        ...
        unsafe { bindings::refcount_inc(self.ptr.as_ref().refcount.get()) };
        ...
    }
}

また、Drop トレイトの drop() メソッドの実装により、スコープから外れるなど Arc 型のインスタンスが不要になったタイミングで、C の refcount_dec_and_test() 関数により参照カウントを1つ下げ、その結果が 0 になったときに、Arc が保持する ArcInner の領域を解放(kfree)しています。

rust/kernel/sync/arc.rs

impl<T: ?Sized> Drop for Arc<T> {
    fn drop(&mut self) {
        ...
        let refcount = unsafe { self.ptr.as_ref() }.refcount.get();
        ...
        let is_zero = unsafe { bindings::refcount_dec_and_test(refcount) };
        if is_zero {
            ...
            unsafe { dealloc(self.ptr.cast().as_ptr(), layout) };     // 内部で kfree() を呼び出す
        }
    }
}

UniqueArc

UniqueArc (旧名は UniqueRef) は、Arc 型でも特に、参照カウントが 1 であることが保証されるような型です。構造体の定義としては、以下のように、Arc<T> 型のフィールドを1つ持つだけです。

rust/kernel/sync/arc.rs

pub struct UniqueArc<T: ?Sized> {
    inner: Arc<T>,
}

これを初期化するためのメソッド try_new() は以下のように定義されており、単純に、Arc 型の try_new() を呼び出しているだけです。

rust/kernel/sync/arc.rs

impl<T> UniqueArc<T> {
    /// Tries to allocate a new [`UniqueArc`] instance.
    pub fn try_new(value: T) -> Result<Self> {
        Ok(Self {
            // INVARIANT: The newly-created object has a ref-count of 1.
            inner: Arc::try_new(value)?,
        })
    }
...
}

Arc 型の生成直後は、他にその領域を参照する人がいないことが保証できるため、参照カウントが 1 であり、その状態で、UniqueArc (のフィールド inner) を初期化しています。以下の実装により、UniqueArc は Arc 型に変換可能です。

rust/kernel/sync/arc.rs

impl<T: ?Sized> From<UniqueArc<T>> for Arc<T> {
    fn from(item: UniqueArc<T>) -> Self {
        item.inner
    }
}

UniqueArc は Arc とは異なり Clone トレイトが実装されていないため、UniqueArc の複製はできず、データを参照する人が1人であることが保証されます。また、参照者が一人しかいないということで、安全に可変参照を取得できます。

UniqueArc は Queue::try_spawn() の中で、UniqueArc::<ClosureAdapter<T>>::try_new() のような形で利用されています(★1)。これは、その直後の Work::init() の呼び出し時(★4)の引数として、 UniqueArc 型が要請されているためです。Arc 型があるにも関わらず UniqueArc 型も定義されている理由としては、Rust 側で獲得した領域を C 言語側の関数を呼び出して更新する際に、他に参照する人がいないことを型で強制して、初期化の安全性を保証するためだと考えられます。

`Work` と `WorkAdapter`

Rust の Work 構造体は、C での struct work_struct に対応します。

rust/kernel/workqueue.rs

pub struct Work(Opaque<bindings::work_struct>);

Work 構造体に対して定義されている、workqueue の動作を説明する上で重要なメソッドの実装を以下に載せます。

rust/kernel/workqueue.rs

impl Work {
    pub fn init<T: WorkAdapter<Target = T>>(obj: &UniqueArc<T>, key: &'static LockClassKey) {
        Self::init_with_adapter::<T>(obj, key)           // ★15
    }
    ...
    pub fn init_with_adapter<A: WorkAdapter>(
        obj: &UniqueArc<A::Target>,
        key: &'static LockClassKey,
    ) {
        let ptr = &**obj as *const _ as *const u8;       // ★16
        let field_ptr = ptr.wrapping_offset(A::FIELD_OFFSET) as *mut bindings::work_struct;            // ★17
        ...
        unsafe {
            bindings::__INIT_WORK_WITH_KEY(field_ptr, Some(Self::work_func::<A>), false, key.get())    // ★18
        };
    }
    ...
    unsafe extern "C" fn work_func<A: WorkAdapter>(work: *mut bindings::work_struct) {
        let field_ptr = work as *const _ as *const u8;                                // ★19
        let ptr = field_ptr.wrapping_offset(-A::FIELD_OFFSET) as *const A::Target;    // ★20

        // SAFETY: This callback is only ever used by the `init_with_adapter` method, so it is
        // always the case that the work item is embedded in a `Work` (Self) struct.
        let w = unsafe { Arc::from_raw(ptr) };                                        // ★21
        A::run(w);                                                                    // ★22
    }
    ...
}

C で workqueue を利用する場合、非同期に処理を実行したい処理は、struct work_struct の func フィールド(関数ポインタ)にその処理(関数のアドレス)を設定します。

include/linux/workqueue.h

struct work_struct {
        atomic_long_t data;
        struct list_head entry;
        work_func_t func;  //★
#ifdef CONFIG_LOCKDEP
        struct lockdep_map lockdep_map;
#endif
};
...
typedef void (*work_func_t)(struct work_struct *work);

Rust から C の workqueue を利用する場合、上記の func フィールド(C)には、 Work 構造体の work_func() メソッド (Rust) が固定で登録されます。具体的には、★18 の、struct work_struct 初期化用のマクロ __INIT_WORK_WITH_KEY の呼び出しによって、func フィールドには、work_func::<A> メソッドが設定されます。

本来、workqueue の利用者は、非同期に実行したい処理を func フィールドに設定したいはずのところを、Rust 実装では work_func メソッド固定で設定しているので、利用者の処理はどこに書けばよいのか、が気になりますが、それは、★22 に現れる、A::run() メソッドに記述します。 work_func メソッドの定義箇所で <A: WorkAdapter> と指定されている通り、A は WorkAdapter トレイトを実装した任意の型です。WorkAdapter トレイトの run() メソッドを workqueue の利用者が実装する、という使い方です。

WorkAdapter は run() メソッド実装以外にも、もう一つ役割があり、「Work (struct work_struct) をフィールドに持つ構造体」に対応するような概念になっています。 WorkAdapter の具体的な定義は以下の通りです。

rust/kernel/workqueue.rs

pub unsafe trait WorkAdapter {
    /// The type that this work adapter is meant to use.
    type Target;

    /// The offset, in bytes, from the beginning of [`Self::Target`] to the instance of [`Work`].
    const FIELD_OFFSET: isize;

    /// Runs when the work item is picked up for execution after it has been enqueued to some work
    /// queue.
    fn run(w: Arc<Self::Target>);
}

Target には、「Work (struct work_struct) をフィールドに持つ構造体」が指定され、その構造体の先頭から見た Work 型のフィールドの位置(オフセット)が FIELD_OFFSET に入ります。

例えば、

#[repr(C)]
struct Example {
    work: Work,
    ...
}

に対して、WorkAdapter トレイトを実装する場合は、Target に Example, FIELD_OFFSET に 0が入ります。また、 run() メソッドは、Target として指定した Example (を Arc の参照カウンタで管理しているもの)を受け取る関数として、workqueue 利用者が具体的に記述します。

Work::work_func() の実装に戻ると、このメソッドでは、Work のアドレスを (引数で受け取った Work への参照に対するキャストにより) 求め(★19)、そのアドレスおよび WorkAdapter が持つ情報(FIELD_OFFSET, Target) を利用して、Work をフィールドとして持つ構造体(上の例では Example 構造体)のアドレスを ptr に保存しています(★20)。その後、ptr に対して、Arc::from_raw() を呼び出した(★21)後に、workqueue で実際に実行したい処理が書かれた WorkAdapter::run() 関数を呼び出しています(★22)。

★21 の Arc::from_raw() は Arc::into_raw() と対になって使うものです。後のコードにも現れますが、 workqueue で実行したい処理が含まれる Work (をフィールドとして持つ構造体) をキューに投入する前には、Work に対して Arc::into_raw() を呼び出しています(★23)。キューに投入する操作は C 言語側の実装に Rust 側で作成した Work を預けることになりますので、(C 言語側が参照している間に Rust 側で壊さないように) Rust 側での管理を一旦外し、C言語側でポインタとして扱えるように、Arc::into_raw() を呼び出しています。外した管理をまた元に戻して Rust 側のインスタンスとして構築しなおすのが Arc::from_raw() の役割であり、★21 でそれを実行しています。Arc のインスタンスが持つ参照カウントの情報も復元され、★22 の直後で (特に他に Arc を複製していなければ) Arc がスコープから外れたことによりその参照カウントが 1 から　0 に変化し、同時に、Arc が参照していたデータ領域も解放されることになります。 Rust を使わない場合は、使用済みの struct work_struct の領域解放処理は workqueue の利用者が自分で記述する必要ありましたが、Rust では Arc の管理のおかげで、利用者が明に解放処理を呼び出す必要はありません。

Work を初期化するメソッドとして init() および init_with_adapter() メソッドがあります。init() は init_with_adapter() を呼び出すだけです。

init_with_adapter() の中身を見ると、★16 の obj は、「参照カウントで管理される、Work をフィールドに持つ構造体、への参照」です。その参照を as *const _ as *const u8 によるキャストでポインタに変換した後、 WorkAdapter が持っている Work フィールドのオフセット情報 (FIELD_OFFSET) を使って、Work のアドレスも取得しています(★17)。

ということで、Rust における Work 型のインスタンスの作成には、以下の 1., 2. が必要となります。

Work をフィールドに持つ構造体
WorkAdapter トレイトを実装した構造体 (1.の構造体を Target に指定し、かつ、非同期に実行したい関数 run() を実装する)

1.の構造体および2のrun()の実装を使って、Work を作成(初期化)し、その
Work を Queue 型のインスタンスに投入して、workqueue の処理が実行される、という流れです。

上記の 1.,2. は別の構造体でも、同じ構造体を使うこともできます。次に説明する ClosureAdapter 構造体は、上記の 1., 2. を両方満たすような構造体です。

ClosureAdapter

ClosureAdapter は以下のように、「Work をフィールドに持つ構造体」です。

rust/kernel/workqueue.rs

struct ClosureAdapter<T: Fn() + Send> {
    work: Work,
    func: T,
}

また、以下のように、「WorkAdapter トレイトを実装した構造体」でもあります。

rust/kernel/workqueue.rs

unsafe impl<T: Fn() + Send> WorkAdapter for ClosureAdapter<T> {
    type Target = Self;
    const FIELD_OFFSET: isize = crate::offset_of!(Self, work);

    fn run(w: Arc<Self::Target>) {
        (w.func)();
    }
}

ClosureAdapter は、引数0、戻り値無しのクロージャ(あるいは関数)によって初期化し、run() の中ではそのクロージャをそのまま呼び出すだけの、WorkAdapter です。

Queue::try_spawn() の続き

UniqueArc や ClosureAdapter について詳細が分かりましたので、Queue::try_spawn() の実装を見ていきます。

try_spawn() の最初に UniqueArc::<ClosureAdapter<T>> 型のインスタンスの初期化をしています(★1)。Work は未初期化状態で作成しておき(★2)、また、Queue::try_spawn() の引数として受け取ったクロージャ func を使って(★3)、ClosureAdapter 構造体のインスタンスを作成し、同時に、UniqueArc としても初期化しています。

その後 Work の中身を、UniqueArc::<ClosureAdapter<T>> (および key) のインスタンスを使って、Work::init() の呼び出しにより、初期化しています(★4)。

このように Work を初期化した後に、 Queue 構造体に実装されている enqueue() メソッドを使って workqueue に Work (をフィールドに持つ構造体のインスタンス) を投入しています(★5)。

最後に enqueue() の中身を説明して、workqueue の Rust での実現方法の確認は終わりになります。

Queue::enqueue()

Queue::enqueue() は workqueue への Work の投入を行うメソッドです。投入には C 側で定義されている queue_work_on 関数を呼び出しています。

実装は以下の通りです。

rust/kernel/workqueue.rs

impl Queue {
    ...
    /// Enqueues a work item.
    ///
    /// Returns `true` if the work item was successfully enqueue; returns `false` if it had already
    /// been (and continued to be) enqueued.
    pub fn enqueue<T: WorkAdapter<Target = T>>(&self, w: Arc<T>) -> bool {
        self.enqueue_adapter::<T>(w)
    }

    /// Enqueues a work item with an explicit adapter.
    ///
    /// Returns `true` if the work item was successfully enqueue; returns `false` if it had already
    /// been (and continued to be) enqueued.
    pub fn enqueue_adapter<A: WorkAdapter + ?Sized>(&self, w: Arc<A::Target>) -> bool {
        let ptr = Arc::into_raw(w);       // ★23
        let field_ptr =
            (ptr as *const u8).wrapping_offset(A::FIELD_OFFSET) as *mut bindings::work_struct;     // ★24

        // SAFETY: Having a shared reference to work queue guarantees that it remains valid, while
        // the work item remains valid because we called `into_raw` and only call `from_raw` again
        // if the object was already queued (so a previous call already guarantees it remains
        // alive), when the work item runs, or when the work item is canceled.
        let ret = unsafe {
            bindings::queue_work_on(bindings::WORK_CPU_UNBOUND as _, self.0.get(), field_ptr)      // ★25
        };

        if !ret {
            // SAFETY: `ptr` comes from a previous call to `into_raw`. Additionally, given that
            // `queue_work_on` returned `false`, we know that no-one is going to use the result of
            // `into_raw`, so we must drop it here to avoid a reference leak.
            unsafe { Arc::from_raw(ptr) };
        }

        ret
    }
    ...
}

enqueue() は enqueue_adapter() を呼び出すだけです。

enqueue_adapter() では、まず、C 言語に Work を渡す前の準備として、Arc::into_raw() を呼び出しています(★23)。少し前にも説明しました通り、 workqueue への投入で Work の管理が一旦 C に移りますので、Rust の管理から外しています。

enqueue_adapter() の引数で受け取るインスタンスは Work をフィールドに持つ構造体ですので、WorkAdapter の持つ FIELD_OFFSET 情報から Work のアドレスを特定します(★24)。

その後は C で定義された queue_work_on() 関数を呼び出しています(★25)。その引数として、Queue (C での struct workqueue_struct) のアドレスや、Work (C での struct work_struct) のアドレスを渡しています。

queue_work_on() が呼ばれれば、後は、C 側の workqueue の実装に従って、非同期に実行したかった処理 (本記事の実行例では pr_info!("hello\n")) が呼ばれます。また、Arc の参照カウントによるメモリ管理により、Work で使用していた領域は、hello メッセージ出力後に自動的に解放されます。

まとめ

カーネルの workqueue の Rust による実装を少し詳しく追ってみました。Rust を使うと、API 利用者から見ると workqueue は簡単に扱えるのですが、API 内部の実装は思ったよりも複雑でした。

API利用者から見ると、(APIの実装にバグがない前提で) Rust を使えば C よりも安全にかつ簡単にカーネルのドライバが実装できる、というのはその通りだと思います。ただ、API の実装に関しては、C 側の挙動を意識してメモリ安全性やスレッド安全性を保証しながら、Rust の unsafe 操作を駆使して実装が必要になりますので、バグが無いように実装するのは大変そうだと感じました。

参考