Skip to main content

ax_task/sync/membarrier/
operations.rs

1//! Scheduler-owned Linux membarrier protocol.
2
3use crate::{
4    runtime::{
5        RuntimeStatus,
6        context::{
7            RuntimeIrqGuard, runtime_current_cpu_mut, runtime_task_system, validate_task_context,
8        },
9        cpu::RuntimeCpuId,
10        lock::PreemptScope,
11        resource::{MembarrierRegistration, RuntimeMembarrierAction},
12        task_runtime,
13    },
14    sched::{
15        CpuSet,
16        system::{MembarrierCpuTargets, MembarrierTarget},
17    },
18    thread::TaskError,
19};
20
21/// Memory-barrier target semantics implemented by the scheduler.
22#[derive(Clone, Copy, Debug, Eq, PartialEq)]
23pub enum MembarrierCommand {
24    /// Orders every online CPU through one synchronous scheduler rendezvous.
25    Global,
26    /// Orders CPUs currently executing any globally registered user `mm`.
27    GlobalExpedited,
28    /// Orders CPUs currently executing the caller's shared `mm`.
29    PrivateExpedited,
30}
31
32fn synchronize_targets(targets: &CpuSet, action: RuntimeMembarrierAction) -> Result<(), TaskError> {
33    for cpu in targets.iter() {
34        match task_runtime::synchronize_membarrier_cpu(RuntimeCpuId::new(cpu.as_u32()), action) {
35            RuntimeStatus::Success => {}
36            status => return Err(TaskError::RuntimeFailure(status as u32)),
37        }
38    }
39    Ok(())
40}
41
42/// Registers one irreversible expedited facility for the caller's shared `mm`.
43pub fn register_current_membarrier(registration: MembarrierRegistration) -> Result<(), TaskError> {
44    validate_task_context()?;
45    let system = runtime_task_system()?;
46    let targets = MembarrierCpuTargets::new(system.cpu_topology_len());
47    let _pin = PreemptScope::enter();
48    let plan = {
49        let mut irq = RuntimeIrqGuard::enter();
50        let mut cpu = runtime_current_cpu_mut(&mut irq)?;
51        system.begin_current_membarrier_registration(cpu.as_mut(), registration, targets)?
52    };
53    synchronize_targets(plan.targets(), RuntimeMembarrierAction::RefreshRunQueue)?;
54    system.complete_membarrier_registration(plan);
55    Ok(())
56}
57
58/// Executes one Linux-style membarrier command.
59pub fn membarrier(
60    command: MembarrierCommand,
61) -> Result<(), crate::sync::membarrier::MembarrierError> {
62    validate_task_context()?;
63    let system = runtime_task_system()?;
64    let targets = MembarrierCpuTargets::new(system.cpu_topology_len());
65    let _pin = PreemptScope::enter();
66
67    // Matches the full barrier before Linux scans rq->curr.
68    core::sync::atomic::fence(core::sync::atomic::Ordering::SeqCst);
69    let targets = {
70        let mut irq = RuntimeIrqGuard::enter();
71        let mut cpu = runtime_current_cpu_mut(&mut irq)?;
72        let target = match command {
73            MembarrierCommand::Global => MembarrierTarget::Global,
74            MembarrierCommand::GlobalExpedited => MembarrierTarget::GlobalExpedited,
75            MembarrierCommand::PrivateExpedited => {
76                system.current_private_membarrier_target(cpu.as_mut())?
77            }
78        };
79        system.current_membarrier_targets(cpu.as_mut(), target, targets)?
80    };
81    synchronize_targets(targets.cpus(), RuntimeMembarrierAction::MemoryBarrier)?;
82    // Matches the full barrier after Linux waits for the final target IPI.
83    core::sync::atomic::fence(core::sync::atomic::Ordering::SeqCst);
84    Ok(())
85}
86
87/// Refreshes the calling CPU's `rq->membarrier_state` from its current task.
88///
89/// This fixed entry is invoked only by the runtime's registration hard-call.
90#[doc(hidden)]
91pub fn refresh_current_membarrier_run_queue() -> Result<(), TaskError> {
92    let mut irq = RuntimeIrqGuard::enter();
93    let mut cpu = runtime_current_cpu_mut(&mut irq)?;
94    runtime_task_system()?.refresh_current_membarrier_run_queue(cpu.as_mut())
95}