Add test_task_priority

Cargo.toml

@@ -88,6 +88,7 @@ exclude = [
 	"applications/test_mlx5",
 	"applications/test_panic",
 	"applications/test_preemption_counter",
+	"applications/test_task_priority",
 	"applications/test_restartable",
 	"applications/test_scheduler",
 	"applications/test_std_fs",

applications/test_task_priority/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "test_task_priority"
+version = "0.1.0"
+authors = ["Klim Tsoutsman <[email protected]>"]
+description = "Test for priority schedulers"
+edition = "2021"
+
+[dependencies]
+cpu = { path = "../../kernel/cpu" }
+fastrand = { version = "2.0.1", default-features = false }
+log = "0.4.8"
+preemption = { path = "../../kernel/preemption" }
+random = { path = "../../kernel/random" }
+spawn = { path = "../../kernel/spawn" }
+task = { path = "../../kernel/task" }

applications/test_task_priority/src/lib.rs

@@ -0,0 +1,90 @@
+//! Test for schedulers supporting priorities i.e. epoch and priority
+//! schedulers.
+//!
+//! The test ensures that tasks are run in order of priority for at least one
+//! time slice.
+
+#![no_std]
+
+extern crate alloc;
+
+use alloc::{string::String, vec::Vec};
+use core::sync::atomic::{AtomicUsize, Ordering};
+
+static CURRENT_PRIORITY: AtomicUsize = AtomicUsize::new(MAX_PRIORITY);
+
+const MAX_PRIORITY: usize = 63;
+
+fn worker(priority: usize) {
+    // Add a bit of chaos.
+    //
+    // NOTE: When using the epoch scheduler, the test relies on the fact that the
+    // worker runs in less than one time slice, and so we can't yield.
+    #[cfg(priority_scheduler)]
+    task::schedule();
+
+    let previous = CURRENT_PRIORITY.fetch_sub(1, Ordering::Relaxed);
+    assert_eq!(previous, priority);
+}
+
+fn spawner(_: ()) {
+    if !task::scheduler::supports_priority() {
+        log::warn!("scheduler does not support priorities");
+        return;
+    }
+
+    let current_cpu = cpu::current_cpu();
+
+    let priorities = priorities();
+    let mut tasks = Vec::with_capacity(MAX_PRIORITY);
+
+    // We hold preemption here so that when the scheduler next runs, all the worker
+    // tasks are unblocked and in a random order on the run queue. Holding
+    // preemption is sufficient as we pin the worker threads to the same core as
+    // the spawner thread.
+    let guard = preemption::hold_preemption();
+
+    for priority in priorities {
+        let task = spawn::new_task_builder(worker, priority)
+            .pin_on_cpu(current_cpu)
+            .block()
+            .spawn()
+            .unwrap();
+        assert!(task::scheduler::set_priority(
+            &task,
+            priority.try_into().unwrap()
+        ));
+        tasks.push(task);
+    }
+
+    for task in tasks.iter() {
+        task.unblock().unwrap();
+    }
+
+    drop(guard);
+
+    for task in tasks {
+        matches!(task.join().unwrap(), task::ExitValue::Completed(_));
+    }
+}
+
+/// Returns a shuffled list of priorities.
+fn priorities() -> Vec<usize> {
+    let mut priorities = (0..=MAX_PRIORITY).collect::<Vec<_>>();
+
+    let mut rng = fastrand::Rng::with_seed(random::next_u64());
+    rng.shuffle(&mut priorities);
+
+    priorities
+}
+
+pub fn main(_: Vec<String>) -> isize {
+    let current_cpu = cpu::current_cpu();
+    // The spawning thread must be pinned to the same CPU as the worker threads.
+    let task = spawn::new_task_builder(spawner, ())
+        .pin_on_cpu(current_cpu)
+        .spawn()
+        .unwrap();
+    matches!(task.join().unwrap(), task::ExitValue::Completed(_));
+    0
+}

kernel/task/src/scheduler.rs

@@ -267,6 +267,11 @@ pub fn inherit_priority(task: &TaskRef) -> PriorityInheritanceGuard<'_> {
     }
 }
 
+/// Returns whether the current scheduler supports priorities.
+pub fn supports_priority() -> bool {
+    SCHEDULER.update(|scheduler| scheduler.as_ref().unwrap().lock().as_priority_scheduler().is_some())
+}
+
 /// A guard that lowers a task's priority back to its previous value when dropped.
 pub struct PriorityInheritanceGuard<'a> {
     inner: Option<(&'a TaskRef, u8)>,

theseus_features/Cargo.toml

@@ -67,6 +67,7 @@ test_scheduler = { path = "../applications/test_scheduler", optional = true }
 test_std_fs = { path = "../applications/test_std_fs", optional = true }
 test_sync_block = { path = "../applications/test_sync_block", optional = true }
 test_task_cancel = { path = "../applications/test_task_cancel", optional = true }
+test_task_priority = { path = "../applications/test_task_priority", optional = true }
 test_tls = { path = "../applications/test_tls", optional = true }
 test_wait_queue = { path = "../applications/test_wait_queue", optional = true }
 test_wasmtime = { path = "../applications/test_wasmtime", optional = true }
@@ -164,6 +165,7 @@ theseus_tests = [
     "test_std_fs",
     "test_sync_block",
     "test_task_cancel",
+    "test_task_priority",
     "test_tls",
     "test_wait_queue",
     "test_wasmtime",