1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
//! The cross-core drain + connection-reap half of the io_uring reactor.
//! Split out of [`crate::uring_reactor`] to keep that file under the
//! 500-LOC house rule — every method here is on the same
//! `impl<C: Commands> Shard<C>` and only ever called from `run_uring`.
use crate::Commands;
use crate::shard::Shard;
use crate::uring_reactor::UringConn;
use core::sync::atomic::Ordering;
use kevy_map::KevyMap;
impl<C: Commands> Shard<C> {
/// Drain cross-core rings: execute forwarded requests, fold replies into
/// their connection's output (no direct write — the io_uring arm/write
/// loop flushes it). The message handling itself is
/// [`Shard::drain_inbound_core_slow`], shared with the epoll reactor.
///
/// Fast-path split: a perf
/// diagnostic showed this at 3.59 % self — almost all from the per-iter
/// fn call overhead despite the cheap Acquire load inside. Now the
/// Acquire load lives here in a tiny `#[inline]` wrapper that LLVM
/// folds into the reactor loop body; the cold drain body is
/// `#[inline(never)]` so its bulk stays off the hot iTLB pages.
#[inline]
pub(crate) fn uring_drain_inbound(&mut self) -> usize {
let me = self.id;
if self.inbound_dirty[me].load(Ordering::Acquire) == 0 {
return 0;
}
self.drain_inbound_core_slow::<false>()
.expect("DIRECT_FLUSH=false drain has no fallible step")
}
/// Close connections that are done: EOF/QUIT seen, all output flushed, no
/// SQE in flight. Dropping the `Conn` closes the fd.
///
/// An earlier attempt tried a two-`any()`-scan fast-path bail (skip
/// the Vec collect when no conn carries a closing flag) and reverted —
/// at c100 the 2×N pre-scan added more cost than the avoided alloc
/// saved (measured -2.9 % on the bench box's c100 SET shape), and
/// the only sound way to use a
/// single scan is to keep io.closing + conn.closing in sync (which
/// requires plumbing the io map down into the dispatch QUIT path).
/// Left for a future iteration that's willing to take that plumb.
// LOC-WAIVER: closing ready-set reap state machine (classify / requeue /
// shared teardown) — io_uring-only path, unverifiable on darwin;
// waived rather than split without a runnable test surface.
pub(crate) fn uring_reap_closed(&mut self, io: &mut KevyMap<u64, UringConn>) {
// Drain the closing ready-set instead of
// walking the whole io map. perf-record-dwarf at c=10 000 -P 1
// SET sustained showed the prior `io.iter().filter(...).map(
// |(cid, _)| (cid, self.conns.get(cid))).collect::<Vec<u64>>()`
// body at 36.74 % of CPU — pure O(N) scan + per-entry second
// hash lookup into `self.conns`. With the ready-set populated
// by `uring_mark_closing` + the QUIT dispatch sites, this is
// O(closing) per reap pass — typically 0-few entries at any
// moment.
//
// A conn that is not yet quiet goes back on the set's tail.
let candidates: Vec<u64> = std::mem::take(&mut self.closing_uring_conns);
let mut done: Vec<u64> = Vec::with_capacity(candidates.len());
let mut requeue: Vec<u64> = Vec::new();
for cid in candidates {
// Already reaped (e.g. dedup on a doubly-pushed cid)?
let Some(uc) = io.get(&cid) else { continue };
let conn = self.conns.get(&cid);
// Sanity: cid was pushed because something flipped closing — but
// accept-fail / EOF races could land it without `closing == true`.
// Skip non-closing rather than reap.
if !(uc.closing || conn.is_some_and(|c| c.closing)) {
continue;
}
if closing_conn_is_quiet(uc, conn) {
done.push(cid);
} else {
requeue.push(cid);
}
}
// Restore retries for the next reap pass.
self.closing_uring_conns.append(&mut requeue);
for cid in done {
// Use the shared teardown (not a local conns.remove): it also
// cancels block waiters (local + cross-shard arbiter) and drops
// pub/sub + pattern subscriptions. Skipping it leaked a parked
// BLPOP/XREAD waiter and psub registrations on every io_uring
// disconnect — a waiter left behind could consume a later push
// meant for a live client. The epoll-only `poller.delete` /
// `fd_to_conn` steps inside are harmless no-ops here (io_uring
// never registered the fd with the readiness poller).
self.close_conn(cid);
io.remove(&cid);
// No per-conn list to maintain. A stale
// entry in `arm_pending` for `cid` is a no-op next iter
// (the arm loop bails when both `conns.get_mut(&cid)` and
// `io.get_mut(&cid)` return None).
}
}
}
/// Whether a closing conn is finished with the ring and can be torn
/// down: nothing of it still in flight, nothing of it still unsent.
///
/// The recv term is the one that was missing. [`Shard::uring_arm_conns`]
/// cancels a closing conn's multishot recv precisely so that `close(fd)`
/// sends a FIN, and its comment ends "the next reap closes cleanly" —
/// but the next reap did not look at `recv_armed`. A reap landing in the
/// window between the cancel being submitted and its terminal CQE
/// arriving tore the conn down with the multishot still armed; the
/// socket stayed pinned in the kernel, `close(fd)` sent nothing, and a
/// client the server had decided to disconnect waited on a live socket
/// forever.
///
/// Measured rather than reasoned. The query-buffer cell failed 5-7 times
/// in 100 on Linux/io_uring, and in every one of those the reactor's own
/// stall dump reported `conns=0` for all 121 heartbeats spanning the
/// client's 30-second wait: the conn was already fully reaped while the
/// client still saw the socket open. That rules out everything upstream
/// of the reap and leaves the teardown itself.
///
/// Waiting here is bounded. The arm loop re-issues the cancel on every
/// visit to a closing conn (idempotent — a redundant one returns
/// `-ENOENT`) and keeps closing conns queued, and `recv_armed` clears on
/// either that cancel's `-ECANCELED` or a multishot that stops
/// delivering. If it somehow did not clear, the conn stays in
/// `self.conns` and the stall dump names it — which is the failure worth
/// having, the alternative being the silent half-open leak this fixes.
fn closing_conn_is_quiet(uc: &UringConn, conn: Option<&crate::conn::Conn>) -> bool {
let drained =
conn.is_none_or(|c| c.output.is_empty() && c.pending.is_empty() && c.write_pos == 0);
let writes_quiet = !uc.write_inflight && uc.write_buf.is_empty();
let recv_quiet = !uc.recv_armed;
// The kernel-direct big-arg read holds a raw pointer into the body
// Vec this conn owns. Reaping while it is in flight frees that Vec
// under the kernel — a use-after-free the Rust side cannot see and
// no test would report as anything but corruption somewhere else.
let big_read_quiet = !uc.big_read_inflight;
writes_quiet && recv_quiet && big_read_quiet && drained
}
#[cfg(test)]
mod tests {
use super::closing_conn_is_quiet;
use crate::uring_conn::UringConn;
/// A `Conn` for the cases that need one. `Conn::new` reads only
/// `peer_addr`, which is allowed to fail, so a listener on an
/// ephemeral port is the cheapest socket that will do — nothing
/// below touches the socket itself.
fn a_conn() -> crate::conn::Conn {
crate::conn::Conn::new(kevy_sys::tcp_listen([127, 0, 0, 1], 0, 1).unwrap())
}
fn a_pending_slot() -> crate::message_agg::PendingSlot {
crate::message_agg::PendingSlot {
remaining: 0,
agg: crate::message_agg::Agg::SumInt(0),
done: None,
proto: kevy_resp::RespVersion::default(),
}
}
/// A `None` conn is a conn already gone from `self.conns`, which the
/// reap treats as drained — so these cases isolate the three terms
/// that live on the `UringConn` side.
#[test]
fn a_fresh_conn_with_nothing_outstanding_is_quiet() {
assert!(closing_conn_is_quiet(&UringConn::new(), None));
}
/// The term this function was extracted to add. An armed multishot
/// recv pins the socket in the kernel, so reaping here closes the
/// descriptor without a FIN ever reaching the client — the
/// query-buffer disconnect that was decided and never landed.
#[test]
fn an_armed_recv_is_not_quiet() {
let mut uc = UringConn::new();
uc.recv_armed = true;
assert!(!closing_conn_is_quiet(&uc, None), "reaped with the recv still armed");
}
/// The term added after `recv_armed`, and found the same way: by
/// asking what else the kernel could still be holding.
///
/// A kernel-direct big-arg read has a raw pointer into the body Vec
/// that `pending_big_arg` owns. Reaping frees it, and the kernel
/// then writes the client's SET body into freed memory. `CLIENT KILL`
/// against a connection stalled mid-body reaches it, and the window
/// is as long as the client cares to hold it.
///
/// Note which flag does NOT protect this: `big_arg_read_pending`
/// means "queue an SQE next pass" and is cleared on submit, so it is
/// false for exactly the dangerous window.
#[test]
fn a_kernel_direct_big_read_in_flight_is_not_quiet() {
let mut uc = UringConn::new();
uc.big_read_inflight = true;
assert!(!closing_conn_is_quiet(&uc, None), "reaped with a read in the kernel");
}
/// And the flag that looks like it should have covered it does not,
/// stated as an assertion rather than left to a reader: a conn
/// waiting to QUEUE a read owns its body outright and is safe to
/// reap. Only a submitted one is not.
#[test]
fn a_big_read_merely_wanted_is_still_quiet() {
let mut uc = UringConn::new();
uc.big_arg_read_pending = true;
assert!(closing_conn_is_quiet(&uc, None));
}
#[test]
fn a_write_in_flight_is_not_quiet() {
let mut uc = UringConn::new();
uc.write_inflight = true;
assert!(!closing_conn_is_quiet(&uc, None));
}
#[test]
fn unsent_bytes_in_write_buf_are_not_quiet() {
let mut uc = UringConn::new();
uc.write_buf.push(b'x');
assert!(!closing_conn_is_quiet(&uc, None));
}
/// The `Some` side of the same question, and the reason it is not
/// covered by the cases above: `drained` is three terms of its own,
/// and the reap tears a conn down on all three.
#[test]
fn a_conn_with_nothing_left_to_send_is_quiet() {
assert!(closing_conn_is_quiet(&UringConn::new(), Some(&a_conn())));
}
#[test]
fn unsent_output_is_not_quiet() {
let mut c = a_conn();
c.output.push(b'x');
assert!(!closing_conn_is_quiet(&UringConn::new(), Some(&c)));
}
#[test]
fn an_unemitted_reply_is_not_quiet() {
let mut c = a_conn();
c.pending.push_back(a_pending_slot());
assert!(!closing_conn_is_quiet(&UringConn::new(), Some(&c)));
}
/// A half-written reply: `output` has been consumed up to
/// `write_pos`, so emptiness alone would call this drained.
#[test]
fn a_partly_written_reply_is_not_quiet() {
let mut c = a_conn();
c.write_pos = 1;
assert!(!closing_conn_is_quiet(&UringConn::new(), Some(&c)));
}
}