Add LogicalCpuSet as a sched affinity mask.

src/context/context.rs

@@ -10,7 +10,7 @@ use alloc::{
 };
 use spin::RwLock;
 
-use crate::LogicalCpuId;
+use crate::{LogicalCpuId, LogicalCpuSet};
 use crate::arch::{interrupt::InterruptStack, paging::PAGE_SIZE};
 use crate::common::aligned_box::AlignedBox;
 use crate::common::unique::Unique;
@@ -147,7 +147,7 @@ pub struct ContextSnapshot {
     pub running: bool,
     pub cpu_id: Option<LogicalCpuId>,
     pub cpu_time: u128,
-    pub sched_affinity: Option<LogicalCpuId>,
+    pub sched_affinity: LogicalCpuSet,
     pub syscall: Option<(usize, usize, usize, usize, usize, usize)>,
     // Clone fields
     //TODO: is there a faster way than allocation?
@@ -239,7 +239,7 @@ pub struct Context {
     /// Scheduler CPU affinity. If set, [`cpu_id`] can except [`None`] never be anything else than
     /// this value.
     // TODO: bitmask (selection of multiple allowed CPUs)?
-    pub sched_affinity: Option<LogicalCpuId>,
+    pub sched_affinity: LogicalCpuSet,
     /// Current system call
     pub syscall: Option<(usize, usize, usize, usize, usize, usize)>,
     /// Head buffer to use when system call buffers are not page aligned
@@ -317,7 +317,7 @@ impl Context {
             cpu_id: None,
             switch_time: 0,
             cpu_time: 0,
-            sched_affinity: None,
+            sched_affinity: LogicalCpuSet::all(),
             syscall: None,
             syscall_head: Some(RaiiFrame::allocate()?),
             syscall_tail: Some(RaiiFrame::allocate()?),

src/context/mod.rs

@@ -7,6 +7,7 @@ use alloc::sync::Arc;
 
 use spin::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 
+use crate::LogicalCpuSet;
 use crate::paging::{RmmA, RmmArch, TableKind};
 use crate::percpu::PercpuBlock;
 use crate::syscall::error::{Error, ESRCH, Result};
@@ -66,7 +67,7 @@ pub fn init() {
     let id = ContextId::from(crate::cpu_id().get() as usize + 1);
     let context_lock = contexts.insert_context_raw(id).expect("could not initialize first context");
     let mut context = context_lock.write();
-    context.sched_affinity = Some(crate::cpu_id());
+    context.sched_affinity = LogicalCpuSet::single(crate::cpu_id());
     context.name = Cow::Borrowed("kmain");
 
     self::arch::EMPTY_CR3.call_once(|| unsafe { RmmA::table(TableKind::User) });

src/context/switch.rs

@@ -28,7 +28,7 @@ unsafe fn update_runnable(context: &mut Context, cpu_id: LogicalCpuId) -> bool {
 
     // Take ownership if not already owned
     // TODO: Support unclaiming context, while still respecting the CPU affinity.
-    if context.cpu_id == None && context.sched_affinity.map_or(true, |id| id == crate::cpu_id()) {
+    if context.cpu_id == None && context.sched_affinity.contains(crate::cpu_id()) {
         context.cpu_id = Some(cpu_id);
         // println!("{}: take {} {}", cpu_id, context.id, *context.name.read());
     }

src/lib.rs

@@ -290,7 +290,7 @@ pub mod kernel_executable_offsets {
 /// This is usually but not necessarily the same as the APIC ID.
 
 // TODO: Differentiate between logical CPU IDs and hardware CPU IDs (e.g. APIC IDs)
-#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
+#[derive(Clone, Copy, Eq, Ord, PartialEq, PartialOrd)]
 // TODO: NonMaxUsize?
 // TODO: Optimize away this type if not cfg!(feature = "multi_core")
 pub struct LogicalCpuId(u32);
@@ -302,8 +302,32 @@ impl LogicalCpuId {
     pub const fn get(self) -> u32 { self.0 }
 }
 
-impl core::fmt::Display for LogicalCpuId {
+impl core::fmt::Debug for LogicalCpuId {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         write!(f, "[logical cpu #{}]", self.0)
     }
 }
+impl core::fmt::Display for LogicalCpuId {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        write!(f, "#{}", self.0)
+    }
+}
+
+// TODO: Support more than 128 CPUs.
+// The maximum number of CPUs on Linux is configurable, and the type for LogicalCpuSet and
+// LogicalCpuId may be optimized accordingly. In that case, box the mask if it's larger than some
+// base size (probably 256 bytes).
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct LogicalCpuSet(u128);
+
+impl LogicalCpuSet {
+    pub const fn new(inner: u128) -> Self { Self(inner) }
+    pub const fn get(self) -> u128 { self.0 }
+
+    pub const fn empty() -> Self { Self::new(0) }
+    pub const fn all() -> Self { Self::new(!0) }
+    pub const fn single(id: LogicalCpuId) -> Self { Self::new(1 << id.get()) }
+    pub const fn contains(&self, id: LogicalCpuId) -> bool {
+        self.0 & (1 << id.get()) != 0
+    }
+}

src/scheme/proc.rs

@@ -13,7 +13,7 @@ use crate::{
         flag::*,
         scheme::{CallerCtx, Scheme},
         self, usercopy::{UserSliceWo, UserSliceRo},
-    }, LogicalCpuId,
+    }, LogicalCpuId, LogicalCpuSet,
 };
 
 use alloc::{
@@ -862,7 +862,15 @@ impl KernelScheme for ProcScheme {
                 Ok(mem::size_of::<usize>())
             }
             Operation::SchedAffinity => {
-                let id = context::contexts().get(info.pid).ok_or(Error::new(EBADFD))?.read().sched_affinity.map_or(u32::MAX, |a| a.get() % crate::cpu_count());
+                // TODO: Improve the sched_affinity interface to allow a full mask.
+                let set = context::contexts().get(info.pid).ok_or(Error::new(EBADFD))?.read().sched_affinity;
+
+                let id = if set == LogicalCpuSet::empty() {
+                    usize::MAX
+                } else {
+                    set.get().trailing_zeros() as usize
+                };
+
                 buf.write_usize(id as usize)?;
                 Ok(mem::size_of::<usize>())
             }
@@ -1074,9 +1082,9 @@ impl KernelScheme for ProcScheme {
                 context::contexts().get(info.pid)
                     .ok_or(Error::new(EBADFD))?.write()
                     .sched_affinity = if val == u32::MAX {
-                        None
+                        LogicalCpuSet::all()
                     } else {
-                        Some(LogicalCpuId::new(val % crate::cpu_count()))
+                        LogicalCpuSet::single(LogicalCpuId::new(val % crate::cpu_count()))
                     };
                 Ok(mem::size_of::<usize>())
             }

src/scheme/sys/context.rs

@@ -64,11 +64,7 @@ pub fn resource() -> Result<Vec<u8>> {
             } else {
                 format!("?")
             };
-            let affinity = if let Some(aff) = context.sched_affinity {
-                format!("{}", aff)
-            } else {
-                format!("?")
-            };
+            let affinity = format!("{:#x}", context.sched_affinity.get() & ((1 << crate::cpu_count()) - 1));
 
             let cpu_time_s = context.cpu_time / crate::time::NANOS_PER_SEC;
             let cpu_time_ns = context.cpu_time % crate::time::NANOS_PER_SEC;