ax_net/queue_runtime/
state.rs

1use alloc::sync::Arc;
2use core::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, AtomicUsize, Ordering};
3
4use super::{
5    QueueNotification, STATE_DISABLED, STATE_IDLE, STATE_MASK, STATE_MISSED, STATE_POLLING,
6    STATE_SCHEDULED,
7};
8
9/// Observable queue statistics used by SMP contract tests.
10#[derive(Clone, Copy, Debug, Eq, PartialEq)]
11pub struct NetQueueStats {
12    pub irq: u64,
13    pub schedule: u64,
14    pub missed: u64,
15    pub poll_batches: u64,
16    pub budget_exhaustion: u64,
17    pub spurious: u64,
18    pub probe_deferred: u64,
19    pub rearm_race: u64,
20    pub owner_cpu: usize,
21    pub last_irq_cpu: Option<usize>,
22    pub last_poll_cpu: Option<usize>,
23    pub irq_to_poll_remote_wake: u64,
24}
25
26pub(super) struct QueueStatsAtomic {
27    pub(super) irq: AtomicU64,
28    pub(super) schedule: AtomicU64,
29    pub(super) missed: AtomicU64,
30    pub(super) poll_batches: AtomicU64,
31    pub(super) budget_exhaustion: AtomicU64,
32    pub(super) spurious: AtomicU64,
33    pub(super) probe_deferred: AtomicU64,
34    pub(super) rearm_race: AtomicU64,
35    pub(super) last_irq_cpu: AtomicUsize,
36    pub(super) last_poll_cpu: AtomicUsize,
37    pub(super) irq_to_poll_remote_wake: AtomicU64,
38}
39
40impl QueueStatsAtomic {
41    const fn new() -> Self {
42        Self {
43            irq: AtomicU64::new(0),
44            schedule: AtomicU64::new(0),
45            missed: AtomicU64::new(0),
46            poll_batches: AtomicU64::new(0),
47            budget_exhaustion: AtomicU64::new(0),
48            spurious: AtomicU64::new(0),
49            probe_deferred: AtomicU64::new(0),
50            rearm_race: AtomicU64::new(0),
51            last_irq_cpu: AtomicUsize::new(usize::MAX),
52            last_poll_cpu: AtomicUsize::new(usize::MAX),
53            irq_to_poll_remote_wake: AtomicU64::new(0),
54        }
55    }
56
57    pub(super) fn snapshot(&self, owner_cpu: usize) -> NetQueueStats {
58        let optional_cpu = |cpu| (cpu != usize::MAX).then_some(cpu);
59        NetQueueStats {
60            irq: self.irq.load(Ordering::Relaxed),
61            schedule: self.schedule.load(Ordering::Relaxed),
62            missed: self.missed.load(Ordering::Relaxed),
63            poll_batches: self.poll_batches.load(Ordering::Relaxed),
64            budget_exhaustion: self.budget_exhaustion.load(Ordering::Relaxed),
65            spurious: self.spurious.load(Ordering::Relaxed),
66            probe_deferred: self.probe_deferred.load(Ordering::Relaxed),
67            rearm_race: self.rearm_race.load(Ordering::Relaxed),
68            owner_cpu,
69            last_irq_cpu: optional_cpu(self.last_irq_cpu.load(Ordering::Acquire)),
70            last_poll_cpu: optional_cpu(self.last_poll_cpu.load(Ordering::Acquire)),
71            irq_to_poll_remote_wake: self.irq_to_poll_remote_wake.load(Ordering::Relaxed),
72        }
73    }
74}
75
76/// Shared atomic state for one poll group.
77pub(super) struct PollGroupState {
78    pub(super) state: AtomicU8,
79    startup_absent: AtomicBool,
80    pub(super) owner_cpu: usize,
81    notify: Arc<QueueNotification>,
82    pub(super) stats: QueueStatsAtomic,
83    rx_drops: AtomicU64,
84}
85
86impl PollGroupState {
87    pub(super) fn new(owner_cpu: usize, notify: Arc<QueueNotification>) -> Self {
88        Self {
89            state: AtomicU8::new(STATE_DISABLED),
90            startup_absent: AtomicBool::new(false),
91            owner_cpu,
92            notify,
93            stats: QueueStatsAtomic::new(),
94            rx_drops: AtomicU64::new(0),
95        }
96    }
97
98    pub(super) fn mark_startup_absent(&self) {
99        // The owner publishes this only after startup cancellation or shutdown has proved
100        // that the unpublished group can be released. The builder's acquire
101        // load precedes IRQ synchronization and removal of protocol endpoints.
102        self.startup_absent.store(true, Ordering::Release);
103    }
104
105    pub(super) fn startup_absent(&self) -> bool {
106        self.startup_absent.load(Ordering::Acquire)
107    }
108
109    pub(super) fn record_rx_drop(&self) {
110        // Statistics only; packet ownership is published through the queues.
111        self.rx_drops.fetch_add(1, Ordering::Relaxed);
112    }
113
114    pub(super) fn take_rx_drops(&self) -> u64 {
115        self.rx_drops.swap(0, Ordering::Relaxed)
116    }
117
118    pub(super) fn activate(&self, pending: bool) {
119        self.state.store(STATE_IDLE, Ordering::Release);
120        if pending {
121            self.schedule_task();
122        }
123    }
124
125    pub(super) fn schedule_irq(&self) {
126        let cpu = ax_hal::percpu::this_cpu_id();
127        self.stats.irq.fetch_add(1, Ordering::Relaxed);
128        self.stats.last_irq_cpu.store(cpu, Ordering::Release);
129        if self.startup_absent() {
130            return;
131        }
132        if cpu != self.owner_cpu {
133            self.stats
134                .irq_to_poll_remote_wake
135                .fetch_add(1, Ordering::Relaxed);
136            self.disable();
137            return;
138        }
139        if self.is_disabled() {
140            // During owner startup queues stay disabled, but the startup
141            // state machine still needs the IRQ notification to advance.
142            self.notify.notify();
143        } else if self.publish_schedule() {
144            self.notify.notify();
145        }
146    }
147
148    pub(super) fn wait_startup_irq(&self, waiter: &ax_task::sync::irq::IrqWorkerWaiter) {
149        self.notify.wait(waiter);
150    }
151
152    pub(super) fn wait_startup_deadline(
153        &self,
154        waiter: &ax_task::sync::irq::IrqWorkerWaiter,
155        deadline_nanos: u64,
156    ) {
157        let now = ax_hal::time::monotonic_time_nanos();
158        if deadline_nanos > now {
159            let duration = core::time::Duration::from_nanos(deadline_nanos - now);
160            self.notify.wait_timeout(waiter, duration);
161        }
162    }
163
164    pub(super) fn schedule_task(&self) {
165        self.publish_schedule();
166        // A task-side publication can be what releases a queue executor that
167        // stopped on RX/TX ring backpressure. In that case the state is
168        // POLLING|MISSED rather than a fresh IDLE->SCHEDULED transition, but
169        // the sleeping owner still needs a precise wakeup.
170        if !self.is_disabled() {
171            self.notify.notify();
172        }
173    }
174
175    fn publish_schedule(&self) -> bool {
176        loop {
177            let old = self.state.load(Ordering::Acquire);
178            match old & STATE_MASK {
179                STATE_DISABLED => return false,
180                STATE_IDLE => {
181                    if self
182                        .state
183                        .compare_exchange(old, STATE_SCHEDULED, Ordering::AcqRel, Ordering::Acquire)
184                        .is_ok()
185                    {
186                        self.stats.schedule.fetch_add(1, Ordering::Relaxed);
187                        return true;
188                    }
189                }
190                STATE_SCHEDULED | STATE_POLLING => {
191                    if old & STATE_MISSED != 0 {
192                        return false;
193                    }
194                    if self
195                        .state
196                        .compare_exchange(
197                            old,
198                            old | STATE_MISSED,
199                            Ordering::AcqRel,
200                            Ordering::Acquire,
201                        )
202                        .is_ok()
203                    {
204                        self.stats.missed.fetch_add(1, Ordering::Relaxed);
205                        return false;
206                    }
207                }
208                _ => return false,
209            }
210        }
211    }
212
213    pub(super) fn claim(&self) -> bool {
214        let current_cpu = ax_hal::percpu::this_cpu_id();
215        if current_cpu != self.owner_cpu {
216            self.disable();
217            return false;
218        }
219        loop {
220            let old = self.state.load(Ordering::Acquire);
221            let claimable = (old & STATE_MASK == STATE_SCHEDULED)
222                || (old & STATE_MASK == STATE_POLLING && old & STATE_MISSED != 0);
223            if !claimable {
224                return false;
225            }
226            if self
227                .state
228                .compare_exchange(old, STATE_POLLING, Ordering::AcqRel, Ordering::Acquire)
229                .is_ok()
230            {
231                self.stats
232                    .last_poll_cpu
233                    .store(current_cpu, Ordering::Release);
234                self.stats.poll_batches.fetch_add(1, Ordering::Relaxed);
235                return true;
236            }
237        }
238    }
239
240    pub(super) fn finish_more(&self) {
241        loop {
242            let old = self.state.load(Ordering::Acquire);
243            if old & STATE_MASK != STATE_POLLING {
244                return;
245            }
246            if self
247                .state
248                .compare_exchange(old, STATE_SCHEDULED, Ordering::AcqRel, Ordering::Acquire)
249                .is_ok()
250            {
251                return;
252            }
253        }
254    }
255
256    pub(super) fn begin_rearm(&self) -> bool {
257        loop {
258            let old = self.state.load(Ordering::Acquire);
259            if old & STATE_MASK != STATE_POLLING {
260                return false;
261            }
262            if old & STATE_MISSED != 0 {
263                if self
264                    .state
265                    .compare_exchange(old, STATE_SCHEDULED, Ordering::AcqRel, Ordering::Acquire)
266                    .is_ok()
267                {
268                    return false;
269                }
270                continue;
271            }
272            if self
273                .state
274                .compare_exchange(old, STATE_IDLE, Ordering::AcqRel, Ordering::Acquire)
275                .is_ok()
276            {
277                return true;
278            }
279        }
280    }
281
282    pub(super) fn disable(&self) {
283        self.state.store(STATE_DISABLED, Ordering::Release);
284        self.notify.notify();
285    }
286
287    pub(super) fn is_disabled(&self) -> bool {
288        self.state.load(Ordering::Acquire) & STATE_MASK == STATE_DISABLED
289    }
290}
ax_net/queue_runtime/state.rs

ax_net/queue_runtime/
state.rs
