feat(sched): add CPU core affinity via ClusteredEDF scheduler

applications/tests/test_partitioned_edf/Cargo.toml → applications/tests/test_clustered_edf/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "test_partitioned_edf"
+name = "test_clustered_edf"
 version = "0.1.0"
 edition = "2021"
 

applications/tests/test_partitioned_edf/src/lib.rs → applications/tests/test_clustered_edf/src/lib.rs

@@ -7,7 +7,7 @@ use alloc::string::ToString;
 use awkernel_async_lib::sleep;
 use awkernel_async_lib::{scheduler::SchedulerType, spawn};
 use awkernel_lib::{
-    cpu::{cpu_id, num_cpu},
+    cpu::{cpu_id, num_cpu, CpuSet},
     delay::{uptime, wait_microsec},
 };
 use core::time::Duration;
@@ -16,24 +16,24 @@ pub async fn run() {
     wait_microsec(2_000_000);
 
     if num_cpu() < 2 {
-        log::warn!("test_partitioned_edf: requires at least 2 CPUs, skipping");
+        log::warn!("test_clustered_edf: requires at least 2 CPUs, skipping");
         return;
     }
 
     test_core_pinning().await;
+    test_cluster_affinity().await;
     test_edf_preemption().await;
     test_multi_core().await;
 
     wait_microsec(100_000_000);
 }
 
 /// Test 1: Verify tasks are pinned to their assigned cores.
-/// Spawns one periodic task per worker core (1..num_cpu()).
-/// Each task checks cpu_id() matches the assigned core.
+/// Spawns one periodic task per worker core (1..num_cpu()) with a
+/// single-core CPU set. Each task checks cpu_id() matches the assigned core.
 async fn test_core_pinning() {
     log::info!("=== test_core_pinning start ===");
     for core in 1..num_cpu() {
-        let core_u16 = core as u16;
         spawn(
             alloc::format!("pinning_core{core}").into(),
             async move {
@@ -49,21 +49,52 @@ async fn test_core_pinning() {
                     awkernel_async_lib::r#yield().await;
                 }
             },
-            SchedulerType::PartitionedEDF(1_000_000, core_u16),
+            SchedulerType::ClusteredEDF(1_000_000, CpuSet::empty().insert(core)),
         )
         .await;
     }
 }
 
-/// Test 2: Verify EDF preemption on a single core.
+/// Test 2: Verify a task with a multi-core CPU set only runs on cores
+/// within the set.
+async fn test_cluster_affinity() {
+    if num_cpu() < 4 {
+        log::info!("test_cluster_affinity: skipped (num_cpu < 4)");
+        return;
+    }
+    log::info!("=== test_cluster_affinity start ===");
+
+    // Cluster {1, 2}: the task must never run on other cores.
+    let cluster = CpuSet::empty().insert(1).insert(2);
+    spawn(
+        "cluster_1_2".into(),
+        async move {
+            for _ in 0..5 {
+                let actual = cpu_id();
+                if cluster.contains(actual) {
+                    log::info!("cluster_affinity: task ran on cpu {actual} [OK]");
+                } else {
+                    log::error!("cluster_affinity: task ran on cpu {actual} [FAIL]");
+                }
+                wait_microsec(100_000);
+                sleep(Duration::from_millis(200)).await;
+                awkernel_async_lib::r#yield().await;
+            }
+        },
+        SchedulerType::ClusteredEDF(1_000_000, cluster),
+    )
+    .await;
+}
+
+/// Test 3: Verify EDF preemption on a single core.
 /// heavy (deadline=9900ms) is preempted by light (deadline=990ms) on core 1.
 async fn test_edf_preemption() {
     log::info!("=== test_edf_preemption start ===");
     spawn_periodic_task("heavy".to_string(), 9_600_000, 10_000_000, 9_900_000, 1).await;
     spawn_periodic_task("light".to_string(), 900_000, 1_000_000, 990_000, 1).await;
 }
 
