Improve readability on SchedPolicy storage & reduce memory foot print on sched entities

kernel/src/sched/sched_class/fair.rs

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: MPL-2.0
 
-use alloc::collections::binary_heap::BinaryHeap;
+use alloc::{collections::binary_heap::BinaryHeap, sync::Arc};
 use core::{
     cmp::{self, Reverse},
     sync::atomic::{AtomicU64, Ordering::*},
@@ -8,14 +8,20 @@ use core::{
 
 use ostd::{
     cpu::{num_cpus, CpuId},
-    task::scheduler::{EnqueueFlags, UpdateFlags},
+    task::{
+        scheduler::{EnqueueFlags, UpdateFlags},
+        Task,
+    },
 };
 
 use super::{
     time::{base_slice_clocks, min_period_clocks},
-    CurrentRuntime, SchedAttr, SchedClassRq, SchedEntity,
+    CurrentRuntime, SchedAttr, SchedClassRq,
+};
+use crate::{
+    sched::priority::{Nice, NiceRange},
+    thread::AsThread,
 };
-use crate::sched::priority::{Nice, NiceRange};
 
 const WEIGHT_0: u64 = 1024;
 pub const fn nice_to_weight(nice: Nice) -> u64 {
@@ -117,7 +123,7 @@ impl FairAttr {
 ///
 /// This structure is used to provide the capability for keying in the
 /// run queue implemented by `BTreeSet` in the `FairClassRq`.
-struct FairQueueItem(SchedEntity);
+struct FairQueueItem(Arc<Task>, u64);
 
 impl core::fmt::Debug for FairQueueItem {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
@@ -127,7 +133,7 @@ impl core::fmt::Debug for FairQueueItem {
 
 impl FairQueueItem {
     fn key(&self) -> u64 {
-        self.0 .1.sched_attr().fair.vruntime.load(Relaxed)
+        self.1
     }
 }
 
@@ -214,16 +220,19 @@ impl FairClassRq {
 }
 
 impl SchedClassRq for FairClassRq {
-    fn enqueue(&mut self, entity: SchedEntity, flags: Option<EnqueueFlags>) {
+    fn enqueue(&mut self, entity: Arc<Task>, flags: Option<EnqueueFlags>) {
-        let fair_attr = &entity.1.sched_attr().fair;
+        let fair_attr = &entity.as_thread().unwrap().sched_attr().fair;
         let vruntime = match flags {
             Some(EnqueueFlags::Spawn) => self.min_vruntime + self.vtime_slice(),
             _ => self.min_vruntime,
         };
-        fair_attr.vruntime.fetch_max(vruntime, Relaxed);
+        let vruntime = fair_attr
+            .vruntime
+            .fetch_max(vruntime, Relaxed)
+            .max(vruntime);
 
         self.total_weight += fair_attr.weight.load(Relaxed);
-        self.entities.push(Reverse(FairQueueItem(entity)));
+        self.entities.push(Reverse(FairQueueItem(entity, vruntime)));
     }
 
     fn len(&mut self) -> usize {
@@ -234,10 +243,11 @@ impl SchedClassRq for FairClassRq {
         self.entities.is_empty()
     }
 
-    fn pick_next(&mut self) -> Option<SchedEntity> {
+    fn pick_next(&mut self) -> Option<Arc<Task>> {
-        let Reverse(FairQueueItem(entity)) = self.entities.pop()?;
+        let Reverse(FairQueueItem(entity, _)) = self.entities.pop()?;
 
-        self.total_weight -= entity.1.sched_attr().fair.weight.load(Relaxed);
+        let sched_attr = entity.as_thread().unwrap().sched_attr();
+        self.total_weight -= sched_attr.fair.weight.load(Relaxed);
 
         Some(entity)
     }

kernel/src/sched/sched_class/idle.rs

@@ -6,7 +6,7 @@ use super::*;
 ///
 /// This run queue is used for the per-cpu idle entity, if any.
 pub(super) struct IdleClassRq {
-    entity: Option<SchedEntity>,
+    entity: Option<Arc<Task>>,
 }
 
 impl IdleClassRq {
@@ -27,10 +27,10 @@ impl core::fmt::Debug for IdleClassRq {
 }
 
 impl SchedClassRq for IdleClassRq {
-    fn enqueue(&mut self, entity: SchedEntity, _: Option<EnqueueFlags>) {
+    fn enqueue(&mut self, entity: Arc<Task>, _: Option<EnqueueFlags>) {
-        let ptr = Arc::as_ptr(&entity.0);
+        let ptr = Arc::as_ptr(&entity);
         if let Some(t) = self.entity.replace(entity)
-            && ptr != Arc::as_ptr(&t.0)
+            && ptr != Arc::as_ptr(&t)
         {
             panic!("Multiple `idle` entities spawned")
         }
@@ -44,7 +44,7 @@ impl SchedClassRq for IdleClassRq {
         self.entity.is_none()
     }
 
-    fn pick_next(&mut self) -> Option<SchedEntity> {
+    fn pick_next(&mut self) -> Option<Arc<Task>> {
         self.entity.clone()
     }
 

kernel/src/sched/sched_class/mod.rs

@@ -3,13 +3,7 @@
 #![warn(unused)]
 
 use alloc::{boxed::Box, sync::Arc};
-use core::{
+use core::{fmt, sync::atomic::AtomicU64};
-    fmt,
-    sync::atomic::{
-        AtomicU64,
-        Ordering::{Relaxed, SeqCst},
-    },
-};
 
 use ostd::{
     cpu::{all_cpus, AtomicCpuSet, CpuId, PinCurrentCpu},
@@ -34,8 +28,8 @@ mod stop;
 
 use ostd::arch::read_tsc as sched_clock;
 
-use self::policy::SchedPolicyKind;
+pub use self::policy::*;
-pub use self::{policy::SchedPolicy, real_time::RealTimePolicy};
+use self::policy::{SchedPolicyKind, SchedPolicyState};
 use super::{
     priority::{Nice, RangedU8},
     stats::SchedulerStats,
@@ -100,7 +94,7 @@ impl CurrentRuntime {
 /// should implement this trait to function as expected.
 trait SchedClassRq: Send + fmt::Debug {
     /// Enqueues a task into the run queue.
-    fn enqueue(&mut self, entity: SchedEntity, flags: Option<EnqueueFlags>);
+    fn enqueue(&mut self, task: Arc<Task>, flags: Option<EnqueueFlags>);
 
     /// Returns the number of threads in the run queue.
     fn len(&mut self) -> usize;
@@ -111,7 +105,7 @@ trait SchedClassRq: Send + fmt::Debug {
     }
 
     /// Picks the next task for running.
-    fn pick_next(&mut self) -> Option<SchedEntity>;
+    fn pick_next(&mut self) -> Option<Arc<Task>>;
 
     /// Update the information of the current task.
     fn update_current(&mut self, rt: &CurrentRuntime, attr: &SchedAttr, flags: UpdateFlags)
@@ -124,7 +118,7 @@ trait SchedClassRq: Send + fmt::Debug {
 /// scheduling class.
 #[derive(Debug)]
 pub struct SchedAttr {
-    policy: AtomicU64,
+    policy: SchedPolicyState,
 
     real_time: real_time::RealTimeAttr,
     fair: fair::FairAttr,
@@ -134,7 +128,7 @@ impl SchedAttr {
     /// Constructs a new `SchedAttr` with the given scheduling policy.
     pub fn new(policy: SchedPolicy) -> Self {
         Self {
-            policy: SchedPolicy::into_raw(policy).into(),
+            policy: SchedPolicyState::new(policy),
             real_time: {
                 let (prio, policy) = match policy {
                     SchedPolicy::RealTime { rt_prio, rt_policy } => (rt_prio.get(), rt_policy),
@@ -151,47 +145,24 @@ impl SchedAttr {
 
     /// Retrieves the current scheduling policy of the thread.
     pub fn policy(&self) -> SchedPolicy {
-        SchedPolicy::from_raw(self.policy.load(Relaxed))
+        self.policy.get()
     }
 
     fn policy_kind(&self) -> SchedPolicyKind {
-        SchedPolicyKind::from_raw(self.policy.load(Relaxed))
+        self.policy.kind()
     }
 
     /// Updates the scheduling policy of the thread.
     ///
     /// Specifically for real-time policies, if the new policy doesn't
     /// specify a base slice factor for RR, the old one will be kept.
-    pub fn set_policy(&self, mut policy: SchedPolicy) {
+    pub fn set_policy(&self, policy: SchedPolicy) {
-        let _ = self.policy.fetch_update(SeqCst, SeqCst, |raw| {
+        self.policy.set(policy, |policy| match policy {
-            let current = SchedPolicy::from_raw(raw);
-            match policy {
             SchedPolicy::RealTime { rt_prio, rt_policy } => {
                 self.real_time.update(rt_prio.get(), rt_policy);
             }
             SchedPolicy::Fair(nice) => self.fair.update(nice),
             _ => {}
-            }
-
-            // Keep the old base slice factor if the new policy doesn't specify one.
-            if let (
-                SchedPolicy::RealTime {
-                    rt_policy:
-                        RealTimePolicy::RoundRobin {
-                            base_slice_factor: slot,
-                        },
-                    ..
-                },
-                SchedPolicy::RealTime {
-                    rt_policy: RealTimePolicy::RoundRobin { base_slice_factor },
-                    ..
-                },
-            ) = (current, &mut policy)
-            {
-                *base_slice_factor = slot.or(*base_slice_factor);
-            }
-
-            Some(SchedPolicy::into_raw(policy))
         });
     }
 }
@@ -269,14 +240,18 @@ impl PerCpuClassRqSet {
             .or_else(|| self.real_time.pick_next())
             .or_else(|| self.fair.pick_next())
             .or_else(|| self.idle.pick_next())
+            .and_then(|task| {
+                let thread = task.as_thread()?.clone();
+                Some((task, thread))
+            })
     }
 
-    fn enqueue_entity(&mut self, entity: SchedEntity, flags: Option<EnqueueFlags>) {
+    fn enqueue_entity(&mut self, (task, thread): SchedEntity, flags: Option<EnqueueFlags>) {
-        match entity.1.sched_attr().policy_kind() {
+        match thread.sched_attr().policy_kind() {
-            SchedPolicyKind::Stop => self.stop.enqueue(entity, flags),
+            SchedPolicyKind::Stop => self.stop.enqueue(task, flags),
-            SchedPolicyKind::RealTime => self.real_time.enqueue(entity, flags),
+            SchedPolicyKind::RealTime => self.real_time.enqueue(task, flags),
-            SchedPolicyKind::Fair => self.fair.enqueue(entity, flags),
+            SchedPolicyKind::Fair => self.fair.enqueue(task, flags),
-            SchedPolicyKind::Idle => self.idle.enqueue(entity, flags),
+            SchedPolicyKind::Idle => self.idle.enqueue(task, flags),
         }
     }
 
@@ -294,9 +269,9 @@ impl LocalRunQueue for PerCpuClassRqSet {
 
     fn pick_next_current(&mut self) -> Option<&Arc<Task>> {
         self.pick_next_entity().and_then(|next| {
-            let next_ptr = Arc::as_ptr(&next.1);
+            let next_ptr = Arc::as_ptr(&next.0);
             if let Some((old, _)) = self.current.replace((next, CurrentRuntime::new())) {
-                if Arc::as_ptr(&old.1) == next_ptr {
+                if Arc::as_ptr(&old.0) == next_ptr {
                     return None;
                 }
                 self.enqueue_entity(old, None);

kernel/src/sched/sched_class/policy.rs

@@ -1,9 +1,13 @@
 // SPDX-License-Identifier: MPL-2.0
 
-use core::num::NonZero;
+use core::sync::atomic::{AtomicU8, Ordering::Relaxed};
 
-use super::real_time::RealTimePolicy;
+use atomic_integer_wrapper::define_atomic_version_of_integer_like_type;
-use crate::sched::priority::{Nice, NiceRange, Priority, RangedU8};
+use int_to_c_enum::TryFromInt;
+use ostd::sync::SpinLock;
+
+pub use super::real_time::RealTimePolicy;
+use crate::sched::priority::{Nice, Priority, RangedU8};
 
 /// The User-chosen scheduling policy.
 ///
@@ -19,12 +23,13 @@ pub enum SchedPolicy {
     Idle,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, TryFromInt)]
+#[repr(u8)]
 pub(super) enum SchedPolicyKind {
-    Stop,
+    Stop = 0,
-    RealTime,
+    RealTime = 1,
-    Fair,
+    Fair = 2,
-    Idle,
+    Idle = 3,
 }
 
 impl From<Priority> for SchedPolicy {
@@ -41,104 +46,73 @@ impl From<Priority> for SchedPolicy {
     }
 }
 
-const TYPE_MASK: u64 = 0x0000_0000_0000_ffff;
-const TYPE_SHIFT: u32 = 0;
-
-const TYPE_STOP: u64 = 0;
-const TYPE_REAL_TIME: u64 = 1;
-const TYPE_FAIR: u64 = 2;
-const TYPE_IDLE: u64 = 3;
-
-const SUBTYPE_MASK: u64 = 0x0000_0000_00ff_0000;
-const SUBTYPE_SHIFT: u32 = 16;
-
-const RT_PRIO_MASK: u64 = SUBTYPE_MASK;
-const RT_PRIO_SHIFT: u32 = SUBTYPE_SHIFT;
-
-const FAIR_NICE_MASK: u64 = SUBTYPE_MASK;
-const FAIR_NICE_SHIFT: u32 = SUBTYPE_SHIFT;
-
-const RT_TYPE_MASK: u64 = 0x0000_0000_ff00_0000;
-const RT_TYPE_SHIFT: u32 = 24;
-
-const RT_TYPE_FIFO: u64 = 0;
-const RT_TYPE_RR: u64 = 1;
-
-const RT_FACTOR_MASK: u64 = 0xffff_ffff_0000_0000;
-const RT_FACTOR_SHIFT: u32 = 32;
-
-fn get(raw: u64, mask: u64, shift: u32) -> u64 {
-    (raw & mask) >> shift
-}
-
-fn set(value: u64, mask: u64, shift: u32) -> u64 {
-    (value << shift) & mask
-}
-
 impl SchedPolicy {
-    pub(super) fn from_raw(raw: u64) -> Self {
+    pub(super) fn kind(&self) -> SchedPolicyKind {
-        match get(raw, TYPE_MASK, TYPE_SHIFT) {
+        match self {
-            TYPE_STOP => SchedPolicy::Stop,
+            SchedPolicy::Stop => SchedPolicyKind::Stop,
-            TYPE_REAL_TIME => SchedPolicy::RealTime {
+            SchedPolicy::RealTime { .. } => SchedPolicyKind::RealTime,
-                rt_prio: RangedU8::new(get(raw, RT_PRIO_MASK, RT_PRIO_SHIFT) as u8),
+            SchedPolicy::Fair(_) => SchedPolicyKind::Fair,
-                rt_policy: match get(raw, RT_TYPE_MASK, RT_TYPE_SHIFT) {
+            SchedPolicy::Idle => SchedPolicyKind::Idle,
-                    RT_TYPE_FIFO => RealTimePolicy::Fifo,
-                    RT_TYPE_RR => RealTimePolicy::RoundRobin {
-                        base_slice_factor: NonZero::new(
-                            get(raw, RT_FACTOR_MASK, RT_FACTOR_SHIFT) as u32
-                        ),
-                    },
-                    _ => unreachable!(),
-                },
-            },
-            TYPE_FAIR => {
-                SchedPolicy::Fair(Nice::new(NiceRange::new(
-                    get(raw, FAIR_NICE_MASK, FAIR_NICE_SHIFT) as i8,
-                )))
-            }
-            TYPE_IDLE => SchedPolicy::Idle,
-            _ => unreachable!(),
-        }
-    }
-
-    pub(super) fn into_raw(this: Self) -> u64 {
-        match this {
-            SchedPolicy::Stop => set(TYPE_STOP, TYPE_MASK, TYPE_SHIFT),
-            SchedPolicy::RealTime { rt_prio, rt_policy } => {
-                let ty = set(TYPE_REAL_TIME, TYPE_MASK, TYPE_SHIFT);
-                let rt_prio = set(rt_prio.get() as u64, RT_PRIO_MASK, RT_PRIO_SHIFT);
-                let rt_policy = match rt_policy {
-                    RealTimePolicy::Fifo => set(RT_TYPE_FIFO, RT_TYPE_MASK, RT_TYPE_SHIFT),
-                    RealTimePolicy::RoundRobin { base_slice_factor } => {
-                        let rt_type = set(RT_TYPE_RR, RT_TYPE_MASK, RT_TYPE_SHIFT);
-                        let rt_factor = set(
-                            base_slice_factor.map_or(0, NonZero::get) as u64,
-                            RT_FACTOR_MASK,
-                            RT_FACTOR_SHIFT,
-                        );
-                        rt_type | rt_factor
-                    }
-                };
-                ty | rt_prio | rt_policy
-            }
-            SchedPolicy::Fair(nice) => {
-                let ty = set(TYPE_FAIR, TYPE_MASK, TYPE_SHIFT);
-                let nice = set(nice.range().get() as u64, FAIR_NICE_MASK, FAIR_NICE_SHIFT);
-                ty | nice
-            }
-            SchedPolicy::Idle => set(TYPE_IDLE, TYPE_MASK, TYPE_SHIFT),
         }
     }
 }
 
-impl SchedPolicyKind {
+define_atomic_version_of_integer_like_type!(SchedPolicyKind, try_from = true, {
-    pub fn from_raw(raw: u64) -> Self {
+    #[derive(Debug)]
-        match get(raw, TYPE_MASK, TYPE_SHIFT) {
+    pub struct AtomicSchedPolicyKind(AtomicU8);
-            TYPE_STOP => SchedPolicyKind::Stop,
+});
-            TYPE_REAL_TIME => SchedPolicyKind::RealTime,
+
-            TYPE_FAIR => SchedPolicyKind::Fair,
+impl From<SchedPolicyKind> for u8 {
-            TYPE_IDLE => SchedPolicyKind::Idle,
+    fn from(value: SchedPolicyKind) -> Self {
-            _ => unreachable!(),
+        value as _
     }
 }
+
+#[derive(Debug)]
+pub(super) struct SchedPolicyState {
+    kind: AtomicSchedPolicyKind,
+    policy: SpinLock<SchedPolicy>,
+}
+
+impl SchedPolicyState {
+    pub fn new(policy: SchedPolicy) -> Self {
+        Self {
+            kind: AtomicSchedPolicyKind::new(policy.kind()),
+            policy: SpinLock::new(policy),
+        }
+    }
+
+    pub fn kind(&self) -> SchedPolicyKind {
+        self.kind.load(Relaxed)
+    }
+
+    pub fn get(&self) -> SchedPolicy {
+        *self.policy.disable_irq().lock()
+    }
+
+    pub fn set(&self, mut policy: SchedPolicy, update: impl FnOnce(SchedPolicy)) {
+        let mut this = self.policy.disable_irq().lock();
+
+        // Keep the old base slice factor if the new policy doesn't specify one.
+        if let (
+            SchedPolicy::RealTime {
+                rt_policy:
+                    RealTimePolicy::RoundRobin {
+                        base_slice_factor: slot,
+                    },
+                ..
+            },
+            SchedPolicy::RealTime {
+                rt_policy: RealTimePolicy::RoundRobin { base_slice_factor },
+                ..
+            },
+        ) = (*this, &mut policy)
+        {
+            *base_slice_factor = slot.or(*base_slice_factor);
+        }
+
+        update(policy);
+        self.kind.store(policy.kind(), Relaxed);
+        *this = policy;
+    }
 }

kernel/src/sched/sched_class/real_time.rs

@@ -77,7 +77,7 @@ impl RealTimeAttr {
 
 struct PrioArray {
     map: BitArr![for 100],
-    queue: [VecDeque<SchedEntity>; 100],
+    queue: [VecDeque<Arc<Task>>; 100],
 }
 
 impl core::fmt::Debug for PrioArray {
@@ -88,7 +88,10 @@ impl core::fmt::Debug for PrioArray {
             })
             .field_with("queue", |f| {
                 f.debug_list()
-                    .entries((self.queue.iter().flatten()).map(|(_, thread)| thread.sched_attr()))
+                    .entries(
+                        (self.queue.iter().flatten())
+                            .map(|task| task.as_thread().unwrap().sched_attr()),
+                    )
                     .finish()
             })
             .finish()
@@ -96,7 +99,7 @@ impl core::fmt::Debug for PrioArray {
 }
 
 impl PrioArray {
-    fn enqueue(&mut self, thread: SchedEntity, prio: u8) {
+    fn enqueue(&mut self, thread: Arc<Task>, prio: u8) {
         let queue = &mut self.queue[usize::from(prio)];
         let is_empty = queue.is_empty();
         queue.push_back(thread);
@@ -105,7 +108,7 @@ impl PrioArray {
         }
     }
 
-    fn pop(&mut self) -> Option<SchedEntity> {
+    fn pop(&mut self) -> Option<Arc<Task>> {
         let mut iter = self.map.iter_ones();
         let prio = iter.next()? as u8;
 
@@ -166,8 +169,9 @@ impl RealTimeClassRq {
 }
 
 impl SchedClassRq for RealTimeClassRq {
-    fn enqueue(&mut self, entity: SchedEntity, _: Option<EnqueueFlags>) {
+    fn enqueue(&mut self, entity: Arc<Task>, _: Option<EnqueueFlags>) {
-        let prio = entity.1.sched_attr().real_time.prio.load(Relaxed);
+        let sched_attr = entity.as_thread().unwrap().sched_attr();
+        let prio = sched_attr.real_time.prio.load(Relaxed);
         self.inactive_array().enqueue(entity, prio);
         self.nr_running += 1;
     }
@@ -180,20 +184,17 @@ impl SchedClassRq for RealTimeClassRq {
         self.nr_running == 0
     }
 
-    fn pick_next(&mut self) -> Option<SchedEntity> {
+    fn pick_next(&mut self) -> Option<Arc<Task>> {
         if self.nr_running == 0 {
             return None;
         }
 
-        let res = self.active_array().pop().or_else(|| {
+        (self.active_array().pop())
+            .or_else(|| {
                 self.swap_arrays();
                 self.active_array().pop()
-        });
+            })
-        if res.is_some() {
+            .inspect(|_| self.nr_running -= 1)
-            self.nr_running -= 1;
-        }
-
-        res
     }
 
     fn update_current(

kernel/src/sched/sched_class/stop.rs

@@ -1,7 +1,5 @@
 // SPDX-License-Identifier: MPL-2.0
 
-use core::sync::atomic::AtomicBool;
-
 use super::*;
 
 /// The per-cpu run queue for the STOP scheduling class.
@@ -9,7 +7,7 @@ use super::*;
 /// This is a singleton class, meaning that only one thread can be in this class at a time.
 /// This is used for the most critical tasks, such as powering off and rebooting.
 pub(super) struct StopClassRq {
-    thread: Option<Arc<Thread>>,
+    thread: Option<Arc<Task>>,
 }
 
 impl StopClassRq {
@@ -30,11 +28,10 @@ impl core::fmt::Debug for StopClassRq {
 }
 
 impl SchedClassRq for StopClassRq {
-    fn enqueue(&mut self, thread: Arc<Thread>, _: Option<EnqueueFlags>) {
+    fn enqueue(&mut self, thread: Arc<Task>, _: Option<EnqueueFlags>) {
         if self.thread.replace(thread).is_some() {
             panic!("Multiple `stop` threads spawned")
         }
-        self.has_value.store(true, Relaxed);
     }
 
     fn len(&mut self) -> usize {
@@ -45,7 +42,7 @@ impl SchedClassRq for StopClassRq {
         self.thread.is_none()
     }
 
-    fn pick_next(&mut self) -> Option<Arc<Thread>> {
+    fn pick_next(&mut self) -> Option<Arc<Task>> {
         self.thread.take()
     }