Allow specifying orderings when loading/storing AtomicCpuSets

kernel/src/sched/sched_class/mod.rs

@@ -3,7 +3,7 @@
 #![warn(unused)]
 
 use alloc::{boxed::Box, sync::Arc};
-use core::fmt;
+use core::{fmt, sync::atomic::Ordering};
 
 use ostd::{
     arch::read_tsc as sched_clock,
@@ -263,7 +263,7 @@ impl ClassScheduler {
         }
         debug_assert!(flags == EnqueueFlags::Spawn);
         let guard = disable_local();
-        let affinity = thread.atomic_cpu_affinity().load();
+        let affinity = thread.atomic_cpu_affinity().load(Ordering::Relaxed);
         let mut selected = guard.current_cpu();
         let mut minimum_load = u32::MAX;
         let last_chosen = match self.last_chosen_cpu.get() {

kernel/src/syscall/sched_affinity.rs

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: MPL-2.0
 
-use core::{cmp, mem};
+use core::{cmp, mem, sync::atomic::Ordering};
 
 use ostd::cpu::{num_cpus, CpuId, CpuSet};
 
@@ -14,9 +14,9 @@ pub fn sys_sched_getaffinity(
     ctx: &Context,
 ) -> Result<SyscallReturn> {
     let cpu_set = match tid {
-        0 => ctx.thread.atomic_cpu_affinity().load(),
+        0 => ctx.thread.atomic_cpu_affinity().load(Ordering::Relaxed),
         _ => match thread_table::get_thread(tid) {
-            Some(thread) => thread.atomic_cpu_affinity().load(),
+            Some(thread) => thread.atomic_cpu_affinity().load(Ordering::Relaxed),
             None => return Err(Error::with_message(Errno::ESRCH, "thread does not exist")),
         },
     };
@@ -39,10 +39,15 @@ pub fn sys_sched_setaffinity(
     let user_cpu_set = read_cpu_set_from(ctx.user_space(), cpuset_size, cpu_set_ptr)?;
 
     match tid {
-        0 => ctx.thread.atomic_cpu_affinity().store(&user_cpu_set),
+        0 => ctx
+            .thread
+            .atomic_cpu_affinity()
+            .store(&user_cpu_set, Ordering::Relaxed),
         _ => match thread_table::get_thread(tid) {
             Some(thread) => {
-                thread.atomic_cpu_affinity().store(&user_cpu_set);
+                thread
+                    .atomic_cpu_affinity()
+                    .store(&user_cpu_set, Ordering::Relaxed);
             }
             None => return Err(Error::with_message(Errno::ESRCH, "thread does not exist")),
         },

ostd/src/cpu/set.rs

@@ -168,26 +168,37 @@ impl AtomicCpuSet {
         Self { bits }
     }
 
-    /// Loads the value of the set.
+    /// Loads the value of the set with the given ordering.
     ///
-    /// This operation can only be done in the [`Ordering::Relaxed`] memory
-    /// order. It cannot be used to synchronize anything between CPUs.
-    pub fn load(&self) -> CpuSet {
+    /// This operation is not atomic. When racing with a [`Self::store`]
+    /// operation, this load may return a set that contains a portion of the
+    /// new value and a portion of the old value. Load on each specific
+    /// word is atomic, and follows the specified ordering.
+    ///
+    /// Note that load with [`Ordering::Release`] is a valid operation, which
+    /// is different from the normal atomic operations. When coupled with
+    /// [`Ordering::Release`], it actually performs `fetch_or(0, Release)`.
+    pub fn load(&self, ordering: Ordering) -> CpuSet {
         let bits = self
             .bits
             .iter()
-            .map(|part| part.load(Ordering::Relaxed))
+            .map(|part| match ordering {
+                Ordering::Release => part.fetch_or(0, ordering),
+                _ => part.load(ordering),
+            })
             .collect();
         CpuSet { bits }
     }
 
-    /// Stores a new value to the set.
+    /// Stores a new value to the set with the given ordering.
     ///
-    /// This operation can only be done in the [`Ordering::Relaxed`] memory
-    /// order. It cannot be used to synchronize anything between CPUs.
-    pub fn store(&self, value: &CpuSet) {
+    /// This operation is not atomic. When racing with a [`Self::load`]
+    /// operation, that load may return a set that contains a portion of the
+    /// new value and a portion of the old value. Load on each specific
+    /// word is atomic, and follows the specified ordering.
+    pub fn store(&self, value: &CpuSet, ordering: Ordering) {
         for (part, new_part) in self.bits.iter().zip(value.bits.iter()) {
-            part.store(*new_part, Ordering::Relaxed);
+            part.store(*new_part, ordering);
         }
     }
 
@@ -271,7 +282,7 @@ mod test {
                 }
             }
 
-            let loaded = atomic_set.load();
+            let loaded = atomic_set.load(Ordering::Relaxed);
             for cpu_id in loaded.iter() {
                 if cpu_id.as_usize() % 3 == 0 {
                     assert!(loaded.contains(cpu_id));
@@ -280,7 +291,10 @@ mod test {
                 }
             }
 
-            atomic_set.store(&CpuSet::with_capacity_val(test_num_cpus, 0));
+            atomic_set.store(
+                &CpuSet::with_capacity_val(test_num_cpus, 0),
+                Ordering::Relaxed,
+            );
 
             for cpu_id in test_all_iter() {
                 assert!(!atomic_set.contains(cpu_id, Ordering::Relaxed));

ostd/src/mm/vm_space.rs

@@ -95,7 +95,7 @@ impl VmSpace {
             .try_write()
             .ok_or(VmSpaceClearError::CursorsAlive)?;
 
-        let cpus = self.cpus.load();
+        let cpus = self.cpus.load(Ordering::Relaxed);
         let cpu = preempt_guard.current_cpu();
         let cpus_set_is_empty = cpus.is_empty();
         let cpus_set_is_single_self = cpus.count() == 1 && cpus.contains(cpu);
@@ -142,7 +142,7 @@ impl VmSpace {
             CursorMut {
                 pt_cursor,
                 activation_lock,
-                flusher: TlbFlusher::new(self.cpus.load(), disable_preempt()),
+                flusher: TlbFlusher::new(self.cpus.load(Ordering::Relaxed), disable_preempt()),
             }
         })?)
     }

ostd/src/sync/rcu/monitor.rs

@@ -129,7 +129,7 @@ impl GracePeriod {
     unsafe fn finish_grace_period(&mut self, this_cpu: CpuId) {
         self.cpu_mask.add(this_cpu, Ordering::Relaxed);
 
-        if self.cpu_mask.load().is_full() {
+        if self.cpu_mask.load(Ordering::Relaxed).is_full() {
             self.is_complete = true;
         }
     }
@@ -140,7 +140,7 @@ impl GracePeriod {
 
     pub fn restart(&mut self, callbacks: Callbacks) {
         self.is_complete = false;
-        self.cpu_mask.store(&CpuSet::new_empty());
+        self.cpu_mask.store(&CpuSet::new_empty(), Ordering::Relaxed);
         self.callbacks = callbacks;
     }
 }