Remove cpu_affinity field from OSTD Task struct

kernel/src/process/posix_thread/builder.rs

@@ -2,7 +2,7 @@
 
 #![allow(dead_code)]
 
-use ostd::{task::Task, user::UserSpace};
+use ostd::{cpu::CpuSet, task::Task, user::UserSpace};
 
 use super::{thread_table, PosixThread};
 use crate::{
@@ -113,11 +113,13 @@ impl PosixThreadBuilder {
 
             let status = ThreadStatus::Init;
             let priority = Priority::default();
+            let cpu_affinity = CpuSet::new_full();
             let thread = Arc::new(Thread::new(
                 weak_task.clone(),
                 posix_thread,
                 status,
                 priority,
+                cpu_affinity,
             ));
 
             thread_table::add_thread(tid, thread.clone());

kernel/src/sched/priority_scheduler.rs

@@ -2,6 +2,7 @@
 
 use ostd::{
     cpu::{num_cpus, CpuSet, PinCurrentCpu},
+    sync::PreemptDisabled,
     task::{
         scheduler::{inject_scheduler, EnqueueFlags, LocalRunQueue, Scheduler, UpdateFlags},
         AtomicCpuId, Task,
@@ -259,8 +260,11 @@ impl PreemptSchedInfo for Task {
         &self.schedule_info().cpu
     }
 
-    fn cpu_affinity(&self) -> &CpuSet {
+    fn cpu_affinity(&self) -> SpinLockGuard<CpuSet, PreemptDisabled> {
-        &self.schedule_info().cpu_affinity
+        self.data()
+            .downcast_ref::<Arc<Thread>>()
+            .unwrap()
+            .lock_cpu_affinity()
     }
 }
 
@@ -272,7 +276,7 @@ trait PreemptSchedInfo {
 
     fn cpu(&self) -> &AtomicCpuId;
 
-    fn cpu_affinity(&self) -> &CpuSet;
+    fn cpu_affinity(&self) -> SpinLockGuard<CpuSet, PreemptDisabled>;
 
     fn is_real_time(&self) -> bool {
         self.priority() < Self::REAL_TIME_TASK_PRIORITY

kernel/src/thread/kernel_thread.rs

@@ -55,19 +55,17 @@ pub fn create_new_kernel_task(mut thread_options: ThreadOptions) -> Arc<Task> {
             let kernel_thread = KernelThread;
             let status = ThreadStatus::Init;
             let priority = thread_options.priority;
+            let cpu_affinity = thread_options.cpu_affinity;
             Arc::new(Thread::new(
                 weak_task.clone(),
                 kernel_thread,
                 status,
                 priority,
+                cpu_affinity,
             ))
         };
 
-        TaskOptions::new(thread_fn)
+        TaskOptions::new(thread_fn).data(thread).build().unwrap()
-            .data(thread)
-            .cpu_affinity(thread_options.cpu_affinity)
-            .build()
-            .unwrap()
     })
 }
 

kernel/src/thread/mod.rs

@@ -4,7 +4,7 @@
 
 use core::sync::atomic::Ordering;
 
-use ostd::task::Task;
+use ostd::{cpu::CpuSet, sync::PreemptDisabled, task::Task};
 
 use self::status::{AtomicThreadStatus, ThreadStatus};
 use crate::{
@@ -33,6 +33,8 @@ pub struct Thread {
     status: AtomicThreadStatus,
     /// Thread priority
     priority: AtomicPriority,
+    /// Thread cpu affinity
+    cpu_affinity: SpinLock<CpuSet>,
 }
 
 impl Thread {
@@ -42,12 +44,14 @@ impl Thread {
         data: impl Send + Sync + Any,
         status: ThreadStatus,
         priority: Priority,
+        cpu_affinity: CpuSet,
     ) -> Self {
         Thread {
             task,
             data: Box::new(data),
             status: AtomicThreadStatus::new(status),
             priority: AtomicPriority::new(priority),
+            cpu_affinity: SpinLock::new(cpu_affinity),
         }
     }
 
@@ -111,6 +115,16 @@ impl Thread {
         self.priority.store(new_priority, Ordering::Relaxed)
     }
 
+    /// Acquires the lock of cpu affinity.
+    pub fn lock_cpu_affinity(&self) -> SpinLockGuard<CpuSet, PreemptDisabled> {
+        self.cpu_affinity.lock()
+    }
+
+    /// Updates the cpu affinity with the new value.
+    pub fn set_cpu_affinity(&self, new_cpu_affinity: CpuSet) {
+        *self.cpu_affinity.lock() = new_cpu_affinity;
+    }
+
     pub fn yield_now() {
         Task::yield_now()
     }

ostd/src/task/mod.rs

@@ -19,7 +19,7 @@ pub use self::{
     scheduler::info::{AtomicCpuId, TaskScheduleInfo},
 };
 pub(crate) use crate::arch::task::{context_switch, TaskContext};
-use crate::{cpu::CpuSet, prelude::*, user::UserSpace};
+use crate::{prelude::*, user::UserSpace};
 
 /// A task that executes a function to the end.
 ///
@@ -124,7 +124,6 @@ pub struct TaskOptions {
     func: Option<Box<dyn Fn() + Send + Sync>>,
     data: Option<Box<dyn Any + Send + Sync>>,
     user_space: Option<Arc<UserSpace>>,
-    cpu_affinity: CpuSet,
 }
 
 impl TaskOptions {
@@ -137,7 +136,6 @@ impl TaskOptions {
             func: Some(Box::new(func)),
             data: None,
             user_space: None,
-            cpu_affinity: CpuSet::new_full(),
         }
     }
 
@@ -165,15 +163,6 @@ impl TaskOptions {
         self
     }
 
-    /// Sets the CPU affinity mask for the task.
-    ///
-    /// The `cpu_affinity` parameter represents
-    /// the desired set of CPUs to run the task on.
-    pub fn cpu_affinity(mut self, cpu_affinity: CpuSet) -> Self {
-        self.cpu_affinity = cpu_affinity;
-        self
-    }
-
     /// Builds a new task without running it immediately.
     pub fn build(self) -> Result<Task> {
         /// all task will entering this function
@@ -212,7 +201,6 @@ impl TaskOptions {
             kstack,
             schedule_info: TaskScheduleInfo {
                 cpu: AtomicCpuId::default(),
-                cpu_affinity: self.cpu_affinity,
             },
         };
 

ostd/src/task/scheduler/info.rs

@@ -4,8 +4,6 @@
 
 use core::sync::atomic::{AtomicU32, Ordering};
 
-use crate::cpu::CpuSet;
-
 /// Fields of a task that OSTD will never touch.
 ///
 /// The type ought to be defined by the OSTD user and injected into the task.
@@ -17,8 +15,6 @@ use crate::cpu::CpuSet;
 pub struct TaskScheduleInfo {
     /// The CPU that the task would like to be running on.
     pub cpu: AtomicCpuId,
-    /// The CPUs that this task can run on.
-    pub cpu_affinity: CpuSet,
 }
 
 /// An atomic CPUID container.