「ゼロからのOS自作入門」を Rust でやる (第14章～第16章)

Rust ゼロからのOS自作入門

ブログの更新は間が空いてしまいましたが、OSの移植じたいは細々と続けていました。 今回はまとめて3章分です。

第14章

第13章に引き続きプリエンプティブマルチタスクを実装する章です。 第14章ではタスクのスリープや優先度を実装します。

タスクのスリープ (day14a)

タスクのスリープを実装します。また、 s キー、 w キーの入力でタスクBをスリープ状態にしたり実行状態にしたりします。

github.com

実装は C++ 版とおおよそ同じです。 Rust 版固有の作り込みとして、コンテキストスイッチ時の Arc の取り扱いがあります。

static TASK_MANAGER: OnceCell<Mutex<TaskManager>> = OnceCell::uninit();

#[derive(Debug)]
struct TaskManager {
    tasks: BTreeMap<TaskId, Arc<Task>>,
    wake_queue: VecDeque<TaskId>,
}

impl SwitchTask {
    fn switch(self) {
        assert!(Arc::strong_count(&self.next_task) > 1);
        assert!(Arc::strong_count(&self.current_task) > 1);
        unsafe {
            let next_task_ptr = Arc::as_ptr(&self.next_task);
            let current_task_ptr = Arc::as_ptr(&self.current_task);
            drop(self.next_task);
            drop(self.current_task);
            #[allow(clippy::unwrap_used)]
            let next_task = next_task_ptr.as_ref().unwrap();
            #[allow(clippy::unwrap_used)]
            let current_task = current_task_ptr.as_ref().unwrap();

            Task::switch(next_task, current_task)
        }
    }
}

Task 構造体は TaskManager 構造体の tasks フィールドに Arc<Task> という形で格納されています。 TaskManager 構造体は複数のコンテキストからアクセスするために static 変数 TASK_MANAGER に保持されており、アクセスするためには Mutex ロックの取得が必要です。 このため、コンテキストスイッチ時にロックを取得したままだと次回のコンテキストスイッチ時にロックを取得しようとしてデッドロックになってしまうため、以下のような順序で処理するようにしました。

  1. TASK_MANAGER のロックを取得
  2. 現在実行中のタスクと次にコンテキストスイッチするタスクの Arc<Task> のクローンを取得
  3. TASK_MANAGER のロックを解除
  4. コンテキストスイッチ対象の Arc<Task> から *const Task を取得し、 Arc<Task> は drop する
  5. コンテキストスイッチを実行

4 で *const Task を取得しているのは、 Arc<Task> の drop (参照カウントのデクリメント) をコンテキストスイッチ前に行いたいためです。 4~5 の時点では TaskManager が Arc<Task> を保持しているため、ポインタの指す先の領域が解放されることはないはずです。

イベント発生時にタスクを起床させる & アイドル時にタスクをスリープさせる (day14b)

イベント待ち状態になったタスクをスリープ状態にする & 他タスクに対してイベント通知をする際に当該タスクを起床状態にするという節です。 これにより、OS の応答性が向上します。

C++ 版ではタスク間でイベントをやりとりするためのイベントキュー関連処理でタスクの起床/スリープを行っていました。 Rust 版では async-await の仕組みを採用しているため、 async-await の waker の仕組みと連携できると良さそうです。

まずは、新たに生成されたタスクで async-await が使えるよう、タスクごとにランタイム (executor) を持つようにしました。

github.com

従来は Task::new の引数にタスクのエントリーポイントとなる関数ポインタを渡していました。 修正後は impl Future<Output=()> を渡すようにしています。 エントリポイントは全タスクで共通化し、第一引数 rax に格納されたポインタから Box<EntryPointArg> を復元し、その中に含まれる executor の run メソッドを実行することで CoTask が実行できるようになりました。

次に、 async-wait の waker をタスクの wake/sleep に対応させます。

github.com

waker が CoTask を起床させる際に、当該 CoTask が所属するタスクも起床させるようにしました。 また、各タスクの Executor で処理可能な CoTask が存在しなくなったら自タスクをスリープさせるようにしました。 これだけでイベント発生時のタスク起床 & アイドル時のタスクスリープが実現できます。 簡単ですね！

タスクに優先度をつける (day14c)

個々のタスクに優先度 (レベル) を設け、優先度の高いタスクが優先的に実装されるようにします。

github.com

C++版とは少し実装方法が異なりますが、簡単に実装することができました。

アイドルタスクを追加 (day14d)

すべてのタスクがスリープ状態になった時にCPU消費を抑えられるよう、 hlt を実行するだけの低優先度タスクであるアイドルタスクを追加します。

github.com

バグ修正

Rust 版のタスク優先度設定にはバグがあり、高優先度のタスクが実行可能状態の場合にタイマー割り込み契機のコンテキストスイッチ要求が発生すると、低優先度のタスクにスイッチしてしまうという問題がありましたので修正しました。

github.com

もうひとつのバグとして、割り込みコンテキストでのコンテキストスイッチ時にメモリ獲得を行ってしまう可能性がありました。 タスクの spawn 時に必要な領域獲得を行うようにして、割り込みコンテキストではメモリ獲得が行われないようにしました。