-/// Test 3: Verify core independence when num_cpu >= 3.
+/// Test 4: Verify core independence when num_cpu >= 3.
 /// Tasks on core 1 and core 2 run in parallel without interfering.
 async fn test_multi_core() {
     if num_cpu() < 3 {
@@ -81,7 +112,7 @@ async fn spawn_periodic_task(
     exe_time: u64,
     period: u64,
     relative_deadline: u64,
-    core: u16,
+    core: usize,
 ) {
     let task_name_clone = task_name.clone();
     spawn(
@@ -102,7 +133,7 @@ async fn spawn_periodic_task(
                 awkernel_async_lib::r#yield().await;
             }
         },
-        SchedulerType::PartitionedEDF(relative_deadline, core),
+        SchedulerType::ClusteredEDF(relative_deadline, CpuSet::empty().insert(core)),
     )
     .await;
 }

awkernel_async_lib/Cargo.toml

@@ -12,6 +12,7 @@ array-macro = "2.1"
 fixedbitset = "0.5.7"
 async-trait = "0.1"
 async-recursion = "1.1"
+affinity_btree_queue = "0.1"
 
 [dependencies.futures]
 version = "0.3"

awkernel_async_lib/src/scheduler.rs

@@ -10,16 +10,16 @@ use alloc::collections::{binary_heap::BinaryHeap, btree_map::BTreeMap};
 use alloc::sync::Arc;
 use awkernel_async_lib_verified::delta_list::DeltaList;
 use awkernel_lib::{
-    cpu::num_cpu,
+    cpu::{num_cpu, CpuSet},
     sync::mutex::{MCSNode, Mutex},
 };
 
 #[cfg(not(feature = "std"))]
 use alloc::boxed::Box;
 
+mod clustered_edf;
 pub mod gedf;
 pub(super) mod panicked;
-mod partitioned_edf;
 mod prioritized_fifo;
 mod prioritized_rr;
 
@@ -73,8 +73,8 @@ pub fn move_preemption_pending(cpu_id: usize) -> Option<BinaryHeap<Arc<Task>>> {
 /// 0 is the lowest priority and 31 is the highest priority.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum SchedulerType {
-    PartitionedEDF(u64, u16), // relative deadline and partitioned core
-    GEDF(u64),                // relative deadline
+    ClusteredEDF(u64, CpuSet), // relative deadline and CPU affinity set
+    GEDF(u64),                 // relative deadline
     PrioritizedFIFO(u8),
     PrioritizedRR(u8),
     Panicked,
@@ -86,8 +86,8 @@ impl SchedulerType {
             (self, other),
             (SchedulerType::GEDF(_), SchedulerType::GEDF(_))
                 | (
-                    SchedulerType::PartitionedEDF(_, _),
-                    SchedulerType::PartitionedEDF(_, _)
+                    SchedulerType::ClusteredEDF(_, _),
+                    SchedulerType::ClusteredEDF(_, _)
                 )
                 | (
                     SchedulerType::PrioritizedFIFO(_),
@@ -101,13 +101,13 @@ impl SchedulerType {
         )
     }
 
-    /// Return the partitioned core index if this is a [`SchedulerType::PartitionedEDF`] scheduler.
+    /// Return the CPU affinity set if this is a [`SchedulerType::ClusteredEDF`] scheduler.
     ///
-    /// Returns `Some(n)` where `n` is the CPU core index (`1..num_cpu()`) assigned to the
-    /// partitioned EDF scheduler. Returns `None` for all other scheduler variants.
-    pub const fn partitioned_core(&self) -> Option<u16> {
+    /// Returns `Some(set)` where `set` is the set of CPU cores (`1..num_cpu()`) the task
+    /// may run on. Returns `None` for all other scheduler variants.
+    pub const fn cpu_set(&self) -> Option<CpuSet> {
         match self {
-            SchedulerType::PartitionedEDF(_, n) => Some(*n),
+            SchedulerType::ClusteredEDF(_, set) => Some(*set),
             _ => None,
         }
     }
@@ -118,7 +118,7 @@ impl SchedulerType {
 /// `priority()` returns the priority of the scheduler for preemption.
 ///
 /// - The highest priority.
-///   - Partitioned EDF scheduler.
+///   - Clustered EDF scheduler.
 /// - The second highest priority.
 ///   - GEDF scheduler.
 /// - The third highest priority.
@@ -128,19 +128,19 @@ impl SchedulerType {
 /// - The lowest priority.
 ///   - Panicked scheduler.
 static PRIORITY_LIST: [SchedulerType; 5] = [
-    SchedulerType::PartitionedEDF(0, 0),
+    SchedulerType::ClusteredEDF(0, CpuSet::empty()),
     SchedulerType::GEDF(0),
     SchedulerType::PrioritizedFIFO(0),
     SchedulerType::PrioritizedRR(0),
     SchedulerType::Panicked,
 ];
 
-/// Return the number of partitioned schedulers in `PRIORITY_LIST`.
-/// Update this function if you add a new partitioned scheduler to `PRIORITY_LIST`.
-const fn get_num_partitioned_schedulers() -> usize {
+/// Return the number of clustered schedulers in `PRIORITY_LIST`.
+/// Update this function if you add a new clustered scheduler to `PRIORITY_LIST`.
+const fn get_num_clustered_schedulers() -> usize {
     let mut count = 0;
     while count < PRIORITY_LIST.len() {
-        if matches!(PRIORITY_LIST[count], SchedulerType::PartitionedEDF(_, _)) {
+        if matches!(PRIORITY_LIST[count], SchedulerType::ClusteredEDF(_, _)) {
             count += 1;
         } else {
             break;
@@ -177,27 +177,29 @@ pub(crate) fn get_next_task(execution_ensured: bool) -> Option<Arc<Task>> {
     let mut node = MCSNode::new();
     let _guard = GLOBAL_WAKE_GET_MUTEX.lock(&mut node);
 
-    let num_partitioned_tasks =
-        crate::task::NUM_PARTITIONED_TASKS_IN_QUEUE[cpu_id].load(Ordering::Relaxed);
+    let num_clustered_tasks =
+        crate::task::NUM_CLUSTERED_TASKS_IN_QUEUE[cpu_id].load(Ordering::Relaxed);
 
-    if num_partitioned_tasks > 0 {
-        let task = PRIORITY_LIST[..get_num_partitioned_schedulers()]
-            .iter()
-            .find_map(|&scheduler_type| get_scheduler(scheduler_type).get_next(execution_ensured));
-
-        // Decrement is handled by PartitionedTask::Drop inside get_next().
-        task
-    } else {
-        let task = PRIORITY_LIST
+    if num_clustered_tasks > 0 {
+        let task = PRIORITY_LIST[..get_num_clustered_schedulers()]
             .iter()
             .find_map(|&scheduler_type| get_scheduler(scheduler_type).get_next(execution_ensured));
 
+        // Decrement is handled by ClusteredTask::Drop inside get_next().
         if task.is_some() {
-            crate::task::NUM_TASK_IN_QUEUE.fetch_sub(1, Ordering::Relaxed);
+            return task;
         }
+    }
+
+    let task = PRIORITY_LIST
+        .iter()
+        .find_map(|&scheduler_type| get_scheduler(scheduler_type).get_next(execution_ensured));
 
-        task
+    if task.is_some() {
+        crate::task::NUM_TASK_IN_QUEUE.fetch_sub(1, Ordering::Relaxed);
     }
+
+    task
 }
 
 /// Get a scheduler.
@@ -206,7 +208,7 @@ pub(crate) fn get_scheduler(sched_type: SchedulerType) -> &'static dyn Scheduler
         SchedulerType::PrioritizedFIFO(_) => &prioritized_fifo::SCHEDULER,
         SchedulerType::PrioritizedRR(_) => &prioritized_rr::SCHEDULER,
         SchedulerType::GEDF(_) => &gedf::SCHEDULER,
-        SchedulerType::PartitionedEDF(_, _) => &partitioned_edf::SCHEDULER,
+        SchedulerType::ClusteredEDF(_, _) => &clustered_edf::SCHEDULER,
         SchedulerType::Panicked => &panicked::SCHEDULER,
     }
 }
@@ -222,72 +224,57 @@ pub const fn get_priority(sched_type: SchedulerType) -> u8 {
     panic!("Scheduler type not registered in PRIORITY_LIST or equals()")
 }
 
-/// RAII wrapper representing one slot in `NUM_PARTITIONED_TASKS_IN_QUEUE[cpu_id]`.
+/// RAII wrapper representing one slot in `NUM_CLUSTERED_TASKS_IN_QUEUE` for every
+/// CPU in `cpu_set`.
 ///
-/// Constructing a `PartitionedTask` increments the counter for `cpu_id`.
-/// The counter is decremented exactly once — either via [`take`] (explicit
+/// Constructing a `ClusteredTask` increments the counter of each CPU in `cpu_set`.
+/// The counters are decremented exactly once — either via [`take`] (explicit
 /// ownership transfer) or via `Drop` (e.g. when a terminated task is
 /// discarded). If `take` has already been called, `Drop` is a no-op.
 ///
-/// [`take`]: PartitionedTask::take
-pub(crate) struct PartitionedTask<T> {
+/// [`take`]: ClusteredTask::take
+pub(crate) struct ClusteredTask<T> {
     inner: Option<T>,
-    cpu_id: usize,
+    cpu_set: CpuSet,
 }
 
-impl<T> PartitionedTask<T> {
-    pub(crate) fn new(inner: T, cpu_id: usize) -> Self {
-        crate::task::NUM_PARTITIONED_TASKS_IN_QUEUE[cpu_id].fetch_add(1, Ordering::Relaxed);
+impl<T> ClusteredTask<T> {
+    pub(crate) fn new(inner: T, cpu_set: CpuSet) -> Self {
+        for cpu_id in cpu_set.iter() {
+            crate::task::NUM_CLUSTERED_TASKS_IN_QUEUE[cpu_id].fetch_add(1, Ordering::Relaxed);
+        }
         Self {
             inner: Some(inner),
-            cpu_id,
+            cpu_set,
         }
     }
 
-    /// Take the inner value and decrement the counter.
+    /// Take the inner value and decrement the counters.
     ///
     /// Returns `None` if the value has already been taken.
     /// After this call `Drop` will be a no-op.
     pub(crate) fn take(&mut self) -> Option<T> {
         let val = self.inner.take();
         if val.is_some() {
-            crate::task::NUM_PARTITIONED_TASKS_IN_QUEUE[self.cpu_id]
-                .fetch_sub(1, Ordering::Relaxed);
+            for cpu_id in self.cpu_set.iter() {
+                crate::task::NUM_CLUSTERED_TASKS_IN_QUEUE[cpu_id].fetch_sub(1, Ordering::Relaxed);
+            }
         }
         val
     }
 }
 
-impl<T> Drop for PartitionedTask<T> {
+impl<T> Drop for ClusteredTask<T> {
     fn drop(&mut self) {
         // Decrement only if take() has not been called yet.
         if self.inner.is_some() {
-            crate::task::NUM_PARTITIONED_TASKS_IN_QUEUE[self.cpu_id]
-                .fetch_sub(1, Ordering::Relaxed);
+            for cpu_id in self.cpu_set.iter() {
+                crate::task::NUM_CLUSTERED_TASKS_IN_QUEUE[cpu_id].fetch_sub(1, Ordering::Relaxed);
+            }
         }
     }
 }
 
-impl<T: PartialEq> PartialEq for PartitionedTask<T> {
-    fn eq(&self, other: &Self) -> bool {
-        self.inner == other.inner
-    }
-}
-
-impl<T: Eq> Eq for PartitionedTask<T> {}
-
-impl<T: PartialOrd> PartialOrd for PartitionedTask<T> {
-    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
-        self.inner.partial_cmp(&other.inner)
-    }
-}
-
-impl<T: Ord> Ord for PartitionedTask<T> {
-    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
-        self.inner.cmp(&other.inner)
-    }
-}
-
 /// Maintain sleeping tasks by a delta list.
 struct SleepingTasks {
     delta_list: DeltaList<Box<dyn FnOnce() + Send>>,