Skip to main content

strat9_kernel/async_io/
complete.rs

1//! Completion routing for async I/O.
2//!
3//! Provides a [`CompletionRouter`] that maps ring ids to their CQ buffers
4//! and pushes [`AsyncCqe`] entries atomically, waking any task waiting
5//! on the ring via `SYS_ASYNC_ENTER`.
6
7use super::{ops::AsyncCqe, ring::find_ring};
8
9fn cq_space_available(ring: &super::ring::Ring) -> u32 {
10    let tail = ring
11        .cq_meta()
12        .tail
13        .load(core::sync::atomic::Ordering::Relaxed);
14    let head = ring.cq_head();
15    ring.entries.saturating_sub(tail.wrapping_sub(head))
16}
17
18/// Write a single CQE into the ring's visible CQ and advance tail.
19///
20/// # Safety : caller must hold `ring.cq_lock`.
21fn write_cqe_at_tail(ring: &super::ring::Ring, cqe: AsyncCqe) {
22    let tail = ring
23        .cq_meta()
24        .tail
25        .load(core::sync::atomic::Ordering::Relaxed);
26    let mask = ring
27        .cq_meta()
28        .mask
29        .load(core::sync::atomic::Ordering::Relaxed);
30    let idx = (tail & mask) as usize;
31    unsafe {
32        let slot = ring.cqe_at(idx as u32);
33        core::ptr::write_volatile(slot, cqe);
34    }
35    ring.cq_advance_tail(1);
36}
37
38/// Drain as many entries from the backlog into the visible CQ as space allows.
39///
40/// # Safety : caller must hold `ring.cq_lock`.
41fn flush_backlog_locked(ring: &super::ring::Ring) -> u32 {
42    let mut backlog = ring.completion_backlog.lock();
43    let flush_count = core::cmp::min(backlog.len(), cq_space_available(ring) as usize);
44    if flush_count == 0 {
45        return 0;
46    }
47
48    for cqe in backlog.drain(..flush_count) {
49        write_cqe_at_tail(ring, cqe);
50    }
51
52    flush_count as u32
53}
54
55// ====  ring-id-based public API  ===============================================================
56
57/// Push a completion event to the ring identified by `ring_id`.
58///
59/// Returns `true` on success, `false` if the ring was not found or destroyed.
60pub fn push_completion(ring_id: u64, user_data: u64, result: i32, flags: u32) -> bool {
61    let Some(ring) = find_ring(ring_id) else {
62        return false;
63    };
64    push_completion_for_ring(&ring, user_data, result, flags)
65}
66
67/// Variant of [`push_completion`] that takes a resolved `&Ring` to avoid a
68/// redundant `find_ring` lookup in the dispatch hot path.
69pub fn push_completion_for_ring(
70    ring: &super::ring::Ring,
71    user_data: u64,
72    result: i32,
73    flags: u32,
74) -> bool {
75    // Reject if destroyed
76    if ring.destroyed.load(core::sync::atomic::Ordering::Acquire) != 0 {
77        return false;
78    }
79
80    let _guard = ring.cq_lock.lock();
81
82    let cqe = AsyncCqe {
83        user_data,
84        res: result,
85        flags,
86    };
87
88    // Single acquisition of the backlog lock.
89    let mut backlog = ring.completion_backlog.lock();
90
91    if !backlog.is_empty() {
92        // Backlog already has entries: always append to it.
93        backlog.push(cqe);
94        drop(backlog);
95        flush_backlog_locked(ring);
96    } else if cq_space_available(ring) == 0 {
97        // CQ visible to userspace is full so push to backlog.
98        backlog.push(cqe);
99        drop(backlog);
100    } else {
101        drop(backlog);
102        write_cqe_at_tail(ring, cqe);
103    }
104
105    // Decrement the in-flight counter
106    let _ = ring.in_flight.fetch_update(
107        core::sync::atomic::Ordering::Relaxed,
108        core::sync::atomic::Ordering::Relaxed,
109        |value| value.checked_sub(1),
110    );
111
112    // Wake any task waiting on this ring
113    ring.wq.wake_all();
114
115    true
116}
117
118/// Drain all pending completions for a ring and return them.
119/// Called by `SYS_ASYNC_ENTER` when userspace wants to consume CQEs.
120pub fn drain_completions(ring_id: u64, max: u32) -> u32 {
121    let Some(ring) = find_ring(ring_id) else {
122        return 0;
123    };
124
125    let _guard = ring.cq_lock.lock();
126
127    let head = ring.cq_head();
128    let tail = ring
129        .cq_meta()
130        .tail
131        .load(core::sync::atomic::Ordering::Acquire);
132    let available = tail.wrapping_sub(head).min(max);
133    if available == 0 {
134        return 0;
135    }
136
137    // No need to copy CQEs : userspace reads them directly from the
138    // shared CQ page. We just advance head to acknowledge consumption.
139    ring.cq_meta().head.store(
140        head.wrapping_add(available),
141        core::sync::atomic::Ordering::Release,
142    );
143
144    if flush_backlog_locked(&ring) > 0 {
145        ring.wq.wake_all();
146    }
147
148    available
149}
150
151/// Number of in-flight operations tracked for a ring.
152pub fn in_flight_count(ring_id: u64) -> u32 {
153    let Some(ring) = find_ring(ring_id) else {
154        return 0;
155    };
156    ring.in_flight.load(core::sync::atomic::Ordering::Relaxed)
157}
158
159#[inline]
160pub fn inc_in_flight(ring_id: u64) {
161    if let Some(ring) = find_ring(ring_id) {
162        inc_in_flight_for_ring(&ring);
163    }
164}
165
166/// Variant of [`inc_in_flight`] that takes a resolved `&Ring`.
167#[inline]
168pub fn inc_in_flight_for_ring(ring: &super::ring::Ring) {
169    ring.in_flight
170        .fetch_add(1, core::sync::atomic::Ordering::Relaxed);
171}
172
173#[inline]
174pub fn dec_in_flight(ring_id: u64) {
175    if let Some(ring) = find_ring(ring_id) {
176        ring.in_flight
177            .fetch_sub(1, core::sync::atomic::Ordering::Relaxed);
178    }
179}