github.com

第15章

端末を実装する章です。

ウインドウ描画をメインタスクで行うようにする (day15a)

Rust 実装では元々描画処理はメインタスクで行うようにしていたため、この節の対応は不要でした。

アクティブウインドウの追加 (day15b)

アクティブウインドウの概念を導入し、タイトルバーの色を変えたり、キー入力イベントの送信先タスクを限定したりします。

まずは、ウインドウ描画処理を共通化します。 具体的には、タイトルバーや枠のあるウインドウを意味する FramedWindow 構造体を追加しました。

github.com

C++ 版では ToplevelWindow という名前でしたが、あまりしっくり来ない名前だったので FramedWindow としています。 また、C++版では ToplevelWindow は Window を継承していましたが、Rust 版では継承ではなくコンポジションでコードの再利用を実現しています (Rust 版では Layer が Window を保持しないため Window 型とのサブタイピング関係が不要であるという事情もあります)。

次にアクティブウインドウを管理する ActiveLayer 構造体を追加します。

github.com

だいたいC++版と同じ実装です。

続いて、ウインドウの状態に応じてタイトルバーの色を変更します。

github.com

ウインドウごとにイベント通知用のキューを保持するようにし、 ActiveLayer が各 Window の active/inactive を変更した時に、当該キュー経由でイベント通知するようにしました。

更にこのキューを利用して、キーボード入力イベントはアクティブウインドウに送信するようにしました。

github.com

これにて最初の節はおしまいです。

ターミナルの追加 (day15c)

ターミナルのウインドウを追加します。

github.com

C++版とだいたい同じです。

描画速度の高速化 (day15d)

従来はウインドウの一部に更新があるとウインドウ全体を再描画していました。 再描画時に必要な範囲だけ描画するよう描画範囲を指定することで描画処理を高速化します。

github.com

C++ 版ではターミナルのカーソル点滅の処理のみ描画範囲指定して高速化していました。 描画が必要な範囲は Window に描画する側のコードで計算する必要があります。 Rust 版では描画処理において再描画が必要な範囲も記憶するようにすることで、描画する側のコードで特殊な考慮をする必要がなくなり、カーソルの点滅以外のすべての描画処理で高速化の恩恵を受けられます。

バグ修正

第15章まで実装したあたりで、並列処理関連のバグ (デッドロックやクラッシュ) が高頻度で発生するようになってしまったため、関連バグを修正しました。

デッドロックが発生した原因は、スピンロックとタスクの優先度の実装方法にあります。 スピンロックは、ロック取得に失敗した場合タスクをスリープさせるのではなく、無限ループで他のタスクがロックを解放するのを待ちます。 また、現時点では高優先度のタスクが実行可能状態で存在する限り、低優先度のタスクは実行されません。 結果として、以下のようなデッドロックが発生する可能性があります。

  1. 低優先度のタスクがスピンロックを取得する
  2. 低優先度のタスクがロックを解放する前に高優先度のタスクにコンテキストスイッチする
  3. 高優先度のタスクが当該ロックを取得しようとすると、いつまで経ってもロックが取得できず待ち続けてしまう

この現象を回避するため、ロック取得失敗時にタスクをスリープさせるような Mutex を実装しました。 (スケジューラーのロジックを変更するのでも良いのですが、 Mutex はいずれにせよ必要になるため実装しました)

github.com

Mutex ロック取得失敗時、タスクIDを Mutex の保持するキューに追加します。 Mutex ロック取得に成功したタスクは、アンロック時にこのキューに含まれるタスクIDのタスクを起床させます。 (上記コミットのソースにはレースコンディションがあったため、後のコミットで修正しています。)

Mutex の実装に SegQue を利用したため、 allocator の排他制御に Mutex は利用できなくなってしまいました。 このため、メモリ割り当て処理中は割り込みを抑止するようにしました。

github.com

Mutex のキューにヒープの領域ではなくスタックの領域を利用するなどすれば allocator の排他制御にも Mutex が利用できるかもしれないですね。

また、複数タスクからロックが取得される可能性がある箇所については Mutex を利用するようにしました。

github.com

ここまでの修正でデッドロックやパニックの発生頻度はかなり低下しました。 ArrayQueue 関連処理が wait-free でないため同様のデッドロックが発生する可能性は残っているのですが、ひとまずはこれでヨシとしました。

ついでに、割り込みコンテキストから allocator が呼び出されたことを検知できるようアサーションも追加しています。

github.com

並列処理が関連すると問題のあるコードがたまたま動いてしまうことが多く、バグがあっても再現性がなくて原因調査が大変だったりするので、アサーションを入れるに越したことはないでしょう。

第16章

ターミナル上での入力やコマンド実行を実装する章です。 全体的にC++実装と同じで難しいことはあまりないので、コメントは省略します。

github.com github.com github.com github.com github.com github.com

まとめ

プリエンプティブマルチタスクと協調的マルチタスク (async-await) をうまく連携させることができました。 Rustらしい実装ができたのではないかと思います。 また、並列処理に関するバグも修正でき、動作の安定度も上がっています。

次章ではファイルシステムを取り扱います。どうなることか。お楽しみに。