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
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
//! Bareset local-shard cancel/single-shot/re-arm
//! cycle handlers. Split out of [`crate::uring_bigbulk`] so that file
//! stays under the 500-LOC house rule; every method here is on the
//! same `impl<C: Commands> Shard<C>`.
//!
//! Flow per local-shard bare-`SET key <BIG>`:
//!
//! 1. `try_promote_bigbulk` (in `uring_bigbulk.rs`) installs
//! `BigArgState::BareSetCancelling` + sets `big_arg_cancel_pending`.
//! 2. `uring_arm_conns` (in `uring_arm.rs`) submits the
//! `IORING_OP_ASYNC_CANCEL` SQE targeting the multishot recv.
//! 3. Two CQEs flip the cancel flags (any order):
//! - `OP_BIG_CANCEL` → [`Shard::uring_on_big_arg_cancel`] sets
//! `cancel_acked`.
//! - Terminal `OP_RECV` with `res = -ECANCELED` →
//! [`Shard::uring_on_big_arg_target_canceled`] sets
//! `target_canceled`.
//! 4. Both flags set → [`Shard::transition_to_reading`] flips the
//! state to `BareSetReading` + sets `big_arg_read_pending`.
//! 5. `uring_arm_conns` submits the single-shot `prep_read` SQE
//! pointing at `body.as_mut_ptr().add(body.len())` for the
//! remaining bytes. Kernel writes recv bytes directly into the
//! body Vec — no userspace memcpy.
//! 6. `OP_BIG_READ` CQE → [`Shard::uring_on_big_arg_read`] advances
//! body via `set_len`. If incomplete, re-schedule prep_read; if
//! complete, [`Shard::dispatch_bareset_owned`] runs the
//! SET + sets `big_arg_rearm_recv`.
//! 7. `uring_arm_conns` re-arms the multishot for the next request.
use crate::Commands;
use crate::shard::Shard;
use crate::uring_conn::{BigArgState, UringConn};
use kevy_map::KevyMap;
impl<C: Commands> Shard<C> {
/// Handler for `OP_BIG_READ` CQE: extend the
/// body Vec by the kernel-reported byte count (the kernel wrote
/// directly into `body.as_mut_ptr().add(body.len())` for `res`
/// bytes). If body still incomplete, mark the conn for another
/// `prep_read` on the next arm pass; if complete, dispatch + mark
/// for multishot re-arm.
// LOC-WAIVER: kernel-direct big-arg read completion state machine
// — body fill / CRLF account / dispatch+re-arm, one unit.
pub(crate) fn uring_on_big_arg_read(
&mut self,
cid: u64,
res: i32,
io: &mut KevyMap<u64, UringConn>,
) {
let Some(uc) = io.get_mut(&cid) else { return };
// The completion IS the kernel handing the buffer back, on every
// path below including error and EOF — so this clears first and
// unconditionally rather than once per branch.
uc.big_read_inflight = false;
if res <= 0 {
// EOF or error mid-body — drop the conn (mirrors
// `uring_on_recv` semantics; partial-body state is
// unrecoverable here).
uc.pending_big_arg = None;
uc.big_arg_read_pending = false;
uc.big_arg_rearm_recv = false;
self.uring_mark_closing(cid, io);
return;
}
let Some(state) = uc.pending_big_arg.as_mut() else { return };
let BigArgState::BareSetReading { body, body_len, crlf_seen, .. } = state.as_mut() else {
// Not in reading phase — defensive ignore.
return;
};
// The kernel-direct read landed `n` bytes; route into body
// first (preserving `body.len() ≤ body.capacity() == body_len`),
// then bump `crlf_seen` for any trailing CRLF bytes that
// arrived in the same CQE.
let n = res as usize;
let body_room = *body_len - body.len();
let body_n = n.min(body_room);
if body_n > 0 {
// SAFETY: kernel wrote into `body.as_mut_ptr().add(body.len())`
// for at most `body_room` bytes (the arm-pass `prep_read`
// submission caps the SQE length).
unsafe {
body.set_len(body.len() + body_n);
}
}
let crlf_n = ((n - body_n).min(2 - *crlf_seen as usize)) as u8;
*crlf_seen += crlf_n;
if body.len() == *body_len {
// Body fully received. Trailing CRLF (if not yet seen) is
// still in the TCP buffer — set `pending_crlf_skip` so the
// re-armed multishot's slab head gets sliced before
// dispatch. Dispatch + re-arm multishot now (body Vec is
// zero-copy-adoptable: len == capacity == body_len).
let crlf_pending_after_dispatch = 2 - *crlf_seen as usize;
if let Some(boxed) = uc.pending_big_arg.take()
&& let BigArgState::BareSetReading { key, body, body_len, .. } = *boxed
{
self.dispatch_bareset_owned(cid, key, body, body_len, io);
}
if let Some(uc) = io.get_mut(&cid) {
uc.big_arg_read_pending = false;
uc.big_arg_rearm_recv = true;
uc.pending_crlf_skip = crlf_pending_after_dispatch as u8;
}
self.mark_arm_pending(cid, io);
} else {
// More body bytes pending — schedule another prep_read.
uc.big_arg_read_pending = true;
self.mark_arm_pending(cid, io);
}
}
/// Handler for `OP_BIG_CANCEL` CQE: mark the
/// cancel side ack'd. If the target ECANCELED has also been seen,
/// transition to `BareSetReading` + schedule the single-shot read.
pub(crate) fn uring_on_big_arg_cancel(
&mut self,
cid: u64,
_res: i32,
io: &mut KevyMap<u64, UringConn>,
) {
// res may be 0 (matched-cancel), -ENOENT (target already gone),
// or -EALREADY (target executing). All three end the cancel
// side — proceed to transition checks.
let Some(uc) = io.get_mut(&cid) else { return };
// The multishot recv can self-terminate (buffer-ring ENOBUFS /
// EOF) in the window between the cancel submission and this ack.
// Its terminal CQE is then NOT -ECANCELED, so the
// `target_canceled` path never fires and the cancel completes
// -ENOENT. `recv_armed == false` is the authoritative "multishot
// is gone" signal — `uring_on_recv` clears it on EVERY terminal
// (cancel or not), and that terminal always precedes the -ENOENT
// ack, so it is already false here. Treat it as the target side
// being done: waiting only on `target_canceled` wedged the conn
// forever in BareSetCancelling under a deep pipeline of big-arg
// SETs (captured: big_arg=true recv_armed=false, target_canceled
// never set).
let multishot_gone = !uc.recv_armed;
let Some(state) = uc.pending_big_arg.as_mut() else {
// The body completed via multishot slabs while the cancel
// was in flight — request a multishot re-arm so the conn
// returns to normal mode.
uc.big_arg_rearm_recv = true;
self.mark_arm_pending(cid, io);
return;
};
let BigArgState::BareSetCancelling { cancel_acked, target_canceled, .. } = state.as_mut()
else {
return;
};
*cancel_acked = true;
if *cancel_acked && (*target_canceled || multishot_gone) {
self.transition_to_reading(cid, io);
}
}
/// Called by `uring_on_recv` when the multishot
/// recv's terminal CQE arrives with `res == -ECANCELED`. Mirrors
/// [`Self::uring_on_big_arg_cancel`] on the target-side flag.
pub(crate) fn uring_on_big_arg_target_canceled(
&mut self,
cid: u64,
io: &mut KevyMap<u64, UringConn>,
) {
let Some(uc) = io.get_mut(&cid) else { return };
let Some(state) = uc.pending_big_arg.as_mut() else {
uc.big_arg_rearm_recv = true;
self.mark_arm_pending(cid, io);
return;
};
let BigArgState::BareSetCancelling { cancel_acked, target_canceled, .. } = state.as_mut()
else {
return;
};
*target_canceled = true;
// Multishot is gone — caller (`uring_on_recv`) already sets
// `recv_armed = false` on !has_more; redundant here for clarity.
uc.recv_armed = false;
if *cancel_acked && *target_canceled {
self.transition_to_reading(cid, io);
}
}
/// `BareSetCancelling` → `BareSetReading`
/// transition: the multishot is fully drained; queue the
/// single-shot `prep_read` for any remaining body bytes. If the
/// body completed via in-flight multishot CQEs BEFORE the
/// transition fired, dispatch immediately and request re-arm.
pub(crate) fn transition_to_reading(&mut self, cid: u64, io: &mut KevyMap<u64, UringConn>) {
let Some(uc) = io.get_mut(&cid) else { return };
let Some(state) = uc.pending_big_arg.take() else { return };
let BigArgState::BareSetCancelling { key, body, body_len, crlf_seen, .. } = *state else {
return;
};
if body.len() == body_len && crlf_seen == 2 {
// Body already complete (last multishot CQE finished it
// before transition fired) — dispatch + re-arm.
self.dispatch_bareset_owned(cid, key, body, body_len, io);
if let Some(uc) = io.get_mut(&cid) {
uc.big_arg_rearm_recv = true;
}
self.mark_arm_pending(cid, io);
return;
}
uc.pending_big_arg =
Some(Box::new(BigArgState::BareSetReading { key, body, body_len, crlf_seen }));
uc.big_arg_read_pending = true;
self.mark_arm_pending(cid, io);
}
/// Dispatch a bare `SET key <BIG>` command with
/// an owned body Vec. Strips the trailing CRLF, runs all post-write
/// hooks (AOF / replication / keyspace notify / BLOCK wake / WATCH
/// bump / Lua wake bridge) on a borrowed three-slice argv view,
/// then hands the Vec to `store.set` (consumed). Reply `+OK\r\n`
/// goes to `conn.output`; caller marks arm-pending for the write
/// SQE.
pub(crate) fn dispatch_bareset_owned(
&mut self,
cid: u64,
key: Vec<u8>,
body: Vec<u8>,
body_len: usize,
io: &mut KevyMap<u64, UringConn>,
) {
// body's capacity invariant: `len == capacity == body_len` (CRLF
// was sunk into `crlf_seen`, never appended). That's the
// requirement for `pick_value_for_set_owned`'s
// `Vec::into_boxed_slice` to be zero-copy.
debug_assert_eq!(body.len(), body_len);
debug_assert_eq!(body.capacity(), body_len);
let view = ThreeSliceView { verb: b"SET", key: &key, body: &body };
// No propagation-override take here (cf. post_write_housekeeping):
// this path records a literal `SET` it built itself — deterministic,
// never SPOP — and none can be pending anyway: every other dispatch
// consumes its own override before this runs.
let w0 = self.always_hold_w0();
if self.aof.is_some() {
self.log_write(&view);
}
if let Some(src) = self.replicate.as_mut().map(|f| f.source_mut())
&& !crate::replication_gate::is_applying_replicated()
{
src.push_mutation(&view);
}
self.maybe_notify_dispatch(&view);
self.wake_key(&key);
let _ok = self.store.set(&key, body, None, false, false);
self.note_key_mutated(&key);
let lua_wakes = crate::lua_wake_bridge::drain_lua_wake_buffer();
for k in lua_wakes {
self.wake_key(&k);
}
if let Some(c) = self.conns.get_mut(&cid) {
c.output.extend_from_slice(b"+OK\r\n");
}
self.uring_stamp_hold(w0, cid, io);
self.mark_arm_pending(cid, io);
}
}
// =====================================================================
// Three-slice borrowed ArgvView for the bareset fast
// path. Implements `kevy_resp::ArgvView` so AOF / replication /
// keyspace-notification hooks accept it without materialising an owned
// `Argv` (which would memcpy the 64 KiB body).
// =====================================================================
pub(crate) struct ThreeSliceView<'a> {
pub(crate) verb: &'a [u8],
pub(crate) key: &'a [u8],
pub(crate) body: &'a [u8],
}
impl<'a> core::ops::Index<usize> for ThreeSliceView<'a> {
type Output = [u8];
#[expect(clippy::panic, reason = "Index's contract is to panic")]
fn index(&self, i: usize) -> &[u8] {
match i {
0 => self.verb,
1 => self.key,
2 => self.body,
// `Index::index` has no fallible form — panicking out of range
// is the trait's contract, the same one `[T]` and `Vec` keep.
_ => panic!("ThreeSliceView index oob: {i}"),
}
}
}
impl<'a> kevy_resp::ArgvView for ThreeSliceView<'a> {
fn len(&self) -> usize {
3
}
fn get(&self, i: usize) -> Option<&[u8]> {
match i {
0 => Some(self.verb),
1 => Some(self.key),
2 => Some(self.body),
_ => None,
}
}
}