Skip to content

feat: axtask set active cpu num #77

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ZR233 merged 2 commits into from
Nov 27, 2025
Merged

feat: axtask set active cpu num #77

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2ee6fd3
fix: update axtask to use CPU_NUM for scheduler initialization and en…
ZR233 Nov 27, 2025
0e0c961
feat: axtask add active cpu num
ZR233 Nov 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 38 additions & 12 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

15 changes: 14 additions & 1 deletion modules/axtask/src/api.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
//! Task APIs for multi-task configuration.

use core::sync::atomic::AtomicUsize;

use alloc::{string::String, sync::Arc};

use kernel_guard::NoPreemptIrqSave;
Expand All @@ -19,6 +21,8 @@ pub type AxTaskRef = Arc<AxTask>;
/// The wrapper type for [`cpumask::CpuMask`] with SMP configuration.
pub type AxCpuMask = cpumask::CpuMask<{ axconfig::plat::CPU_NUM }>;

static CPU_NUM: AtomicUsize = AtomicUsize::new(1);

cfg_if::cfg_if! {
if #[cfg(feature = "sched-rr")] {
const MAX_TIME_SLICE: usize = 5;
Expand Down Expand Up @@ -70,15 +74,24 @@ pub fn current() -> CurrentTask {

/// Initializes the task scheduler (for the primary CPU).
pub fn init_scheduler() {
info!("Initialize scheduling...");
init_scheduler_with_cpu_num(axconfig::plat::CPU_NUM);
}

/// Initializes the task scheduler with cpu_num (for the primary CPU).
Copy link

Copilot AI Nov 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new public API function init_scheduler_with_cpu_num is missing documentation. Public APIs should have doc comments explaining their purpose, parameters, and any safety considerations.

Recommended addition:

/// Initializes the task scheduler with a specific number of active CPUs (for the primary CPU).
///
/// This allows limiting the number of active CPUs to fewer than the compile-time maximum
/// specified by `axconfig::plat::CPU_NUM`.
///
/// # Arguments
///
/// * `cpu_num` - The number of active CPUs to use. Must be <= `axconfig::plat::CPU_NUM`.
pub fn init_scheduler_with_cpu_num(cpu_num: usize) {
Suggested change
/// Initializes the task scheduler with cpu_num (for the primary CPU).
/// Initializes the task scheduler with a specific number of active CPUs (for the primary CPU).
///
/// This allows limiting the number of active CPUs to fewer than the compile-time maximum
/// specified by `axconfig::plat::CPU_NUM`.
///
/// # Arguments
///
/// * `cpu_num` - The number of active CPUs to use. Must be <= `axconfig::plat::CPU_NUM`.

Copilot uses AI. Check for mistakes.
pub fn init_scheduler_with_cpu_num(cpu_num: usize) {
Copy link

Copilot AI Nov 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Missing input validation for cpu_num. The function should validate that cpu_num is within valid bounds (1 <= cpu_num <= axconfig::plat::CPU_NUM).

If cpu_num is 0, tasks would be created with empty CPU masks, leading to a panic in select_run_queue_index (which asserts !cpumask.is_empty()). If cpu_num > axconfig::plat::CPU_NUM, the loop in new_common would attempt to set bits beyond the valid range of the AxCpuMask.

Recommended fix:

pub fn init_scheduler_with_cpu_num(cpu_num: usize) {
    assert!(
        cpu_num > 0 && cpu_num <= axconfig::plat::CPU_NUM,
        "cpu_num must be in range [1, {}], got {}",
        axconfig::plat::CPU_NUM,
        cpu_num
    );
    info!("Initialize scheduling...");
    CPU_NUM.store(cpu_num, core::sync::atomic::Ordering::Release);
    // ...
}
Suggested change
pub fn init_scheduler_with_cpu_num(cpu_num: usize) {
pub fn init_scheduler_with_cpu_num(cpu_num: usize) {
assert!(
cpu_num > 0 && cpu_num <= axconfig::plat::CPU_NUM,
"cpu_num must be in range [1, {}], got {}",
axconfig::plat::CPU_NUM,
cpu_num
);

Copilot uses AI. Check for mistakes.
info!("Initialize scheduling...");
CPU_NUM.store(cpu_num, core::sync::atomic::Ordering::Relaxed);
Copy link

Copilot AI Nov 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The memory ordering used for storing and loading CPU_NUM should be Ordering::Release and Ordering::Acquire respectively (or stronger), not Ordering::Relaxed.

Currently, Relaxed ordering provides no synchronization guarantees. When init_scheduler_with_cpu_num stores the CPU count, tasks created immediately after in crate::run_queue::init() (line 84) call active_cpu_num() to read this value. Without proper ordering, there's no guarantee the write is visible to these reads, potentially causing tasks to be initialized with an incorrect CPU mask.

Recommended fix:

CPU_NUM.store(cpu_num, core::sync::atomic::Ordering::Release);

and

CPU_NUM.load(core::sync::atomic::Ordering::Acquire)

Copilot uses AI. Check for mistakes.
crate::run_queue::init();
#[cfg(feature = "irq")]
crate::timers::init();

info!(" use {} scheduler.", Scheduler::scheduler_name());
}
Comment on lines +81 to 89
Copy link

Copilot AI Nov 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new public API function init_scheduler_with_cpu_num lacks test coverage. Given that the module has comprehensive tests in tests.rs for other scheduler functionality, this new function should also be tested.

Consider adding test cases that verify:

  1. Tasks created after calling init_scheduler_with_cpu_num have the correct CPU mask (matching the specified cpu_num)
  2. The function properly handles edge cases (cpu_num = 1, cpu_num = max)
  3. Invalid inputs are properly rejected (if validation is added)

Copilot uses AI. Check for mistakes.

pub(crate) fn active_cpu_num() -> usize {
CPU_NUM.load(core::sync::atomic::Ordering::Relaxed)
Copy link

Copilot AI Nov 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The memory ordering should be Ordering::Acquire to properly synchronize with the Release store in init_scheduler_with_cpu_num.

Using Relaxed ordering provides no guarantees that the value written by the primary CPU during initialization will be visible to secondary CPUs or to code reading this value later. This could lead to tasks being created with incorrect CPU masks.

Recommended fix:

CPU_NUM.load(core::sync::atomic::Ordering::Acquire)
Suggested change
CPU_NUM.load(core::sync::atomic::Ordering::Relaxed)
CPU_NUM.load(core::sync::atomic::Ordering::Acquire)

Copilot uses AI. Check for mistakes.
}

/// Initializes the task scheduler for secondary CPUs.
pub fn init_scheduler_secondary() {
crate::run_queue::init_secondary();
Expand Down
7 changes: 6 additions & 1 deletion modules/axtask/src/task.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,11 @@ impl TaskInner {
// private methods
impl TaskInner {
fn new_common(id: TaskId, name: String) -> Self {
let mut cpumask = AxCpuMask::new();
for cpu_id in 0..crate::api::active_cpu_num() {
cpumask.set(cpu_id, true);
}

Self {
id,
name,
Expand All @@ -232,7 +237,7 @@ impl TaskInner {
entry: None,
state: AtomicU8::new(TaskState::Ready as u8),
// By default, the task is allowed to run on all CPUs.
cpumask: SpinNoIrq::new(AxCpuMask::full()),
cpumask: SpinNoIrq::new(cpumask),
in_wait_queue: AtomicBool::new(false),
#[cfg(feature = "irq")]
timer_ticket_id: AtomicU64::new(0),
Expand Down
Loading