net/adapter/net/cortex/rpc.rs
1//! nRPC — request/response on top of CortEX folds.
2//!
3//! See `docs/misc/NRPC_DESIGN.md` for the full architectural framing.
4//! In short: an RPC server is a `RedexFold` whose state is the
5//! in-flight call set, whose events are typed `(REQUEST, RESPONSE,
6//! CANCEL, DEADLINE_EXCEEDED)`, whose `EventMeta::seq_or_ts` is the
7//! correlation id, and whose `EventMeta::origin_hash` is the
8//! AEAD-verified caller. The mesh-channel layer's queue-group
9//! subscription mode (see `channel::SubscriptionMode`) does the
10//! one-of-N work distribution across replica servers.
11//!
12//! This module is the **wire codec layer**: dispatch constants for
13//! `EventMeta::dispatch`, payload structs for `RpcRequestPayload` /
14//! `RpcResponsePayload`, and the `RpcStatus` enumeration. The fold
15//! types and the `Mesh::serve_rpc` / `Mesh::call` glue layer build
16//! on top.
17
18use bytes::{Buf, BufMut, Bytes};
19use parking_lot::Mutex;
20use std::collections::HashMap;
21use std::sync::atomic::{AtomicBool, Ordering};
22use std::sync::Arc;
23use tokio::sync::Notify;
24
25use super::super::redex::{RedexError, RedexEvent, RedexFold};
26use super::meta::{EventMeta, EVENT_META_SIZE};
27
28// ============================================================================
29// `EventMeta::dispatch` byte assignments for nRPC.
30//
31// All four values live in the cortex-internal range (`0x00..0x7F`).
32// Application/vendor dispatches stay in `0x80..0xFF`. Adapters that
33// don't care about RPC ignore unknown dispatches as they ignore any
34// other.
35// ============================================================================
36
37/// Caller → server. The first frame of an RPC call. `EventMeta::seq_or_ts`
38/// is the caller-generated `call_id`; `EventMeta::origin_hash` is the
39/// AEAD-verified caller. Payload is an [`RpcRequestPayload`].
40pub const DISPATCH_RPC_REQUEST: u8 = 0x10;
41
42/// Server → caller. The (terminal, for unary) frame of an RPC call.
43/// `EventMeta::seq_or_ts` matches the request's `call_id`. Payload is
44/// an [`RpcResponsePayload`].
45pub const DISPATCH_RPC_RESPONSE: u8 = 0x11;
46
47/// Caller → server. Cancellation signal. `EventMeta::seq_or_ts` matches
48/// the request's `call_id`. Empty payload — the dispatch byte plus
49/// the matching `call_id` is the whole signal. Server's fold removes
50/// the in-flight entry and (if cooperative) flips the handler's
51/// `CancellationToken`.
52pub const DISPATCH_RPC_CANCEL: u8 = 0x12;
53
54/// Server → caller. Deadline-exceeded signal. Emitted when the
55/// server's fold sees `now_ns() > request.deadline_ns` before
56/// starting the handler (or, optionally, when a long-running handler
57/// is aborted by the deadline timer). `EventMeta::seq_or_ts` matches
58/// the request's `call_id`. Empty payload.
59pub const DISPATCH_RPC_DEADLINE_EXCEEDED: u8 = 0x13;
60
61/// Caller → server. Stream credit grant. Carries a 4-byte
62/// big-endian `u32` in the payload after `EventMeta`: the number
63/// of additional response chunks the caller is willing to accept
64/// for the streaming call identified by `EventMeta::seq_or_ts`.
65///
66/// Only meaningful when the caller opted into flow control via
67/// the `nrpc-stream-window-initial` request header
68/// ([`HEADER_NRPC_STREAM_WINDOW_INITIAL`]). On a flow-controlled
69/// stream the server's pump task awaits one credit per chunk; on
70/// a non-flow-controlled stream (no header) the server ignores
71/// every GRANT.
72///
73/// Phase 3.
74pub const DISPATCH_RPC_STREAM_GRANT: u8 = 0x14;
75
76/// Caller → server. Continuation chunk of a client-streaming or
77/// duplex REQUEST. Carries an [`RpcRequestChunkPayload`] after the
78/// `EventMeta` prefix. `EventMeta::seq_or_ts` matches the initial
79/// REQUEST's `call_id`. Non-terminal chunks have
80/// `flags & FLAG_RPC_REQUEST_END == 0`; the terminal upload chunk
81/// sets [`FLAG_RPC_REQUEST_END`].
82///
83/// Only meaningful for calls whose initial REQUEST set
84/// [`FLAG_RPC_CLIENT_STREAMING_REQUEST`]; otherwise the server
85/// silently drops the chunk (caller bug; no observable effect).
86pub const DISPATCH_RPC_REQUEST_CHUNK: u8 = 0x15;
87
88/// Server → caller. Request-direction stream-credit grant. Mirror
89/// of [`DISPATCH_RPC_STREAM_GRANT`] for the upload direction.
90/// Carries an [`RpcRequestGrantPayload`] after `EventMeta`: a
91/// `call_id` plus a `u32` credit count. `EventMeta::seq_or_ts`
92/// matches the call's `call_id` (redundant with the payload, but
93/// kept symmetric with the rest of the dispatch family).
94///
95/// Only meaningful when the caller opted into request-direction
96/// flow control via the `nrpc-request-window-initial` header
97/// ([`HEADER_NRPC_REQUEST_WINDOW_INITIAL`]). Caller's sink
98/// awaits one credit per `REQUEST_CHUNK`; absent header →
99/// unbounded credit (sink emits as fast as the publish path can
100/// take it).
101pub const DISPATCH_RPC_REQUEST_GRANT: u8 = 0x16;
102
103// ============================================================================
104// `RpcRequestPayload::flags` bit assignments.
105// ============================================================================
106
107// Bit 0 (`1 << 0`) is RESERVED — was previously documented as
108// `FLAG_RPC_IDEMPOTENT`, but the server-side replay-cache (LRU of
109// `(origin_hash, call_id) -> RpcResponsePayload`) was never landed,
110// so the flag silently no-op'd despite a load-bearing contract in
111// its doc-string. Removed to avoid shipping a documented behavior
112// the runtime doesn't implement; reservation kept so a future
113// re-add (with the LRU) preserves wire compatibility.
114
115/// Set if the server may emit multiple `DISPATCH_RPC_RESPONSE` events
116/// for this call. Without it, the first response terminates the
117/// call. With it, each response except the terminal one carries
118/// `headers["nrpc-streaming"] = b"continue"`; the terminal response
119/// has either `b"end"` (success) or a non-`Ok` status.
120pub const FLAG_RPC_STREAMING_RESPONSE: u16 = 1 << 1;
121
122/// Set if the request carries W3C Trace Context headers
123/// (`traceparent`, `tracestate`). Server propagates them to its own
124/// span emission. Phase 3.
125pub const FLAG_RPC_PROPAGATE_TRACE: u16 = 1 << 2;
126
127// Bit `1 << 3` reserved — symmetric to the reserved bit 0 above,
128// kept as breathing room for a future protocol-level flag without
129// pushing every existing live bit.
130
131/// Set on the initial REQUEST if the caller will follow up with
132/// one or more [`DISPATCH_RPC_REQUEST_CHUNK`] events. Distinguishes
133/// client-streaming / duplex calls from unary at the very first
134/// frame so the server's fold knows to open a request-side stream
135/// instead of treating the REQUEST as complete.
136///
137/// Combined with [`FLAG_RPC_STREAMING_RESPONSE`] on the same
138/// REQUEST: full duplex.
139///
140/// Bidi streaming plan (Phase A).
141pub const FLAG_RPC_CLIENT_STREAMING_REQUEST: u16 = 1 << 4;
142
143/// Set on a [`DISPATCH_RPC_REQUEST_CHUNK`] (or on the initial
144/// REQUEST itself) to signal the terminal upload frame for a
145/// client-streaming or duplex call. After receiving this, the
146/// server's request-side stream yields `None` and the handler
147/// proceeds to its terminal response.
148///
149/// Setting this on the initial REQUEST is the degenerate "client-
150/// streaming with exactly one item" path — saves a round-trip
151/// for the trivial case.
152///
153/// Bidi streaming plan (Phase A).
154pub const FLAG_RPC_REQUEST_END: u16 = 1 << 5;
155
156// Bits `6..=15` reserved; producers MUST write zero, consumers MUST
157// ignore unknown bits (forward-compat with future flags).
158
159// ============================================================================
160// `RpcResponsePayload::status` enumeration.
161// ============================================================================
162
163/// Outcome of an nRPC call. Net-native numbering with documented
164/// gRPC equivalents (see comments). Numeric stability: callers and
165/// servers across versions agree on `0x0000..=0x7FFF`; the
166/// application-defined range is `0x8000..=0xFFFF`.
167#[derive(Debug, Clone, Copy, PartialEq, Eq)]
168#[repr(u16)]
169pub enum RpcStatus {
170 /// Success. Payload carries the application response. Terminal
171 /// (or, for streaming responses, may be one of many — see the
172 /// streaming flag).
173 /// gRPC equivalent: `OK` (0).
174 Ok = 0x0000,
175 /// No service registered with the requested name on the server.
176 /// gRPC equivalent: `NOT_FOUND` (5).
177 NotFound = 0x0001,
178 /// Caller's token doesn't list the requested service in scope,
179 /// or the channel-level capability check failed.
180 /// gRPC equivalent: `PERMISSION_DENIED` (7).
181 Unauthorized = 0x0002,
182 /// Server observed `now_ns() > deadline_ns` before starting work.
183 /// (For the in-flight case after the handler started, see
184 /// [`DISPATCH_RPC_DEADLINE_EXCEEDED`].)
185 /// gRPC equivalent: `DEADLINE_EXCEEDED` (4).
186 Timeout = 0x0003,
187 /// Server's per-service queue is at `max_in_flight` capacity.
188 /// gRPC equivalent: `RESOURCE_EXHAUSTED` (8).
189 Backpressure = 0x0004,
190 /// Caller emitted `DISPATCH_RPC_CANCEL` before the server
191 /// completed.
192 /// gRPC equivalent: `CANCELLED` (1).
193 Cancelled = 0x0005,
194 /// Handler panicked or returned an error not classified as one
195 /// of the above. Payload carries a UTF-8 diagnostic.
196 /// gRPC equivalent: `INTERNAL` (13).
197 Internal = 0x0006,
198 /// Request payload version not supported by the server. Should
199 /// normally be caught earlier by subprotocol-version
200 /// negotiation; the in-payload guard is the floor.
201 /// gRPC equivalent: `UNIMPLEMENTED` (12).
202 UnknownVersion = 0x0007,
203 /// v0.4 capability-auth: the target's `CapabilityAnnouncement`
204 /// either does not list the requested `nrpc:<service>` tag, or
205 /// lists it with allow-lists the caller does not match. See
206 /// `docs/plans/CAPABILITY_AUTH_PLAN.md` §3. Distinct from
207 /// `Unauthorized` (channel-auth / token-scope failures) so
208 /// operators can tell the two enforcement surfaces apart in
209 /// audit logs.
210 /// gRPC equivalent: `PERMISSION_DENIED` (7) — same outward
211 /// shape as `Unauthorized` but a separate substrate code.
212 CapabilityDenied = 0x0008,
213 /// Application-defined status. The wire carries the raw u16;
214 /// callers / servers agree on the meaning out of band.
215 Application(u16),
216}
217
218impl RpcStatus {
219 /// Encode to the wire `u16`.
220 pub fn to_wire(self) -> u16 {
221 match self {
222 Self::Ok => 0x0000,
223 Self::NotFound => 0x0001,
224 Self::Unauthorized => 0x0002,
225 Self::Timeout => 0x0003,
226 Self::Backpressure => 0x0004,
227 Self::Cancelled => 0x0005,
228 Self::Internal => 0x0006,
229 Self::UnknownVersion => 0x0007,
230 Self::CapabilityDenied => 0x0008,
231 Self::Application(v) => v,
232 }
233 }
234
235 /// Decode from the wire `u16`. Reserved values
236 /// (`0x0009..=0x7FFF`) decode as `Application(v)` rather than
237 /// failing — forward-compat with future status assignments.
238 pub fn from_wire(v: u16) -> Self {
239 match v {
240 0x0000 => Self::Ok,
241 0x0001 => Self::NotFound,
242 0x0002 => Self::Unauthorized,
243 0x0003 => Self::Timeout,
244 0x0004 => Self::Backpressure,
245 0x0005 => Self::Cancelled,
246 0x0006 => Self::Internal,
247 0x0007 => Self::UnknownVersion,
248 0x0008 => Self::CapabilityDenied,
249 other => Self::Application(other),
250 }
251 }
252
253 /// True iff `self == Ok`. Convenience for the hot caller-side
254 /// success-or-error branch.
255 #[inline]
256 pub fn is_ok(self) -> bool {
257 matches!(self, Self::Ok)
258 }
259}
260
261// ============================================================================
262// Request / response payloads.
263//
264// These ride in the bytes AFTER the 24-byte `EventMeta` prefix on a
265// CortEX-adapted event. The cortex adapter handles meta + tail
266// concatenation; this codec produces only the tail.
267// ============================================================================
268
269/// Header name + value pair. Used for trace-context propagation,
270/// idempotency-key carriage, content-type hints. Names are
271/// case-sensitive UTF-8; values are opaque bytes.
272pub type RpcHeader = (String, Vec<u8>);
273
274/// Maximum service-name length on the wire (matches
275/// `MAX_CHANNEL_NAME_LEN` upstream; reasonable upper bound for a
276/// human-readable identifier).
277pub const MAX_RPC_SERVICE_NAME_LEN: usize = 255;
278
279/// Maximum number of headers in a single request or response.
280/// Prevents pathological `headers.len()` reads from a malformed
281/// peer; legitimate callers stay well below this.
282pub const MAX_RPC_HEADERS: usize = 32;
283
284/// Maximum length of a single header name (UTF-8 bytes).
285pub const MAX_RPC_HEADER_NAME_LEN: usize = 64;
286
287/// Maximum length of a single header value (bytes).
288pub const MAX_RPC_HEADER_VALUE_LEN: usize = 4096;
289
290/// Maximum length of a request or response body. Larger payloads
291/// must use streaming responses (Phase 3) or chunk at the
292/// application layer. Comparable to gRPC's default `max_message_size`
293/// of 4 MiB; tuned downward to match RedEX's
294/// `MAX_REDEX_HEAP_PAYLOAD` ceiling.
295pub const MAX_RPC_BODY_LEN: usize = 4 * 1024 * 1024;
296
297/// nRPC request payload. Lives after the 24-byte `EventMeta` prefix
298/// in a `DISPATCH_RPC_REQUEST` event.
299#[derive(Debug, Clone, PartialEq, Eq)]
300pub struct RpcRequestPayload {
301 /// Service-name dispatch key. The server's fold looks this up
302 /// in its `serve_rpc` registry and routes to the registered
303 /// handler.
304 pub service: String,
305 /// Absolute deadline (unix nanos). `0` means no deadline; the
306 /// caller will cancel via `DISPATCH_RPC_CANCEL` if it changes
307 /// its mind.
308 pub deadline_ns: u64,
309 /// Bitfield of `FLAG_RPC_*` constants.
310 pub flags: u16,
311 /// Headers (trace context, idempotency key, content-type, etc.).
312 /// Capped at `MAX_RPC_HEADERS` entries, name <= `MAX_RPC_HEADER_NAME_LEN`,
313 /// value <= `MAX_RPC_HEADER_VALUE_LEN`.
314 pub headers: Vec<RpcHeader>,
315 /// Application-defined request body. Caller and server agree on
316 /// the codec out-of-band; nRPC doesn't interpret these bytes.
317 ///
318 /// Held as [`Bytes`] so [`Self::decode`] can zero-copy `slice_ref`
319 /// the body out of the inbound event's `Bytes` payload — pre-fix
320 /// perf #84 in `docs/performance/net-perf-analysis.md` this was
321 /// `Vec<u8>` and every decode did a `data[body_start..body_end].to_vec()`
322 /// (a memcpy per frame). For high-RPS systems doing 100K+ RPCs/sec
323 /// with 1 KB+ bodies that was 100+ MB/sec of pure memcpy.
324 pub body: Bytes,
325}
326
327/// Continuation chunk for a client-streaming or duplex REQUEST.
328/// Lives after the 24-byte `EventMeta` prefix in a
329/// [`DISPATCH_RPC_REQUEST_CHUNK`] event.
330///
331/// Unlike the initial [`RpcRequestPayload`] there is no
332/// `service` field (server already routed by service at the
333/// initial REQUEST) and no `deadline_ns` (the initial REQUEST's
334/// deadline applies to the whole call). The `call_id` field is
335/// redundant with `EventMeta::seq_or_ts` but kept on the
336/// payload so the codec is self-contained — a reader handed a
337/// chunk's bytes without the meta header can still recover its
338/// correlation id.
339///
340/// Bidi streaming plan (Phase A).
341#[derive(Debug, Clone, PartialEq, Eq)]
342pub struct RpcRequestChunkPayload {
343 /// Matches `EventMeta::seq_or_ts` and the original REQUEST's
344 /// `call_id`. Kept on the payload so the codec round-trips
345 /// in isolation.
346 pub call_id: u64,
347 /// Bitfield of `FLAG_RPC_*` constants. The only flag that
348 /// makes sense on a chunk today is [`FLAG_RPC_REQUEST_END`];
349 /// other flags MUST be zero on the wire so future protocol
350 /// extensions can claim them without colliding with
351 /// existing chunks.
352 pub flags: u16,
353 /// Per-chunk metadata. Typically empty; reserved for
354 /// trace-span continuity across long uploads, content-type
355 /// changes mid-stream, or other rare per-chunk concerns.
356 /// Capped at `MAX_RPC_HEADERS` entries with the same
357 /// per-field caps as `RpcRequestPayload::headers`.
358 pub headers: Vec<RpcHeader>,
359 /// Application-defined chunk body. Cap is `MAX_RPC_BODY_LEN`
360 /// (4 MiB), same as the initial REQUEST body — clients that
361 /// need >4 MiB total payload chunk their upload across
362 /// multiple `REQUEST_CHUNK` events.
363 ///
364 /// See [`RpcRequestPayload::body`] for the `Bytes`-vs-`Vec<u8>`
365 /// rationale.
366 pub body: Bytes,
367}
368
369/// Request-direction credit grant. Lives after the 24-byte
370/// `EventMeta` prefix in a [`DISPATCH_RPC_REQUEST_GRANT`] event.
371/// Mirror of the response-direction [`encode_stream_grant`] /
372/// [`decode_stream_grant`] pair, but with an explicit `call_id`
373/// in the payload (instead of relying solely on
374/// `EventMeta::seq_or_ts`) so the codec is self-contained — same
375/// rationale as [`RpcRequestChunkPayload::call_id`].
376///
377/// Bidi streaming plan (Phase A).
378#[derive(Debug, Clone, Copy, PartialEq, Eq)]
379pub struct RpcRequestGrantPayload {
380 /// Matches the call's `call_id`.
381 pub call_id: u64,
382 /// Additional REQUEST_CHUNK frames the server is willing to
383 /// admit beyond the current credit. Server's incoming-credit
384 /// counter is capped defensively (see PHASE-B server fold)
385 /// so a misbehaving grant can't overflow.
386 pub credits: u32,
387}
388
389/// nRPC response payload. Lives after the 24-byte `EventMeta`
390/// prefix in a `DISPATCH_RPC_RESPONSE` event.
391#[derive(Debug, Clone, PartialEq, Eq)]
392pub struct RpcResponsePayload {
393 /// Outcome of the call. Decoded on the caller side via
394 /// [`RpcStatus::from_wire`].
395 pub status: RpcStatus,
396 /// Headers (trace context, content-type, content-encoding,
397 /// etc.). Same caps as `RpcRequestPayload::headers`.
398 pub headers: Vec<RpcHeader>,
399 /// For `status == Ok`: the application response body.
400 /// For non-`Ok` statuses: a UTF-8 diagnostic string (callers
401 /// `String::from_utf8_lossy` for display; the bytes are not
402 /// guaranteed to be valid UTF-8 against a malicious server).
403 ///
404 /// See [`RpcRequestPayload::body`] for the `Bytes`-vs-`Vec<u8>`
405 /// rationale.
406 pub body: Bytes,
407}
408
409// ============================================================================
410// Codec.
411//
412// All wire integers are little-endian. Lengths are u32_le where the
413// upper bound exceeds u16, u16_le where it fits, u8 where it fits.
414// ============================================================================
415
416/// Errors from the request / response codecs.
417#[derive(Debug, thiserror::Error)]
418pub enum RpcCodecError {
419 /// Buffer ended mid-field.
420 #[error("truncated payload at {0}")]
421 Truncated(&'static str),
422 /// Length prefix exceeds the configured maximum.
423 #[error("length {actual} exceeds limit {limit} for {field}")]
424 TooLarge {
425 /// Field name whose declared length exceeded the cap (e.g.
426 /// `"body"`, `"headers"`, `"service"`). Stable strings —
427 /// callers may match on them for diagnostics.
428 field: &'static str,
429 /// The length the wire claimed for the field.
430 actual: usize,
431 /// The maximum the codec accepts (one of the `MAX_RPC_*`
432 /// constants).
433 limit: usize,
434 },
435 /// String field contains non-UTF-8 bytes.
436 #[error("non-utf8 string in {0}")]
437 InvalidUtf8(&'static str),
438}
439
440impl RpcRequestPayload {
441 /// Compute the encoded byte length WITHOUT actually encoding.
442 /// Used by [`request_wire_size`] and any caller that needs to
443 /// budget event size at the bus layer (e.g., to refuse a
444 /// request that wouldn't fit in the configured packet budget)
445 /// without paying the encode cost.
446 pub fn encoded_len(&self) -> usize {
447 // service: u8 length + bytes
448 1 + self.service.len()
449 // deadline_ns: u64
450 + 8
451 // flags: u16
452 + 2
453 // headers: u8 count + per-header (u8 name_len + name + u16 value_len + value)
454 + 1
455 + self
456 .headers
457 .iter()
458 .map(|(n, v)| 1 + n.len() + 2 + v.len())
459 .sum::<usize>()
460 // body: u32 length + bytes
461 + 4
462 + self.body.len()
463 }
464
465 /// Encode to the wire format. The result is the bytes that
466 /// follow the 24-byte `EventMeta` prefix in the RedEX payload.
467 ///
468 /// **Encoder bounds:** every field that has a `MAX_RPC_*` cap
469 /// is asserted against that cap. In debug builds an oversize
470 /// field panics with a useful diagnostic so the programmer
471 /// notices in tests; in release builds the assert is dropped
472 /// (the decoder side still enforces the cap, so a malformed
473 /// frame would be rejected by the receiver — but constructing
474 /// one is always a caller bug).
475 pub fn encode(&self) -> Vec<u8> {
476 let mut buf = Vec::with_capacity(self.encoded_len());
477 self.encode_into(&mut buf);
478 buf
479 }
480
481 /// Encode directly into `buf`, appending the wire bytes (audit T2.2).
482 /// Callers that already hold an `EventMeta`-prefixed buffer use this to
483 /// skip `encode()`'s intermediate `Vec` allocation + copy. Produces the
484 /// identical bytes `encode()` returns. Pre-`reserve(encoded_len())` for an
485 /// exact-fit single allocation when starting from an empty buffer.
486 pub fn encode_into(&self, buf: &mut Vec<u8>) {
487 // service
488 let svc = self.service.as_bytes();
489 debug_assert!(
490 svc.len() <= MAX_RPC_SERVICE_NAME_LEN,
491 "service name {} exceeds MAX_RPC_SERVICE_NAME_LEN ({})",
492 svc.len(),
493 MAX_RPC_SERVICE_NAME_LEN,
494 );
495 buf.put_u8(svc.len() as u8);
496 buf.extend_from_slice(svc);
497 // deadline_ns
498 buf.put_u64_le(self.deadline_ns);
499 // flags
500 buf.put_u16_le(self.flags);
501 // headers
502 encode_headers(&self.headers, buf);
503 // body
504 debug_assert!(
505 self.body.len() <= MAX_RPC_BODY_LEN,
506 "body length {} exceeds MAX_RPC_BODY_LEN ({})",
507 self.body.len(),
508 MAX_RPC_BODY_LEN,
509 );
510 buf.put_u32_le(self.body.len() as u32);
511 buf.extend_from_slice(&self.body);
512 }
513
514 /// Decode from the wire bytes following the `EventMeta` prefix.
515 /// All length fields are bounded by the `MAX_RPC_*` constants;
516 /// over-cap inputs error rather than allocate unbounded
517 /// buffers.
518 ///
519 /// Takes [`Bytes`] (not `&[u8]`) so the decoded `body` field
520 /// can be a zero-copy `data.slice(..)` instead of an owned
521 /// `to_vec` — see perf #84.
522 pub fn decode(data: Bytes) -> Result<Self, RpcCodecError> {
523 let mut cur = std::io::Cursor::new(data.as_ref());
524 // service
525 if cur.remaining() < 1 {
526 return Err(RpcCodecError::Truncated("service length"));
527 }
528 let svc_len = cur.get_u8() as usize;
529 if svc_len == 0 {
530 return Err(RpcCodecError::Truncated("empty service name"));
531 }
532 if svc_len > MAX_RPC_SERVICE_NAME_LEN {
533 return Err(RpcCodecError::TooLarge {
534 field: "service",
535 actual: svc_len,
536 limit: MAX_RPC_SERVICE_NAME_LEN,
537 });
538 }
539 if cur.remaining() < svc_len {
540 return Err(RpcCodecError::Truncated("service bytes"));
541 }
542 let svc_start = cur.position() as usize;
543 let svc_end = svc_start + svc_len;
544 let service = std::str::from_utf8(&data[svc_start..svc_end])
545 .map_err(|_| RpcCodecError::InvalidUtf8("service"))?
546 .to_string();
547 cur.set_position(svc_end as u64);
548 // deadline_ns
549 if cur.remaining() < 8 {
550 return Err(RpcCodecError::Truncated("deadline_ns"));
551 }
552 let deadline_ns = cur.get_u64_le();
553 // flags
554 if cur.remaining() < 2 {
555 return Err(RpcCodecError::Truncated("flags"));
556 }
557 let flags = cur.get_u16_le();
558 // headers
559 let headers = decode_headers(&mut cur, &data)?;
560 // body
561 if cur.remaining() < 4 {
562 return Err(RpcCodecError::Truncated("body length"));
563 }
564 let body_len = cur.get_u32_le() as usize;
565 if body_len > MAX_RPC_BODY_LEN {
566 return Err(RpcCodecError::TooLarge {
567 field: "body",
568 actual: body_len,
569 limit: MAX_RPC_BODY_LEN,
570 });
571 }
572 if cur.remaining() < body_len {
573 return Err(RpcCodecError::Truncated("body bytes"));
574 }
575 let body_start = cur.position() as usize;
576 let body_end = body_start + body_len;
577 // Zero-copy slice over the input — refcount bump only.
578 let body = data.slice(body_start..body_end);
579 Ok(Self {
580 service,
581 deadline_ns,
582 flags,
583 headers,
584 body,
585 })
586 }
587}
588
589impl RpcRequestChunkPayload {
590 /// Compute the encoded byte length WITHOUT actually encoding.
591 /// See [`RpcRequestPayload::encoded_len`] for the rationale.
592 pub fn encoded_len(&self) -> usize {
593 // call_id: u64
594 8
595 // flags: u16
596 + 2
597 // headers: u8 count + per-header (u8 name_len + name + u16 value_len + value)
598 + 1
599 + self
600 .headers
601 .iter()
602 .map(|(n, v)| 1 + n.len() + 2 + v.len())
603 .sum::<usize>()
604 // body: u32 length + bytes
605 + 4
606 + self.body.len()
607 }
608
609 /// Encode to the wire bytes that follow the 24-byte `EventMeta`
610 /// prefix in a [`DISPATCH_RPC_REQUEST_CHUNK`] event. Same
611 /// encoder-bounds policy as [`RpcRequestPayload::encode`]:
612 /// oversize fields panic in debug, the decoder enforces in
613 /// release.
614 pub fn encode(&self) -> Vec<u8> {
615 let mut buf = Vec::with_capacity(self.encoded_len());
616 self.encode_into(&mut buf);
617 buf
618 }
619
620 /// Encode directly into `buf`, appending the wire bytes (audit T2.2).
621 /// See [`RpcRequestPayload::encode_into`].
622 pub fn encode_into(&self, buf: &mut Vec<u8>) {
623 // call_id
624 buf.put_u64_le(self.call_id);
625 // flags
626 buf.put_u16_le(self.flags);
627 // headers
628 encode_headers(&self.headers, buf);
629 // body
630 debug_assert!(
631 self.body.len() <= MAX_RPC_BODY_LEN,
632 "body length {} exceeds MAX_RPC_BODY_LEN ({})",
633 self.body.len(),
634 MAX_RPC_BODY_LEN,
635 );
636 buf.put_u32_le(self.body.len() as u32);
637 buf.extend_from_slice(&self.body);
638 }
639
640 /// Decode from the wire bytes following the `EventMeta` prefix.
641 /// Bounded by the same `MAX_RPC_*` caps as the initial REQUEST.
642 /// Takes [`Bytes`] for zero-copy `body` slicing — see perf #84.
643 pub fn decode(data: Bytes) -> Result<Self, RpcCodecError> {
644 let mut cur = std::io::Cursor::new(data.as_ref());
645 // call_id
646 if cur.remaining() < 8 {
647 return Err(RpcCodecError::Truncated("call_id"));
648 }
649 let call_id = cur.get_u64_le();
650 // flags
651 if cur.remaining() < 2 {
652 return Err(RpcCodecError::Truncated("flags"));
653 }
654 let flags = cur.get_u16_le();
655 // headers
656 let headers = decode_headers(&mut cur, &data)?;
657 // body
658 if cur.remaining() < 4 {
659 return Err(RpcCodecError::Truncated("body length"));
660 }
661 let body_len = cur.get_u32_le() as usize;
662 if body_len > MAX_RPC_BODY_LEN {
663 return Err(RpcCodecError::TooLarge {
664 field: "body",
665 actual: body_len,
666 limit: MAX_RPC_BODY_LEN,
667 });
668 }
669 if cur.remaining() < body_len {
670 return Err(RpcCodecError::Truncated("body bytes"));
671 }
672 let body_start = cur.position() as usize;
673 let body_end = body_start + body_len;
674 let body = data.slice(body_start..body_end);
675 Ok(Self {
676 call_id,
677 flags,
678 headers,
679 body,
680 })
681 }
682}
683
684impl RpcResponsePayload {
685 /// Compute the encoded byte length WITHOUT actually encoding.
686 /// See [`RpcRequestPayload::encoded_len`].
687 pub fn encoded_len(&self) -> usize {
688 // status: u16
689 2
690 // headers: u8 count + per-header
691 + 1
692 + self
693 .headers
694 .iter()
695 .map(|(n, v)| 1 + n.len() + 2 + v.len())
696 .sum::<usize>()
697 // body: u32 length + bytes
698 + 4
699 + self.body.len()
700 }
701
702 /// Encode to the wire format. The result is the bytes that
703 /// follow the 24-byte `EventMeta` prefix in the RedEX payload.
704 /// Same encoder-bounds policy as
705 /// [`RpcRequestPayload::encode`] — see that method's doc.
706 pub fn encode(&self) -> Vec<u8> {
707 let mut buf = Vec::with_capacity(self.encoded_len());
708 self.encode_into(&mut buf);
709 buf
710 }
711
712 /// Encode directly into `buf`, appending the wire bytes (audit T2.2).
713 /// See [`RpcRequestPayload::encode_into`].
714 pub fn encode_into(&self, buf: &mut Vec<u8>) {
715 buf.put_u16_le(self.status.to_wire());
716 encode_headers(&self.headers, buf);
717 debug_assert!(
718 self.body.len() <= MAX_RPC_BODY_LEN,
719 "body length {} exceeds MAX_RPC_BODY_LEN ({})",
720 self.body.len(),
721 MAX_RPC_BODY_LEN,
722 );
723 buf.put_u32_le(self.body.len() as u32);
724 buf.extend_from_slice(&self.body);
725 }
726
727 /// Decode from the wire bytes following the `EventMeta` prefix.
728 /// Takes [`Bytes`] for zero-copy `body` slicing — see perf #84.
729 pub fn decode(data: Bytes) -> Result<Self, RpcCodecError> {
730 let mut cur = std::io::Cursor::new(data.as_ref());
731 if cur.remaining() < 2 {
732 return Err(RpcCodecError::Truncated("status"));
733 }
734 let status = RpcStatus::from_wire(cur.get_u16_le());
735 let headers = decode_headers(&mut cur, &data)?;
736 if cur.remaining() < 4 {
737 return Err(RpcCodecError::Truncated("body length"));
738 }
739 let body_len = cur.get_u32_le() as usize;
740 if body_len > MAX_RPC_BODY_LEN {
741 return Err(RpcCodecError::TooLarge {
742 field: "body",
743 actual: body_len,
744 limit: MAX_RPC_BODY_LEN,
745 });
746 }
747 if cur.remaining() < body_len {
748 return Err(RpcCodecError::Truncated("body bytes"));
749 }
750 let body_start = cur.position() as usize;
751 let body_end = body_start + body_len;
752 let body = data.slice(body_start..body_end);
753 Ok(Self {
754 status,
755 headers,
756 body,
757 })
758 }
759}
760
761/// Pull `traceparent` / `tracestate` out of `headers` if present.
762/// Caller-side helper: callers building an `RpcRequestPayload`
763/// with a `TraceContext` use [`build_trace_headers`] to emit the
764/// matching headers; this is the inverse on the server side.
765///
766/// Returns `Some(TraceContext)` if `traceparent` is present;
767/// `None` otherwise. `tracestate` defaults to empty when absent
768/// — W3C says tracestate is optional even when traceparent is
769/// set.
770pub fn extract_trace_context(headers: &[RpcHeader]) -> Option<TraceContext> {
771 let mut traceparent: Option<String> = None;
772 let mut tracestate = String::new();
773 for (name, value) in headers {
774 // Header names are case-insensitive (matches W3C and HTTP
775 // convention) — same comparison style as `parse_stream_
776 // window_initial` for consistency. The wire spec doesn't
777 // mandate case so a peer that emits `Traceparent` (capital
778 // T) shouldn't be silently ignored.
779 if name.eq_ignore_ascii_case("traceparent") {
780 if let Ok(s) = std::str::from_utf8(value) {
781 traceparent = Some(s.to_string());
782 }
783 } else if name.eq_ignore_ascii_case("tracestate") {
784 if let Ok(s) = std::str::from_utf8(value) {
785 tracestate = s.to_string();
786 }
787 }
788 }
789 traceparent.map(|tp| TraceContext {
790 traceparent: tp,
791 tracestate,
792 })
793}
794
795/// Build the headers a caller appends to its
796/// `RpcRequestPayload::headers` to propagate the trace context
797/// across the call. Set `RpcRequestPayload::flags |= FLAG_RPC_PROPAGATE_TRACE`
798/// alongside this so the server's fold knows to extract them.
799///
800/// Always emits `traceparent`. Emits `tracestate` only when
801/// non-empty (matches the W3C convention of skipping empty
802/// tracestate values on the wire).
803pub fn build_trace_headers(ctx: &TraceContext) -> Vec<RpcHeader> {
804 let mut headers = Vec::with_capacity(2);
805 headers.push((
806 "traceparent".to_string(),
807 ctx.traceparent.clone().into_bytes(),
808 ));
809 if !ctx.tracestate.is_empty() {
810 headers.push((
811 "tracestate".to_string(),
812 ctx.tracestate.clone().into_bytes(),
813 ));
814 }
815 headers
816}
817
818fn encode_headers(headers: &[RpcHeader], buf: &mut Vec<u8>) {
819 debug_assert!(
820 headers.len() <= MAX_RPC_HEADERS,
821 "headers count {} exceeds MAX_RPC_HEADERS ({})",
822 headers.len(),
823 MAX_RPC_HEADERS,
824 );
825 buf.put_u8(headers.len() as u8);
826 for (name, value) in headers {
827 let nbytes = name.as_bytes();
828 debug_assert!(
829 nbytes.len() <= MAX_RPC_HEADER_NAME_LEN,
830 "header name {} exceeds MAX_RPC_HEADER_NAME_LEN ({})",
831 nbytes.len(),
832 MAX_RPC_HEADER_NAME_LEN,
833 );
834 debug_assert!(
835 value.len() <= MAX_RPC_HEADER_VALUE_LEN,
836 "header value {} exceeds MAX_RPC_HEADER_VALUE_LEN ({})",
837 value.len(),
838 MAX_RPC_HEADER_VALUE_LEN,
839 );
840 buf.put_u8(nbytes.len() as u8);
841 buf.extend_from_slice(nbytes);
842 buf.put_u16_le(value.len() as u16);
843 buf.extend_from_slice(value);
844 }
845}
846
847fn decode_headers(
848 cur: &mut std::io::Cursor<&[u8]>,
849 data: &[u8],
850) -> Result<Vec<RpcHeader>, RpcCodecError> {
851 if cur.remaining() < 1 {
852 return Err(RpcCodecError::Truncated("headers count"));
853 }
854 let count = cur.get_u8() as usize;
855 if count > MAX_RPC_HEADERS {
856 return Err(RpcCodecError::TooLarge {
857 field: "headers",
858 actual: count,
859 limit: MAX_RPC_HEADERS,
860 });
861 }
862 let mut headers = Vec::with_capacity(count);
863 for _ in 0..count {
864 if cur.remaining() < 1 {
865 return Err(RpcCodecError::Truncated("header name length"));
866 }
867 let name_len = cur.get_u8() as usize;
868 if name_len == 0 {
869 return Err(RpcCodecError::Truncated("empty header name"));
870 }
871 if name_len > MAX_RPC_HEADER_NAME_LEN {
872 return Err(RpcCodecError::TooLarge {
873 field: "header name",
874 actual: name_len,
875 limit: MAX_RPC_HEADER_NAME_LEN,
876 });
877 }
878 if cur.remaining() < name_len {
879 return Err(RpcCodecError::Truncated("header name bytes"));
880 }
881 let nstart = cur.position() as usize;
882 let nend = nstart + name_len;
883 let name = std::str::from_utf8(&data[nstart..nend])
884 .map_err(|_| RpcCodecError::InvalidUtf8("header name"))?
885 .to_string();
886 cur.set_position(nend as u64);
887
888 if cur.remaining() < 2 {
889 return Err(RpcCodecError::Truncated("header value length"));
890 }
891 let value_len = cur.get_u16_le() as usize;
892 if value_len > MAX_RPC_HEADER_VALUE_LEN {
893 return Err(RpcCodecError::TooLarge {
894 field: "header value",
895 actual: value_len,
896 limit: MAX_RPC_HEADER_VALUE_LEN,
897 });
898 }
899 if cur.remaining() < value_len {
900 return Err(RpcCodecError::Truncated("header value bytes"));
901 }
902 let vstart = cur.position() as usize;
903 let vend = vstart + value_len;
904 let value = data[vstart..vend].to_vec();
905 cur.set_position(vend as u64);
906 headers.push((name, value));
907 }
908 Ok(headers)
909}
910
911/// Convenience: the byte layout of an `RpcRequestPayload` that lands
912/// after the `EventMeta` prefix in a `DISPATCH_RPC_REQUEST` event.
913/// Exposed so callers can budget the total event size at the bus
914/// layer without doing the encode first.
915pub fn request_wire_size(payload: &RpcRequestPayload) -> usize {
916 EVENT_META_SIZE + payload.encoded_len()
917}
918
919/// Same for `RpcResponsePayload` after the `EventMeta` prefix in a
920/// `DISPATCH_RPC_RESPONSE` event.
921pub fn response_wire_size(payload: &RpcResponsePayload) -> usize {
922 EVENT_META_SIZE + payload.encoded_len()
923}
924
925// ============================================================================
926// Mesh inbound dispatch hook.
927//
928// `MeshNode::dispatch_packet` normally pushes inbound channel
929// events onto a per-shard `inbound` queue keyed by `shard_id`. The
930// channel name / hash is stripped on the way in — by the time the
931// event lands in the queue, only the payload remains.
932//
933// RPC needs per-channel routing (events for `<service>.requests`
934// drive the server fold; events for `<service>.replies.<origin>`
935// drive the client fold). Without channel info on the queued
936// event, we can't filter from the consumer side.
937//
938// The hook below adds a per-channel-hash dispatcher map that the
939// mesh's inbound dispatch consults BEFORE pushing to the shard
940// queue. If a dispatcher is registered for the event's
941// canonical [`ChannelHash`], the event is routed there directly
942// (bypassing the shard queue); otherwise the existing shard-queue
943// path runs.
944//
945// **Collision posture.** The dispatch event carries the canonical
946// 32-bit [`ChannelHash`] (joint-collision threshold ~65 K
947// channels, well above realistic deployment); the wire
948// `NetHeader::channel_hash` is `u16` and may bucket-collide at
949// scale, so the mesh's inbound dispatch indexes by the wire `u16`
950// and dispatches every canonical entry registered in that bucket
951// (the canonical match resolves on the dispatcher side). At
952// typical sizing this is a single entry per bucket.
953// ============================================================================
954
955/// One inbound event delivered to a registered RPC dispatcher.
956#[derive(Debug, Clone)]
957pub struct RpcInboundEvent {
958 /// Canonical [`ChannelHash`](crate::adapter::net::channel::ChannelHash)
959 /// (u32) of the channel this event arrived on — widened from the
960 /// per-packet wire `u16` `NetHeader::channel_hash` via the
961 /// registered-dispatcher table at receive time.
962 /// Collision-resistant at realistic scale; the wire `u16` may
963 /// bucket-collide but the canonical hash uniquely identifies the
964 /// registered dispatcher target.
965 pub channel_hash: super::super::channel::ChannelHash,
966 /// Caller's `origin_hash` from the packet header — the full
967 /// 64-bit `EntityKeypair::origin_hash()` mirroring the wire
968 /// field's width post-`WIRE_ORIGIN_HASH_64BIT`. The dispatcher
969 /// should treat this as routing metadata, not identity
970 /// authentication (the AEAD-verified `session_node` field
971 /// below carries that).
972 pub origin_hash: u64,
973 /// Wire-session peer's `NodeId` resolved at packet receive
974 /// time from the AEAD-verified session_id. Distinct from
975 /// `origin_hash`: this is the full 64-bit network identity
976 /// of the peer that delivered the packet. Used by
977 /// `RpcClientPending::deliver`
978 /// to reject spoofed RESPONSE frames whose call_id happens
979 /// to match an in-flight request but whose session peer
980 /// isn't the recorded target.
981 ///
982 /// Set to `0` on test / loopback paths that don't have a
983 /// session to resolve against; callers that register
984 /// pending entries with `target_node = 0` opt out of the
985 /// binding gate (and trust the call_id randomness alone).
986 ///
987 /// **Production wire-path invariant**: real over-the-wire
988 /// inbound delivery MUST NOT produce `from_node = 0`. The
989 /// dispatcher in `mesh.rs` (`handle_inbound_user_payload`)
990 /// drops the event when the wire session has no resolvable
991 /// `NodeId`, rather than forwarding under the sentinel — see
992 /// the explicit drop + warn at the
993 /// `dropping cortex-RPC event: wire session has no resolvable NodeId`
994 /// log site. The v0.4 capability-auth callee-side gate in
995 /// `MeshNode::serve_rpc`'s bridge relies on this: it skips
996 /// permissively when `from_node == 0` (loopback compat), so
997 /// a wire-path leak of the sentinel would silently re-open
998 /// the gate. If you change the dispatcher to fall back to 0
999 /// instead of dropping, you ALSO have to teach the bridge
1000 /// to deny on the sentinel.
1001 pub from_node: super::super::behavior::placement::NodeId,
1002 /// Event payload bytes — the same bytes that would have been
1003 /// pushed onto the shard inbound queue. For RPC events these
1004 /// start with a 24-byte `EventMeta` followed by the
1005 /// `RpcRequestPayload` / `RpcResponsePayload` encoding.
1006 pub payload: bytes::Bytes,
1007}
1008
1009/// Type-erased callback fired by the mesh's inbound dispatch
1010/// when an event arrives for a registered `channel_hash`. The
1011/// callback runs on the mesh's dispatch task, so the body should
1012/// be quick (push the event onto an mpsc / fold consumer rather
1013/// than do real work).
1014pub type RpcInboundDispatcher = Arc<dyn Fn(RpcInboundEvent) + Send + Sync + 'static>;
1015
1016// ============================================================================
1017// Streaming-response protocol markers.
1018//
1019// When a caller sets `FLAG_RPC_STREAMING_RESPONSE` on the request,
1020// the server emits multiple `DISPATCH_RPC_RESPONSE` events for the
1021// same `call_id`. Non-terminal chunks carry the
1022// `nrpc-streaming = continue` header; the terminal chunk carries
1023// `nrpc-streaming = end` (or any non-`Ok` status, which is also
1024// terminal). The client-side stream collects chunks until it sees
1025// a terminal marker.
1026// ============================================================================
1027
1028/// Header name nRPC uses to mark streaming-response chunks.
1029/// Present on every chunk of a streaming response, with one of two
1030/// values defined below.
1031pub const HEADER_NRPC_STREAMING: &str = "nrpc-streaming";
1032
1033/// `nrpc-streaming` value on a non-terminal chunk. The client-side
1034/// stream yields the chunk's body and continues waiting for more.
1035pub const HEADER_NRPC_STREAMING_CONTINUE: &[u8] = b"continue";
1036
1037/// `nrpc-streaming` value on the terminal chunk. The client-side
1038/// stream yields the chunk's body (if non-empty) and then closes.
1039/// A non-`Ok` status is also terminal, regardless of header — the
1040/// stream yields the error and closes.
1041pub const HEADER_NRPC_STREAMING_END: &[u8] = b"end";
1042
1043/// Header on a streaming REQUEST that opts into flow control with
1044/// the given initial credit window. Value is the ASCII decimal
1045/// representation of a `u32` (e.g. `"32"`). When present, the
1046/// server's streaming fold creates a per-call semaphore initialized
1047/// to that count and the pump awaits one credit per emitted chunk.
1048/// The caller refills via [`DISPATCH_RPC_STREAM_GRANT`] events.
1049///
1050/// Absent → unbounded credit (the pump emits chunks as fast as
1051/// the publish path can take them). Long-running streams that
1052/// could outpace a slow consumer SHOULD opt into flow control —
1053/// without it, the server's sink mpsc grows unbounded under a
1054/// stalled caller.
1055pub const HEADER_NRPC_STREAM_WINDOW_INITIAL: &str = "nrpc-stream-window-initial";
1056
1057/// Encode a stream-grant payload — 4 bytes big-endian `u32`
1058/// representing additional credit. Pair with [`decode_stream_grant`]
1059/// on the server side.
1060pub fn encode_stream_grant(amount: u32) -> Vec<u8> {
1061 amount.to_be_bytes().to_vec()
1062}
1063
1064/// Decode a stream-grant payload. Returns `None` if the slice is
1065/// not exactly 4 bytes — defends the server fold against
1066/// malformed grants without killing the cortex adapter.
1067pub fn decode_stream_grant(payload: &[u8]) -> Option<u32> {
1068 if payload.len() != 4 {
1069 return None;
1070 }
1071 let mut bytes = [0u8; 4];
1072 bytes.copy_from_slice(payload);
1073 Some(u32::from_be_bytes(bytes))
1074}
1075
1076/// Parse the `nrpc-stream-window-initial` header from a request's
1077/// header list. Returns `Some(window)` if a valid u32 ASCII-decimal
1078/// value is present, else `None` (no header / malformed value /
1079/// non-utf8 — all treated as "no flow control").
1080pub fn parse_stream_window_initial(headers: &[RpcHeader]) -> Option<u32> {
1081 for (name, value) in headers {
1082 if name.eq_ignore_ascii_case(HEADER_NRPC_STREAM_WINDOW_INITIAL) {
1083 return std::str::from_utf8(value).ok()?.parse::<u32>().ok();
1084 }
1085 }
1086 None
1087}
1088
1089/// Header on the initial REQUEST of a client-streaming or duplex
1090/// call that opts the upload direction into flow control with the
1091/// given initial credit window. Value is the ASCII decimal
1092/// representation of a `u32`. When present, the server's
1093/// streaming-request fold creates a per-call semaphore and the
1094/// caller's sink awaits one credit per `REQUEST_CHUNK`. The server
1095/// refills via [`DISPATCH_RPC_REQUEST_GRANT`] events.
1096///
1097/// Absent → unbounded credit (caller's sink emits as fast as the
1098/// publish path can take it). Long client-streaming calls that
1099/// could outpace a slow handler SHOULD opt into flow control —
1100/// without it, the server's chunk mpsc grows unbounded under a
1101/// stalled handler.
1102///
1103/// Bidi streaming plan (Phase A).
1104pub const HEADER_NRPC_REQUEST_WINDOW_INITIAL: &str = "nrpc-request-window-initial";
1105
1106/// Encode a request-grant payload — `call_id` (u64 little-endian)
1107/// followed by additional credit (u32 big-endian). Big-endian on
1108/// the credit field matches [`encode_stream_grant`]; little-endian
1109/// on `call_id` matches the rest of the RPC codec's u64 fields.
1110///
1111/// Pair with [`decode_request_grant`] on the caller side.
1112pub fn encode_request_grant(call_id: u64, credits: u32) -> Vec<u8> {
1113 let mut buf = Vec::with_capacity(12);
1114 buf.put_u64_le(call_id);
1115 buf.extend_from_slice(&credits.to_be_bytes());
1116 buf
1117}
1118
1119/// Decode a request-grant payload. Returns `None` if the slice is
1120/// not exactly 12 bytes — defends the caller's fold against
1121/// malformed grants without killing the cortex adapter.
1122pub fn decode_request_grant(payload: &[u8]) -> Option<RpcRequestGrantPayload> {
1123 if payload.len() != 12 {
1124 return None;
1125 }
1126 let mut cid = [0u8; 8];
1127 cid.copy_from_slice(&payload[..8]);
1128 let call_id = u64::from_le_bytes(cid);
1129 let mut credits = [0u8; 4];
1130 credits.copy_from_slice(&payload[8..]);
1131 Some(RpcRequestGrantPayload {
1132 call_id,
1133 credits: u32::from_be_bytes(credits),
1134 })
1135}
1136
1137/// Parse the `nrpc-request-window-initial` header from a request's
1138/// header list. Same semantics as [`parse_stream_window_initial`]
1139/// but for the upload direction.
1140pub fn parse_request_window_initial(headers: &[RpcHeader]) -> Option<u32> {
1141 for (name, value) in headers {
1142 if name.eq_ignore_ascii_case(HEADER_NRPC_REQUEST_WINDOW_INITIAL) {
1143 return std::str::from_utf8(value).ok()?.parse::<u32>().ok();
1144 }
1145 }
1146 None
1147}
1148
1149/// Inspect a `RpcResponsePayload`'s headers and decide whether
1150/// it's a non-terminal streaming chunk (`continue`), a terminal
1151/// streaming chunk (`end` OR non-`Ok` status), OR a unary
1152/// response (no streaming header at all). Used by the client-side
1153/// fold to demux streaming vs unary responses without needing a
1154/// separate flag.
1155#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1156pub enum StreamingChunkKind {
1157 /// Non-terminal chunk — yield body, continue waiting.
1158 Continue,
1159 /// Terminal chunk — yield body (if any), close stream.
1160 Terminal,
1161 /// Not a streaming response — unary semantics apply.
1162 Unary,
1163}
1164
1165/// Classify a response per the streaming-protocol markers.
1166pub fn classify_streaming_chunk(resp: &RpcResponsePayload) -> StreamingChunkKind {
1167 // Non-Ok status is always terminal regardless of header — the
1168 // stream surfaces the error and closes.
1169 if !resp.status.is_ok() {
1170 return StreamingChunkKind::Terminal;
1171 }
1172 // Walk headers for the streaming marker. Absence = unary
1173 // semantics (caller used `call`, not `call_streaming`).
1174 for (name, value) in &resp.headers {
1175 if name == HEADER_NRPC_STREAMING {
1176 return if value.as_slice() == HEADER_NRPC_STREAMING_END {
1177 StreamingChunkKind::Terminal
1178 } else if value.as_slice() == HEADER_NRPC_STREAMING_CONTINUE {
1179 StreamingChunkKind::Continue
1180 } else {
1181 // Unknown marker value — be defensive, treat as
1182 // terminal so a misbehaving server doesn't keep
1183 // a stream open forever.
1184 StreamingChunkKind::Terminal
1185 };
1186 }
1187 }
1188 StreamingChunkKind::Unary
1189}
1190
1191// ============================================================================
1192// Server-side fold.
1193//
1194// `RpcServerFold` is the `RedexFold` half of the server. It sees
1195// REQUEST events on the channel the cortex adapter is opened against,
1196// spawns the user handler in a tokio task, and emits the RESPONSE
1197// via a callback the `Mesh::serve_rpc` glue layer wires up. The
1198// fold itself is small and pure — all I/O happens in the spawned
1199// task and the emitter callback.
1200//
1201// Cancellation: each in-flight call gets an `RpcCancellationToken`
1202// that the handler can `select!` on. CANCEL events flip the
1203// matching token; the handler observes `cancellation.cancelled()`
1204// firing and aborts cooperatively.
1205// ============================================================================
1206
1207/// Cancellation signal for an in-flight RPC handler.
1208///
1209/// Created when the fold dispatches a REQUEST; cloned into the
1210/// handler's `RpcContext` and held in the fold's in-flight map. A
1211/// matching CANCEL event flips the token; handlers observe via
1212/// either [`Self::is_cancelled`] (synchronous probe) or
1213/// [`Self::cancelled`] (await for the signal).
1214#[derive(Clone, Default)]
1215pub struct RpcCancellationToken {
1216 inner: Arc<RpcCancellationInner>,
1217}
1218
1219#[derive(Default)]
1220struct RpcCancellationInner {
1221 fired: AtomicBool,
1222 notify: Notify,
1223}
1224
1225impl RpcCancellationToken {
1226 /// Construct a fresh, un-fired token.
1227 pub fn new() -> Self {
1228 Self::default()
1229 }
1230
1231 /// Flip the token. Idempotent — repeated calls are no-ops.
1232 /// Wakes any task currently in [`Self::cancelled`].
1233 pub fn cancel(&self) {
1234 // Release pairs with the Acquire load in `is_cancelled`
1235 // so a handler that observes `is_cancelled() == true` is
1236 // guaranteed to see every prior write the canceller did.
1237 self.inner.fired.store(true, Ordering::Release);
1238 self.inner.notify.notify_waiters();
1239 }
1240
1241 /// Synchronous probe. `true` once `cancel()` has been called.
1242 #[inline]
1243 pub fn is_cancelled(&self) -> bool {
1244 self.inner.fired.load(Ordering::Acquire)
1245 }
1246
1247 /// Await the cancellation. Returns immediately if already
1248 /// cancelled. Otherwise registers as a waiter and returns when
1249 /// `cancel()` is called.
1250 ///
1251 /// Race-safe: registering the `notified()` future BEFORE the
1252 /// `is_cancelled` check means a `cancel()` racing this method
1253 /// either (a) is observed by the post-register check and we
1254 /// return immediately, OR (b) lands after we register and wakes
1255 /// our future. Either way we don't sleep past a cancellation.
1256 pub async fn cancelled(&self) {
1257 let notified = self.inner.notify.notified();
1258 if self.is_cancelled() {
1259 return;
1260 }
1261 notified.await;
1262 }
1263}
1264
1265/// W3C Trace Context — `traceparent` and `tracestate` headers
1266/// propagated through nRPC for distributed-tracing systems.
1267///
1268/// `traceparent` carries the trace id, parent span id, and flags;
1269/// `tracestate` carries vendor-specific tracing extensions. nRPC
1270/// is **transport-only** for these — it doesn't parse or generate
1271/// IDs, doesn't emit spans, doesn't talk to any tracing backend.
1272/// Application code (typically via `tracing-opentelemetry` or a
1273/// Datadog client) reads these on the server side and continues
1274/// the trace.
1275///
1276/// See <https://www.w3.org/TR/trace-context/> for the wire format
1277/// of each field.
1278#[derive(Debug, Clone, Default, PartialEq, Eq)]
1279pub struct TraceContext {
1280 /// `traceparent` header value (e.g.
1281 /// `"00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"`).
1282 /// Required by the W3C spec; nRPC treats it as opaque bytes.
1283 pub traceparent: String,
1284 /// `tracestate` header value — vendor-specific extensions.
1285 /// Optional in W3C; empty string when absent.
1286 pub tracestate: String,
1287}
1288
1289/// Context handed to a `RpcHandler::call`. Carries everything the
1290/// handler needs to fulfill the request: the AEAD-verified caller
1291/// identity, the request payload, and a cancellation token.
1292pub struct RpcContext {
1293 /// AEAD-verified caller `origin_hash`. The bus sets this from
1294 /// the verified peer; not self-claimable from the request body.
1295 pub caller_origin: u64,
1296 /// Caller-generated correlation id. Same value on the matching
1297 /// CANCEL or RESPONSE.
1298 pub call_id: u64,
1299 /// Decoded request payload.
1300 pub payload: RpcRequestPayload,
1301 /// Cancellation signal. Handlers should `select!` on
1302 /// `cancellation.cancelled()` if their work is async-cancellable;
1303 /// long-running synchronous handlers should periodically check
1304 /// `cancellation.is_cancelled()`.
1305 pub cancellation: RpcCancellationToken,
1306 /// W3C Trace Context propagated from the caller, if the
1307 /// caller set `FLAG_RPC_PROPAGATE_TRACE` and supplied
1308 /// `traceparent` / `tracestate` headers in the request. The
1309 /// server's handler reads this to continue the distributed
1310 /// trace. `None` for calls that didn't propagate trace
1311 /// context.
1312 pub trace_context: Option<TraceContext>,
1313}
1314
1315/// Handler-side error that doesn't fit the application's normal
1316/// `Ok(RpcResponsePayload)` channel. The fold maps these onto a
1317/// failure-status `RpcResponsePayload` for the caller.
1318#[derive(Debug, thiserror::Error)]
1319pub enum RpcHandlerError {
1320 /// Application-defined error. The fold encodes this as
1321 /// `RpcStatus::Application(code)` with `message` as the body.
1322 #[error("application error {code:#06x}: {message}")]
1323 Application {
1324 /// Application error code; surfaces as `RpcStatus::Application(code)`
1325 /// to the caller. Use `0x8000..=0xFFFF` to avoid the
1326 /// reserved canonical range.
1327 code: u16,
1328 /// Diagnostic. Becomes the response body (UTF-8 bytes).
1329 message: String,
1330 },
1331 /// Catch-all for handler-internal failures. The fold encodes this
1332 /// as `RpcStatus::Internal` with `message` as the body.
1333 #[error("internal: {0}")]
1334 Internal(String),
1335}
1336
1337/// User-supplied handler. Implementors typically wrap their state
1338/// (or an `Arc<Mutex<State>>`) and route to the appropriate logic
1339/// based on `ctx.payload.service` or per-handler dispatch.
1340///
1341/// Multiple `Mesh::serve_rpc` registrations on different services
1342/// each install their own handler; a single handler typically
1343/// services one service.
1344#[async_trait::async_trait]
1345pub trait RpcHandler: Send + Sync + 'static {
1346 /// Process one request and return the response payload. The
1347 /// fold spawns this in a tokio task; the fold itself doesn't
1348 /// block on it. Handlers should respect `ctx.cancellation` for
1349 /// cooperative early-abort.
1350 async fn call(&self, ctx: RpcContext) -> Result<RpcResponsePayload, RpcHandlerError>;
1351}
1352
1353/// Callback the fold invokes to publish a response back to the
1354/// caller. Wired up by `Mesh::serve_rpc` to publish on
1355/// `<service>.replies.<caller_origin>`. Type-erased so the fold
1356/// doesn't depend on the mesh layer directly.
1357///
1358/// Arguments: `(caller_origin, call_id, response_payload)`.
1359pub type RpcResponseEmitter = Arc<dyn Fn(u64, u64, RpcResponsePayload) + Send + Sync + 'static>;
1360
1361/// Async counterpart of [`RpcResponseEmitter`] used by the
1362/// streaming fold's pump task to serialize per-call publishes.
1363///
1364/// The streaming pump awaits each emit before reading the next
1365/// chunk from the sink — this guarantees that chunks for one
1366/// `call_id` reach the network publish path in the order the
1367/// handler emitted them. (The unary fold has no such requirement
1368/// — it emits exactly one RESPONSE per call — so it sticks with
1369/// the simpler sync `RpcResponseEmitter`.)
1370pub type RpcAsyncResponseEmitter = Arc<
1371 dyn Fn(u64, u64, RpcResponsePayload) -> futures::future::BoxFuture<'static, ()>
1372 + Send
1373 + Sync
1374 + 'static,
1375>;
1376
1377/// Server-side fold. Sees REQUEST events on the configured channel,
1378/// dispatches to the user-supplied handler, emits RESPONSE events
1379/// via the supplied emitter. CANCEL events flip the matching
1380/// in-flight token.
1381///
1382/// State `()` — the user's state lives on whatever the `RpcHandler`
1383/// captures (typically `Arc<Mutex<S>>`). The fold's own state (the
1384/// in-flight map) lives on `&mut self` and is shared with spawned
1385/// handler tasks via `Arc<Mutex<...>>` so the task can self-clean
1386/// on completion.
1387pub struct RpcServerFold {
1388 handler: Arc<dyn RpcHandler>,
1389 emit: RpcResponseEmitter,
1390 /// (caller_origin, call_id) → cancellation token for the
1391 /// in-flight handler. Inserted on REQUEST, removed by either
1392 /// the spawned handler task on completion or by the fold on
1393 /// CANCEL. Wrapped in `Arc<Mutex<...>>` so spawned tasks can
1394 /// remove their own entries without going back through the
1395 /// fold.
1396 in_flight: Arc<Mutex<HashMap<(u64, u64), RpcCancellationToken>>>,
1397 /// Optional per-service metrics handle. When `Some`, the
1398 /// spawned handler task bumps `handler_invocations_total` /
1399 /// `handler_in_flight` / `handler_panics_total` and records
1400 /// per-task wall-clock durations. `None` → no metrics
1401 /// (test-only path; production `Mesh::serve_rpc` always
1402 /// supplies one).
1403 metrics: Option<Arc<crate::adapter::net::mesh_rpc_metrics::ServiceMetricsAtomic>>,
1404 /// Optional clock override for tests. `None` → real wall-clock
1405 /// `unix_nanos`. `Some(...)` → fixed value, lets tests pin
1406 /// deadline-already-passed behavior without sleeping.
1407 #[cfg(test)]
1408 test_now_ns: Option<u64>,
1409}
1410
1411impl RpcServerFold {
1412 /// Construct a server fold around `handler`. `emit` is the
1413 /// callback that publishes RESPONSE events to the caller's
1414 /// reply channel — `Mesh::serve_rpc` wires this to the
1415 /// publisher for `<service>.replies.<caller_origin>`.
1416 /// Constructed without a metrics handle; production callers
1417 /// chain `.with_metrics(...)` to opt into per-service
1418 /// counters.
1419 pub fn new(handler: Arc<dyn RpcHandler>, emit: RpcResponseEmitter) -> Self {
1420 Self {
1421 handler,
1422 emit,
1423 in_flight: Arc::new(Mutex::new(HashMap::new())),
1424 metrics: None,
1425 #[cfg(test)]
1426 test_now_ns: None,
1427 }
1428 }
1429
1430 /// Attach a per-service metrics handle. Hooks the spawned
1431 /// handler task to bump `handler_invocations_total`, balance
1432 /// `handler_in_flight`, count panics, and record handler
1433 /// duration into the histogram.
1434 pub fn with_metrics(
1435 mut self,
1436 metrics: Arc<crate::adapter::net::mesh_rpc_metrics::ServiceMetricsAtomic>,
1437 ) -> Self {
1438 self.metrics = Some(metrics);
1439 self
1440 }
1441
1442 /// Test-only: pin the clock the fold uses for deadline
1443 /// short-circuit. Lets a unit test exercise the
1444 /// deadline-already-passed branch without waiting for wall
1445 /// time.
1446 #[cfg(test)]
1447 pub fn with_test_now_ns(mut self, now_ns: u64) -> Self {
1448 self.test_now_ns = Some(now_ns);
1449 self
1450 }
1451
1452 /// Test-only: snapshot of the in-flight call set.
1453 #[cfg(test)]
1454 pub fn in_flight_keys(&self) -> Vec<(u64, u64)> {
1455 self.in_flight.lock().keys().copied().collect()
1456 }
1457
1458 fn now_ns(&self) -> u64 {
1459 #[cfg(test)]
1460 if let Some(t) = self.test_now_ns {
1461 return t;
1462 }
1463 std::time::SystemTime::now()
1464 .duration_since(std::time::UNIX_EPOCH)
1465 .map(|d| d.as_nanos() as u64)
1466 .unwrap_or(0)
1467 }
1468
1469 /// `true` if the request's deadline has already elapsed at
1470 /// the server's current wall-clock — accounting for a small
1471 /// tolerance window that absorbs clock skew between caller
1472 /// and server. Without the tolerance, a request from a peer
1473 /// whose clock is a few hundred ms ahead of the server's
1474 /// would be timed out before the handler even saw it. Matches
1475 /// gRPC's default deadline-clock-skew tolerance shape (gRPC
1476 /// uses ~10 s).
1477 fn deadline_already_passed(&self, deadline_ns: u64) -> bool {
1478 if deadline_ns == 0 {
1479 return false;
1480 }
1481 self.now_ns().saturating_sub(DEADLINE_SKEW_TOLERANCE_NS) > deadline_ns
1482 }
1483}
1484
1485/// Tolerance for clock skew between caller and server when the
1486/// server short-circuits a request whose `deadline_ns` looks like
1487/// it has already elapsed. The check is
1488/// `now_ns - SKEW > deadline_ns`, so a request from a peer whose
1489/// clock is up to `SKEW` nanoseconds ahead of ours never hits the
1490/// short-circuit path. 10 s matches gRPC's default and is well
1491/// within the threshold an NTP-disciplined cluster ever drifts to.
1492pub const DEADLINE_SKEW_TOLERANCE_NS: u64 = 10_000_000_000; // 10 seconds
1493
1494impl RedexFold<()> for RpcServerFold {
1495 fn apply(&mut self, ev: &RedexEvent, _state: &mut ()) -> Result<(), RedexError> {
1496 // Decode the meta header. A garbled meta means the event
1497 // doesn't even claim to be an RPC packet — log and skip
1498 // rather than killing the fold. Returning `Err(Decode)`
1499 // here would stop the entire cortex adapter for one
1500 // malformed event, which is wrong for an RPC server that
1501 // needs to keep serving.
1502 let Some(meta) = (if ev.payload.len() >= EVENT_META_SIZE {
1503 EventMeta::from_bytes(&ev.payload[..EVENT_META_SIZE])
1504 } else {
1505 None
1506 }) else {
1507 tracing::warn!(
1508 payload_len = ev.payload.len(),
1509 "rpc server fold: event payload too short for EventMeta; skipping",
1510 );
1511 return Ok(());
1512 };
1513 let key = (meta.origin_hash, meta.seq_or_ts);
1514 match meta.dispatch {
1515 DISPATCH_RPC_REQUEST => {
1516 let payload = match RpcRequestPayload::decode(ev.payload.slice(EVENT_META_SIZE..)) {
1517 Ok(p) => p,
1518 Err(e) => {
1519 // Malformed request payload. Surface as
1520 // `UnknownVersion` to the caller — they sent
1521 // bytes we couldn't parse, which usually
1522 // means a wire-format mismatch (the most
1523 // common cause). Log so operators can
1524 // diagnose.
1525 tracing::warn!(
1526 error = %e,
1527 caller_origin = format!("{:#x}", meta.origin_hash),
1528 call_id = meta.seq_or_ts,
1529 "rpc server fold: malformed request payload",
1530 );
1531 let resp = RpcResponsePayload {
1532 status: RpcStatus::UnknownVersion,
1533 headers: vec![],
1534 body: Bytes::from(format!("malformed request: {e}")),
1535 };
1536 (self.emit)(meta.origin_hash, meta.seq_or_ts, resp);
1537 return Ok(());
1538 }
1539 };
1540 // Fast deadline-already-passed short-circuit.
1541 // Server-side `Timeout` without invoking the
1542 // handler. Includes a clock-skew tolerance window
1543 // so a peer with a slightly-fast clock isn't
1544 // prematurely timed out — see
1545 // `deadline_already_passed`.
1546 if self.deadline_already_passed(payload.deadline_ns) {
1547 let resp = RpcResponsePayload {
1548 status: RpcStatus::Timeout,
1549 headers: vec![],
1550 body: Bytes::from_static(b"deadline already passed when request landed"),
1551 };
1552 (self.emit)(meta.origin_hash, meta.seq_or_ts, resp);
1553 return Ok(());
1554 }
1555 // Refuse a duplicate REQUEST with the same
1556 // `(origin_hash, call_id)` — see streaming fold for
1557 // the full rationale. For the unary fold this would
1558 // spawn a second handler under the same key, and
1559 // whichever handler completes first removes the
1560 // in-flight entry — leaving the second handler's
1561 // CANCEL handling broken (CANCEL events look up
1562 // the now-missing key and no-op). Cleaner to refuse.
1563 {
1564 let in_flight = self.in_flight.lock();
1565 if in_flight.contains_key(&key) {
1566 drop(in_flight);
1567 tracing::warn!(
1568 caller_origin = format!("{:#x}", meta.origin_hash),
1569 call_id = meta.seq_or_ts,
1570 "rpc server fold: duplicate REQUEST for in-flight call_id; refusing",
1571 );
1572 let resp = RpcResponsePayload {
1573 status: RpcStatus::Internal,
1574 headers: vec![],
1575 body: Bytes::from_static(
1576 b"duplicate REQUEST for already-in-flight call_id",
1577 ),
1578 };
1579 (self.emit)(meta.origin_hash, meta.seq_or_ts, resp);
1580 return Ok(());
1581 }
1582 }
1583 let cancellation = RpcCancellationToken::new();
1584 self.in_flight.lock().insert(key, cancellation.clone());
1585 let handler = self.handler.clone();
1586 let emit = self.emit.clone();
1587 let in_flight = self.in_flight.clone();
1588 let caller_origin = meta.origin_hash;
1589 let call_id = meta.seq_or_ts;
1590 // Decode the W3C Trace Context if the caller
1591 // signaled it via `FLAG_RPC_PROPAGATE_TRACE` and
1592 // included the `traceparent` / `tracestate`
1593 // headers. nRPC is transport-only — application
1594 // code reads `ctx.trace_context` to continue the
1595 // trace via whatever backend it has wired up.
1596 let trace_context = if payload.flags & FLAG_RPC_PROPAGATE_TRACE != 0 {
1597 extract_trace_context(&payload.headers)
1598 } else {
1599 None
1600 };
1601 let metrics = self.metrics.clone();
1602 // Keep a probe handle so the spawned task can detect
1603 // a CANCEL that fired during handler execution and
1604 // override its response with `RpcStatus::Cancelled`.
1605 let cancel_probe = cancellation.clone();
1606 tokio::spawn(async move {
1607 // Server-side metrics: count this invocation;
1608 // bump in_flight; time the handler; tally
1609 // panics. Only fires when a metrics handle was
1610 // attached via `with_metrics(...)` — test-only
1611 // folds construct without one.
1612 if let Some(m) = metrics.as_ref() {
1613 m.handler_invocations_total
1614 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1615 m.handler_in_flight
1616 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1617 }
1618 let handler_started = std::time::Instant::now();
1619 let ctx = RpcContext {
1620 caller_origin,
1621 call_id,
1622 payload,
1623 cancellation,
1624 trace_context,
1625 };
1626 // Catch panics so a misbehaving handler can't
1627 // take down the runtime. `AssertUnwindSafe` is
1628 // load-bearing because `RpcHandler::call`
1629 // returns a future that may borrow non-
1630 // `UnwindSafe` types from the handler; we
1631 // accept the assertion because the handler's
1632 // state is untouched on panic (we just don't
1633 // observe its in-progress mutations).
1634 let outcome = futures::FutureExt::catch_unwind(std::panic::AssertUnwindSafe(
1635 handler.call(ctx),
1636 ))
1637 .await;
1638 if let Some(m) = metrics.as_ref() {
1639 m.handler_in_flight
1640 .fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
1641 m.record_handler_duration(handler_started.elapsed());
1642 if outcome.is_err() {
1643 m.handler_panics_total
1644 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1645 }
1646 }
1647 // CANCEL-wins ordering: if the cancellation
1648 // token fired at any point during handler
1649 // execution, override the handler's outcome
1650 // with `RpcStatus::Cancelled` so the caller
1651 // (or hedge primary, retry layer, etc.) sees
1652 // the documented `Cancelled` status code rather
1653 // than whatever the handler happened to return
1654 // before / despite cancellation. A cooperative
1655 // handler that observes the token and bails
1656 // early gets the same Cancelled framing as a
1657 // handler that ignored cancellation and ran to
1658 // completion — the caller's view is uniform.
1659 let resp = if cancel_probe.is_cancelled() {
1660 RpcResponsePayload {
1661 status: RpcStatus::Cancelled,
1662 headers: vec![],
1663 body: Bytes::from_static(
1664 b"server observed CANCEL during handler execution",
1665 ),
1666 }
1667 } else {
1668 match outcome {
1669 Ok(Ok(payload)) => payload,
1670 Ok(Err(RpcHandlerError::Application { code, message })) => {
1671 RpcResponsePayload {
1672 status: RpcStatus::Application(code),
1673 headers: vec![],
1674 body: Bytes::from(message),
1675 }
1676 }
1677 Ok(Err(RpcHandlerError::Internal(message))) => RpcResponsePayload {
1678 status: RpcStatus::Internal,
1679 headers: vec![],
1680 body: Bytes::from(message),
1681 },
1682 Err(panic) => {
1683 let panic_msg = panic
1684 .downcast_ref::<&'static str>()
1685 .map(|s| s.to_string())
1686 .or_else(|| panic.downcast_ref::<String>().cloned())
1687 .unwrap_or_else(|| "<non-string panic>".into());
1688 tracing::error!(
1689 caller_origin = format!("{:#x}", caller_origin),
1690 call_id,
1691 panic = %panic_msg,
1692 "rpc server handler panicked",
1693 );
1694 RpcResponsePayload {
1695 status: RpcStatus::Internal,
1696 headers: vec![],
1697 body: Bytes::from(format!("handler panicked: {panic_msg}")),
1698 }
1699 }
1700 }
1701 };
1702 in_flight.lock().remove(&key);
1703 emit(caller_origin, call_id, resp);
1704 });
1705 }
1706 DISPATCH_RPC_CANCEL => {
1707 if let Some(token) = self.in_flight.lock().remove(&key) {
1708 token.cancel();
1709 }
1710 // Idempotent — CANCEL for an unknown call_id (e.g.
1711 // a CANCEL that races the handler's completion) is
1712 // a no-op rather than an error. The spawned handler
1713 // task observes `cancel_probe.is_cancelled()` after
1714 // its future resolves and overrides the response
1715 // with `RpcStatus::Cancelled` so the caller sees a
1716 // documented status code rather than the handler's
1717 // accidental Ok / Internal payload.
1718 }
1719 // RESPONSE / DEADLINE_EXCEEDED are server-emitted; if
1720 // the server's own fold sees them (e.g. from a replay)
1721 // there's nothing to do.
1722 _ => {}
1723 }
1724 Ok(())
1725 }
1726}
1727
1728// ============================================================================
1729// Streaming server-side: handler trait + sink + fold.
1730// ============================================================================
1731
1732/// Sink the handler writes to in order to emit streaming-response
1733/// chunks. Each `send` produces one non-terminal `RESPONSE` event
1734/// to the caller. The terminal frame is emitted automatically when
1735/// the sink is dropped — the handler returning `Ok(())` drops the
1736/// sink, which closes the stream cleanly. Returning
1737/// `Err(RpcHandlerError)` drops the sink and emits the error as a
1738/// terminal non-`Ok` RESPONSE.
1739///
1740/// `send` is best-effort and infallible: the underlying mpsc is
1741/// **bounded** at [`STREAMING_PUMP_CAPACITY`] chunks. If the pump
1742/// can't keep up (publish path is congested, caller hasn't granted
1743/// flow-control credits), `send` discards on overflow — same
1744/// observable shape as a closed receiver (caller cancelled mid-
1745/// stream). Counts the drop in `streaming_chunks_dropped_total` so
1746/// operators can see backpressure occurring. Cooperative
1747/// cancellation via `ctx.cancellation` is the right way for the
1748/// handler to notice the consumer is gone; opt-in flow control via
1749/// `CallOptions::stream_window_initial` is the right way to
1750/// throttle a fast handler against a slow consumer.
1751pub struct RpcResponseSink {
1752 inner: tokio::sync::mpsc::Sender<bytes::Bytes>,
1753 /// Optional metrics handle so a dropped-on-full chunk bumps the
1754 /// `streaming_chunks_dropped_total` counter. `None` for unit-
1755 /// test folds that construct without metrics.
1756 metrics: Option<Arc<crate::adapter::net::mesh_rpc_metrics::ServiceMetricsAtomic>>,
1757}
1758
1759impl RpcResponseSink {
1760 /// Emit one non-terminal chunk. Cheap (`try_send` on a
1761 /// [`STREAMING_PUMP_CAPACITY`]-bounded mpsc); never blocks. On
1762 /// overflow OR receiver-closed, the chunk is dropped and (when
1763 /// metrics are wired) `streaming_chunks_dropped_total` is
1764 /// incremented for the service.
1765 pub fn send(&self, body: impl Into<bytes::Bytes>) {
1766 if self.inner.try_send(body.into()).is_err() {
1767 if let Some(m) = self.metrics.as_ref() {
1768 m.streaming_chunks_dropped_total
1769 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1770 }
1771 }
1772 }
1773}
1774
1775/// Bounded capacity for the streaming pump's internal mpsc. A
1776/// runaway handler that produces chunks faster than the publish
1777/// path can drain them stops blocking the runtime past this many
1778/// queued chunks — additional chunks are dropped (and counted via
1779/// `streaming_chunks_dropped_total`). 1024 is generous for typical
1780/// streaming patterns; opt-in flow control via
1781/// `CallOptions::stream_window_initial` is the right primitive for
1782/// strict throttling.
1783pub const STREAMING_PUMP_CAPACITY: usize = 1024;
1784
1785/// Bounded capacity for the client-streaming server fold's
1786/// per-call request mpsc. Mirror of [`STREAMING_PUMP_CAPACITY`]
1787/// for the upload direction. A runaway caller that emits
1788/// REQUEST_CHUNKs faster than the handler can drain stops
1789/// queueing past this many chunks — additional chunks are
1790/// dropped (and counted via `streaming_request_chunks_dropped_total`
1791/// when metrics are wired). Opt-in flow control via the
1792/// `nrpc-request-window-initial` header is the right primitive
1793/// for strict throttling on the upload side.
1794///
1795/// Bidi streaming plan (Phase B).
1796pub const STREAMING_REQUEST_PUMP_CAPACITY: usize = 1024;
1797
1798// ============================================================================
1799// Phase B — server-side primitives for client-streaming.
1800// ============================================================================
1801
1802/// Context handed to an [`RpcClientStreamingHandler::call`]. Same
1803/// shape as [`RpcContext`] minus the eager `payload` (the request
1804/// stream delivers chunk bodies on the fly) and plus the
1805/// per-call `deadline_ns` (which would otherwise have ridden on
1806/// the eager payload).
1807///
1808/// Bidi streaming plan (Phase B).
1809pub struct RpcStreamingContext {
1810 /// AEAD-verified caller `origin_hash`. Same source as
1811 /// [`RpcContext::caller_origin`].
1812 pub caller_origin: u64,
1813 /// Caller-generated correlation id. Matches the initial
1814 /// REQUEST's `call_id` and every subsequent REQUEST_CHUNK /
1815 /// CANCEL / REQUEST_GRANT for this call.
1816 pub call_id: u64,
1817 /// Absolute deadline (unix nanos) from the initial REQUEST.
1818 /// `0` means no deadline; the fold does NOT auto-cancel on
1819 /// deadline (handlers self-supervise via tokio timers, same
1820 /// contract as the unary fold).
1821 pub deadline_ns: u64,
1822 /// Per-chunk metadata headers from the initial REQUEST.
1823 /// Per-REQUEST_CHUNK headers are NOT surfaced at the substrate
1824 /// layer — the typed SDK veneer (Phase E) is where header
1825 /// inspection across chunks lives (if it lands at all; the
1826 /// plan defers per-chunk-headers as opt-in raw-path access).
1827 pub headers: Vec<RpcHeader>,
1828 /// Cancellation signal. Flipped by the fold when a
1829 /// `DISPATCH_RPC_CANCEL` arrives for this call's `call_id`.
1830 /// Long-running handlers should `select!` on
1831 /// `cancellation.cancelled()`; the request stream also
1832 /// terminates on cancellation, but the token is the
1833 /// authoritative signal (the stream's terminator is shared
1834 /// with REQUEST_END).
1835 pub cancellation: RpcCancellationToken,
1836 /// W3C Trace Context propagated from the caller's initial
1837 /// REQUEST. Same semantics as [`RpcContext::trace_context`].
1838 pub trace_context: Option<TraceContext>,
1839}
1840
1841/// Callback the fold invokes to publish a [`DISPATCH_RPC_REQUEST_GRANT`]
1842/// event back to the caller. Wired up by the `Mesh` glue (Phase C)
1843/// to publish on the caller's reply channel. Type-erased so the
1844/// fold doesn't depend on the mesh layer directly.
1845///
1846/// Arguments: `(caller_origin, call_id, credits)`. Synchronous —
1847/// the publish itself is non-blocking (the underlying transport
1848/// has its own internal queueing); the fold fires-and-forgets
1849/// every grant, so dropped grants are at worst a latency wobble,
1850/// not a correctness issue (the caller's send sink will retry
1851/// when its credit budget refills via the next grant or via the
1852/// initial window).
1853///
1854/// Bidi streaming plan (Phase B).
1855pub type RpcRequestGrantEmitter = Arc<dyn Fn(u64, u64, u32) + Send + Sync + 'static>;
1856
1857/// Server-side stream of inbound request chunk bodies for one
1858/// client-streaming (or duplex) call. Yields one `Bytes` per
1859/// `DISPATCH_RPC_REQUEST` / `DISPATCH_RPC_REQUEST_CHUNK` frame
1860/// (including empty bodies — the semantics of "empty bytes" are
1861/// the application's concern, not the substrate's). Closes on
1862/// `FLAG_RPC_REQUEST_END` or on CANCEL.
1863///
1864/// **Stream item ordering convention**: the first item this
1865/// stream yields corresponds to the initial REQUEST body; every
1866/// subsequent item corresponds to a REQUEST_CHUNK body, in the
1867/// order the chunks were received from the wire. The substrate
1868/// does not tag items with their frame kind — the SDK veneer
1869/// (Phase E) is responsible for the Init / Data classification
1870/// via its `Chunk<T>` enum.
1871///
1872/// **Auto-grant behavior**: when the caller opted into
1873/// request-direction flow control via
1874/// [`HEADER_NRPC_REQUEST_WINDOW_INITIAL`], every successful
1875/// `poll_next()` fires one credit back to the caller via the
1876/// captured `grant_emitter`. This keeps the in-flight window
1877/// at the caller's initial value as the handler drains the
1878/// stream. When the caller did NOT opt in (no header), the
1879/// `grant_emitter` is `None` and the auto-grant path is a no-op
1880/// (caller is on the unbounded-credit fast path).
1881///
1882/// Bidi streaming plan (Phase B).
1883pub struct RequestStream {
1884 inner: tokio::sync::mpsc::Receiver<bytes::Bytes>,
1885 grant_emitter: Option<RpcRequestGrantEmitter>,
1886 caller_origin: u64,
1887 call_id: u64,
1888}
1889
1890impl RequestStream {
1891 /// Visible to the fold (and only the fold) for constructing
1892 /// a stream tied to a specific receiver + caller. The
1893 /// `grant_emitter` is `None` when the caller didn't opt into
1894 /// flow control; `Some(...)` when they did.
1895 pub(crate) fn new(
1896 inner: tokio::sync::mpsc::Receiver<bytes::Bytes>,
1897 grant_emitter: Option<RpcRequestGrantEmitter>,
1898 caller_origin: u64,
1899 call_id: u64,
1900 ) -> Self {
1901 Self {
1902 inner,
1903 grant_emitter,
1904 caller_origin,
1905 call_id,
1906 }
1907 }
1908}
1909
1910impl futures::Stream for RequestStream {
1911 type Item = bytes::Bytes;
1912
1913 fn poll_next(
1914 mut self: std::pin::Pin<&mut Self>,
1915 cx: &mut std::task::Context<'_>,
1916 ) -> std::task::Poll<Option<Self::Item>> {
1917 match self.inner.poll_recv(cx) {
1918 std::task::Poll::Ready(Some(bytes)) => {
1919 // Auto-grant fires on every successful pull when
1920 // flow control was opted into. Cheap and
1921 // fire-and-forget; missed grants are recovered
1922 // by subsequent pulls.
1923 if let Some(emit) = self.grant_emitter.as_ref() {
1924 emit(self.caller_origin, self.call_id, 1);
1925 }
1926 std::task::Poll::Ready(Some(bytes))
1927 }
1928 other => other,
1929 }
1930 }
1931}
1932
1933/// User-supplied handler for a client-streaming RPC. Receives an
1934/// [`RpcStreamingContext`] (caller identity, deadline, cancellation,
1935/// trace context, initial REQUEST headers) plus a [`RequestStream`]
1936/// of chunk bodies. Returns one terminal [`RpcResponsePayload`] —
1937/// the fold publishes it as the call's single RESPONSE frame.
1938///
1939/// **Cancellation contract.** Long-running handlers should
1940/// `select!` on `ctx.cancellation.cancelled()` so a caller-side
1941/// drop / deadline correctly stops the handler. The request
1942/// stream also terminates on cancellation (yields `None`), but
1943/// the token is the authoritative signal — the stream's `None`
1944/// is shared with the clean REQUEST_END path, so handlers can't
1945/// distinguish "caller finished cleanly" from "caller cancelled"
1946/// without consulting the token.
1947///
1948/// **Auto-grant.** When the caller opted into request-direction
1949/// flow control via [`HEADER_NRPC_REQUEST_WINDOW_INITIAL`], every
1950/// `stream.next().await` that yields `Some` fires one
1951/// REQUEST_GRANT back to the caller, maintaining the in-flight
1952/// window at the caller's initial value. Handlers don't need to
1953/// think about credit management for the common case.
1954///
1955/// Bidi streaming plan (Phase B).
1956#[async_trait::async_trait]
1957pub trait RpcClientStreamingHandler: Send + Sync + 'static {
1958 /// Process a client-streaming call. Drain the request stream,
1959 /// produce one terminal response payload (or an
1960 /// [`RpcHandlerError`] for failure mapping).
1961 async fn call(
1962 &self,
1963 ctx: RpcStreamingContext,
1964 requests: RequestStream,
1965 ) -> Result<RpcResponsePayload, RpcHandlerError>;
1966}
1967
1968/// User-supplied handler for a duplex RPC — many requests in,
1969/// many responses out, interleaved. Receives an [`RpcStreamingContext`]
1970/// plus a [`RequestStream`] of chunk bodies plus an
1971/// [`RpcResponseSink`] for emitting response chunks. The handler's
1972/// return value is its terminal status, NOT a final payload:
1973/// `Ok(())` closes the response stream cleanly with a terminal
1974/// `Ok` frame, `Err(RpcHandlerError)` closes with the matching
1975/// error status.
1976///
1977/// **Composition.** A duplex handler is a hybrid of an
1978/// [`RpcClientStreamingHandler`] (drains request chunks) and an
1979/// [`RpcStreamingHandler`] (emits response chunks). The two
1980/// directions are independent — a handler can finish emitting
1981/// responses before reading all requests, or vice versa. The
1982/// server fold serializes RESPONSE chunk publishes per call_id
1983/// so wire order matches handler order.
1984///
1985/// **Cancellation contract.** Identical to
1986/// [`RpcClientStreamingHandler`]: long-running work should
1987/// `select!` on `ctx.cancellation.cancelled()`.
1988///
1989/// **Auto-grant.** Identical to [`RpcClientStreamingHandler`]:
1990/// every successful `requests.next().await` emits one
1991/// REQUEST_GRANT back to the caller (when the caller opted in).
1992///
1993/// Bidi streaming plan (Phase D).
1994#[async_trait::async_trait]
1995pub trait RpcDuplexHandler: Send + Sync + 'static {
1996 /// Process one duplex call. Drain inbound chunks via
1997 /// `requests.next().await`; emit outbound chunks via
1998 /// `responses.send(...)`. Return `Ok(())` for clean close,
1999 /// `Err(RpcHandlerError)` for failure mapping.
2000 async fn call(
2001 &self,
2002 ctx: RpcStreamingContext,
2003 requests: RequestStream,
2004 responses: RpcResponseSink,
2005 ) -> Result<(), RpcHandlerError>;
2006}
2007
2008/// User-supplied streaming handler. Receives the same `RpcContext`
2009/// as a unary handler plus a `RpcResponseSink` for emitting chunks.
2010/// Returning `Ok(())` closes the stream cleanly with a terminal
2011/// `Ok` RESPONSE; `Err(RpcHandlerError)` closes the stream with a
2012/// terminal non-`Ok` RESPONSE carrying the diagnostic.
2013///
2014/// **Cancellation contract.** Long-running streams should
2015/// `select!` on `ctx.cancellation.cancelled()` so a caller-side
2016/// drop / deadline correctly stops the handler. Continuing to
2017/// `send` after cancellation is harmless (sink discards) but
2018/// wastes work.
2019#[async_trait::async_trait]
2020pub trait RpcStreamingHandler: Send + Sync + 'static {
2021 /// Process one streaming request. Emit chunks via `sink.send(...)`.
2022 /// Drop the sink (or return) to close the stream.
2023 async fn call(&self, ctx: RpcContext, sink: RpcResponseSink) -> Result<(), RpcHandlerError>;
2024}
2025
2026/// Per-call flow-control map type. Keyed on
2027/// `(caller_origin_hash, call_id)`; value is a tokio
2028/// `Semaphore` shared between the pump task (which awaits
2029/// permits) and the fold's `apply()` method handling
2030/// STREAM_GRANT events (which add permits).
2031type FlowControlMap = Arc<Mutex<HashMap<(u64, u64), Arc<tokio::sync::Semaphore>>>>;
2032
2033/// Server-side fold for streaming RPC. Parallel to `RpcServerFold`
2034/// but multi-fire emit: each handler invocation may produce many
2035/// `RESPONSE` events for the same `call_id`, marked
2036/// non-terminal/terminal via the `nrpc-streaming` header.
2037///
2038/// State `()` — like the unary fold, the handler owns user state
2039/// via captured `Arc<Mutex<S>>`. The fold's own state (in-flight
2040/// cancellation tokens) lives on `&mut self`.
2041pub struct RpcServerStreamingFold {
2042 handler: Arc<dyn RpcStreamingHandler>,
2043 emit: RpcAsyncResponseEmitter,
2044 in_flight: Arc<Mutex<HashMap<(u64, u64), RpcCancellationToken>>>,
2045 /// Per-call flow-control semaphore (when the caller opted in).
2046 /// `Some(sem)` means "pump must `acquire().await` one permit
2047 /// per chunk before emitting; STREAM_GRANT events
2048 /// `add_permits(n)`". Absence of an entry for a `(origin,
2049 /// call_id)` key means unbounded credit (no flow control —
2050 /// pump emits as fast as the publish path can take chunks).
2051 flow_control: FlowControlMap,
2052 /// Optional per-service metrics handle. Same shape as
2053 /// `RpcServerFold::metrics`; the streaming fold ALSO bumps
2054 /// `streaming_chunks_emitted_total` from the pump task on
2055 /// every chunk.
2056 metrics: Option<Arc<crate::adapter::net::mesh_rpc_metrics::ServiceMetricsAtomic>>,
2057}
2058
2059impl RpcServerStreamingFold {
2060 /// Construct a streaming server fold. `emit` publishes
2061 /// individual chunks (and the terminal frame) on the caller's
2062 /// reply channel.
2063 ///
2064 /// Uses the **async** emitter variant so the pump task can
2065 /// serialize per-call publishes — without that ordering
2066 /// guarantee, two chunks emitted in succession can race into
2067 /// the publish path and arrive at the caller out of order
2068 /// (or be eclipsed by the terminal frame and lost entirely).
2069 pub fn new(handler: Arc<dyn RpcStreamingHandler>, emit: RpcAsyncResponseEmitter) -> Self {
2070 Self {
2071 handler,
2072 emit,
2073 in_flight: Arc::new(Mutex::new(HashMap::new())),
2074 flow_control: Arc::new(Mutex::new(HashMap::new())),
2075 metrics: None,
2076 }
2077 }
2078
2079 /// Attach a per-service metrics handle. Hooks the spawned
2080 /// handler task to bump `handler_invocations_total` /
2081 /// `handler_in_flight` / `handler_panics_total` /
2082 /// `handler_duration_*`, and the pump task to bump
2083 /// `streaming_chunks_emitted_total` per emitted chunk.
2084 pub fn with_metrics(
2085 mut self,
2086 metrics: Arc<crate::adapter::net::mesh_rpc_metrics::ServiceMetricsAtomic>,
2087 ) -> Self {
2088 self.metrics = Some(metrics);
2089 self
2090 }
2091
2092 /// Test-only: snapshot of the in-flight call set.
2093 #[cfg(test)]
2094 pub fn in_flight_keys(&self) -> Vec<(u64, u64)> {
2095 self.in_flight.lock().keys().copied().collect()
2096 }
2097}
2098
2099impl RedexFold<()> for RpcServerStreamingFold {
2100 fn apply(&mut self, ev: &RedexEvent, _state: &mut ()) -> Result<(), RedexError> {
2101 let Some(meta) = (if ev.payload.len() >= EVENT_META_SIZE {
2102 EventMeta::from_bytes(&ev.payload[..EVENT_META_SIZE])
2103 } else {
2104 None
2105 }) else {
2106 tracing::warn!(
2107 payload_len = ev.payload.len(),
2108 "rpc streaming server fold: event payload too short for EventMeta",
2109 );
2110 return Ok(());
2111 };
2112 let key = (meta.origin_hash, meta.seq_or_ts);
2113 match meta.dispatch {
2114 DISPATCH_RPC_REQUEST => {
2115 let payload = match RpcRequestPayload::decode(ev.payload.slice(EVENT_META_SIZE..)) {
2116 Ok(p) => p,
2117 Err(e) => {
2118 tracing::warn!(
2119 error = %e,
2120 caller_origin = format!("{:#x}", meta.origin_hash),
2121 call_id = meta.seq_or_ts,
2122 "rpc streaming server fold: malformed request payload",
2123 );
2124 // Surface as a terminal error chunk. Spawn
2125 // because the apply method is sync and the
2126 // emit is async; this is a one-shot publish
2127 // so ordering doesn't matter here.
2128 let resp = RpcResponsePayload {
2129 status: RpcStatus::UnknownVersion,
2130 headers: vec![(
2131 HEADER_NRPC_STREAMING.to_string(),
2132 HEADER_NRPC_STREAMING_END.to_vec(),
2133 )],
2134 body: Bytes::from(format!("malformed request: {e}")),
2135 };
2136 let emit = self.emit.clone();
2137 let caller_origin = meta.origin_hash;
2138 let call_id = meta.seq_or_ts;
2139 tokio::spawn(async move {
2140 emit(caller_origin, call_id, resp).await;
2141 });
2142 return Ok(());
2143 }
2144 };
2145 // Refuse a duplicate REQUEST with the same
2146 // `(origin_hash, call_id)`. Without this, a retry
2147 // that arrives while the first attempt's pump is
2148 // still draining will overwrite the prior
2149 // semaphore Arc in `flow_control`, leaving the
2150 // first pump awaiting an orphaned semaphore (the
2151 // terminal cleanup keys on `key` and removes the
2152 // *new* entry, so the orphan never gets dropped
2153 // and the first handler hangs forever).
2154 //
2155 // Idempotent for the caller: we emit a terminal
2156 // `Internal` chunk so the duplicate sender sees a
2157 // clean refusal rather than waiting on a stream
2158 // that will never produce output.
2159 {
2160 let in_flight = self.in_flight.lock();
2161 if in_flight.contains_key(&key) {
2162 drop(in_flight);
2163 tracing::warn!(
2164 caller_origin = format!("{:#x}", meta.origin_hash),
2165 call_id = meta.seq_or_ts,
2166 "rpc streaming server fold: duplicate REQUEST for in-flight call_id; refusing",
2167 );
2168 let resp = RpcResponsePayload {
2169 status: RpcStatus::Internal,
2170 headers: vec![(
2171 HEADER_NRPC_STREAMING.to_string(),
2172 HEADER_NRPC_STREAMING_END.to_vec(),
2173 )],
2174 body: Bytes::from_static(
2175 b"duplicate REQUEST for already-in-flight call_id",
2176 ),
2177 };
2178 let emit = self.emit.clone();
2179 let caller_origin = meta.origin_hash;
2180 let call_id = meta.seq_or_ts;
2181 tokio::spawn(async move {
2182 emit(caller_origin, call_id, resp).await;
2183 });
2184 return Ok(());
2185 }
2186 }
2187 // Cancellation token + in-flight bookkeeping —
2188 // identical to the unary fold's pattern.
2189 let cancellation = RpcCancellationToken::new();
2190 self.in_flight.lock().insert(key, cancellation.clone());
2191 // Flow-control opt-in: parse the
2192 // `nrpc-stream-window-initial` header. When
2193 // present, install a per-call semaphore the pump
2194 // task will await per chunk; subsequent
2195 // STREAM_GRANT events refill it. When absent, no
2196 // entry → pump skips the await (back-compat).
2197 let flow_sem = parse_stream_window_initial(&payload.headers).map(|n| {
2198 let sem = Arc::new(tokio::sync::Semaphore::new(n as usize));
2199 self.flow_control.lock().insert(key, sem.clone());
2200 sem
2201 });
2202 let handler = self.handler.clone();
2203 let emit = self.emit.clone();
2204 let in_flight = self.in_flight.clone();
2205 let flow_control = self.flow_control.clone();
2206 let caller_origin = meta.origin_hash;
2207 let call_id = meta.seq_or_ts;
2208 let trace_context = if payload.flags & FLAG_RPC_PROPAGATE_TRACE != 0 {
2209 extract_trace_context(&payload.headers)
2210 } else {
2211 None
2212 };
2213 let metrics = self.metrics.clone();
2214 // See unary fold for rationale — clone the
2215 // cancellation handle so the spawned task can probe
2216 // it after the handler returns and override the
2217 // terminal frame with `RpcStatus::Cancelled`.
2218 let cancel_probe = cancellation.clone();
2219 tokio::spawn(async move {
2220 if let Some(m) = metrics.as_ref() {
2221 m.handler_invocations_total
2222 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2223 m.handler_in_flight
2224 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2225 }
2226 let handler_started = std::time::Instant::now();
2227 let ctx = RpcContext {
2228 caller_origin,
2229 call_id,
2230 payload,
2231 cancellation,
2232 trace_context,
2233 };
2234 // Build the sink + receive end. Spawn a
2235 // pump that forwards each chunk to the emit
2236 // closure. The handler's `sink.send(...)`
2237 // calls show up here as items on the receiver.
2238 // **Bounded** at STREAMING_PUMP_CAPACITY: a
2239 // runaway handler that produces chunks faster
2240 // than the publish path can drain stops
2241 // blocking the runtime past this many queued
2242 // chunks; additional chunks are dropped and
2243 // counted via streaming_chunks_dropped_total.
2244 let (tx, mut rx) =
2245 tokio::sync::mpsc::channel::<bytes::Bytes>(STREAMING_PUMP_CAPACITY);
2246 let sink = RpcResponseSink {
2247 inner: tx,
2248 metrics: metrics.clone(),
2249 };
2250 let pump_emit = emit.clone();
2251 let pump_metrics = metrics.clone();
2252 let pump_flow = flow_sem.clone();
2253 let pump = tokio::spawn(async move {
2254 while let Some(chunk) = rx.recv().await {
2255 // Flow control: when the caller opted
2256 // in, await one semaphore permit per
2257 // chunk before publishing. The semaphore
2258 // starts at the caller's `initial_window`
2259 // and refills when the caller sends
2260 // STREAM_GRANT events. `forget()`
2261 // consumes the slot — each chunk uses
2262 // exactly one credit, never returned.
2263 // No-op when `pump_flow` is None
2264 // (back-compat path).
2265 if let Some(sem) = pump_flow.as_ref() {
2266 let permit = match sem.clone().acquire_owned().await {
2267 Ok(p) => p,
2268 Err(_) => {
2269 // Semaphore was closed —
2270 // shouldn't happen during
2271 // normal operation; bail.
2272 break;
2273 }
2274 };
2275 permit.forget();
2276 }
2277 if let Some(m) = pump_metrics.as_ref() {
2278 m.streaming_chunks_emitted_total
2279 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2280 }
2281 let resp = RpcResponsePayload {
2282 status: RpcStatus::Ok,
2283 headers: vec![(
2284 HEADER_NRPC_STREAMING.to_string(),
2285 HEADER_NRPC_STREAMING_CONTINUE.to_vec(),
2286 )],
2287 body: chunk.clone(),
2288 };
2289 // Await per-chunk publish so chunks for
2290 // one call_id reach the network in send
2291 // order. Without this, two chunks emitted
2292 // in tight succession can race into the
2293 // publish path and arrive out of order
2294 // (or be eclipsed by the terminal frame
2295 // and lost entirely on the caller side).
2296 pump_emit(caller_origin, call_id, resp).await;
2297 }
2298 });
2299 // Run the handler. Catch panics so a
2300 // misbehaving handler can't take down the
2301 // runtime — same shape as the unary fold.
2302 let outcome = futures::FutureExt::catch_unwind(std::panic::AssertUnwindSafe(
2303 handler.call(ctx, sink),
2304 ))
2305 .await;
2306 // The handler dropped the sink (either by
2307 // returning or by panicking through the
2308 // catch_unwind). Wait for the pump to drain
2309 // any final in-flight chunks.
2310 let _ = pump.await;
2311 if let Some(m) = metrics.as_ref() {
2312 m.handler_in_flight
2313 .fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
2314 m.record_handler_duration(handler_started.elapsed());
2315 if outcome.is_err() {
2316 m.handler_panics_total
2317 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2318 }
2319 }
2320 // Emit the terminal frame. CANCEL-wins ordering
2321 // matches the unary fold: if the cancellation
2322 // token fired during execution, override the
2323 // handler's terminal with `RpcStatus::Cancelled`.
2324 let terminal = if cancel_probe.is_cancelled() {
2325 RpcResponsePayload {
2326 status: RpcStatus::Cancelled,
2327 headers: vec![],
2328 body: Bytes::from_static(
2329 b"server observed CANCEL during streaming handler execution",
2330 ),
2331 }
2332 } else {
2333 match outcome {
2334 Ok(Ok(())) => RpcResponsePayload {
2335 status: RpcStatus::Ok,
2336 headers: vec![(
2337 HEADER_NRPC_STREAMING.to_string(),
2338 HEADER_NRPC_STREAMING_END.to_vec(),
2339 )],
2340 body: Bytes::new(),
2341 },
2342 Ok(Err(RpcHandlerError::Application { code, message })) => {
2343 RpcResponsePayload {
2344 status: RpcStatus::Application(code),
2345 headers: vec![],
2346 body: Bytes::from(message),
2347 }
2348 }
2349 Ok(Err(RpcHandlerError::Internal(message))) => RpcResponsePayload {
2350 status: RpcStatus::Internal,
2351 headers: vec![],
2352 body: Bytes::from(message),
2353 },
2354 Err(panic) => {
2355 let panic_msg = panic
2356 .downcast_ref::<&'static str>()
2357 .map(|s| s.to_string())
2358 .or_else(|| panic.downcast_ref::<String>().cloned())
2359 .unwrap_or_else(|| "<non-string panic>".into());
2360 tracing::error!(
2361 caller_origin = format!("{:#x}", caller_origin),
2362 call_id,
2363 panic = %panic_msg,
2364 "rpc streaming server handler panicked",
2365 );
2366 RpcResponsePayload {
2367 status: RpcStatus::Internal,
2368 headers: vec![],
2369 body: Bytes::from(format!("handler panicked: {panic_msg}")),
2370 }
2371 }
2372 }
2373 };
2374 in_flight.lock().remove(&key);
2375 // Drop the per-call flow-control semaphore
2376 // (if any) so a stale GRANT arriving after
2377 // termination is silently dropped — the entry
2378 // is gone, lookup misses.
2379 flow_control.lock().remove(&key);
2380 // Await the terminal frame's publish too so it
2381 // arrives strictly AFTER the last chunk on the
2382 // wire (the pump has already drained, but the
2383 // emit itself is still async and we must await
2384 // it before the spawned task ends).
2385 emit(caller_origin, call_id, terminal).await;
2386 });
2387 }
2388 DISPATCH_RPC_CANCEL => {
2389 if let Some(token) = self.in_flight.lock().remove(&key) {
2390 token.cancel();
2391 }
2392 // Also drop the flow-control entry — the spawned
2393 // task's terminal cleanup will run too, but doing
2394 // it here makes the CANCEL path immediately stop
2395 // refilling the pump (the pending `acquire().await`
2396 // will resolve once the semaphore is dropped or
2397 // when the task exits).
2398 self.flow_control.lock().remove(&key);
2399 }
2400 DISPATCH_RPC_STREAM_GRANT => {
2401 // Add credit to the per-call semaphore. Silently
2402 // drop GRANT events for unknown / non-flow-
2403 // controlled calls — server can't tell whether
2404 // the caller is racing a terminal vs. sending a
2405 // grant for a non-flow-controlled stream, and
2406 // both are harmless to ignore.
2407 let amount = match decode_stream_grant(&ev.payload[EVENT_META_SIZE..]) {
2408 Some(n) => n,
2409 None => {
2410 tracing::debug!(
2411 caller_origin = format!("{:#x}", meta.origin_hash),
2412 call_id = meta.seq_or_ts,
2413 "rpc streaming server fold: malformed STREAM_GRANT payload",
2414 );
2415 return Ok(());
2416 }
2417 };
2418 if amount == 0 {
2419 return Ok(());
2420 }
2421 if let Some(sem) = self.flow_control.lock().get(&key).cloned() {
2422 // Tokio's `Semaphore::add_permits` is bounded
2423 // by `MAX_PERMITS = usize::MAX >> 3`. A
2424 // misbehaving caller flooding huge grants
2425 // would eventually saturate; cap defensively.
2426 let safe = (amount as usize).min(usize::MAX >> 4);
2427 sem.add_permits(safe);
2428 }
2429 }
2430 _ => {}
2431 }
2432 Ok(())
2433 }
2434}
2435
2436// ============================================================================
2437// Phase B — server-side fold for client-streaming.
2438//
2439// `RpcStreamingRequestFold` mirrors `RpcServerStreamingFold` but
2440// flipped on the data-direction axis: the SERVER consumes a
2441// stream of REQUEST_CHUNK events and the handler produces ONE
2442// terminal RESPONSE (vs. the response-side fold where one REQUEST
2443// drives many RESPONSE chunks).
2444//
2445// Wire shape it handles:
2446// DISPATCH_RPC_REQUEST (FLAG_RPC_CLIENT_STREAMING_REQUEST)
2447// DISPATCH_RPC_REQUEST_CHUNK (zero or more)
2448// DISPATCH_RPC_REQUEST_CHUNK (FLAG_RPC_REQUEST_END)
2449// DISPATCH_RPC_CANCEL (any time; flips token + closes stream)
2450//
2451// Wire shape it EMITS (via callbacks):
2452// DISPATCH_RPC_RESPONSE (one terminal frame; via RpcResponseEmitter)
2453// DISPATCH_RPC_REQUEST_GRANT (one per consumed chunk when flow
2454// control is opted in; via
2455// RpcRequestGrantEmitter)
2456//
2457// Each service binds to exactly one fold shape (unary, server-
2458// streaming, or client-streaming) at `serve_rpc*` registration.
2459// A REQUEST without FLAG_RPC_CLIENT_STREAMING_REQUEST that lands
2460// on the client-streaming fold is a caller bug — the fold emits a
2461// terminal `Internal` and drops the call.
2462// ============================================================================
2463
2464/// Per-call request-direction sender map type. Keyed on
2465/// `(caller_origin_hash, call_id)`; value is the bounded mpsc
2466/// sender the fold's `apply()` pushes REQUEST_CHUNK bodies into.
2467/// The matching receiver lives inside the handler's
2468/// [`RequestStream`]; dropping the sender (on REQUEST_END or
2469/// CANCEL) closes the stream.
2470type RequestChunkSenders = Arc<Mutex<HashMap<(u64, u64), tokio::sync::mpsc::Sender<bytes::Bytes>>>>;
2471
2472/// Shared REQUEST_CHUNK handling used by both
2473/// [`RpcStreamingRequestFold`] and [`RpcDuplexFold`]. Decodes the
2474/// payload, validates the call_id agreement, looks up the per-call
2475/// sender, pushes the body (skipping the empty-body FLAG_END
2476/// terminator), and removes the sender on FLAG_END so the
2477/// handler's stream observes EOF.
2478///
2479/// `diag_tag` selects the log prefix ("client-streaming" or
2480/// "duplex") so the two call sites surface identically-shaped
2481/// diagnostics with the correct fold name. The behavior is
2482/// otherwise identical — both folds carry the same wire format
2483/// and the same per-call mpsc + sender-map contract.
2484fn apply_request_chunk_to_senders(
2485 payload_bytes: Bytes,
2486 meta: &EventMeta,
2487 senders: &RequestChunkSenders,
2488 diag_tag: &'static str,
2489) {
2490 let payload = match RpcRequestChunkPayload::decode(payload_bytes) {
2491 Ok(p) => p,
2492 Err(e) => {
2493 tracing::warn!(
2494 error = %e,
2495 caller_origin = format!("{:#x}", meta.origin_hash),
2496 call_id = meta.seq_or_ts,
2497 tag = diag_tag,
2498 "rpc server fold: malformed REQUEST_CHUNK payload",
2499 );
2500 return;
2501 }
2502 };
2503 if payload.call_id != meta.seq_or_ts {
2504 tracing::warn!(
2505 caller_origin = format!("{:#x}", meta.origin_hash),
2506 meta_call_id = meta.seq_or_ts,
2507 payload_call_id = payload.call_id,
2508 tag = diag_tag,
2509 "rpc server fold: REQUEST_CHUNK payload call_id does not match EventMeta",
2510 );
2511 return;
2512 }
2513 let key = (meta.origin_hash, meta.seq_or_ts);
2514 let is_end = payload.flags & FLAG_RPC_REQUEST_END != 0;
2515 let sender = senders.lock().get(&key).cloned();
2516 let Some(sender) = sender else {
2517 // Unknown call — either the initial REQUEST hasn't
2518 // arrived yet (out-of-order delivery is possible on the
2519 // bus) or the handler already completed and the entry is
2520 // gone. Drop silently.
2521 tracing::debug!(
2522 caller_origin = format!("{:#x}", meta.origin_hash),
2523 call_id = meta.seq_or_ts,
2524 tag = diag_tag,
2525 "rpc server fold: REQUEST_CHUNK for unknown call_id; dropping",
2526 );
2527 return;
2528 };
2529 let is_pure_terminator = is_end && payload.body.is_empty();
2530 if !is_pure_terminator && sender.try_send(payload.body).is_err() {
2531 tracing::debug!(
2532 caller_origin = format!("{:#x}", meta.origin_hash),
2533 call_id = meta.seq_or_ts,
2534 tag = diag_tag,
2535 "rpc server fold: request-chunk mpsc full or closed; dropping",
2536 );
2537 }
2538 if is_end {
2539 // Drop the sender from the map → its clone here goes out
2540 // of scope at end of function → the receiver in the
2541 // handler's RequestStream sees EOF on the next poll.
2542 senders.lock().remove(&key);
2543 }
2544}
2545
2546/// Server-side fold for client-streaming RPC. Parallel to
2547/// [`RpcServerStreamingFold`] but consumes REQUEST_CHUNK on the
2548/// input side and produces one terminal RESPONSE on the output
2549/// side (vs. one REQUEST in / many RESPONSE chunks out).
2550///
2551/// State `()` — like the other folds, application state lives in
2552/// the handler's captured `Arc<Mutex<S>>`. The fold's own state
2553/// (in-flight cancellation tokens + per-call request-chunk
2554/// senders) lives on `&mut self` via `Arc<Mutex<...>>` so spawned
2555/// handler tasks can self-clean on completion.
2556///
2557/// Bidi streaming plan (Phase B).
2558pub struct RpcStreamingRequestFold {
2559 handler: Arc<dyn RpcClientStreamingHandler>,
2560 emit: RpcResponseEmitter,
2561 /// Optional request-direction grant emitter. `Some(...)`
2562 /// when the surrounding mesh glue is wired to publish
2563 /// REQUEST_GRANT events; `None` in unit tests / contexts
2564 /// without a real publish path. When `None`, the auto-grant
2565 /// path on every `RequestStream::poll_next` becomes a no-op
2566 /// (callers that opted into flow control will see no
2567 /// refill and stall once their initial window is exhausted —
2568 /// honest behavior for a fold not wired up for grants).
2569 grant_emit: Option<RpcRequestGrantEmitter>,
2570 in_flight: Arc<Mutex<HashMap<(u64, u64), RpcCancellationToken>>>,
2571 senders: RequestChunkSenders,
2572 /// Optional per-service metrics handle. Same shape as the
2573 /// other folds. Reuses the response-side counters where they
2574 /// apply (handler_invocations / handler_panics / etc.) and
2575 /// would gain request-side counters (e.g.
2576 /// `streaming_request_chunks_dropped_total`) in a follow-up.
2577 metrics: Option<Arc<crate::adapter::net::mesh_rpc_metrics::ServiceMetricsAtomic>>,
2578}
2579
2580impl RpcStreamingRequestFold {
2581 /// Construct a client-streaming server fold. `emit` publishes
2582 /// the terminal RESPONSE on the caller's reply channel.
2583 ///
2584 /// Use the sync [`RpcResponseEmitter`] here — there's only
2585 /// one RESPONSE per call (the terminal frame), so the
2586 /// per-call serialization the async emitter buys for the
2587 /// response-side fold is not needed here.
2588 pub fn new(handler: Arc<dyn RpcClientStreamingHandler>, emit: RpcResponseEmitter) -> Self {
2589 Self {
2590 handler,
2591 emit,
2592 grant_emit: None,
2593 in_flight: Arc::new(Mutex::new(HashMap::new())),
2594 senders: Arc::new(Mutex::new(HashMap::new())),
2595 metrics: None,
2596 }
2597 }
2598
2599 /// Attach the request-direction grant emitter. Hands every
2600 /// `RequestStream::poll_next` a hook to fire one REQUEST_GRANT
2601 /// back to the caller after a chunk is consumed. Optional —
2602 /// folds constructed without it still work, callers that
2603 /// opted into flow control just won't be refilled.
2604 pub fn with_grant_emitter(mut self, grant_emit: RpcRequestGrantEmitter) -> Self {
2605 self.grant_emit = Some(grant_emit);
2606 self
2607 }
2608
2609 /// Attach a per-service metrics handle. Hooks the spawned
2610 /// handler task to bump `handler_invocations_total` /
2611 /// `handler_in_flight` / `handler_panics_total` /
2612 /// `handler_duration_*`. Symmetric with `RpcServerFold` and
2613 /// `RpcServerStreamingFold`.
2614 pub fn with_metrics(
2615 mut self,
2616 metrics: Arc<crate::adapter::net::mesh_rpc_metrics::ServiceMetricsAtomic>,
2617 ) -> Self {
2618 self.metrics = Some(metrics);
2619 self
2620 }
2621
2622 /// Test-only: snapshot of the in-flight call set.
2623 #[cfg(test)]
2624 pub fn in_flight_keys(&self) -> Vec<(u64, u64)> {
2625 self.in_flight.lock().keys().copied().collect()
2626 }
2627
2628 /// Test-only: snapshot of the in-flight per-call senders.
2629 /// Useful for tests that need to assert a call's sender has
2630 /// been dropped after REQUEST_END / CANCEL.
2631 #[cfg(test)]
2632 pub fn sender_keys(&self) -> Vec<(u64, u64)> {
2633 self.senders.lock().keys().copied().collect()
2634 }
2635}
2636
2637impl RedexFold<()> for RpcStreamingRequestFold {
2638 fn apply(&mut self, ev: &RedexEvent, _state: &mut ()) -> Result<(), RedexError> {
2639 let Some(meta) = (if ev.payload.len() >= EVENT_META_SIZE {
2640 EventMeta::from_bytes(&ev.payload[..EVENT_META_SIZE])
2641 } else {
2642 None
2643 }) else {
2644 tracing::warn!(
2645 payload_len = ev.payload.len(),
2646 "rpc client-streaming server fold: event payload too short for EventMeta",
2647 );
2648 return Ok(());
2649 };
2650 let key = (meta.origin_hash, meta.seq_or_ts);
2651 match meta.dispatch {
2652 DISPATCH_RPC_REQUEST => {
2653 let payload = match RpcRequestPayload::decode(ev.payload.slice(EVENT_META_SIZE..)) {
2654 Ok(p) => p,
2655 Err(e) => {
2656 tracing::warn!(
2657 error = %e,
2658 caller_origin = format!("{:#x}", meta.origin_hash),
2659 call_id = meta.seq_or_ts,
2660 "rpc client-streaming server fold: malformed request payload",
2661 );
2662 let resp = RpcResponsePayload {
2663 status: RpcStatus::UnknownVersion,
2664 headers: vec![],
2665 body: Bytes::from(format!("malformed request: {e}")),
2666 };
2667 (self.emit)(meta.origin_hash, meta.seq_or_ts, resp);
2668 return Ok(());
2669 }
2670 };
2671 // A REQUEST without the client-streaming flag on
2672 // this fold is a caller bug — the service was
2673 // registered as client-streaming. Refuse cleanly.
2674 if payload.flags & FLAG_RPC_CLIENT_STREAMING_REQUEST == 0 {
2675 tracing::warn!(
2676 caller_origin = format!("{:#x}", meta.origin_hash),
2677 call_id = meta.seq_or_ts,
2678 flags = format!("{:#06x}", payload.flags),
2679 "rpc client-streaming server fold: REQUEST missing FLAG_RPC_CLIENT_STREAMING_REQUEST",
2680 );
2681 let resp = RpcResponsePayload {
2682 status: RpcStatus::Internal,
2683 headers: vec![],
2684 body: Bytes::from_static(
2685 b"REQUEST on a client-streaming service must set FLAG_RPC_CLIENT_STREAMING_REQUEST",
2686 ),
2687 };
2688 (self.emit)(meta.origin_hash, meta.seq_or_ts, resp);
2689 return Ok(());
2690 }
2691 // Refuse a duplicate REQUEST with the same
2692 // `(origin_hash, call_id)` — same rationale as
2693 // the response-side fold: a retry that arrives
2694 // while the first attempt is still in-flight
2695 // would overwrite the prior sender and orphan the
2696 // existing handler.
2697 {
2698 let in_flight = self.in_flight.lock();
2699 if in_flight.contains_key(&key) {
2700 drop(in_flight);
2701 tracing::warn!(
2702 caller_origin = format!("{:#x}", meta.origin_hash),
2703 call_id = meta.seq_or_ts,
2704 "rpc client-streaming server fold: duplicate REQUEST for in-flight call_id; refusing",
2705 );
2706 let resp = RpcResponsePayload {
2707 status: RpcStatus::Internal,
2708 headers: vec![],
2709 body: Bytes::from_static(
2710 b"duplicate REQUEST for already-in-flight call_id",
2711 ),
2712 };
2713 (self.emit)(meta.origin_hash, meta.seq_or_ts, resp);
2714 return Ok(());
2715 }
2716 }
2717 let cancellation = RpcCancellationToken::new();
2718 self.in_flight.lock().insert(key, cancellation.clone());
2719 // Build the per-call request-chunk mpsc. Bounded
2720 // capacity — overflow on the sender side drops the
2721 // chunk (caller can re-send or, if flow-control is
2722 // wired, will naturally not push past the credit
2723 // window).
2724 let (tx, rx) =
2725 tokio::sync::mpsc::channel::<bytes::Bytes>(STREAMING_REQUEST_PUMP_CAPACITY);
2726 // Terminator-semantics rule: an empty body
2727 // combined with FLAG_REQUEST_END is a pure
2728 // terminator — the caller's `finish()` emits it
2729 // to close the stream without yielding a phantom
2730 // empty item to the handler. A non-empty body on
2731 // a FLAG_END frame IS a final item (used by the
2732 // "single-item degenerate path": initial REQUEST
2733 // with FLAG_END + a real body sends one item +
2734 // closes in a single frame).
2735 let end_on_initial = payload.flags & FLAG_RPC_REQUEST_END != 0;
2736 let is_pure_terminator = end_on_initial && payload.body.is_empty();
2737 if !is_pure_terminator {
2738 // Fresh `mpsc::channel(STREAMING_REQUEST_PUMP_CAPACITY)`
2739 // with a live receiver — try_send cannot fail.
2740 // debug_assert surfaces the invariant break in
2741 // tests; release logs at error level rather than
2742 // silently swallowing the first request body.
2743 if tx.try_send(payload.body).is_err() {
2744 debug_assert!(
2745 false,
2746 "fresh client-streaming request mpsc rejected initial body"
2747 );
2748 tracing::error!(
2749 caller_origin = format!("{:#x}", meta.origin_hash),
2750 call_id = meta.seq_or_ts,
2751 "rpc client-streaming server fold: fresh mpsc rejected initial REQUEST body (invariant break)",
2752 );
2753 }
2754 }
2755 // If the initial REQUEST also set FLAG_REQUEST_END,
2756 // close the stream immediately — degenerate case of
2757 // "one-item upload" where the caller didn't bother
2758 // with a trailing REQUEST_CHUNK. Don't even insert
2759 // the sender into the map; just drop it here.
2760 if !end_on_initial {
2761 self.senders.lock().insert(key, tx);
2762 }
2763 // Build the handler's context + stream. Auto-grant
2764 // is opted into when the caller set the request
2765 // window header AND the fold was wired with a
2766 // grant emitter; both must be present for grants
2767 // to actually fly.
2768 let grant_emitter = if parse_request_window_initial(&payload.headers).is_some() {
2769 self.grant_emit.clone()
2770 } else {
2771 None
2772 };
2773 let request_stream =
2774 RequestStream::new(rx, grant_emitter, meta.origin_hash, meta.seq_or_ts);
2775 let trace_context = if payload.flags & FLAG_RPC_PROPAGATE_TRACE != 0 {
2776 extract_trace_context(&payload.headers)
2777 } else {
2778 None
2779 };
2780 let deadline_ns = payload.deadline_ns;
2781 let ctx = RpcStreamingContext {
2782 caller_origin: meta.origin_hash,
2783 call_id: meta.seq_or_ts,
2784 deadline_ns,
2785 headers: payload.headers,
2786 cancellation: cancellation.clone(),
2787 trace_context,
2788 };
2789 let handler = self.handler.clone();
2790 let emit = self.emit.clone();
2791 let in_flight = self.in_flight.clone();
2792 let senders = self.senders.clone();
2793 let caller_origin = meta.origin_hash;
2794 let call_id = meta.seq_or_ts;
2795 let cancel_probe = cancellation.clone();
2796 let cancel_for_deadline = cancellation.clone();
2797 let metrics = self.metrics.clone();
2798 tokio::spawn(async move {
2799 if let Some(m) = metrics.as_ref() {
2800 m.handler_invocations_total
2801 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2802 m.handler_in_flight
2803 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2804 }
2805 let handler_started = std::time::Instant::now();
2806 // Deadline guard: if the caller declared
2807 // `deadline_ns`, force-drop the handler future
2808 // after it elapses so an orphaned request stream
2809 // (caller-side network partition before
2810 // REQUEST_END arrives) can never hang the call
2811 // indefinitely. `deadline_ns = 0` means "no
2812 // deadline" — caller's responsibility.
2813 let call_fut = futures::FutureExt::catch_unwind(std::panic::AssertUnwindSafe(
2814 handler.call(ctx, request_stream),
2815 ));
2816 let outcome = if deadline_ns > 0 {
2817 let now_ns = std::time::SystemTime::now()
2818 .duration_since(std::time::UNIX_EPOCH)
2819 .map(|d| d.as_nanos() as u64)
2820 .unwrap_or(0);
2821 let remaining = deadline_ns.saturating_sub(now_ns);
2822 if remaining == 0 {
2823 cancel_for_deadline.cancel();
2824 Ok(Err(RpcHandlerError::Internal(
2825 "handler deadline_ns already expired at spawn".to_string(),
2826 )))
2827 } else {
2828 match tokio::time::timeout(
2829 std::time::Duration::from_nanos(remaining),
2830 call_fut,
2831 )
2832 .await
2833 {
2834 Ok(o) => o,
2835 Err(_) => {
2836 cancel_for_deadline.cancel();
2837 Ok(Err(RpcHandlerError::Internal(
2838 "handler deadline_ns exceeded".to_string(),
2839 )))
2840 }
2841 }
2842 }
2843 } else {
2844 call_fut.await
2845 };
2846 if let Some(m) = metrics.as_ref() {
2847 m.handler_in_flight
2848 .fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
2849 m.record_handler_duration(handler_started.elapsed());
2850 if outcome.is_err() {
2851 m.handler_panics_total
2852 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2853 }
2854 }
2855 // CANCEL-wins ordering: if the cancellation
2856 // token fired during execution, override the
2857 // handler's terminal with Cancelled.
2858 let terminal = if cancel_probe.is_cancelled() {
2859 RpcResponsePayload {
2860 status: RpcStatus::Cancelled,
2861 headers: vec![],
2862 body: Bytes::from_static(
2863 b"server observed CANCEL during client-streaming handler execution",
2864 ),
2865 }
2866 } else {
2867 match outcome {
2868 Ok(Ok(resp)) => resp,
2869 Ok(Err(RpcHandlerError::Application { code, message })) => {
2870 RpcResponsePayload {
2871 status: RpcStatus::Application(code),
2872 headers: vec![],
2873 body: Bytes::from(message),
2874 }
2875 }
2876 Ok(Err(RpcHandlerError::Internal(message))) => RpcResponsePayload {
2877 status: RpcStatus::Internal,
2878 headers: vec![],
2879 body: Bytes::from(message),
2880 },
2881 Err(panic) => {
2882 let panic_msg = panic
2883 .downcast_ref::<&'static str>()
2884 .map(|s| s.to_string())
2885 .or_else(|| panic.downcast_ref::<String>().cloned())
2886 .unwrap_or_else(|| "<non-string panic>".into());
2887 tracing::error!(
2888 caller_origin = format!("{:#x}", caller_origin),
2889 call_id,
2890 panic = %panic_msg,
2891 "rpc client-streaming server handler panicked",
2892 );
2893 RpcResponsePayload {
2894 status: RpcStatus::Internal,
2895 headers: vec![],
2896 body: Bytes::from(format!("handler panicked: {panic_msg}")),
2897 }
2898 }
2899 }
2900 };
2901 in_flight.lock().remove(&key);
2902 // Drop the per-call request-chunk sender too
2903 // (idempotent — already gone if REQUEST_END
2904 // arrived; defensive otherwise so a handler
2905 // that returned without consuming all chunks
2906 // doesn't leak the entry).
2907 senders.lock().remove(&key);
2908 (emit)(caller_origin, call_id, terminal);
2909 });
2910 }
2911 DISPATCH_RPC_REQUEST_CHUNK => {
2912 apply_request_chunk_to_senders(
2913 ev.payload.slice(EVENT_META_SIZE..),
2914 &meta,
2915 &self.senders,
2916 "client-streaming",
2917 );
2918 }
2919 DISPATCH_RPC_CANCEL => {
2920 if let Some(token) = self.in_flight.lock().remove(&key) {
2921 token.cancel();
2922 }
2923 // Drop the per-call sender so the handler's
2924 // RequestStream yields None on the next poll
2925 // (handler observes cancel via the token OR via
2926 // the stream's EOF; the cancel_probe in the
2927 // spawned task ensures the terminal RESPONSE is
2928 // Cancelled regardless of which the handler
2929 // checks first).
2930 self.senders.lock().remove(&key);
2931 }
2932 _ => {}
2933 }
2934 Ok(())
2935 }
2936}
2937
2938// ============================================================================
2939// Phase D — server-side fold for full duplex.
2940//
2941// `RpcDuplexFold` is the hybrid of `RpcStreamingRequestFold`
2942// (Phase B — request side) and `RpcServerStreamingFold` (existing
2943// — response side). The handler trait takes BOTH a `RequestStream`
2944// AND an `RpcResponseSink`; the fold spawns one handler task per
2945// REQUEST and one pump task per call_id, then emits a terminal
2946// RESPONSE on handler return.
2947//
2948// Wire shape it consumes:
2949// DISPATCH_RPC_REQUEST (FLAG_CLIENT_STREAMING_REQUEST + FLAG_STREAMING_RESPONSE)
2950// DISPATCH_RPC_REQUEST_CHUNK (zero or more, with FLAG_REQUEST_END on the last)
2951// DISPATCH_RPC_CANCEL (flips token + closes both directions)
2952//
2953// Wire shape it produces:
2954// DISPATCH_RPC_RESPONSE (multi-fire; nrpc-streaming: continue / end)
2955// DISPATCH_RPC_REQUEST_GRANT (one per consumed request-chunk when flow
2956// control is opted in)
2957//
2958// Bidi streaming plan (Phase D).
2959// ============================================================================
2960
2961/// Server-side fold for duplex RPC. Composes Phase B's request
2962/// stream + per-call request-chunk senders with the existing
2963/// response-side pump + multi-fire RESPONSE emit.
2964///
2965/// State `()` — same as the sibling folds.
2966///
2967/// Bidi streaming plan (Phase D).
2968pub struct RpcDuplexFold {
2969 handler: Arc<dyn RpcDuplexHandler>,
2970 /// Async emitter for response chunks (per-call ordering via
2971 /// awaited emits — same rationale as `RpcServerStreamingFold`).
2972 emit: RpcAsyncResponseEmitter,
2973 /// Optional request-direction grant emitter.
2974 grant_emit: Option<RpcRequestGrantEmitter>,
2975 in_flight: Arc<Mutex<HashMap<(u64, u64), RpcCancellationToken>>>,
2976 senders: RequestChunkSenders,
2977 metrics: Option<Arc<crate::adapter::net::mesh_rpc_metrics::ServiceMetricsAtomic>>,
2978}
2979
2980impl RpcDuplexFold {
2981 /// Construct a duplex server fold. `emit` publishes individual
2982 /// response chunks AND the terminal frame on the caller's
2983 /// reply channel (uses the async emitter for per-call
2984 /// ordering, same as `RpcServerStreamingFold`).
2985 pub fn new(handler: Arc<dyn RpcDuplexHandler>, emit: RpcAsyncResponseEmitter) -> Self {
2986 Self {
2987 handler,
2988 emit,
2989 grant_emit: None,
2990 in_flight: Arc::new(Mutex::new(HashMap::new())),
2991 senders: Arc::new(Mutex::new(HashMap::new())),
2992 metrics: None,
2993 }
2994 }
2995
2996 /// Attach the request-direction grant emitter. See
2997 /// [`RpcStreamingRequestFold::with_grant_emitter`] for the
2998 /// auto-grant behavior. When unset, callers that opted into
2999 /// flow control simply won't get refilled.
3000 pub fn with_grant_emitter(mut self, grant_emit: RpcRequestGrantEmitter) -> Self {
3001 self.grant_emit = Some(grant_emit);
3002 self
3003 }
3004
3005 /// Attach a per-service metrics handle. Bumps
3006 /// handler_invocations / handler_in_flight / handler_panics /
3007 /// handler_duration_* + the response pump's
3008 /// streaming_chunks_emitted_total per emitted chunk.
3009 pub fn with_metrics(
3010 mut self,
3011 metrics: Arc<crate::adapter::net::mesh_rpc_metrics::ServiceMetricsAtomic>,
3012 ) -> Self {
3013 self.metrics = Some(metrics);
3014 self
3015 }
3016
3017 /// Test-only: snapshot of the in-flight call set.
3018 #[cfg(test)]
3019 pub fn in_flight_keys(&self) -> Vec<(u64, u64)> {
3020 self.in_flight.lock().keys().copied().collect()
3021 }
3022
3023 /// Test-only: snapshot of the in-flight per-call senders.
3024 #[cfg(test)]
3025 pub fn sender_keys(&self) -> Vec<(u64, u64)> {
3026 self.senders.lock().keys().copied().collect()
3027 }
3028}
3029
3030impl RedexFold<()> for RpcDuplexFold {
3031 fn apply(&mut self, ev: &RedexEvent, _state: &mut ()) -> Result<(), RedexError> {
3032 let Some(meta) = (if ev.payload.len() >= EVENT_META_SIZE {
3033 EventMeta::from_bytes(&ev.payload[..EVENT_META_SIZE])
3034 } else {
3035 None
3036 }) else {
3037 tracing::warn!(
3038 payload_len = ev.payload.len(),
3039 "rpc duplex server fold: event payload too short for EventMeta",
3040 );
3041 return Ok(());
3042 };
3043 let key = (meta.origin_hash, meta.seq_or_ts);
3044 match meta.dispatch {
3045 DISPATCH_RPC_REQUEST => {
3046 let payload = match RpcRequestPayload::decode(ev.payload.slice(EVENT_META_SIZE..)) {
3047 Ok(p) => p,
3048 Err(e) => {
3049 tracing::warn!(
3050 error = %e,
3051 caller_origin = format!("{:#x}", meta.origin_hash),
3052 call_id = meta.seq_or_ts,
3053 "rpc duplex server fold: malformed request payload",
3054 );
3055 let resp = RpcResponsePayload {
3056 status: RpcStatus::UnknownVersion,
3057 headers: vec![(
3058 HEADER_NRPC_STREAMING.to_string(),
3059 HEADER_NRPC_STREAMING_END.to_vec(),
3060 )],
3061 body: Bytes::from(format!("malformed request: {e}")),
3062 };
3063 let emit = self.emit.clone();
3064 let caller_origin = meta.origin_hash;
3065 let call_id = meta.seq_or_ts;
3066 tokio::spawn(async move {
3067 emit(caller_origin, call_id, resp).await;
3068 });
3069 return Ok(());
3070 }
3071 };
3072 // Caller-bug guard: a duplex REQUEST must set
3073 // BOTH the client-streaming flag (we'll receive
3074 // request chunks) AND the streaming-response flag
3075 // (we'll emit response chunks). Missing flags →
3076 // refuse cleanly.
3077 let required = FLAG_RPC_CLIENT_STREAMING_REQUEST | FLAG_RPC_STREAMING_RESPONSE;
3078 if payload.flags & required != required {
3079 tracing::warn!(
3080 caller_origin = format!("{:#x}", meta.origin_hash),
3081 call_id = meta.seq_or_ts,
3082 flags = format!("{:#06x}", payload.flags),
3083 "rpc duplex server fold: REQUEST missing required flags",
3084 );
3085 let resp = RpcResponsePayload {
3086 status: RpcStatus::Internal,
3087 headers: vec![(
3088 HEADER_NRPC_STREAMING.to_string(),
3089 HEADER_NRPC_STREAMING_END.to_vec(),
3090 )],
3091 body: Bytes::from_static(
3092 b"REQUEST on a duplex service must set FLAG_RPC_CLIENT_STREAMING_REQUEST and FLAG_RPC_STREAMING_RESPONSE",
3093 ),
3094 };
3095 let emit = self.emit.clone();
3096 let caller_origin = meta.origin_hash;
3097 let call_id = meta.seq_or_ts;
3098 tokio::spawn(async move {
3099 emit(caller_origin, call_id, resp).await;
3100 });
3101 return Ok(());
3102 }
3103 // Duplicate-REQUEST refusal.
3104 {
3105 let in_flight = self.in_flight.lock();
3106 if in_flight.contains_key(&key) {
3107 drop(in_flight);
3108 tracing::warn!(
3109 caller_origin = format!("{:#x}", meta.origin_hash),
3110 call_id = meta.seq_or_ts,
3111 "rpc duplex server fold: duplicate REQUEST for in-flight call_id; refusing",
3112 );
3113 let resp = RpcResponsePayload {
3114 status: RpcStatus::Internal,
3115 headers: vec![(
3116 HEADER_NRPC_STREAMING.to_string(),
3117 HEADER_NRPC_STREAMING_END.to_vec(),
3118 )],
3119 body: Bytes::from_static(
3120 b"duplicate REQUEST for already-in-flight call_id",
3121 ),
3122 };
3123 let emit = self.emit.clone();
3124 let caller_origin = meta.origin_hash;
3125 let call_id = meta.seq_or_ts;
3126 tokio::spawn(async move {
3127 emit(caller_origin, call_id, resp).await;
3128 });
3129 return Ok(());
3130 }
3131 }
3132 let cancellation = RpcCancellationToken::new();
3133 self.in_flight.lock().insert(key, cancellation.clone());
3134
3135 // Build per-call request-side mpsc (Phase B
3136 // pattern).
3137 let (req_tx, req_rx) =
3138 tokio::sync::mpsc::channel::<bytes::Bytes>(STREAMING_REQUEST_PUMP_CAPACITY);
3139 let end_on_initial = payload.flags & FLAG_RPC_REQUEST_END != 0;
3140 let is_pure_terminator = end_on_initial && payload.body.is_empty();
3141 if !is_pure_terminator {
3142 // Same invariant as the client-streaming fold:
3143 // fresh bounded mpsc with a live receiver cannot
3144 // reject the first send.
3145 if req_tx.try_send(payload.body).is_err() {
3146 debug_assert!(false, "fresh duplex request mpsc rejected initial body");
3147 tracing::error!(
3148 caller_origin = format!("{:#x}", meta.origin_hash),
3149 call_id = meta.seq_or_ts,
3150 "rpc duplex server fold: fresh mpsc rejected initial REQUEST body (invariant break)",
3151 );
3152 }
3153 }
3154 if !end_on_initial {
3155 self.senders.lock().insert(key, req_tx);
3156 }
3157 // Hand the handler an auto-granting RequestStream
3158 // when the caller opted into request-direction
3159 // flow control AND the fold was wired with a
3160 // grant emitter.
3161 let grant_emitter = if parse_request_window_initial(&payload.headers).is_some() {
3162 self.grant_emit.clone()
3163 } else {
3164 None
3165 };
3166 let request_stream =
3167 RequestStream::new(req_rx, grant_emitter, meta.origin_hash, meta.seq_or_ts);
3168
3169 // Build the per-call response-side mpsc (existing
3170 // server-streaming-response pattern). The handler
3171 // writes chunks to the sink; the pump task drains
3172 // the receiver and publishes RESPONSE events.
3173 let (resp_tx, mut resp_rx) =
3174 tokio::sync::mpsc::channel::<bytes::Bytes>(STREAMING_PUMP_CAPACITY);
3175 let response_sink = RpcResponseSink {
3176 inner: resp_tx,
3177 metrics: self.metrics.clone(),
3178 };
3179
3180 let trace_context = if payload.flags & FLAG_RPC_PROPAGATE_TRACE != 0 {
3181 extract_trace_context(&payload.headers)
3182 } else {
3183 None
3184 };
3185 let deadline_ns = payload.deadline_ns;
3186 let ctx = RpcStreamingContext {
3187 caller_origin: meta.origin_hash,
3188 call_id: meta.seq_or_ts,
3189 deadline_ns,
3190 headers: payload.headers,
3191 cancellation: cancellation.clone(),
3192 trace_context,
3193 };
3194 let handler = self.handler.clone();
3195 let emit = self.emit.clone();
3196 let in_flight = self.in_flight.clone();
3197 let senders = self.senders.clone();
3198 let caller_origin = meta.origin_hash;
3199 let call_id = meta.seq_or_ts;
3200 let cancel_probe = cancellation.clone();
3201 let cancel_for_deadline = cancellation.clone();
3202 let metrics = self.metrics.clone();
3203
3204 // Pump: drains resp_rx, emits per-chunk RESPONSE
3205 // events with `nrpc-streaming: continue`.
3206 let pump_emit = emit.clone();
3207 let pump_metrics = metrics.clone();
3208 let pump = tokio::spawn(async move {
3209 while let Some(chunk) = resp_rx.recv().await {
3210 if let Some(m) = pump_metrics.as_ref() {
3211 m.streaming_chunks_emitted_total
3212 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3213 }
3214 let resp = RpcResponsePayload {
3215 status: RpcStatus::Ok,
3216 headers: vec![(
3217 HEADER_NRPC_STREAMING.to_string(),
3218 HEADER_NRPC_STREAMING_CONTINUE.to_vec(),
3219 )],
3220 body: chunk.clone(),
3221 };
3222 pump_emit(caller_origin, call_id, resp).await;
3223 }
3224 });
3225
3226 tokio::spawn(async move {
3227 if let Some(m) = metrics.as_ref() {
3228 m.handler_invocations_total
3229 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3230 m.handler_in_flight
3231 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3232 }
3233 let handler_started = std::time::Instant::now();
3234 // Same deadline guard as the client-streaming
3235 // fold: force-drop the handler future at
3236 // deadline_ns so an orphaned request stream
3237 // can't hang the call. `0` means no deadline.
3238 let call_fut = futures::FutureExt::catch_unwind(std::panic::AssertUnwindSafe(
3239 handler.call(ctx, request_stream, response_sink),
3240 ));
3241 let outcome = if deadline_ns > 0 {
3242 let now_ns = std::time::SystemTime::now()
3243 .duration_since(std::time::UNIX_EPOCH)
3244 .map(|d| d.as_nanos() as u64)
3245 .unwrap_or(0);
3246 let remaining = deadline_ns.saturating_sub(now_ns);
3247 if remaining == 0 {
3248 cancel_for_deadline.cancel();
3249 Ok(Err(RpcHandlerError::Internal(
3250 "duplex handler deadline_ns already expired at spawn".to_string(),
3251 )))
3252 } else {
3253 match tokio::time::timeout(
3254 std::time::Duration::from_nanos(remaining),
3255 call_fut,
3256 )
3257 .await
3258 {
3259 Ok(o) => o,
3260 Err(_) => {
3261 cancel_for_deadline.cancel();
3262 Ok(Err(RpcHandlerError::Internal(
3263 "duplex handler deadline_ns exceeded".to_string(),
3264 )))
3265 }
3266 }
3267 }
3268 } else {
3269 call_fut.await
3270 };
3271 // Handler dropped the sink — let the pump
3272 // drain any final in-flight chunks before we
3273 // emit the terminal frame.
3274 let _ = pump.await;
3275 if let Some(m) = metrics.as_ref() {
3276 m.handler_in_flight
3277 .fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
3278 m.record_handler_duration(handler_started.elapsed());
3279 if outcome.is_err() {
3280 m.handler_panics_total
3281 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3282 }
3283 }
3284 let terminal = if cancel_probe.is_cancelled() {
3285 RpcResponsePayload {
3286 status: RpcStatus::Cancelled,
3287 headers: vec![],
3288 body: Bytes::from_static(
3289 b"server observed CANCEL during duplex handler execution",
3290 ),
3291 }
3292 } else {
3293 match outcome {
3294 Ok(Ok(())) => RpcResponsePayload {
3295 status: RpcStatus::Ok,
3296 headers: vec![(
3297 HEADER_NRPC_STREAMING.to_string(),
3298 HEADER_NRPC_STREAMING_END.to_vec(),
3299 )],
3300 body: Bytes::new(),
3301 },
3302 Ok(Err(RpcHandlerError::Application { code, message })) => {
3303 RpcResponsePayload {
3304 status: RpcStatus::Application(code),
3305 headers: vec![],
3306 body: Bytes::from(message),
3307 }
3308 }
3309 Ok(Err(RpcHandlerError::Internal(message))) => RpcResponsePayload {
3310 status: RpcStatus::Internal,
3311 headers: vec![],
3312 body: Bytes::from(message),
3313 },
3314 Err(panic) => {
3315 let panic_msg = panic
3316 .downcast_ref::<&'static str>()
3317 .map(|s| s.to_string())
3318 .or_else(|| panic.downcast_ref::<String>().cloned())
3319 .unwrap_or_else(|| "<non-string panic>".into());
3320 tracing::error!(
3321 caller_origin = format!("{:#x}", caller_origin),
3322 call_id,
3323 panic = %panic_msg,
3324 "rpc duplex server handler panicked",
3325 );
3326 RpcResponsePayload {
3327 status: RpcStatus::Internal,
3328 headers: vec![],
3329 body: Bytes::from(format!("handler panicked: {panic_msg}")),
3330 }
3331 }
3332 }
3333 };
3334 in_flight.lock().remove(&key);
3335 senders.lock().remove(&key);
3336 emit(caller_origin, call_id, terminal).await;
3337 });
3338 }
3339 DISPATCH_RPC_REQUEST_CHUNK => {
3340 apply_request_chunk_to_senders(
3341 ev.payload.slice(EVENT_META_SIZE..),
3342 &meta,
3343 &self.senders,
3344 "duplex",
3345 );
3346 }
3347 DISPATCH_RPC_CANCEL => {
3348 if let Some(token) = self.in_flight.lock().remove(&key) {
3349 token.cancel();
3350 }
3351 self.senders.lock().remove(&key);
3352 }
3353 _ => {}
3354 }
3355 Ok(())
3356 }
3357}
3358
3359// ============================================================================
3360// Client-side fold.
3361//
3362// `RpcClientFold` is the symmetric companion of `RpcServerFold`.
3363// It sees RESPONSE events on the caller's reply channel
3364// (`<service>.replies.<self_origin>`) and routes each one to the
3365// matching call's awaiting `oneshot::Receiver` keyed on `call_id`
3366// (the `EventMeta::seq_or_ts`).
3367//
3368// The fold's mutable state (the pending-senders map) is shared
3369// with the `Mesh::call` API via a clone of the same Arc — so the
3370// publisher side can `register(call_id)` to stage a receiver
3371// before publishing the REQUEST, and the fold side can `deliver`
3372// when the matching RESPONSE arrives.
3373// ============================================================================
3374
3375/// One pending entry — unary oneshot, server-streaming mpsc, or
3376/// client-streaming (one terminal oneshot + a separate grant
3377/// mpsc). The fold dispatches to the right variant based on
3378/// what's registered for the `call_id`.
3379enum PendingEntry {
3380 /// Unary call — exactly one RESPONSE expected. Completes the
3381 /// oneshot with the decoded payload.
3382 Unary(tokio::sync::oneshot::Sender<RpcResponsePayload>),
3383 /// Server-streaming call — multiple non-terminal `Continue`
3384 /// chunks followed by one terminal frame. Each non-terminal
3385 /// chunk pushes a `StreamItem::Chunk(body)` onto the mpsc;
3386 /// the terminal frame pushes `StreamItem::End` (Ok) or
3387 /// `StreamItem::Error(payload)` (non-Ok status) and the
3388 /// pending entry is removed.
3389 Streaming(tokio::sync::mpsc::UnboundedSender<StreamItem>),
3390 /// Client-streaming or duplex call. Two sender halves:
3391 ///
3392 /// - `terminal_tx`: oneshot that completes when the server's
3393 /// single terminal RESPONSE arrives. Response shape and
3394 /// delivery semantics are identical to the unary variant —
3395 /// the caller awaits one payload, success or failure status.
3396 /// - `grant_tx`: mpsc that ferries REQUEST_GRANT credit values
3397 /// from the client fold to the caller's send sink. Each
3398 /// `DISPATCH_RPC_REQUEST_GRANT` event for this call_id
3399 /// pushes one `u32` credit onto the mpsc; the caller's send
3400 /// sink consumes credits to gate `send(...).await`.
3401 ///
3402 /// Bidi streaming plan (Phase C). Used for pure client-
3403 /// streaming (one terminal RESPONSE closes the call). Duplex
3404 /// calls use the [`PendingEntry::Duplex`] variant instead,
3405 /// since they receive many response chunks rather than one
3406 /// terminal payload.
3407 ClientStreaming {
3408 terminal_tx: tokio::sync::oneshot::Sender<RpcResponsePayload>,
3409 grant_tx: tokio::sync::mpsc::UnboundedSender<u32>,
3410 },
3411 /// Duplex call — many request chunks out, many response
3412 /// chunks in. Two senders, same shape as `ClientStreaming`
3413 /// except the terminal slot is an mpsc instead of a oneshot
3414 /// because the response side is multi-chunk (terminator is
3415 /// implicit in `StreamItem::End` / `StreamItem::Error` on
3416 /// the chunks_tx mpsc, same as `PendingEntry::Streaming`).
3417 ///
3418 /// - `chunks_tx`: response-chunk mpsc — fed by `deliver`
3419 /// when RESPONSE events arrive on the reply channel.
3420 /// `StreamItem::Chunk` for non-terminal, `StreamItem::End`
3421 /// / `StreamItem::Error` terminates and removes the entry.
3422 /// - `grant_tx`: request-direction credit mpsc — fed by
3423 /// `deliver_grant` when REQUEST_GRANT events arrive.
3424 ///
3425 /// Bidi streaming plan (Phase D).
3426 Duplex {
3427 chunks_tx: tokio::sync::mpsc::UnboundedSender<StreamItem>,
3428 grant_tx: tokio::sync::mpsc::UnboundedSender<u32>,
3429 },
3430}
3431
3432/// One item delivered to a streaming caller. The caller's
3433/// `RpcStream` translates these into `Stream::Item =
3434/// Result<Bytes, RpcError>` plus stream termination.
3435#[derive(Debug, Clone)]
3436pub enum StreamItem {
3437 /// Non-terminal chunk — a body slice from the server.
3438 Chunk(bytes::Bytes),
3439 /// Terminal frame, server signaled clean stream end.
3440 End,
3441 /// Terminal frame with a non-`Ok` status. Body is the
3442 /// server's diagnostic; status is the wire `RpcStatus` value.
3443 Error(RpcResponsePayload),
3444}
3445
3446/// Shared pending-call state. Held by both the `RpcClientFold`
3447/// (writer side: completes oneshot senders / pushes streaming
3448/// chunks on RESPONSE arrival) and the `Mesh::call*` APIs (reader
3449/// side: registers entries before publishing the REQUEST).
3450/// Concurrent access is mediated by `DashMap`.
3451///
3452/// Multiplexes unary AND streaming calls in a single map keyed
3453/// on `call_id` — the entry's enum variant tells the fold how
3454/// to dispatch incoming RESPONSE events.
3455pub struct RpcClientPending {
3456 /// Map keyed on `call_id`, value carries `(expected_target,
3457 /// PendingEntry)`. `expected_target` is the `NodeId` of the
3458 /// peer the request was dispatched to; `deliver` rejects
3459 /// frames whose wire `from_node` doesn't match. A
3460 /// `expected_target == 0` entry opts out of the binding
3461 /// (loopback tests + paths with no session).
3462 senders: dashmap::DashMap<u64, (super::super::behavior::placement::NodeId, PendingEntry)>,
3463}
3464
3465impl RpcClientPending {
3466 /// Construct an empty pending-call store.
3467 pub fn new() -> Self {
3468 Self {
3469 senders: dashmap::DashMap::new(),
3470 }
3471 }
3472
3473 /// Register a oneshot for a unary `call_id`. Returns the
3474 /// receiver the caller awaits. The caller MUST publish the
3475 /// REQUEST after registration (and not before) so the
3476 /// matching RESPONSE can't arrive while the pending entry is
3477 /// missing.
3478 ///
3479 /// `target_node` is the wire-session peer the request will
3480 /// be sent to; `deliver` rejects RESPONSE frames whose
3481 /// `from_node` doesn't match. Pass `0` for loopback / no-
3482 /// session test paths to opt out of the binding gate.
3483 ///
3484 /// If a sender already exists for `call_id` (improperly reused
3485 /// id), it is replaced and the old receiver gets a
3486 /// `RecvError::Closed` — surfacing the misuse as a hard error
3487 /// at the caller rather than silently delivering the response
3488 /// to the wrong waiter.
3489 pub fn register(
3490 &self,
3491 call_id: u64,
3492 target_node: super::super::behavior::placement::NodeId,
3493 ) -> tokio::sync::oneshot::Receiver<RpcResponsePayload> {
3494 let (tx, rx) = tokio::sync::oneshot::channel();
3495 self.senders
3496 .insert(call_id, (target_node, PendingEntry::Unary(tx)));
3497 rx
3498 }
3499
3500 /// Register a streaming entry for `call_id`. Returns the
3501 /// receive end of an mpsc the fold will push chunks onto.
3502 /// Same registration ordering rules as `register` —
3503 /// publisher must call this BEFORE publishing the REQUEST.
3504 pub fn register_streaming(
3505 &self,
3506 call_id: u64,
3507 target_node: super::super::behavior::placement::NodeId,
3508 ) -> tokio::sync::mpsc::UnboundedReceiver<StreamItem> {
3509 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
3510 self.senders
3511 .insert(call_id, (target_node, PendingEntry::Streaming(tx)));
3512 rx
3513 }
3514
3515 /// Register a client-streaming (or duplex) entry for
3516 /// `call_id`. Returns BOTH the terminal-response receiver
3517 /// (the caller awaits on this for the single terminal
3518 /// RESPONSE that ends the call) AND a grant receiver (the
3519 /// caller's send sink consumes this to gate `send().await`
3520 /// when the caller opted into request-direction flow
3521 /// control).
3522 ///
3523 /// Same registration ordering rules as `register` /
3524 /// `register_streaming` — publisher must call this BEFORE
3525 /// publishing the REQUEST so a fast server's RESPONSE /
3526 /// REQUEST_GRANT can't arrive while no pending entry exists.
3527 ///
3528 /// Bidi streaming plan (Phase C).
3529 pub fn register_client_streaming(
3530 &self,
3531 call_id: u64,
3532 target_node: super::super::behavior::placement::NodeId,
3533 ) -> (
3534 tokio::sync::oneshot::Receiver<RpcResponsePayload>,
3535 tokio::sync::mpsc::UnboundedReceiver<u32>,
3536 ) {
3537 let (terminal_tx, terminal_rx) = tokio::sync::oneshot::channel();
3538 let (grant_tx, grant_rx) = tokio::sync::mpsc::unbounded_channel();
3539 self.senders.insert(
3540 call_id,
3541 (
3542 target_node,
3543 PendingEntry::ClientStreaming {
3544 terminal_tx,
3545 grant_tx,
3546 },
3547 ),
3548 );
3549 (terminal_rx, grant_rx)
3550 }
3551
3552 /// Register a duplex entry for `call_id`. Returns BOTH a
3553 /// response-chunk receiver (yields `StreamItem` per inbound
3554 /// RESPONSE chunk; terminator is `End` / `Error`) AND a
3555 /// grant receiver (yields `u32` credits per inbound
3556 /// REQUEST_GRANT).
3557 ///
3558 /// Same registration ordering rules as the other `register_*`
3559 /// methods: publisher must call this BEFORE publishing the
3560 /// REQUEST so the server's response chunks / grants can't
3561 /// arrive while no pending entry exists.
3562 ///
3563 /// Bidi streaming plan (Phase D).
3564 pub fn register_duplex(
3565 &self,
3566 call_id: u64,
3567 target_node: super::super::behavior::placement::NodeId,
3568 ) -> (
3569 tokio::sync::mpsc::UnboundedReceiver<StreamItem>,
3570 tokio::sync::mpsc::UnboundedReceiver<u32>,
3571 ) {
3572 let (chunks_tx, chunks_rx) = tokio::sync::mpsc::unbounded_channel();
3573 let (grant_tx, grant_rx) = tokio::sync::mpsc::unbounded_channel();
3574 self.senders.insert(
3575 call_id,
3576 (
3577 target_node,
3578 PendingEntry::Duplex {
3579 chunks_tx,
3580 grant_tx,
3581 },
3582 ),
3583 );
3584 (chunks_rx, grant_rx)
3585 }
3586
3587 /// Drop the pending entry for `call_id`. Called by the
3588 /// caller-side cancellation path (e.g. `Mesh::call`'s future
3589 /// being dropped, the stream being dropped, or a deadline
3590 /// timer firing). The matching RESPONSE(s) that may still
3591 /// arrive afterwards are silently discarded by `deliver`.
3592 pub fn cancel(&self, call_id: u64) {
3593 self.senders.remove(&call_id);
3594 }
3595
3596 /// Deliver `resp` to the waiter for `call_id`, if any.
3597 ///
3598 /// `from_node` is the wire-session peer of the inbound
3599 /// RESPONSE. If the pending entry's recorded `target_node`
3600 /// is non-zero and does not match `from_node`, the frame is
3601 /// dropped with a trace log and the pending entry stays
3602 /// intact — a forged response on a shared reply channel
3603 /// can't resolve a victim's call. A recorded `target_node
3604 /// == 0` opts the call out of the binding (loopback paths).
3605 ///
3606 /// For a unary entry: completes the oneshot and removes the
3607 /// entry.
3608 ///
3609 /// For a streaming entry: examines the response's headers to
3610 /// decide whether it's a non-terminal chunk (`Continue` —
3611 /// push `StreamItem::Chunk`, keep the entry) or terminal
3612 /// (`End` / non-`Ok` — push `StreamItem::End` or `Error`,
3613 /// remove the entry).
3614 ///
3615 /// Idempotent on subsequent deliveries to a removed entry.
3616 fn deliver(
3617 &self,
3618 call_id: u64,
3619 from_node: super::super::behavior::placement::NodeId,
3620 resp: RpcResponsePayload,
3621 ) {
3622 // Look up the entry — but DON'T remove it yet, because for
3623 // streaming we may want to keep it for non-terminal chunks.
3624 // The remove decision is per-variant.
3625 let entry = self.senders.get(&call_id);
3626 let Some(entry) = entry else { return };
3627 // S-4 part 2 gate. The pending registry binds each call
3628 // to the AEAD-verified `target_node` the request was
3629 // dispatched to; any other session peer publishing on the
3630 // shared reply channel with a guessed call_id is dropped
3631 // here without touching the waiter. `0` opts out — used
3632 // by loopback paths that have no session peer.
3633 let (target_node, _entry_value) = entry.value();
3634 if *target_node != 0 && *target_node != from_node {
3635 tracing::trace!(
3636 call_id,
3637 from_node,
3638 expected = *target_node,
3639 "rpc client: dropping RESPONSE from non-target session peer"
3640 );
3641 return;
3642 }
3643 match entry.value() {
3644 (_, PendingEntry::Unary(_)) => {
3645 drop(entry);
3646 if let Some((_, (_, PendingEntry::Unary(tx)))) = self.senders.remove(&call_id) {
3647 let _ = tx.send(resp);
3648 }
3649 }
3650 (_, PendingEntry::ClientStreaming { .. }) => {
3651 // Terminal RESPONSE for a client-streaming /
3652 // duplex call. Same delivery shape as Unary —
3653 // complete the oneshot, remove the entry. The
3654 // grant_tx half drops with the entry, which is
3655 // fine (no more grants will arrive after the
3656 // terminal frame).
3657 drop(entry);
3658 if let Some((
3659 _,
3660 (
3661 _,
3662 PendingEntry::ClientStreaming {
3663 terminal_tx,
3664 grant_tx: _,
3665 },
3666 ),
3667 )) = self.senders.remove(&call_id)
3668 {
3669 let _ = terminal_tx.send(resp);
3670 }
3671 }
3672 (_, PendingEntry::Streaming(tx)) => {
3673 let tx = tx.clone();
3674 drop(entry);
3675 self.dispatch_streaming_chunk(&tx, resp, call_id);
3676 }
3677 (_, PendingEntry::Duplex { chunks_tx, .. }) => {
3678 // Same dispatch logic as Streaming — duplex
3679 // response side IS a multi-chunk stream.
3680 let tx = chunks_tx.clone();
3681 drop(entry);
3682 self.dispatch_streaming_chunk(&tx, resp, call_id);
3683 }
3684 }
3685 }
3686
3687 /// Shared response-chunk dispatch used by both
3688 /// `PendingEntry::Streaming` and `PendingEntry::Duplex`. The
3689 /// caller has already verified the target-binding gate and
3690 /// dropped its `entry` ref; this helper does the classify-
3691 /// and-push and removes the entry from the senders map on
3692 /// terminal frames.
3693 fn dispatch_streaming_chunk(
3694 &self,
3695 tx: &tokio::sync::mpsc::UnboundedSender<StreamItem>,
3696 resp: RpcResponsePayload,
3697 call_id: u64,
3698 ) {
3699 let kind = classify_streaming_chunk(&resp);
3700 match kind {
3701 StreamingChunkKind::Continue => {
3702 let _ = tx.send(StreamItem::Chunk(resp.body));
3703 }
3704 StreamingChunkKind::Terminal => {
3705 let item = if resp.status.is_ok() {
3706 if !resp.body.is_empty() {
3707 let _ = tx.send(StreamItem::Chunk(resp.body));
3708 }
3709 StreamItem::End
3710 } else {
3711 StreamItem::Error(resp)
3712 };
3713 let _ = tx.send(item);
3714 self.senders.remove(&call_id);
3715 }
3716 StreamingChunkKind::Unary => {
3717 tracing::warn!(
3718 call_id,
3719 body_len = resp.body.len(),
3720 "rpc client: streaming / duplex consumer received unary-shaped \
3721 response (no nrpc-streaming header); server may have bridged a \
3722 unary path. Bridging to single-chunk + EOF.",
3723 );
3724 if !resp.body.is_empty() {
3725 let _ = tx.send(StreamItem::Chunk(resp.body));
3726 }
3727 let _ = tx.send(StreamItem::End);
3728 self.senders.remove(&call_id);
3729 }
3730 }
3731 }
3732
3733 /// Deliver a request-direction grant credit to the waiter
3734 /// for `call_id`, if it's a client-streaming / duplex entry.
3735 /// Silently no-op for unknown call_ids, for unary entries
3736 /// (caller bug — grant for a unary call makes no sense),
3737 /// and for server-streaming entries (grants apply only to
3738 /// the upload direction).
3739 ///
3740 /// `from_node` is gated by the same target-binding check
3741 /// as `deliver`: a grant from a non-target session peer is
3742 /// dropped (a forged grant on a shared reply channel can't
3743 /// inject credit into a victim's call).
3744 ///
3745 /// Bidi streaming plan (Phase C).
3746 fn deliver_grant(
3747 &self,
3748 call_id: u64,
3749 from_node: super::super::behavior::placement::NodeId,
3750 credits: u32,
3751 ) {
3752 let entry = self.senders.get(&call_id);
3753 let Some(entry) = entry else { return };
3754 let (target_node, _entry_value) = entry.value();
3755 if *target_node != 0 && *target_node != from_node {
3756 tracing::trace!(
3757 call_id,
3758 from_node,
3759 expected = *target_node,
3760 "rpc client: dropping REQUEST_GRANT from non-target session peer"
3761 );
3762 return;
3763 }
3764 match entry.value() {
3765 (_, PendingEntry::ClientStreaming { grant_tx, .. })
3766 | (_, PendingEntry::Duplex { grant_tx, .. }) => {
3767 let _ = grant_tx.send(credits);
3768 }
3769 // Unary / Streaming entries silently ignore — see
3770 // method docs for the rationale.
3771 _ => {}
3772 }
3773 }
3774
3775 /// Test-only: how many pending calls are registered. Used by
3776 /// integration tests to confirm cleanup after happy-path / cancel.
3777 #[cfg(test)]
3778 pub fn pending_count(&self) -> usize {
3779 self.senders.len()
3780 }
3781}
3782
3783impl Default for RpcClientPending {
3784 fn default() -> Self {
3785 Self::new()
3786 }
3787}
3788
3789/// Client-side fold. Decodes RESPONSE events and routes them to
3790/// awaiting oneshots in the shared [`RpcClientPending`].
3791///
3792/// `Mesh::call` clones the same `Arc<RpcClientPending>` to register
3793/// oneshots before publishing REQUESTs.
3794pub struct RpcClientFold {
3795 pending: Arc<RpcClientPending>,
3796}
3797
3798impl RpcClientFold {
3799 /// Construct a client fold that delivers responses through
3800 /// `pending`. Typical pattern:
3801 ///
3802 /// ```ignore
3803 /// let pending = Arc::new(RpcClientPending::new());
3804 /// let fold = RpcClientFold::new(pending.clone());
3805 /// let adapter = CortexAdapter::open(..., fold, ())?;
3806 /// // `pending` is still usable for register / cancel.
3807 /// ```
3808 pub fn new(pending: Arc<RpcClientPending>) -> Self {
3809 Self { pending }
3810 }
3811
3812 /// Production-path entry point. Mesh dispatch calls this with
3813 /// the AEAD-verified session peer's `NodeId` in
3814 /// `ev.from_node`; the pending registry's S-4 binding gate
3815 /// uses it to reject responses from the wrong target.
3816 pub fn apply_inbound(&mut self, ev: &RpcInboundEvent) {
3817 let Some(meta) = (if ev.payload.len() >= EVENT_META_SIZE {
3818 EventMeta::from_bytes(&ev.payload[..EVENT_META_SIZE])
3819 } else {
3820 None
3821 }) else {
3822 tracing::warn!(
3823 payload_len = ev.payload.len(),
3824 "rpc client fold: event payload too short for EventMeta; skipping",
3825 );
3826 return;
3827 };
3828 match meta.dispatch {
3829 DISPATCH_RPC_RESPONSE => {
3830 match RpcResponsePayload::decode(ev.payload.slice(EVENT_META_SIZE..)) {
3831 Ok(resp) => self.pending.deliver(meta.seq_or_ts, ev.from_node, resp),
3832 Err(e) => {
3833 tracing::warn!(
3834 error = %e,
3835 call_id = meta.seq_or_ts,
3836 "rpc client fold: malformed response payload",
3837 );
3838 }
3839 }
3840 }
3841 DISPATCH_RPC_REQUEST_GRANT => {
3842 // Server granted upload credit for a
3843 // client-streaming / duplex call. Route it to the
3844 // matching pending entry's grant mpsc; non-client-
3845 // streaming entries silently ignore (see
3846 // RpcClientPending::deliver_grant docs).
3847 match decode_request_grant(&ev.payload[EVENT_META_SIZE..]) {
3848 Some(grant) => {
3849 // The payload's `call_id` MUST agree with
3850 // the EventMeta's `seq_or_ts`: producer
3851 // encodes both to the same value (see
3852 // `RpcRequestGrantPayload::call_id` docs).
3853 // If they disagree, the frame is malformed
3854 // or forged — drop it. Otherwise a peer
3855 // could publish a GRANT whose meta names
3856 // one call but whose payload credits a
3857 // different in-flight call_id.
3858 if grant.call_id != meta.seq_or_ts {
3859 tracing::debug!(
3860 meta_call_id = meta.seq_or_ts,
3861 payload_call_id = grant.call_id,
3862 "rpc client fold: REQUEST_GRANT meta/payload call_id mismatch; dropping",
3863 );
3864 return;
3865 }
3866 if grant.credits == 0 {
3867 return;
3868 }
3869 self.pending
3870 .deliver_grant(grant.call_id, ev.from_node, grant.credits);
3871 }
3872 None => {
3873 tracing::debug!(
3874 call_id = meta.seq_or_ts,
3875 "rpc client fold: malformed REQUEST_GRANT payload"
3876 );
3877 }
3878 }
3879 }
3880 _ => {
3881 // Unknown / unexpected dispatch on the reply
3882 // channel — ignore (a misconfigured publisher
3883 // shouldn't take down the fold).
3884 }
3885 }
3886 }
3887}
3888
3889impl RedexFold<()> for RpcClientFold {
3890 /// Legacy entry point used by loopback / test paths that
3891 /// don't have a session peer to resolve. Calls `deliver`
3892 /// with `from_node = 0`, which the pending registry treats
3893 /// as "no binding" — callers that registered with
3894 /// `target_node = 0` accept it, callers that registered
3895 /// with a real target reject it.
3896 fn apply(&mut self, ev: &RedexEvent, _state: &mut ()) -> Result<(), RedexError> {
3897 let Some(meta) = (if ev.payload.len() >= EVENT_META_SIZE {
3898 EventMeta::from_bytes(&ev.payload[..EVENT_META_SIZE])
3899 } else {
3900 None
3901 }) else {
3902 tracing::warn!(
3903 payload_len = ev.payload.len(),
3904 "rpc client fold: event payload too short for EventMeta; skipping",
3905 );
3906 return Ok(());
3907 };
3908 // Route RESPONSE and REQUEST_GRANT events; ignore other
3909 // dispatches a misconfigured publisher might send. The
3910 // loopback path uses `from_node = 0` which the pending
3911 // registry treats as "no binding" — see the apply_inbound
3912 // production-path counterpart above for the AEAD-verified
3913 // peer routing.
3914 match meta.dispatch {
3915 DISPATCH_RPC_RESPONSE => {
3916 match RpcResponsePayload::decode(ev.payload.slice(EVENT_META_SIZE..)) {
3917 Ok(resp) => self.pending.deliver(meta.seq_or_ts, 0, resp),
3918 Err(e) => {
3919 // Malformed RESPONSE on the reply channel.
3920 // We can't fabricate a synthetic response
3921 // (the call_id might be valid; we just
3922 // can't tell what it was supposed to
3923 // mean). Log and leave the pending entry
3924 // intact — the caller's deadline /
3925 // cancellation path will eventually clean
3926 // it up.
3927 tracing::warn!(
3928 error = %e,
3929 call_id = meta.seq_or_ts,
3930 "rpc client fold: malformed response payload",
3931 );
3932 }
3933 }
3934 }
3935 DISPATCH_RPC_REQUEST_GRANT => {
3936 match decode_request_grant(&ev.payload[EVENT_META_SIZE..]) {
3937 Some(grant) => {
3938 // See `apply_inbound` REQUEST_GRANT arm for
3939 // the meta/payload call_id invariant.
3940 if grant.call_id != meta.seq_or_ts {
3941 tracing::debug!(
3942 meta_call_id = meta.seq_or_ts,
3943 payload_call_id = grant.call_id,
3944 "rpc client fold: REQUEST_GRANT meta/payload call_id mismatch; dropping",
3945 );
3946 return Ok(());
3947 }
3948 if grant.credits == 0 {
3949 return Ok(());
3950 }
3951 self.pending.deliver_grant(grant.call_id, 0, grant.credits);
3952 }
3953 None => {
3954 tracing::debug!(
3955 call_id = meta.seq_or_ts,
3956 "rpc client fold: malformed REQUEST_GRANT payload"
3957 );
3958 }
3959 }
3960 }
3961 _ => {}
3962 }
3963 Ok(())
3964 }
3965}
3966
3967#[cfg(test)]
3968mod tests {
3969 use super::*;
3970
3971 fn header(name: &str, value: &[u8]) -> RpcHeader {
3972 (name.to_string(), value.to_vec())
3973 }
3974
3975 // --------------------------------------------------------------------
3976 // Status code numbering.
3977 // --------------------------------------------------------------------
3978
3979 /// Status codes have stable wire numbers. A regression that
3980 /// renumbered any of the canonical statuses would break
3981 /// every cross-version caller / server pair on the wire — pin
3982 /// the numbers explicitly so the test catches it before the
3983 /// bug ships.
3984 #[test]
3985 fn status_wire_numbers_are_stable() {
3986 for (status, expected) in [
3987 (RpcStatus::Ok, 0x0000u16),
3988 (RpcStatus::NotFound, 0x0001),
3989 (RpcStatus::Unauthorized, 0x0002),
3990 (RpcStatus::Timeout, 0x0003),
3991 (RpcStatus::Backpressure, 0x0004),
3992 (RpcStatus::Cancelled, 0x0005),
3993 (RpcStatus::Internal, 0x0006),
3994 (RpcStatus::UnknownVersion, 0x0007),
3995 (RpcStatus::CapabilityDenied, 0x0008),
3996 ] {
3997 assert_eq!(status.to_wire(), expected, "{status:?}");
3998 assert_eq!(RpcStatus::from_wire(expected), status);
3999 }
4000 }
4001
4002 /// Reserved numeric range (`0x0009..=0x7FFF`) decodes as
4003 /// `Application(v)` for forward-compat with future canonical
4004 /// assignments. A future status numbered `0x0009` would round-
4005 /// trip via `from_wire(0x0009)` until that variant is added,
4006 /// at which point the variant takes precedence.
4007 #[test]
4008 fn reserved_status_range_decodes_as_application_for_forward_compat() {
4009 let decoded = RpcStatus::from_wire(0x0009);
4010 assert_eq!(decoded, RpcStatus::Application(0x0009));
4011 assert_eq!(decoded.to_wire(), 0x0009);
4012 }
4013
4014 /// Application range (`0x8000..=0xFFFF`) encodes / decodes
4015 /// transparently as `Application(v)`.
4016 #[test]
4017 fn application_status_range_roundtrips() {
4018 for v in [0x8000u16, 0x8001, 0xCAFE, 0xFFFF] {
4019 let s = RpcStatus::from_wire(v);
4020 assert_eq!(s, RpcStatus::Application(v));
4021 assert_eq!(s.to_wire(), v);
4022 }
4023 }
4024
4025 // --------------------------------------------------------------------
4026 // Dispatch byte assignments.
4027 // --------------------------------------------------------------------
4028
4029 /// Pin the `dispatch` byte assignments so a renumber surfaces
4030 /// here before it ships on the wire. These also live in the
4031 /// design doc; this test is the source-of-truth check.
4032 #[test]
4033 fn dispatch_byte_assignments_are_stable() {
4034 assert_eq!(DISPATCH_RPC_REQUEST, 0x10);
4035 assert_eq!(DISPATCH_RPC_RESPONSE, 0x11);
4036 assert_eq!(DISPATCH_RPC_CANCEL, 0x12);
4037 assert_eq!(DISPATCH_RPC_DEADLINE_EXCEEDED, 0x13);
4038 assert_eq!(DISPATCH_RPC_STREAM_GRANT, 0x14);
4039 assert_eq!(DISPATCH_RPC_REQUEST_CHUNK, 0x15);
4040 assert_eq!(DISPATCH_RPC_REQUEST_GRANT, 0x16);
4041 }
4042
4043 /// Regression: encoder bounds. Encoding a service name longer
4044 /// than `MAX_RPC_SERVICE_NAME_LEN` panics in debug, catching
4045 /// the programmer error in tests rather than silently writing
4046 /// a truncated `as u8` length that the receiver decodes as
4047 /// garbage. The matching debug_asserts guard body length,
4048 /// header count, header name length, and header value length.
4049 #[cfg(debug_assertions)]
4050 #[test]
4051 #[should_panic(expected = "service name")]
4052 fn request_encode_panics_on_oversize_service_name() {
4053 let p = RpcRequestPayload {
4054 service: "x".repeat(MAX_RPC_SERVICE_NAME_LEN + 1),
4055 deadline_ns: 0,
4056 flags: 0,
4057 headers: vec![],
4058 body: Bytes::new(),
4059 };
4060 let _ = p.encode();
4061 }
4062
4063 #[cfg(debug_assertions)]
4064 #[test]
4065 #[should_panic(expected = "body length")]
4066 fn request_encode_panics_on_oversize_body() {
4067 let p = RpcRequestPayload {
4068 service: "x".to_string(),
4069 deadline_ns: 0,
4070 flags: 0,
4071 headers: vec![],
4072 body: Bytes::from(vec![0; MAX_RPC_BODY_LEN + 1]),
4073 };
4074 let _ = p.encode();
4075 }
4076
4077 #[cfg(debug_assertions)]
4078 #[test]
4079 #[should_panic(expected = "header name")]
4080 fn request_encode_panics_on_oversize_header_name() {
4081 let p = RpcRequestPayload {
4082 service: "x".to_string(),
4083 deadline_ns: 0,
4084 flags: 0,
4085 headers: vec![("a".repeat(MAX_RPC_HEADER_NAME_LEN + 1), vec![])],
4086 body: Bytes::new(),
4087 };
4088 let _ = p.encode();
4089 }
4090
4091 /// `encoded_len()` must agree with `encode().len()` for every
4092 /// payload shape — pin this so a future codec change can't
4093 /// silently desynchronize the size-budgeting helper from the
4094 /// actual wire size.
4095 #[test]
4096 fn encoded_len_matches_encode_len_for_request_and_response() {
4097 let req = RpcRequestPayload {
4098 service: "echo.v1".to_string(),
4099 deadline_ns: 1_700_000_000_000_000_000,
4100 flags: FLAG_RPC_PROPAGATE_TRACE,
4101 headers: vec![
4102 header("traceparent", b"00-aabb"),
4103 header("idempotency-key", &7u64.to_le_bytes()),
4104 ],
4105 body: Bytes::from_static(b"{\"hello\":\"world\"}"),
4106 };
4107 assert_eq!(req.encoded_len(), req.encode().len());
4108
4109 let resp = RpcResponsePayload {
4110 status: RpcStatus::Application(0x8001),
4111 headers: vec![header("content-type", b"application/json")],
4112 body: Bytes::from_static(b"ok"),
4113 };
4114 assert_eq!(resp.encoded_len(), resp.encode().len());
4115
4116 // Empty edge cases.
4117 let empty_req = RpcRequestPayload {
4118 service: "x".to_string(),
4119 deadline_ns: 0,
4120 flags: 0,
4121 headers: vec![],
4122 body: Bytes::new(),
4123 };
4124 assert_eq!(empty_req.encoded_len(), empty_req.encode().len());
4125 let empty_resp = RpcResponsePayload {
4126 status: RpcStatus::Ok,
4127 headers: vec![],
4128 body: Bytes::new(),
4129 };
4130 assert_eq!(empty_resp.encoded_len(), empty_resp.encode().len());
4131 }
4132
4133 /// Bit 0 of `RpcRequestPayload::flags` is reserved (was the
4134 /// removed `FLAG_RPC_IDEMPOTENT`). Pin: live flag constants
4135 /// must NOT collide with bit 0, so a future re-add can safely
4136 /// reuse it without breaking existing senders.
4137 #[test]
4138 fn flag_bit_assignments_leave_idempotent_slot_reserved() {
4139 // Bit 0 (1 << 0) and bit 3 (1 << 3) are reserved; live flags
4140 // occupy other bits. Pinning the exact assignments here so
4141 // a renumber that collides with bit 0 (future `IDEMPOTENT`
4142 // re-add) or bit 3 (held in reserve for a future protocol
4143 // flag) surfaces in the test suite before it ships.
4144 assert_eq!(FLAG_RPC_STREAMING_RESPONSE, 1 << 1);
4145 assert_eq!(FLAG_RPC_PROPAGATE_TRACE, 1 << 2);
4146 assert_eq!(FLAG_RPC_CLIENT_STREAMING_REQUEST, 1 << 4);
4147 assert_eq!(FLAG_RPC_REQUEST_END, 1 << 5);
4148 for flag in [
4149 FLAG_RPC_STREAMING_RESPONSE,
4150 FLAG_RPC_PROPAGATE_TRACE,
4151 FLAG_RPC_CLIENT_STREAMING_REQUEST,
4152 FLAG_RPC_REQUEST_END,
4153 ] {
4154 assert_eq!(
4155 flag & (1 << 0),
4156 0,
4157 "flag {flag:#06x} collides with reserved bit 0"
4158 );
4159 assert_eq!(
4160 flag & (1 << 3),
4161 0,
4162 "flag {flag:#06x} collides with reserved bit 3"
4163 );
4164 }
4165 }
4166
4167 // --------------------------------------------------------------------
4168 // Bidi streaming (Phase A) — RpcRequestChunkPayload and
4169 // RpcRequestGrantPayload wire-stability tests.
4170 // --------------------------------------------------------------------
4171
4172 /// 1/5 — RequestChunk round-trip with realistic header set and
4173 /// 1 KiB body. Pins the encode/decode loop on the full shape.
4174 #[test]
4175 fn request_chunk_roundtrip_with_headers_and_body() {
4176 let mut headers = Vec::new();
4177 for i in 0..10u8 {
4178 headers.push(header(&format!("x-chunk-meta-{i}"), &[0xAA, 0xBB, i, !i]));
4179 }
4180 let body: Vec<u8> = (0..1024u32).map(|n| (n & 0xFF) as u8).collect();
4181 let p = RpcRequestChunkPayload {
4182 call_id: 0xCAFE_F00D_DEAD_BEEF,
4183 flags: FLAG_RPC_REQUEST_END | FLAG_RPC_PROPAGATE_TRACE,
4184 headers,
4185 body: Bytes::from(body),
4186 };
4187 let bytes = p.encode();
4188 assert_eq!(
4189 p.encoded_len(),
4190 bytes.len(),
4191 "encoded_len must agree with encode().len()"
4192 );
4193 let decoded = RpcRequestChunkPayload::decode(Bytes::from(bytes)).expect("decode");
4194 assert_eq!(decoded, p);
4195 }
4196
4197 /// 2/5 — truncation rejection at every field boundary. The
4198 /// codec must error rather than panic / allocate-unbounded on
4199 /// any short slice.
4200 #[test]
4201 fn request_chunk_decode_rejects_truncation_at_every_boundary() {
4202 let p = RpcRequestChunkPayload {
4203 call_id: 0x1234,
4204 flags: 0,
4205 headers: vec![header("x", b"y")],
4206 body: Bytes::from_static(b"hello"),
4207 };
4208 let full = p.encode();
4209 // Walk every prefix shorter than the full encoding; every
4210 // one must produce a Truncated / TooLarge / InvalidUtf8
4211 // error, not panic.
4212 for n in 0..full.len() {
4213 let prefix = &full[..n];
4214 let result = RpcRequestChunkPayload::decode(Bytes::copy_from_slice(prefix));
4215 assert!(result.is_err(), "n={n}: expected Err, got Ok({:?})", result);
4216 }
4217 // Full length must decode cleanly.
4218 assert!(RpcRequestChunkPayload::decode(Bytes::from(full)).is_ok());
4219 }
4220
4221 /// 3/5 — body length cap rejection. A wire-claimed body length
4222 /// over `MAX_RPC_BODY_LEN` must error rather than try to
4223 /// allocate 4+ MiB of garbage.
4224 #[test]
4225 fn request_chunk_decode_rejects_oversized_body_length() {
4226 // Build a synthetic encoding by hand: small valid prefix
4227 // up to body_len, then claim body_len = MAX_RPC_BODY_LEN + 1.
4228 let mut buf = Vec::new();
4229 buf.put_u64_le(0x42); // call_id
4230 buf.put_u16_le(0); // flags
4231 buf.put_u8(0); // headers count = 0
4232 buf.put_u32_le((MAX_RPC_BODY_LEN + 1) as u32);
4233 // (no body bytes follow — we want the decoder to reject at
4234 // the length check before it even tries to read body bytes)
4235 let err = RpcRequestChunkPayload::decode(Bytes::from(buf))
4236 .expect_err("oversized body length must reject");
4237 match err {
4238 RpcCodecError::TooLarge {
4239 field,
4240 actual,
4241 limit,
4242 } => {
4243 assert_eq!(field, "body");
4244 assert_eq!(actual, MAX_RPC_BODY_LEN + 1);
4245 assert_eq!(limit, MAX_RPC_BODY_LEN);
4246 }
4247 other => panic!("expected TooLarge {{ field=body }}, got {other:?}"),
4248 }
4249 }
4250
4251 /// 4/5 — header count cap rejection. A header count over
4252 /// `MAX_RPC_HEADERS` must error before the per-header decode
4253 /// loop even starts.
4254 #[test]
4255 fn request_chunk_decode_rejects_oversized_header_count() {
4256 let mut buf = Vec::new();
4257 buf.put_u64_le(0x42); // call_id
4258 buf.put_u16_le(0); // flags
4259 buf.put_u8((MAX_RPC_HEADERS + 1) as u8); // over the cap
4260 let err = RpcRequestChunkPayload::decode(Bytes::from(buf))
4261 .expect_err("oversized header count must reject");
4262 match err {
4263 RpcCodecError::TooLarge {
4264 field,
4265 actual,
4266 limit,
4267 } => {
4268 // The shared `decode_headers` helper reports this
4269 // field as "headers".
4270 assert_eq!(field, "headers");
4271 assert_eq!(actual, MAX_RPC_HEADERS + 1);
4272 assert_eq!(limit, MAX_RPC_HEADERS);
4273 }
4274 other => panic!("expected TooLarge {{ field=headers }}, got {other:?}"),
4275 }
4276 }
4277
4278 /// 5/5 — RequestGrant round-trip + truncation rejection. The
4279 /// payload is fixed-size (12 bytes), so the test surface is
4280 /// "exactly 12 bytes decodes" + "any other length errors".
4281 #[test]
4282 fn request_grant_roundtrip_and_truncation_rejection() {
4283 // Round-trip across the full u32 range corners + an
4284 // arbitrary mid-value.
4285 for (call_id, credits) in [
4286 (0u64, 0u32),
4287 (1, 1),
4288 (0xFFFF_FFFF_FFFF_FFFF, 0xFFFF_FFFF),
4289 (0xCAFE_F00D, 0x10203040),
4290 ] {
4291 let bytes = encode_request_grant(call_id, credits);
4292 assert_eq!(bytes.len(), 12, "request grant is always 12 bytes");
4293 let decoded = decode_request_grant(&bytes).expect("decode");
4294 assert_eq!(decoded.call_id, call_id);
4295 assert_eq!(decoded.credits, credits);
4296 }
4297 // Wrong-length payloads must reject (return None), not
4298 // panic. Empty, short, long, off-by-one each get covered.
4299 assert!(decode_request_grant(&[]).is_none());
4300 assert!(decode_request_grant(&[0u8; 11]).is_none());
4301 assert!(decode_request_grant(&[0u8; 13]).is_none());
4302 }
4303
4304 /// Bonus pin: `parse_request_window_initial` extracts a valid
4305 /// u32 ASCII-decimal header and rejects everything else.
4306 /// Same coverage shape as `parse_stream_window_initial`'s
4307 /// implicit contract, made explicit here so the request-side
4308 /// helper doesn't drift away from the response-side one.
4309 #[test]
4310 fn parse_request_window_initial_matches_response_side_semantics() {
4311 // Happy path.
4312 let headers = vec![header(HEADER_NRPC_REQUEST_WINDOW_INITIAL, b"32")];
4313 assert_eq!(parse_request_window_initial(&headers), Some(32));
4314 // Case-insensitive on header name.
4315 let headers = vec![header("Nrpc-Request-Window-Initial", b"7")];
4316 assert_eq!(parse_request_window_initial(&headers), Some(7));
4317 // Absent.
4318 assert_eq!(parse_request_window_initial(&[]), None);
4319 // Malformed value (non-numeric).
4320 let headers = vec![header(HEADER_NRPC_REQUEST_WINDOW_INITIAL, b"twelve")];
4321 assert_eq!(parse_request_window_initial(&headers), None);
4322 // Malformed value (non-utf8 bytes).
4323 let headers = vec![header(HEADER_NRPC_REQUEST_WINDOW_INITIAL, &[0xFF, 0xFE])];
4324 assert_eq!(parse_request_window_initial(&headers), None);
4325 // Empty value.
4326 let headers = vec![header(HEADER_NRPC_REQUEST_WINDOW_INITIAL, b"")];
4327 assert_eq!(parse_request_window_initial(&headers), None);
4328 }
4329
4330 // --------------------------------------------------------------------
4331 // RpcRequestPayload codec.
4332 // --------------------------------------------------------------------
4333
4334 #[test]
4335 fn request_roundtrip_minimal() {
4336 let p = RpcRequestPayload {
4337 service: "hello".to_string(),
4338 deadline_ns: 0,
4339 flags: 0,
4340 headers: vec![],
4341 body: Bytes::new(),
4342 };
4343 let bytes = p.encode();
4344 let decoded = RpcRequestPayload::decode(Bytes::from(bytes)).unwrap();
4345 assert_eq!(decoded, p);
4346 }
4347
4348 #[test]
4349 fn request_roundtrip_full() {
4350 let p = RpcRequestPayload {
4351 service: "echo.v1".to_string(),
4352 deadline_ns: 1_700_000_000_000_000_000,
4353 flags: FLAG_RPC_PROPAGATE_TRACE,
4354 headers: vec![
4355 header("traceparent", b"00-aabb..."),
4356 header("idempotency-key", &7u64.to_le_bytes()),
4357 header("content-type", b"application/json"),
4358 ],
4359 body: Bytes::from_static(b"{\"hello\":\"world\"}"),
4360 };
4361 let bytes = p.encode();
4362 let decoded = RpcRequestPayload::decode(Bytes::from(bytes)).unwrap();
4363 assert_eq!(decoded, p);
4364 }
4365
4366 #[test]
4367 fn request_decode_rejects_empty_service() {
4368 let bytes = vec![0x00];
4369 let err = RpcRequestPayload::decode(Bytes::from(bytes)).unwrap_err();
4370 assert!(matches!(err, RpcCodecError::Truncated(_)));
4371 }
4372
4373 #[test]
4374 fn request_decode_rejects_oversize_body_length() {
4375 // Forge: service "x", deadline 0, flags 0, no headers,
4376 // body length = MAX_RPC_BODY_LEN + 1 (no body bytes).
4377 let mut bytes = vec![1u8, b'x'];
4378 bytes.extend_from_slice(&0u64.to_le_bytes()); // deadline
4379 bytes.extend_from_slice(&0u16.to_le_bytes()); // flags
4380 bytes.push(0); // 0 headers
4381 bytes.extend_from_slice(&((MAX_RPC_BODY_LEN as u32) + 1).to_le_bytes());
4382 let err = RpcRequestPayload::decode(Bytes::from(bytes)).unwrap_err();
4383 assert!(
4384 matches!(err, RpcCodecError::TooLarge { field, .. } if field == "body"),
4385 "got {err:?}",
4386 );
4387 }
4388
4389 #[test]
4390 fn request_decode_rejects_oversize_headers_count() {
4391 // Forge: service "x", deadline 0, flags 0, headers count =
4392 // MAX_RPC_HEADERS + 1 (no header bytes).
4393 let mut bytes = vec![1u8, b'x'];
4394 bytes.extend_from_slice(&0u64.to_le_bytes());
4395 bytes.extend_from_slice(&0u16.to_le_bytes());
4396 bytes.push((MAX_RPC_HEADERS as u8).wrapping_add(1));
4397 let err = RpcRequestPayload::decode(Bytes::from(bytes)).unwrap_err();
4398 assert!(
4399 matches!(err, RpcCodecError::TooLarge { field, .. } if field == "headers"),
4400 "got {err:?}",
4401 );
4402 }
4403
4404 #[test]
4405 fn request_decode_rejects_truncated_at_each_field() {
4406 // Build a valid payload then truncate at each field
4407 // boundary; every truncation must error rather than silently
4408 // accept partial state.
4409 let p = RpcRequestPayload {
4410 service: "svc".to_string(),
4411 deadline_ns: 1,
4412 flags: 0,
4413 headers: vec![header("h", b"v")],
4414 body: Bytes::from_static(b"body"),
4415 };
4416 let bytes = p.encode();
4417 // Try each prefix length up to but not including the full
4418 // length — every one must be a decode error.
4419 for trim_to in 0..bytes.len() {
4420 let truncated = &bytes[..trim_to];
4421 let result = RpcRequestPayload::decode(Bytes::copy_from_slice(truncated));
4422 assert!(
4423 result.is_err(),
4424 "trim_to={trim_to} of {} must error, got {:?}",
4425 bytes.len(),
4426 result,
4427 );
4428 }
4429 // Full length must succeed.
4430 assert!(RpcRequestPayload::decode(Bytes::from(bytes)).is_ok());
4431 }
4432
4433 // --------------------------------------------------------------------
4434 // RpcResponsePayload codec.
4435 // --------------------------------------------------------------------
4436
4437 #[test]
4438 fn response_roundtrip_ok_with_body() {
4439 let p = RpcResponsePayload {
4440 status: RpcStatus::Ok,
4441 headers: vec![header("content-type", b"application/json")],
4442 body: Bytes::from_static(b"{\"answer\":42}"),
4443 };
4444 let bytes = p.encode();
4445 let decoded = RpcResponsePayload::decode(Bytes::from(bytes)).unwrap();
4446 assert_eq!(decoded, p);
4447 }
4448
4449 #[test]
4450 fn response_roundtrip_application_status() {
4451 let p = RpcResponsePayload {
4452 status: RpcStatus::Application(0xBEEF),
4453 headers: vec![],
4454 body: Bytes::from_static(b"app-specific diagnostic"),
4455 };
4456 let bytes = p.encode();
4457 let decoded = RpcResponsePayload::decode(Bytes::from(bytes)).unwrap();
4458 assert_eq!(decoded.status, RpcStatus::Application(0xBEEF));
4459 assert_eq!(decoded.body, p.body);
4460 }
4461
4462 #[test]
4463 fn response_decode_rejects_empty_buffer() {
4464 let err = RpcResponsePayload::decode(Bytes::new()).unwrap_err();
4465 assert!(matches!(err, RpcCodecError::Truncated(_)));
4466 }
4467
4468 // --------------------------------------------------------------------
4469 // Invariant: encoded sizes are reasonable.
4470 // --------------------------------------------------------------------
4471
4472 /// Wire-size budget regression: a tiny request encodes in a
4473 /// small constant number of bytes plus body. Pre-fix the headers
4474 /// or service-length encoding could have grown unbounded; pin
4475 /// the small-case so a regression in either inflates the
4476 /// minimum.
4477 #[test]
4478 fn request_minimum_wire_size_is_bounded() {
4479 let p = RpcRequestPayload {
4480 service: "x".to_string(),
4481 deadline_ns: 0,
4482 flags: 0,
4483 headers: vec![],
4484 body: Bytes::new(),
4485 };
4486 let size = p.encode().len();
4487 // 1 (svc len) + 1 (svc bytes) + 8 (deadline) + 2 (flags) + 1 (headers count) + 4 (body len) = 17
4488 assert_eq!(size, 17, "minimum request encodes in 17 bytes");
4489 assert_eq!(request_wire_size(&p), EVENT_META_SIZE + 17);
4490 }
4491
4492 #[test]
4493 fn response_minimum_wire_size_is_bounded() {
4494 let p = RpcResponsePayload {
4495 status: RpcStatus::Ok,
4496 headers: vec![],
4497 body: Bytes::new(),
4498 };
4499 let size = p.encode().len();
4500 // 2 (status) + 1 (headers count) + 4 (body len) = 7
4501 assert_eq!(size, 7, "minimum response encodes in 7 bytes");
4502 assert_eq!(response_wire_size(&p), EVENT_META_SIZE + 7);
4503 }
4504
4505 // ====================================================================
4506 // RpcServerFold — server-side dispatch behavior.
4507 //
4508 // These tests drive the fold directly with synthetic events
4509 // and observe the emitter callback. The end-to-end story
4510 // (Mesh::serve_rpc + bus + cortex adapter) is integration-
4511 // tested separately once the glue layer lands.
4512 // ====================================================================
4513
4514 use super::super::super::redex::{RedexEntry, RedexEvent};
4515 use std::sync::atomic::AtomicUsize;
4516 use std::time::Duration;
4517
4518 /// Captured-response store. Test-local typedef so the
4519 /// `capturing_emitter` signature stays under the `clippy::
4520 /// type_complexity` lint.
4521 type CapturedResponses = Arc<Mutex<Vec<(u64, u64, RpcResponsePayload)>>>;
4522
4523 /// Build a synthetic RedexEvent carrying an RPC request payload.
4524 /// Tests use this to drive the fold without going through the
4525 /// real ingest/cortex pipeline.
4526 fn rpc_request_event(
4527 caller_origin: u64,
4528 call_id: u64,
4529 payload: RpcRequestPayload,
4530 ) -> RedexEvent {
4531 let meta = EventMeta::new(DISPATCH_RPC_REQUEST, 0, caller_origin, call_id, 0);
4532 let mut buf = Vec::new();
4533 buf.extend_from_slice(&meta.to_bytes());
4534 buf.extend_from_slice(&payload.encode());
4535 RedexEvent {
4536 entry: RedexEntry::new_heap(0, 0, buf.len() as u32, 0, 0),
4537 payload: bytes::Bytes::from(buf),
4538 }
4539 }
4540
4541 fn rpc_cancel_event(caller_origin: u64, call_id: u64) -> RedexEvent {
4542 let meta = EventMeta::new(DISPATCH_RPC_CANCEL, 0, caller_origin, call_id, 0);
4543 let buf = meta.to_bytes().to_vec();
4544 RedexEvent {
4545 entry: RedexEntry::new_heap(0, 0, buf.len() as u32, 0, 0),
4546 payload: bytes::Bytes::from(buf),
4547 }
4548 }
4549
4550 /// Captures responses emitted by the fold for assertion in tests.
4551 fn capturing_emitter() -> (RpcResponseEmitter, CapturedResponses) {
4552 let captured: CapturedResponses = Arc::new(Mutex::new(Vec::new()));
4553 let captured_clone = captured.clone();
4554 let emit: RpcResponseEmitter = Arc::new(move |origin, call_id, resp| {
4555 captured_clone.lock().push((origin, call_id, resp));
4556 });
4557 (emit, captured)
4558 }
4559
4560 /// A handler that just echoes the request body back as the
4561 /// response body, with `RpcStatus::Ok`.
4562 struct EchoHandler;
4563 #[async_trait::async_trait]
4564 impl RpcHandler for EchoHandler {
4565 async fn call(&self, ctx: RpcContext) -> Result<RpcResponsePayload, RpcHandlerError> {
4566 Ok(RpcResponsePayload {
4567 status: RpcStatus::Ok,
4568 headers: vec![],
4569 body: ctx.payload.body,
4570 })
4571 }
4572 }
4573
4574 /// Wait until `pred` is true, polling at 10ms intervals up to
4575 /// `timeout`. Used to await spawned-handler completion in tests
4576 /// without a sleep-and-pray.
4577 async fn wait_until<F: Fn() -> bool>(pred: F, timeout: Duration) -> bool {
4578 let start = std::time::Instant::now();
4579 while start.elapsed() < timeout {
4580 if pred() {
4581 return true;
4582 }
4583 tokio::time::sleep(Duration::from_millis(10)).await;
4584 }
4585 pred()
4586 }
4587
4588 /// Happy path: a REQUEST event triggers the handler; the fold
4589 /// emits a RESPONSE with the handler's payload.
4590 #[tokio::test]
4591 async fn server_fold_request_invokes_handler_and_emits_response() {
4592 let (emit, captured) = capturing_emitter();
4593 let mut fold = RpcServerFold::new(Arc::new(EchoHandler), emit);
4594 let req = RpcRequestPayload {
4595 service: "echo".to_string(),
4596 deadline_ns: 0,
4597 flags: 0,
4598 headers: vec![],
4599 body: Bytes::from_static(b"hello"),
4600 };
4601 let ev = rpc_request_event(0xCAFE, 7, req);
4602 fold.apply(&ev, &mut ()).unwrap();
4603
4604 // Handler runs in tokio::spawn; wait for the emit.
4605 assert!(
4606 wait_until(|| !captured.lock().is_empty(), Duration::from_secs(2)).await,
4607 "expected one emitted response"
4608 );
4609 let captured = captured.lock();
4610 assert_eq!(captured.len(), 1);
4611 let (origin, call_id, resp) = &captured[0];
4612 assert_eq!(*origin, 0xCAFE);
4613 assert_eq!(*call_id, 7);
4614 assert_eq!(resp.status, RpcStatus::Ok);
4615 assert_eq!(resp.body.as_ref(), b"hello");
4616 // In-flight set is cleaned up after the handler completes.
4617 assert!(fold.in_flight_keys().is_empty());
4618 }
4619
4620 /// Application error: handler returns
4621 /// `RpcHandlerError::Application` → fold emits a response with
4622 /// `RpcStatus::Application(code)` and the message as body.
4623 #[tokio::test]
4624 async fn server_fold_application_error_maps_to_application_status() {
4625 struct AppErrHandler;
4626 #[async_trait::async_trait]
4627 impl RpcHandler for AppErrHandler {
4628 async fn call(&self, _ctx: RpcContext) -> Result<RpcResponsePayload, RpcHandlerError> {
4629 Err(RpcHandlerError::Application {
4630 code: 0xBEEF,
4631 message: "bad input".to_string(),
4632 })
4633 }
4634 }
4635 let (emit, captured) = capturing_emitter();
4636 let mut fold = RpcServerFold::new(Arc::new(AppErrHandler), emit);
4637 let req = RpcRequestPayload {
4638 service: "x".to_string(),
4639 deadline_ns: 0,
4640 flags: 0,
4641 headers: vec![],
4642 body: Bytes::new(),
4643 };
4644 fold.apply(&rpc_request_event(1, 1, req), &mut ()).unwrap();
4645 assert!(wait_until(|| !captured.lock().is_empty(), Duration::from_secs(2)).await);
4646 let captured = captured.lock();
4647 let (_, _, resp) = &captured[0];
4648 assert_eq!(resp.status, RpcStatus::Application(0xBEEF));
4649 assert_eq!(resp.body.as_ref(), b"bad input");
4650 }
4651
4652 /// Internal error: handler returns `RpcHandlerError::Internal`
4653 /// → fold emits `RpcStatus::Internal` with the message body.
4654 #[tokio::test]
4655 async fn server_fold_internal_error_maps_to_internal_status() {
4656 struct IntErrHandler;
4657 #[async_trait::async_trait]
4658 impl RpcHandler for IntErrHandler {
4659 async fn call(&self, _ctx: RpcContext) -> Result<RpcResponsePayload, RpcHandlerError> {
4660 Err(RpcHandlerError::Internal("db timeout".to_string()))
4661 }
4662 }
4663 let (emit, captured) = capturing_emitter();
4664 let mut fold = RpcServerFold::new(Arc::new(IntErrHandler), emit);
4665 let req = RpcRequestPayload {
4666 service: "x".to_string(),
4667 deadline_ns: 0,
4668 flags: 0,
4669 headers: vec![],
4670 body: Bytes::new(),
4671 };
4672 fold.apply(&rpc_request_event(1, 1, req), &mut ()).unwrap();
4673 assert!(wait_until(|| !captured.lock().is_empty(), Duration::from_secs(2)).await);
4674 let captured = captured.lock();
4675 let (_, _, resp) = &captured[0];
4676 assert_eq!(resp.status, RpcStatus::Internal);
4677 assert_eq!(resp.body.as_ref(), b"db timeout");
4678 }
4679
4680 /// Handler panic: caught by the fold's `catch_unwind`; surfaces
4681 /// as `RpcStatus::Internal` to the caller. Pre-fix the panic
4682 /// would propagate up the spawned task, log a tokio
4683 /// uncaught-panic message, and silently leave the caller
4684 /// waiting forever.
4685 #[tokio::test]
4686 async fn server_fold_handler_panic_surfaces_as_internal_status() {
4687 struct PanicHandler;
4688 #[async_trait::async_trait]
4689 impl RpcHandler for PanicHandler {
4690 async fn call(&self, _ctx: RpcContext) -> Result<RpcResponsePayload, RpcHandlerError> {
4691 panic!("kaboom");
4692 }
4693 }
4694 let (emit, captured) = capturing_emitter();
4695 let mut fold = RpcServerFold::new(Arc::new(PanicHandler), emit);
4696 let req = RpcRequestPayload {
4697 service: "x".to_string(),
4698 deadline_ns: 0,
4699 flags: 0,
4700 headers: vec![],
4701 body: Bytes::new(),
4702 };
4703 fold.apply(&rpc_request_event(1, 1, req), &mut ()).unwrap();
4704 assert!(wait_until(|| !captured.lock().is_empty(), Duration::from_secs(2)).await);
4705 let captured = captured.lock();
4706 let (_, _, resp) = &captured[0];
4707 assert_eq!(resp.status, RpcStatus::Internal);
4708 assert!(
4709 String::from_utf8_lossy(&resp.body).contains("kaboom"),
4710 "panic message must surface in body, got {}",
4711 String::from_utf8_lossy(&resp.body),
4712 );
4713 }
4714
4715 /// Deadline already passed: server short-circuits with
4716 /// `Timeout` without invoking the handler. Pinned via the
4717 /// `with_test_now_ns` clock override so the test doesn't race
4718 /// wall time.
4719 #[tokio::test]
4720 async fn server_fold_deadline_already_passed_short_circuits_to_timeout() {
4721 let invoked = Arc::new(AtomicBool::new(false));
4722 struct CountingHandler {
4723 invoked: Arc<AtomicBool>,
4724 }
4725 #[async_trait::async_trait]
4726 impl RpcHandler for CountingHandler {
4727 async fn call(&self, _ctx: RpcContext) -> Result<RpcResponsePayload, RpcHandlerError> {
4728 self.invoked.store(true, Ordering::Release);
4729 Ok(RpcResponsePayload {
4730 status: RpcStatus::Ok,
4731 headers: vec![],
4732 body: Bytes::new(),
4733 })
4734 }
4735 }
4736 let (emit, captured) = capturing_emitter();
4737 let mut fold = RpcServerFold::new(
4738 Arc::new(CountingHandler {
4739 invoked: invoked.clone(),
4740 }),
4741 emit,
4742 )
4743 // Use a clock value > DEADLINE_SKEW_TOLERANCE_NS + 1
4744 // (10s + 1ns) so the deadline-passed check fires past the
4745 // skew tolerance window. With now=20s and deadline=1ns,
4746 // (now - 10s) > 1ns.
4747 .with_test_now_ns(20_000_000_000);
4748 let req = RpcRequestPayload {
4749 service: "x".to_string(),
4750 // Deadline well in the past — past the skew tolerance.
4751 deadline_ns: 1_000,
4752 flags: 0,
4753 headers: vec![],
4754 body: Bytes::new(),
4755 };
4756 fold.apply(&rpc_request_event(1, 1, req), &mut ()).unwrap();
4757 // Emit happens synchronously in the deadline-passed branch
4758 // (no handler spawn).
4759 let captured = captured.lock();
4760 assert_eq!(captured.len(), 1);
4761 let (_, _, resp) = &captured[0];
4762 assert_eq!(resp.status, RpcStatus::Timeout);
4763 assert!(
4764 !invoked.load(Ordering::Acquire),
4765 "handler must NOT be invoked when deadline already passed",
4766 );
4767 }
4768
4769 /// Regression: a deadline that has elapsed by less than
4770 /// `DEADLINE_SKEW_TOLERANCE_NS` does NOT short-circuit. A
4771 /// peer with a slightly-fast clock would otherwise be
4772 /// prematurely timed out before the handler ever ran.
4773 #[tokio::test]
4774 async fn server_fold_deadline_within_skew_tolerance_invokes_handler() {
4775 let invoked = Arc::new(AtomicBool::new(false));
4776 struct CountingHandler {
4777 invoked: Arc<AtomicBool>,
4778 }
4779 #[async_trait::async_trait]
4780 impl RpcHandler for CountingHandler {
4781 async fn call(&self, _ctx: RpcContext) -> Result<RpcResponsePayload, RpcHandlerError> {
4782 self.invoked.store(true, Ordering::Release);
4783 Ok(RpcResponsePayload {
4784 status: RpcStatus::Ok,
4785 headers: vec![],
4786 body: Bytes::new(),
4787 })
4788 }
4789 }
4790 let (emit, captured) = capturing_emitter();
4791 let mut fold = RpcServerFold::new(
4792 Arc::new(CountingHandler {
4793 invoked: invoked.clone(),
4794 }),
4795 emit,
4796 )
4797 // now = 100s, deadline = 95s → elapsed = 5s, within the
4798 // 10s skew tolerance.
4799 .with_test_now_ns(100_000_000_000);
4800 let req = RpcRequestPayload {
4801 service: "x".to_string(),
4802 deadline_ns: 95_000_000_000,
4803 flags: 0,
4804 headers: vec![],
4805 body: Bytes::new(),
4806 };
4807 fold.apply(&rpc_request_event(1, 1, req), &mut ()).unwrap();
4808 assert!(
4809 wait_until(|| invoked.load(Ordering::Acquire), Duration::from_secs(1)).await,
4810 "handler must run when deadline is within skew tolerance",
4811 );
4812 let captured = captured.lock();
4813 assert_eq!(captured.len(), 1);
4814 assert_eq!(captured[0].2.status, RpcStatus::Ok);
4815 }
4816
4817 /// CANCEL flips the matching in-flight token. The handler that
4818 /// `select!`s on the cancellation observes the signal and can
4819 /// short-circuit. The fold removes the in-flight entry on
4820 /// CANCEL.
4821 #[tokio::test]
4822 async fn server_fold_cancel_flips_token_and_clears_in_flight() {
4823 let resumed_after_cancel = Arc::new(AtomicBool::new(false));
4824 struct CancelObservingHandler {
4825 resumed: Arc<AtomicBool>,
4826 }
4827 #[async_trait::async_trait]
4828 impl RpcHandler for CancelObservingHandler {
4829 async fn call(&self, ctx: RpcContext) -> Result<RpcResponsePayload, RpcHandlerError> {
4830 tokio::select! {
4831 _ = ctx.cancellation.cancelled() => {
4832 self.resumed.store(true, Ordering::Release);
4833 Err(RpcHandlerError::Internal("cancelled by caller".to_string()))
4834 }
4835 _ = tokio::time::sleep(Duration::from_secs(5)) => {
4836 Ok(RpcResponsePayload {
4837 status: RpcStatus::Ok,
4838 headers: vec![],
4839 body: Bytes::from_static(b"slept the full window"),
4840 })
4841 }
4842 }
4843 }
4844 }
4845 let (emit, captured) = capturing_emitter();
4846 let mut fold = RpcServerFold::new(
4847 Arc::new(CancelObservingHandler {
4848 resumed: resumed_after_cancel.clone(),
4849 }),
4850 emit,
4851 );
4852 let req = RpcRequestPayload {
4853 service: "x".to_string(),
4854 deadline_ns: 0,
4855 flags: 0,
4856 headers: vec![],
4857 body: Bytes::new(),
4858 };
4859 fold.apply(&rpc_request_event(1, 42, req), &mut ()).unwrap();
4860 // Wait until the handler's `select!` is parked; then send
4861 // CANCEL.
4862 assert!(
4863 wait_until(
4864 || fold.in_flight_keys().contains(&(1, 42)),
4865 Duration::from_secs(1)
4866 )
4867 .await
4868 );
4869 fold.apply(&rpc_cancel_event(1, 42), &mut ()).unwrap();
4870 // The cancellation is observed by the handler. Even though
4871 // the handler returns `Internal("cancelled by caller")`,
4872 // the fold's CANCEL-wins ordering overrides the response
4873 // with `RpcStatus::Cancelled` so the caller sees the
4874 // documented status code rather than the handler's
4875 // accidental Internal payload.
4876 assert!(
4877 wait_until(|| !captured.lock().is_empty(), Duration::from_secs(2)).await,
4878 "handler should observe cancellation and emit response"
4879 );
4880 assert!(
4881 resumed_after_cancel.load(Ordering::Acquire),
4882 "handler must observe cancellation"
4883 );
4884 let captured = captured.lock();
4885 assert_eq!(captured.len(), 1);
4886 let (_, _, resp) = &captured[0];
4887 assert_eq!(
4888 resp.status,
4889 RpcStatus::Cancelled,
4890 "CANCEL must override handler outcome with RpcStatus::Cancelled"
4891 );
4892 // CANCEL also removes the in-flight entry directly.
4893 // Handler completion removes it again (idempotent).
4894 assert!(fold.in_flight_keys().is_empty());
4895 }
4896
4897 /// Regression: a duplicate REQUEST for an already-in-flight
4898 /// `(origin_hash, call_id)` must be refused with a synthetic
4899 /// `Internal` response and must NOT spawn a second handler.
4900 /// Without the refusal, two handlers race under the same key
4901 /// and CANCEL handling is broken (CANCEL removes the entry
4902 /// the first handler reinserts, etc.).
4903 #[tokio::test]
4904 async fn server_fold_duplicate_request_refuses_without_double_dispatch() {
4905 let invocations = Arc::new(AtomicUsize::new(0));
4906 struct CountingHandler {
4907 invocations: Arc<AtomicUsize>,
4908 }
4909 #[async_trait::async_trait]
4910 impl RpcHandler for CountingHandler {
4911 async fn call(&self, _ctx: RpcContext) -> Result<RpcResponsePayload, RpcHandlerError> {
4912 self.invocations.fetch_add(1, Ordering::SeqCst);
4913 tokio::time::sleep(Duration::from_millis(80)).await;
4914 Ok(RpcResponsePayload {
4915 status: RpcStatus::Ok,
4916 headers: vec![],
4917 body: Bytes::from_static(b"done"),
4918 })
4919 }
4920 }
4921 let (emit, captured) = capturing_emitter();
4922 let mut fold = RpcServerFold::new(
4923 Arc::new(CountingHandler {
4924 invocations: invocations.clone(),
4925 }),
4926 emit,
4927 );
4928 let req = RpcRequestPayload {
4929 service: "x".to_string(),
4930 deadline_ns: 0,
4931 flags: 0,
4932 headers: vec![],
4933 body: Bytes::new(),
4934 };
4935 // First REQUEST — handler spawns and parks in sleep.
4936 fold.apply(&rpc_request_event(1, 99, req.clone()), &mut ())
4937 .unwrap();
4938 assert!(
4939 wait_until(
4940 || fold.in_flight_keys().contains(&(1, 99)),
4941 Duration::from_secs(1)
4942 )
4943 .await
4944 );
4945 // Second REQUEST with same key — must be refused
4946 // synchronously with a synthetic Internal response.
4947 fold.apply(&rpc_request_event(1, 99, req), &mut ()).unwrap();
4948 // The refusal emit happens synchronously in the fold's
4949 // sync emitter path.
4950 let after_dup = captured.lock().clone();
4951 assert_eq!(
4952 after_dup.len(),
4953 1,
4954 "duplicate REQUEST must emit exactly one synthetic refusal",
4955 );
4956 assert_eq!(after_dup[0].2.status, RpcStatus::Internal);
4957 assert!(String::from_utf8_lossy(&after_dup[0].2.body).contains("duplicate"));
4958 // Wait for the first handler to complete.
4959 assert!(
4960 wait_until(|| captured.lock().len() == 2, Duration::from_secs(2)).await,
4961 "first handler should still complete normally"
4962 );
4963 let captured = captured.lock();
4964 assert_eq!(captured.len(), 2);
4965 // The first handler's response is the second emit (Ok).
4966 assert_eq!(captured[1].2.status, RpcStatus::Ok);
4967 assert_eq!(
4968 invocations.load(Ordering::SeqCst),
4969 1,
4970 "duplicate REQUEST must NOT spawn a second handler",
4971 );
4972 }
4973
4974 /// Regression: a CANCEL that fires while the handler is mid-
4975 /// flight must override the handler's outcome with
4976 /// `RpcStatus::Cancelled` even when the handler ignores
4977 /// cancellation and returns `Ok(...)`. Without this, a caller
4978 /// who cancelled would see the handler's accidental success
4979 /// payload and could not tell whether their CANCEL won.
4980 #[tokio::test]
4981 async fn server_fold_cancel_overrides_handler_ok_with_cancelled_status() {
4982 struct IgnoresCancellation;
4983 #[async_trait::async_trait]
4984 impl RpcHandler for IgnoresCancellation {
4985 async fn call(&self, _ctx: RpcContext) -> Result<RpcResponsePayload, RpcHandlerError> {
4986 // Sleep long enough for the test to send CANCEL,
4987 // then return Ok regardless. This models a handler
4988 // that doesn't `select!` on `ctx.cancellation`.
4989 tokio::time::sleep(Duration::from_millis(80)).await;
4990 Ok(RpcResponsePayload {
4991 status: RpcStatus::Ok,
4992 headers: vec![],
4993 body: Bytes::from_static(b"finished despite cancellation"),
4994 })
4995 }
4996 }
4997 let (emit, captured) = capturing_emitter();
4998 let mut fold = RpcServerFold::new(Arc::new(IgnoresCancellation), emit);
4999 let req = RpcRequestPayload {
5000 service: "x".to_string(),
5001 deadline_ns: 0,
5002 flags: 0,
5003 headers: vec![],
5004 body: Bytes::new(),
5005 };
5006 fold.apply(&rpc_request_event(7, 11, req), &mut ()).unwrap();
5007 // Wait until the handler is parked, then send CANCEL well
5008 // before the handler's sleep elapses.
5009 assert!(
5010 wait_until(
5011 || fold.in_flight_keys().contains(&(7, 11)),
5012 Duration::from_secs(1)
5013 )
5014 .await
5015 );
5016 fold.apply(&rpc_cancel_event(7, 11), &mut ()).unwrap();
5017 assert!(
5018 wait_until(|| !captured.lock().is_empty(), Duration::from_secs(2)).await,
5019 "handler should complete and emit response"
5020 );
5021 let captured = captured.lock();
5022 assert_eq!(captured.len(), 1);
5023 let (_, _, resp) = &captured[0];
5024 assert_eq!(
5025 resp.status,
5026 RpcStatus::Cancelled,
5027 "handler that returned Ok despite CANCEL must surface as Cancelled"
5028 );
5029 assert!(fold.in_flight_keys().is_empty());
5030 }
5031
5032 /// CANCEL for an unknown call_id is a no-op (no panic, no
5033 /// stray emission). This is the case where a CANCEL races a
5034 /// handler completion or a duplicate CANCEL arrives.
5035 #[tokio::test]
5036 async fn server_fold_cancel_for_unknown_call_id_is_no_op() {
5037 let (emit, captured) = capturing_emitter();
5038 let mut fold = RpcServerFold::new(Arc::new(EchoHandler), emit);
5039 // CANCEL with no matching REQUEST.
5040 fold.apply(&rpc_cancel_event(1, 999), &mut ()).unwrap();
5041 assert!(captured.lock().is_empty());
5042 assert!(fold.in_flight_keys().is_empty());
5043 }
5044
5045 /// Malformed request payload: fold emits a
5046 /// `RpcStatus::UnknownVersion` response and continues. A
5047 /// regression that returned `Err` here would kill the cortex
5048 /// adapter's tail-and-fold task on the first malformed event,
5049 /// which is the wrong behavior for an RPC server that needs
5050 /// to keep serving past garbage.
5051 #[tokio::test]
5052 async fn server_fold_malformed_payload_emits_unknown_version_and_keeps_going() {
5053 let (emit, captured) = capturing_emitter();
5054 let mut fold = RpcServerFold::new(Arc::new(EchoHandler), emit);
5055 // Build an event with valid meta but a garbage tail (just
5056 // a single 0x00 byte, which fails the service-len check).
5057 let meta = EventMeta::new(DISPATCH_RPC_REQUEST, 0, 7, 1, 0);
5058 let mut buf = Vec::new();
5059 buf.extend_from_slice(&meta.to_bytes());
5060 buf.push(0x00); // svc_len = 0 → empty service → Truncated
5061 let ev = RedexEvent {
5062 entry: RedexEntry::new_heap(0, 0, buf.len() as u32, 0, 0),
5063 payload: bytes::Bytes::from(buf),
5064 };
5065 let result = fold.apply(&ev, &mut ());
5066 assert!(
5067 result.is_ok(),
5068 "fold must NOT return Err on malformed payload (would kill the adapter); got {result:?}"
5069 );
5070 let captured = captured.lock();
5071 assert_eq!(captured.len(), 1);
5072 let (_, _, resp) = &captured[0];
5073 assert_eq!(resp.status, RpcStatus::UnknownVersion);
5074 }
5075
5076 /// Cancellation token roundtrip: `cancel()` sets `is_cancelled`
5077 /// and wakes a parked `cancelled().await`.
5078 #[tokio::test]
5079 async fn cancellation_token_signals_waiters() {
5080 let token = RpcCancellationToken::new();
5081 assert!(!token.is_cancelled());
5082 let token2 = token.clone();
5083 let waiter = tokio::spawn(async move {
5084 token2.cancelled().await;
5085 });
5086 // Give the waiter a chance to park.
5087 tokio::time::sleep(Duration::from_millis(10)).await;
5088 token.cancel();
5089 // Waiter wakes.
5090 tokio::time::timeout(Duration::from_secs(1), waiter)
5091 .await
5092 .expect("waiter must wake within 1s")
5093 .expect("waiter task must not panic");
5094 assert!(token.is_cancelled());
5095 }
5096
5097 // ====================================================================
5098 // W3C Trace Context propagation.
5099 // ====================================================================
5100
5101 /// `build_trace_headers` + `extract_trace_context` round-trip
5102 /// a typical W3C trace context through the request headers.
5103 #[test]
5104 fn trace_context_round_trips_through_headers() {
5105 let tc = TraceContext {
5106 traceparent: "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".to_string(),
5107 tracestate: "vendor1=opaque-value,vendor2=other".to_string(),
5108 };
5109 let headers = build_trace_headers(&tc);
5110 assert_eq!(headers.len(), 2, "non-empty tracestate emits both headers");
5111 let extracted = extract_trace_context(&headers).expect("must extract");
5112 assert_eq!(extracted, tc);
5113 }
5114
5115 /// Regression for M21: `extract_trace_context` does
5116 /// case-INsensitive matching on the header names, matching the
5117 /// W3C and HTTP conventions. A peer that emits capitalized
5118 /// `Traceparent` or `TRACESTATE` must still be picked up — the
5119 /// previous implementation used `name.as_str() == "traceparent"`
5120 /// and silently dropped any non-lowercase variant.
5121 #[test]
5122 fn extract_trace_context_is_case_insensitive_on_header_names() {
5123 // Capital-T traceparent + uppercase TRACESTATE — both must
5124 // be picked up by the extractor.
5125 let headers = vec![
5126 (
5127 "Traceparent".to_string(),
5128 b"00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".to_vec(),
5129 ),
5130 ("TRACESTATE".to_string(), b"vendor=value".to_vec()),
5131 ];
5132 let extracted =
5133 extract_trace_context(&headers).expect("capital-T traceparent must be recognized");
5134 assert_eq!(
5135 extracted.traceparent,
5136 "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
5137 );
5138 assert_eq!(extracted.tracestate, "vendor=value");
5139
5140 // Mixed-case still works.
5141 let headers = vec![
5142 ("traceParent".to_string(), b"00-aa-bb-01".to_vec()),
5143 ("TraceState".to_string(), b"v=1".to_vec()),
5144 ];
5145 let extracted =
5146 extract_trace_context(&headers).expect("mixed-case traceparent must be recognized");
5147 assert_eq!(extracted.traceparent, "00-aa-bb-01");
5148 assert_eq!(extracted.tracestate, "v=1");
5149 }
5150
5151 /// Empty `tracestate` is omitted on the wire (W3C convention)
5152 /// but extracted as empty on the receive side.
5153 #[test]
5154 fn trace_context_empty_tracestate_omitted_from_wire() {
5155 let tc = TraceContext {
5156 traceparent: "00-aa-bb-01".to_string(),
5157 tracestate: String::new(),
5158 };
5159 let headers = build_trace_headers(&tc);
5160 assert_eq!(
5161 headers.len(),
5162 1,
5163 "empty tracestate must NOT be emitted on the wire",
5164 );
5165 assert_eq!(headers[0].0, "traceparent");
5166 let extracted = extract_trace_context(&headers).expect("must extract");
5167 assert_eq!(extracted.traceparent, "00-aa-bb-01");
5168 assert_eq!(extracted.tracestate, "");
5169 }
5170
5171 /// Headers without `traceparent` decode as `None`. Useful for
5172 /// the FLAG_RPC_PROPAGATE_TRACE-set-but-no-headers misuse
5173 /// case — the server gets `None` rather than a bogus context.
5174 #[test]
5175 fn trace_context_missing_traceparent_returns_none() {
5176 let headers = vec![
5177 ("content-type".to_string(), b"application/json".to_vec()),
5178 ("idempotency-key".to_string(), b"abc".to_vec()),
5179 ];
5180 assert!(extract_trace_context(&headers).is_none());
5181 }
5182
5183 /// Server fold populates `RpcContext::trace_context` only when
5184 /// the caller signals `FLAG_RPC_PROPAGATE_TRACE`. End-to-end
5185 /// through the fold's apply path.
5186 #[tokio::test]
5187 async fn server_fold_propagates_trace_context_via_flag() {
5188 struct CapturingHandler {
5189 captured: Arc<Mutex<Option<Option<TraceContext>>>>,
5190 }
5191 #[async_trait::async_trait]
5192 impl RpcHandler for CapturingHandler {
5193 async fn call(&self, ctx: RpcContext) -> Result<RpcResponsePayload, RpcHandlerError> {
5194 *self.captured.lock() = Some(ctx.trace_context.clone());
5195 Ok(RpcResponsePayload {
5196 status: RpcStatus::Ok,
5197 headers: vec![],
5198 body: Bytes::new(),
5199 })
5200 }
5201 }
5202
5203 // Helper: run one request through a fresh fold and return
5204 // what the handler captured for trace_context.
5205 async fn run(req: RpcRequestPayload) -> Option<TraceContext> {
5206 let captured: Arc<Mutex<Option<Option<TraceContext>>>> = Arc::new(Mutex::new(None));
5207 let (emit, _captured_responses) = capturing_emitter();
5208 let handler = Arc::new(CapturingHandler {
5209 captured: captured.clone(),
5210 });
5211 let mut fold = RpcServerFold::new(handler, emit);
5212 fold.apply(&rpc_request_event(1, 1, req), &mut ()).unwrap();
5213 // Wait for the spawned handler to finish.
5214 assert!(
5215 wait_until(|| captured.lock().is_some(), Duration::from_secs(2)).await,
5216 "handler must run"
5217 );
5218 let observed = captured.lock().take().unwrap();
5219 observed
5220 }
5221
5222 // Case 1: FLAG_RPC_PROPAGATE_TRACE NOT set → trace_context is None.
5223 let req_no_flag = RpcRequestPayload {
5224 service: "x".to_string(),
5225 deadline_ns: 0,
5226 flags: 0,
5227 headers: vec![("traceparent".to_string(), b"00-aa-bb-01".to_vec())],
5228 body: Bytes::new(),
5229 };
5230 assert!(
5231 run(req_no_flag).await.is_none(),
5232 "without the flag, server must NOT extract trace_context"
5233 );
5234
5235 // Case 2: FLAG set + headers present → server gets the context.
5236 let tc = TraceContext {
5237 traceparent: "00-trace-span-01".to_string(),
5238 tracestate: "vendor=value".to_string(),
5239 };
5240 let req_with_flag = RpcRequestPayload {
5241 service: "x".to_string(),
5242 deadline_ns: 0,
5243 flags: FLAG_RPC_PROPAGATE_TRACE,
5244 headers: build_trace_headers(&tc),
5245 body: Bytes::new(),
5246 };
5247 let observed = run(req_with_flag).await.expect("flag set → should be Some");
5248 assert_eq!(observed, tc);
5249
5250 // Case 3: FLAG set but headers missing → None (defensive).
5251 let req_flag_no_headers = RpcRequestPayload {
5252 service: "x".to_string(),
5253 deadline_ns: 0,
5254 flags: FLAG_RPC_PROPAGATE_TRACE,
5255 headers: vec![],
5256 body: Bytes::new(),
5257 };
5258 assert!(
5259 run(req_flag_no_headers).await.is_none(),
5260 "flag set but no headers → server gets None (no synthesis)"
5261 );
5262 }
5263
5264 /// Race: cancel fires AFTER the future is registered but
5265 /// BEFORE the await actually parks. The token's
5266 /// `notified()`-then-check ordering must catch this case
5267 /// without sleeping past the cancellation.
5268 #[tokio::test]
5269 async fn cancellation_token_does_not_miss_cancel_racing_register() {
5270 for _ in 0..50 {
5271 let token = RpcCancellationToken::new();
5272 let token2 = token.clone();
5273 let waiter = tokio::spawn(async move {
5274 token2.cancelled().await;
5275 });
5276 // No sleep — fire cancel as fast as possible against
5277 // the just-spawned waiter. In the worst case the
5278 // waiter has not yet reached `notified()`; it will see
5279 // `is_cancelled() == true` on its first check and
5280 // return immediately. In the other case it parks and
5281 // gets woken by `notify_waiters`.
5282 token.cancel();
5283 tokio::time::timeout(Duration::from_secs(1), waiter)
5284 .await
5285 .expect("waiter must complete within 1s")
5286 .expect("waiter task must not panic");
5287 }
5288 }
5289
5290 // ====================================================================
5291 // RpcClientFold — caller-side response routing.
5292 // ====================================================================
5293
5294 fn rpc_response_event(
5295 caller_origin: u64,
5296 call_id: u64,
5297 payload: RpcResponsePayload,
5298 ) -> RedexEvent {
5299 let meta = EventMeta::new(DISPATCH_RPC_RESPONSE, 0, caller_origin, call_id, 0);
5300 let mut buf = Vec::new();
5301 buf.extend_from_slice(&meta.to_bytes());
5302 buf.extend_from_slice(&payload.encode());
5303 RedexEvent {
5304 entry: RedexEntry::new_heap(0, 0, buf.len() as u32, 0, 0),
5305 payload: bytes::Bytes::from(buf),
5306 }
5307 }
5308
5309 /// Happy path: register a call, drive the matching RESPONSE
5310 /// through the fold, the awaiting receiver gets the payload.
5311 #[tokio::test]
5312 async fn client_fold_routes_response_to_registered_waiter() {
5313 let pending = Arc::new(RpcClientPending::new());
5314 let mut fold = RpcClientFold::new(pending.clone());
5315 let rx = pending.register(42, 0);
5316 assert_eq!(pending.pending_count(), 1);
5317
5318 let resp = RpcResponsePayload {
5319 status: RpcStatus::Ok,
5320 headers: vec![],
5321 body: Bytes::from_static(b"hello back"),
5322 };
5323 fold.apply(&rpc_response_event(0xCAFE, 42, resp.clone()), &mut ())
5324 .unwrap();
5325
5326 // Receiver is completed.
5327 let got = tokio::time::timeout(Duration::from_secs(1), rx)
5328 .await
5329 .expect("receiver must complete within 1s")
5330 .expect("sender must not be dropped");
5331 assert_eq!(got, resp);
5332 // Pending entry cleared after delivery.
5333 assert_eq!(pending.pending_count(), 0);
5334 }
5335
5336 /// RESPONSE for an unknown call_id is a no-op (no panic, no
5337 /// stray side effect). This is the case where a stale RESPONSE
5338 /// arrives after the caller has cancelled or timed out.
5339 #[tokio::test]
5340 async fn client_fold_response_for_unknown_call_id_is_no_op() {
5341 let pending = Arc::new(RpcClientPending::new());
5342 let mut fold = RpcClientFold::new(pending.clone());
5343 let resp = RpcResponsePayload {
5344 status: RpcStatus::Ok,
5345 headers: vec![],
5346 body: Bytes::new(),
5347 };
5348 fold.apply(&rpc_response_event(1, 999, resp), &mut ())
5349 .unwrap();
5350 assert_eq!(pending.pending_count(), 0);
5351 }
5352
5353 /// REQUEST / CANCEL events on the reply channel are ignored
5354 /// rather than producing a stray decode-error or affecting
5355 /// pending state. The reply channel shouldn't carry these in
5356 /// practice (they belong on `<service>.requests`), but a
5357 /// misconfigured publisher must not break the fold.
5358 #[tokio::test]
5359 async fn client_fold_ignores_non_response_dispatches() {
5360 let pending = Arc::new(RpcClientPending::new());
5361 let mut fold = RpcClientFold::new(pending.clone());
5362 let _rx = pending.register(7, 0);
5363
5364 // REQUEST event landing on the caller's reply channel is
5365 // ignored.
5366 let req = RpcRequestPayload {
5367 service: "stray".to_string(),
5368 deadline_ns: 0,
5369 flags: 0,
5370 headers: vec![],
5371 body: Bytes::new(),
5372 };
5373 fold.apply(&rpc_request_event(1, 7, req), &mut ()).unwrap();
5374 // Pending entry untouched.
5375 assert_eq!(pending.pending_count(), 1);
5376
5377 // CANCEL on the reply channel: also ignored.
5378 fold.apply(&rpc_cancel_event(1, 7), &mut ()).unwrap();
5379 assert_eq!(pending.pending_count(), 1);
5380 }
5381
5382 /// `cancel(call_id)` removes the pending entry; a subsequent
5383 /// RESPONSE for that call_id is dropped silently.
5384 #[tokio::test]
5385 async fn client_pending_cancel_drops_subsequent_response() {
5386 let pending = Arc::new(RpcClientPending::new());
5387 let mut fold = RpcClientFold::new(pending.clone());
5388 let rx = pending.register(5, 0);
5389 pending.cancel(5);
5390 assert_eq!(pending.pending_count(), 0);
5391
5392 let resp = RpcResponsePayload {
5393 status: RpcStatus::Ok,
5394 headers: vec![],
5395 body: Bytes::new(),
5396 };
5397 fold.apply(&rpc_response_event(1, 5, resp), &mut ())
5398 .unwrap();
5399
5400 // Receiver was dropped along with the cancel. The previously-
5401 // returned `rx` errors with `Closed`.
5402 let result = tokio::time::timeout(Duration::from_secs(1), rx).await;
5403 let inner = result.expect("must complete within 1s");
5404 assert!(
5405 inner.is_err(),
5406 "receiver after cancel must error (sender dropped)",
5407 );
5408 }
5409
5410 /// Malformed RESPONSE payload: fold returns Ok (does not kill
5411 /// the cortex adapter) and leaves the pending entry intact for
5412 /// the caller's deadline / cancellation path to clean up. Pre-
5413 /// fix a bad payload could either kill the fold or fabricate a
5414 /// synthetic response — both wrong.
5415 #[tokio::test]
5416 async fn client_fold_malformed_response_is_logged_not_fatal() {
5417 let pending = Arc::new(RpcClientPending::new());
5418 let mut fold = RpcClientFold::new(pending.clone());
5419 let rx = pending.register(11, 0);
5420
5421 // Build a malformed RESPONSE: valid meta, garbage tail
5422 // (just `[0xFF]`, which is shorter than the required 2-byte
5423 // status + 1-byte headers count + 4-byte body length).
5424 let meta = EventMeta::new(DISPATCH_RPC_RESPONSE, 0, 1, 11, 0);
5425 let mut buf = Vec::new();
5426 buf.extend_from_slice(&meta.to_bytes());
5427 buf.push(0xFF);
5428 let ev = RedexEvent {
5429 entry: RedexEntry::new_heap(0, 0, buf.len() as u32, 0, 0),
5430 payload: bytes::Bytes::from(buf),
5431 };
5432 let result = fold.apply(&ev, &mut ());
5433 assert!(
5434 result.is_ok(),
5435 "fold must not return Err on malformed response"
5436 );
5437 // Pending entry NOT cleared — the caller's cancellation
5438 // path will eventually clean it up via `cancel(call_id)`.
5439 assert_eq!(pending.pending_count(), 1);
5440 // Receiver is still pending (not delivered, not closed).
5441 assert!(
5442 tokio::time::timeout(Duration::from_millis(50), rx)
5443 .await
5444 .is_err(),
5445 "receiver should still be parked (no delivery, no drop)",
5446 );
5447 }
5448
5449 /// Re-registering the same call_id replaces the prior sender;
5450 /// the prior `Receiver` errors with `RecvError::Closed`. This
5451 /// is the misuse-detection path — call_ids should be unique
5452 /// per (caller, target) for the lifetime of the call, and a
5453 /// clash surfaces as a hard error rather than silently
5454 /// delivering the response to the wrong waiter.
5455 #[tokio::test]
5456 async fn client_pending_re_register_closes_prior_receiver() {
5457 let pending = Arc::new(RpcClientPending::new());
5458 let rx_a = pending.register(99, 0);
5459 let _rx_b = pending.register(99, 0);
5460 // The first receiver is now closed (sender dropped on
5461 // re-insert).
5462 let result = tokio::time::timeout(Duration::from_secs(1), rx_a).await;
5463 let inner = result.expect("must complete within 1s");
5464 assert!(inner.is_err(), "re-register must close prior receiver");
5465 assert_eq!(pending.pending_count(), 1);
5466 }
5467
5468 /// S-4 part 2 regression: a RESPONSE whose wire `from_node`
5469 /// doesn't match the recorded `target_node` must not resolve
5470 /// the call. Without the gate, any peer with publish access
5471 /// to the caller's reply channel could ship a spoofed
5472 /// response (random call_ids from S-4 part 1 narrow the
5473 /// attack surface, but this gate closes the residual case
5474 /// of an attacker who has observed the victim's call_id via
5475 /// some side channel).
5476 #[tokio::test]
5477 async fn client_pending_drops_response_from_wrong_target() {
5478 let pending = Arc::new(RpcClientPending::new());
5479 let rx = pending.register(0xDEAD_BEEF, 0x42);
5480 let resp = RpcResponsePayload {
5481 status: RpcStatus::Ok,
5482 headers: Vec::new(),
5483 body: Bytes::from_static(b"forged"),
5484 };
5485 // Forged from a different session peer — must drop.
5486 pending.deliver(0xDEAD_BEEF, 0x99, resp.clone());
5487 // Receiver is still parked; pending entry is intact.
5488 let parked = tokio::time::timeout(Duration::from_millis(50), rx).await;
5489 assert!(
5490 parked.is_err(),
5491 "forged RESPONSE from wrong target must not resolve the call"
5492 );
5493 assert_eq!(pending.pending_count(), 1);
5494
5495 // Legitimate RESPONSE from the recorded target resolves.
5496 let rx2 = pending.register(0xCAFE, 0x42);
5497 let ok_resp = RpcResponsePayload {
5498 status: RpcStatus::Ok,
5499 headers: Vec::new(),
5500 body: Bytes::from_static(b"ok"),
5501 };
5502 pending.deliver(0xCAFE, 0x42, ok_resp);
5503 let delivered = tokio::time::timeout(Duration::from_millis(50), rx2)
5504 .await
5505 .expect("must complete")
5506 .expect("must receive");
5507 assert_eq!(delivered.body.as_ref(), b"ok");
5508 }
5509
5510 // ====================================================================
5511 // Phase C — RpcClientPending + RpcClientFold for client-streaming.
5512 // ====================================================================
5513
5514 /// Build a REQUEST_GRANT event for tests. Mirror of
5515 /// `rpc_stream_grant_event` for the request direction.
5516 fn rpc_request_grant_event(caller_origin: u64, call_id: u64, credits: u32) -> RedexEvent {
5517 let meta = EventMeta::new(DISPATCH_RPC_REQUEST_GRANT, 0, caller_origin, call_id, 0);
5518 let mut buf = Vec::with_capacity(EVENT_META_SIZE + 12);
5519 buf.extend_from_slice(&meta.to_bytes());
5520 buf.extend_from_slice(&encode_request_grant(call_id, credits));
5521 RedexEvent {
5522 entry: RedexEntry::new_heap(0, 0, buf.len() as u32, 0, 0),
5523 payload: bytes::Bytes::from(buf),
5524 }
5525 }
5526
5527 /// `register_client_streaming` returns two halves: a terminal
5528 /// oneshot and a grant mpsc. A terminal RESPONSE resolves the
5529 /// oneshot (same shape as unary delivery); a REQUEST_GRANT
5530 /// for the same call_id pushes its credit onto the mpsc.
5531 #[tokio::test]
5532 async fn client_pending_client_streaming_routes_terminal_and_grants() {
5533 let pending = Arc::new(RpcClientPending::new());
5534 let (terminal_rx, mut grant_rx) = pending.register_client_streaming(0xCAFE_F00D, 0);
5535 // Push two grants — both should land on the mpsc.
5536 pending.deliver_grant(0xCAFE_F00D, 0, 3);
5537 pending.deliver_grant(0xCAFE_F00D, 0, 7);
5538 assert_eq!(grant_rx.recv().await, Some(3));
5539 assert_eq!(grant_rx.recv().await, Some(7));
5540 // Terminal RESPONSE resolves the oneshot and removes the
5541 // entry. Grant mpsc closes too (its sender drops with
5542 // the entry).
5543 let resp = RpcResponsePayload {
5544 status: RpcStatus::Ok,
5545 headers: vec![],
5546 body: Bytes::from_static(b"done"),
5547 };
5548 pending.deliver(0xCAFE_F00D, 0, resp.clone());
5549 let delivered = tokio::time::timeout(Duration::from_millis(50), terminal_rx)
5550 .await
5551 .expect("terminal must complete")
5552 .expect("terminal must receive");
5553 assert_eq!(delivered.body.as_ref(), b"done");
5554 // Grant mpsc now closed.
5555 assert_eq!(grant_rx.recv().await, None);
5556 assert_eq!(pending.pending_count(), 0);
5557 }
5558
5559 /// REQUEST_GRANT from a non-target session peer is dropped
5560 /// without injecting credit. Same S-4-style binding gate as
5561 /// the RESPONSE delivery path — a forged grant on a shared
5562 /// reply channel can't inflate a victim's credit budget.
5563 #[tokio::test]
5564 async fn client_pending_grant_from_wrong_target_is_dropped() {
5565 let pending = Arc::new(RpcClientPending::new());
5566 let (_terminal_rx, mut grant_rx) = pending.register_client_streaming(0xCAFE_F00D, 0x42);
5567 // Forged grant from a different session peer — must drop.
5568 pending.deliver_grant(0xCAFE_F00D, 0x99, 100);
5569 let parked = tokio::time::timeout(Duration::from_millis(50), grant_rx.recv()).await;
5570 assert!(
5571 parked.is_err(),
5572 "forged REQUEST_GRANT from wrong target must not inject credit"
5573 );
5574 // Legitimate grant from the recorded target lands.
5575 pending.deliver_grant(0xCAFE_F00D, 0x42, 5);
5576 let delivered = tokio::time::timeout(Duration::from_millis(50), grant_rx.recv())
5577 .await
5578 .expect("must complete")
5579 .expect("must receive");
5580 assert_eq!(delivered, 5);
5581 }
5582
5583 /// `deliver_grant` for an unknown call_id is a silent no-op.
5584 /// Same harmless-drop semantics as a STREAM_GRANT for an
5585 /// unknown / non-flow-controlled call (CANCEL/GRANT race is
5586 /// always possible).
5587 #[tokio::test]
5588 async fn client_pending_grant_for_unknown_call_id_is_no_op() {
5589 let pending = Arc::new(RpcClientPending::new());
5590 // No entry registered for this call_id.
5591 pending.deliver_grant(0xDEAD, 0, 42);
5592 // No panics, no entries created.
5593 assert_eq!(pending.pending_count(), 0);
5594 }
5595
5596 /// `deliver_grant` for a unary entry is silently dropped
5597 /// (grants only apply to client-streaming / duplex calls).
5598 #[tokio::test]
5599 async fn client_pending_grant_for_unary_entry_is_no_op() {
5600 let pending = Arc::new(RpcClientPending::new());
5601 let _rx = pending.register(0xDEAD, 0);
5602 pending.deliver_grant(0xDEAD, 0, 42);
5603 // No state changes — entry still pending, no leak.
5604 assert_eq!(pending.pending_count(), 1);
5605 }
5606
5607 /// `RpcClientFold::apply` (legacy / loopback path) routes
5608 /// DISPATCH_RPC_REQUEST_GRANT events through to the matching
5609 /// ClientStreaming entry's grant mpsc. Pins the second
5610 /// dispatch arm the fold gained for Phase C.
5611 #[tokio::test]
5612 async fn client_fold_routes_request_grant_to_registered_waiter() {
5613 let pending = Arc::new(RpcClientPending::new());
5614 let mut fold = RpcClientFold::new(pending.clone());
5615 let (_terminal_rx, mut grant_rx) = pending.register_client_streaming(0xC0DE, 0);
5616 let ev = rpc_request_grant_event(0xCAFE, 0xC0DE, 9);
5617 fold.apply(&ev, &mut ()).unwrap();
5618 let delivered = tokio::time::timeout(Duration::from_millis(50), grant_rx.recv())
5619 .await
5620 .expect("must complete")
5621 .expect("must receive");
5622 assert_eq!(delivered, 9);
5623 }
5624
5625 /// `RpcClientFold::apply` ignores REQUEST_GRANT events whose
5626 /// payload is malformed (wrong length): no panic, no entry
5627 /// state change, fold returns Ok and keeps going. Mirror of
5628 /// the response-side malformed-payload regression.
5629 #[tokio::test]
5630 async fn client_fold_malformed_request_grant_is_logged_not_fatal() {
5631 let pending = Arc::new(RpcClientPending::new());
5632 let mut fold = RpcClientFold::new(pending.clone());
5633 let (_terminal_rx, mut grant_rx) = pending.register_client_streaming(0xC0DE, 0);
5634 // Build a GRANT event whose payload is only 4 bytes
5635 // (truncated — codec needs 12).
5636 let meta = EventMeta::new(DISPATCH_RPC_REQUEST_GRANT, 0, 0xCAFE, 0xC0DE, 0);
5637 let mut buf = Vec::new();
5638 buf.extend_from_slice(&meta.to_bytes());
5639 buf.extend_from_slice(&[0xAA, 0xBB, 0xCC, 0xDD]);
5640 let ev = RedexEvent {
5641 entry: RedexEntry::new_heap(0, 0, buf.len() as u32, 0, 0),
5642 payload: bytes::Bytes::from(buf),
5643 };
5644 let result = fold.apply(&ev, &mut ());
5645 assert!(
5646 result.is_ok(),
5647 "malformed REQUEST_GRANT must NOT kill the fold"
5648 );
5649 // No credit landed on the mpsc.
5650 let parked = tokio::time::timeout(Duration::from_millis(30), grant_rx.recv()).await;
5651 assert!(
5652 parked.is_err(),
5653 "malformed REQUEST_GRANT must not inject credit"
5654 );
5655 }
5656
5657 /// REQUEST_GRANT frames where the payload `call_id` does NOT
5658 /// agree with `EventMeta::seq_or_ts` must be dropped: the
5659 /// producer is contracted to encode both fields to the same
5660 /// value (see `RpcRequestGrantPayload::call_id` doc), so a
5661 /// mismatch is either a malformed frame or an attempted
5662 /// cross-call credit-injection. Without this check, a peer
5663 /// could publish a GRANT whose meta names one call but whose
5664 /// payload credits a different in-flight call_id.
5665 ///
5666 /// Regression: cubic-dev-ai bot P2 review comment on the
5667 /// `nrpc-streaming` branch.
5668 #[tokio::test]
5669 async fn client_fold_drops_request_grant_with_mismatched_call_ids() {
5670 let pending = Arc::new(RpcClientPending::new());
5671 let mut fold = RpcClientFold::new(pending.clone());
5672 let (_terminal_rx_victim, mut grant_rx_victim) =
5673 pending.register_client_streaming(0xC0DE, 0);
5674 let (_terminal_rx_other, mut grant_rx_other) = pending.register_client_streaming(0xBEEF, 0);
5675
5676 // Build a hand-rolled frame: meta names call 0xC0DE,
5677 // payload encodes credit for call 0xBEEF. Either the
5678 // producer is broken or this is a forged frame; the
5679 // consumer must drop, not deliver.
5680 let meta = EventMeta::new(DISPATCH_RPC_REQUEST_GRANT, 0, 0xCAFE, 0xC0DE, 0);
5681 let mut buf = Vec::with_capacity(EVENT_META_SIZE + 12);
5682 buf.extend_from_slice(&meta.to_bytes());
5683 buf.extend_from_slice(&encode_request_grant(0xBEEF, 5));
5684 let ev = RedexEvent {
5685 entry: RedexEntry::new_heap(0, 0, buf.len() as u32, 0, 0),
5686 payload: bytes::Bytes::from(buf),
5687 };
5688 fold.apply(&ev, &mut ()).unwrap();
5689
5690 let parked_victim =
5691 tokio::time::timeout(Duration::from_millis(30), grant_rx_victim.recv()).await;
5692 assert!(
5693 parked_victim.is_err(),
5694 "mismatched REQUEST_GRANT must not credit the call named in meta",
5695 );
5696 let parked_other =
5697 tokio::time::timeout(Duration::from_millis(30), grant_rx_other.recv()).await;
5698 assert!(
5699 parked_other.is_err(),
5700 "mismatched REQUEST_GRANT must not credit the call named in payload either",
5701 );
5702 }
5703
5704 // ====================================================================
5705 // RpcServerStreamingFold — coverage for the multi-fire emit path.
5706 //
5707 // The streaming fold is the most complex code in this file:
5708 // - Per-call cancellation token (same as unary)
5709 // - Pump task that drains an mpsc and awaits each emit to
5710 // enforce per-call ordering
5711 // - Optional flow-control semaphore (caller-set window +
5712 // STREAM_GRANT credit refills)
5713 // - Terminal-frame emission with CANCEL-wins override
5714 //
5715 // These tests pin each branch: ordered chunks + clean EOF;
5716 // application error after partial stream; panic surfacing as
5717 // Internal; CANCEL flipping the cancellation token AND being
5718 // surfaced as the terminal status; STREAM_GRANT permits;
5719 // duplicate-REQUEST refusal.
5720 // ====================================================================
5721
5722 /// Build an async `RpcAsyncResponseEmitter` that captures every
5723 /// emit into a shared Vec. Streaming fold tests use this to
5724 /// inspect the multi-frame emit pattern.
5725 fn capturing_async_emitter() -> (RpcAsyncResponseEmitter, CapturedResponses) {
5726 let captured: CapturedResponses = Arc::new(Mutex::new(Vec::new()));
5727 let captured_clone = captured.clone();
5728 let emit: RpcAsyncResponseEmitter = Arc::new(move |origin, call_id, resp| {
5729 let captured_clone = captured_clone.clone();
5730 Box::pin(async move {
5731 captured_clone.lock().push((origin, call_id, resp));
5732 })
5733 });
5734 (emit, captured)
5735 }
5736
5737 /// Synthesize a STREAM_GRANT event for a `(caller_origin, call_id)`
5738 /// asking for `n` additional credits.
5739 fn rpc_stream_grant_event(caller_origin: u64, call_id: u64, n: u32) -> RedexEvent {
5740 let meta = EventMeta::new(DISPATCH_RPC_STREAM_GRANT, 0, caller_origin, call_id, 0);
5741 let mut buf = Vec::with_capacity(EVENT_META_SIZE + 4);
5742 buf.extend_from_slice(&meta.to_bytes());
5743 buf.extend_from_slice(&encode_stream_grant(n));
5744 RedexEvent {
5745 entry: RedexEntry::new_heap(0, 0, buf.len() as u32, 0, 0),
5746 payload: bytes::Bytes::from(buf),
5747 }
5748 }
5749
5750 /// Streaming handler that emits N chunks and returns Ok. The
5751 /// caller-side test asserts (a) all N chunks arrive in order
5752 /// with the `nrpc-streaming: continue` header, (b) a final
5753 /// terminal frame with `nrpc-streaming: end` follows.
5754 #[tokio::test]
5755 async fn streaming_fold_emits_chunks_in_order_and_clean_terminal() {
5756 struct CountingHandler {
5757 n: usize,
5758 }
5759 #[async_trait::async_trait]
5760 impl RpcStreamingHandler for CountingHandler {
5761 async fn call(
5762 &self,
5763 _ctx: RpcContext,
5764 sink: RpcResponseSink,
5765 ) -> Result<(), RpcHandlerError> {
5766 for i in 0..self.n {
5767 sink.send(format!("chunk-{i}").into_bytes());
5768 }
5769 Ok(())
5770 }
5771 }
5772 let (emit, captured) = capturing_async_emitter();
5773 let mut fold = RpcServerStreamingFold::new(Arc::new(CountingHandler { n: 5 }), emit);
5774 let req = RpcRequestPayload {
5775 service: "stream".to_string(),
5776 deadline_ns: 0,
5777 flags: FLAG_RPC_STREAMING_RESPONSE,
5778 headers: vec![],
5779 body: Bytes::new(),
5780 };
5781 fold.apply(&rpc_request_event(11, 22, req), &mut ())
5782 .unwrap();
5783 // 5 continue chunks + 1 terminal end frame.
5784 assert!(
5785 wait_until(|| captured.lock().len() == 6, Duration::from_secs(2)).await,
5786 "expected 6 frames (5 chunks + terminal end), got {}",
5787 captured.lock().len(),
5788 );
5789 let captured = captured.lock();
5790 for (i, (_, _, resp)) in captured.iter().take(5).enumerate() {
5791 assert_eq!(resp.status, RpcStatus::Ok);
5792 // continue header on every non-terminal chunk
5793 let hdr = resp
5794 .headers
5795 .iter()
5796 .find(|(n, _)| n == HEADER_NRPC_STREAMING)
5797 .expect("streaming header present");
5798 assert_eq!(hdr.1.as_slice(), HEADER_NRPC_STREAMING_CONTINUE);
5799 assert_eq!(resp.body, format!("chunk-{i}").into_bytes());
5800 }
5801 // Terminal frame
5802 let (_, _, term) = captured.last().unwrap();
5803 assert_eq!(term.status, RpcStatus::Ok);
5804 let hdr = term
5805 .headers
5806 .iter()
5807 .find(|(n, _)| n == HEADER_NRPC_STREAMING)
5808 .expect("terminal must have streaming header");
5809 assert_eq!(hdr.1.as_slice(), HEADER_NRPC_STREAMING_END);
5810 assert!(term.body.is_empty());
5811 }
5812
5813 /// Handler returns `Err(Internal)` after sending 2 chunks. Caller
5814 /// must see (a) both chunks with the continue header, (b) a
5815 /// terminal frame carrying `RpcStatus::Internal` (NOT the end
5816 /// marker — the terminal-error path drops the header).
5817 #[tokio::test]
5818 async fn streaming_fold_terminal_error_after_partial_stream() {
5819 struct PartialErrHandler;
5820 #[async_trait::async_trait]
5821 impl RpcStreamingHandler for PartialErrHandler {
5822 async fn call(
5823 &self,
5824 _ctx: RpcContext,
5825 sink: RpcResponseSink,
5826 ) -> Result<(), RpcHandlerError> {
5827 sink.send(b"first".to_vec());
5828 sink.send(b"second".to_vec());
5829 Err(RpcHandlerError::Internal("ran out of fuel".into()))
5830 }
5831 }
5832 let (emit, captured) = capturing_async_emitter();
5833 let mut fold = RpcServerStreamingFold::new(Arc::new(PartialErrHandler), emit);
5834 let req = RpcRequestPayload {
5835 service: "x".to_string(),
5836 deadline_ns: 0,
5837 flags: FLAG_RPC_STREAMING_RESPONSE,
5838 headers: vec![],
5839 body: Bytes::new(),
5840 };
5841 fold.apply(&rpc_request_event(1, 1, req), &mut ()).unwrap();
5842 assert!(
5843 wait_until(|| captured.lock().len() == 3, Duration::from_secs(2)).await,
5844 "expected 2 chunks + 1 terminal error",
5845 );
5846 let captured = captured.lock();
5847 assert_eq!(captured[0].2.body.as_ref(), b"first");
5848 assert_eq!(captured[1].2.body.as_ref(), b"second");
5849 let (_, _, term) = &captured[2];
5850 assert_eq!(term.status, RpcStatus::Internal);
5851 assert!(
5852 String::from_utf8_lossy(&term.body).contains("ran out of fuel"),
5853 "diagnostic must round-trip, got {:?}",
5854 String::from_utf8_lossy(&term.body),
5855 );
5856 }
5857
5858 /// Handler panics. The fold's `catch_unwind` surfaces it as a
5859 /// terminal `RpcStatus::Internal` rather than killing the
5860 /// runtime.
5861 #[tokio::test]
5862 async fn streaming_fold_handler_panic_surfaces_as_internal_terminal() {
5863 struct PanicHandler;
5864 #[async_trait::async_trait]
5865 impl RpcStreamingHandler for PanicHandler {
5866 async fn call(
5867 &self,
5868 _ctx: RpcContext,
5869 _sink: RpcResponseSink,
5870 ) -> Result<(), RpcHandlerError> {
5871 panic!("kaboom in streaming handler");
5872 }
5873 }
5874 let (emit, captured) = capturing_async_emitter();
5875 let mut fold = RpcServerStreamingFold::new(Arc::new(PanicHandler), emit);
5876 let req = RpcRequestPayload {
5877 service: "x".to_string(),
5878 deadline_ns: 0,
5879 flags: FLAG_RPC_STREAMING_RESPONSE,
5880 headers: vec![],
5881 body: Bytes::new(),
5882 };
5883 fold.apply(&rpc_request_event(1, 2, req), &mut ()).unwrap();
5884 assert!(
5885 wait_until(|| !captured.lock().is_empty(), Duration::from_secs(2)).await,
5886 "panic must surface as a terminal frame",
5887 );
5888 let captured = captured.lock();
5889 assert_eq!(captured.len(), 1);
5890 let (_, _, resp) = &captured[0];
5891 assert_eq!(resp.status, RpcStatus::Internal);
5892 assert!(
5893 String::from_utf8_lossy(&resp.body).contains("kaboom"),
5894 "panic message must surface, got {:?}",
5895 String::from_utf8_lossy(&resp.body),
5896 );
5897 }
5898
5899 /// CANCEL during a streaming call overrides the terminal frame
5900 /// with `RpcStatus::Cancelled` — same CANCEL-wins ordering as
5901 /// the unary fold.
5902 #[tokio::test]
5903 async fn streaming_fold_cancel_overrides_terminal_with_cancelled() {
5904 struct CooperativeHandler;
5905 #[async_trait::async_trait]
5906 impl RpcStreamingHandler for CooperativeHandler {
5907 async fn call(
5908 &self,
5909 ctx: RpcContext,
5910 sink: RpcResponseSink,
5911 ) -> Result<(), RpcHandlerError> {
5912 sink.send(b"chunk-0".to_vec());
5913 tokio::select! {
5914 _ = ctx.cancellation.cancelled() => Ok(()),
5915 _ = tokio::time::sleep(Duration::from_secs(5)) => Ok(()),
5916 }
5917 }
5918 }
5919 let (emit, captured) = capturing_async_emitter();
5920 let mut fold = RpcServerStreamingFold::new(Arc::new(CooperativeHandler), emit);
5921 let req = RpcRequestPayload {
5922 service: "x".to_string(),
5923 deadline_ns: 0,
5924 flags: FLAG_RPC_STREAMING_RESPONSE,
5925 headers: vec![],
5926 body: Bytes::new(),
5927 };
5928 fold.apply(&rpc_request_event(7, 13, req), &mut ()).unwrap();
5929 // Wait until at least the first chunk is captured AND the
5930 // handler is parked (in_flight key present), then CANCEL.
5931 assert!(
5932 wait_until(
5933 || !captured.lock().is_empty() && fold.in_flight_keys().contains(&(7, 13)),
5934 Duration::from_secs(2)
5935 )
5936 .await
5937 );
5938 fold.apply(&rpc_cancel_event(7, 13), &mut ()).unwrap();
5939 // Wait for the terminal frame.
5940 assert!(
5941 wait_until(|| captured.lock().len() >= 2, Duration::from_secs(2)).await,
5942 "expected first chunk + terminal frame",
5943 );
5944 let captured = captured.lock();
5945 // First emit was the chunk; the LAST should be the
5946 // Cancelled terminal.
5947 assert_eq!(
5948 captured.last().unwrap().2.status,
5949 RpcStatus::Cancelled,
5950 "CANCEL must override terminal status",
5951 );
5952 }
5953
5954 /// Duplicate REQUEST with the same `(origin, call_id)` is
5955 /// refused with a synthetic Internal terminal frame and does
5956 /// NOT spawn a second handler. Mirror of the unary fold's
5957 /// regression at server_fold_duplicate_request_refuses_*.
5958 #[tokio::test]
5959 async fn streaming_fold_duplicate_request_refuses_without_double_dispatch() {
5960 let invocations = Arc::new(AtomicUsize::new(0));
5961 struct CountingHandler {
5962 invocations: Arc<AtomicUsize>,
5963 }
5964 #[async_trait::async_trait]
5965 impl RpcStreamingHandler for CountingHandler {
5966 async fn call(
5967 &self,
5968 _ctx: RpcContext,
5969 sink: RpcResponseSink,
5970 ) -> Result<(), RpcHandlerError> {
5971 self.invocations.fetch_add(1, Ordering::SeqCst);
5972 tokio::time::sleep(Duration::from_millis(80)).await;
5973 sink.send(b"chunk".to_vec());
5974 Ok(())
5975 }
5976 }
5977 let (emit, captured) = capturing_async_emitter();
5978 let mut fold = RpcServerStreamingFold::new(
5979 Arc::new(CountingHandler {
5980 invocations: invocations.clone(),
5981 }),
5982 emit,
5983 );
5984 let req = RpcRequestPayload {
5985 service: "x".to_string(),
5986 deadline_ns: 0,
5987 flags: FLAG_RPC_STREAMING_RESPONSE,
5988 headers: vec![],
5989 body: Bytes::new(),
5990 };
5991 fold.apply(&rpc_request_event(1, 99, req.clone()), &mut ())
5992 .unwrap();
5993 assert!(
5994 wait_until(
5995 || fold.in_flight_keys().contains(&(1, 99)),
5996 Duration::from_secs(1)
5997 )
5998 .await
5999 );
6000 // Duplicate REQUEST — must emit a synthetic Internal
6001 // terminal and not invoke the handler a second time.
6002 fold.apply(&rpc_request_event(1, 99, req), &mut ()).unwrap();
6003 assert!(
6004 wait_until(|| !captured.lock().is_empty(), Duration::from_secs(1)).await,
6005 "synthetic refusal should be emitted",
6006 );
6007 // First emit (chronologically) is the synthetic refusal.
6008 let refusal = captured.lock()[0].clone();
6009 assert_eq!(refusal.2.status, RpcStatus::Internal);
6010 assert!(String::from_utf8_lossy(&refusal.2.body).contains("duplicate"));
6011 // Wait for the original handler to complete (chunk + terminal).
6012 assert!(
6013 wait_until(|| captured.lock().len() >= 3, Duration::from_secs(2)).await,
6014 "first handler should still complete normally",
6015 );
6016 assert_eq!(
6017 invocations.load(Ordering::SeqCst),
6018 1,
6019 "duplicate REQUEST must NOT spawn a second handler",
6020 );
6021 }
6022
6023 /// STREAM_GRANT for an unknown call_id is silently dropped
6024 /// (no panic, no tracing event escalation). Pin the
6025 /// always-safe behavior so a misbehaving caller (or a CANCEL/
6026 /// GRANT race) can't crash the fold.
6027 #[tokio::test]
6028 async fn streaming_fold_grant_for_unknown_call_id_is_no_op() {
6029 struct NoopHandler;
6030 #[async_trait::async_trait]
6031 impl RpcStreamingHandler for NoopHandler {
6032 async fn call(
6033 &self,
6034 _ctx: RpcContext,
6035 _sink: RpcResponseSink,
6036 ) -> Result<(), RpcHandlerError> {
6037 Ok(())
6038 }
6039 }
6040 let (emit, captured) = capturing_async_emitter();
6041 let mut fold = RpcServerStreamingFold::new(Arc::new(NoopHandler), emit);
6042 let result = fold.apply(&rpc_stream_grant_event(99, 42, 5), &mut ());
6043 assert!(result.is_ok(), "GRANT for unknown call_id must be Ok");
6044 assert!(captured.lock().is_empty(), "no emit for unknown GRANT");
6045 }
6046
6047 /// Regression for M20: the streaming pump's mpsc is bounded
6048 /// at `STREAMING_PUMP_CAPACITY`. A handler that produces
6049 /// chunks faster than the pump drains gets its excess
6050 /// `sink.send(...)` calls silently dropped (matching the
6051 /// "caller cancelled" semantic) — and the metric counter
6052 /// `streaming_chunks_dropped_total` increments.
6053 ///
6054 /// We construct the sink directly with a tiny bounded mpsc
6055 /// (capacity 2) and a metrics handle, then call `send` 5
6056 /// times without a receiver. The first 2 fit in the channel;
6057 /// the next 3 are dropped and counted.
6058 #[tokio::test]
6059 async fn streaming_sink_drops_on_full_and_increments_metric() {
6060 use crate::adapter::net::mesh_rpc_metrics::{RpcMetricsRegistry, ServiceMetricsAtomic};
6061 // Tiny channel to make overflow easy to observe.
6062 let (tx, _rx) = tokio::sync::mpsc::channel::<bytes::Bytes>(2);
6063 let registry = RpcMetricsRegistry::new();
6064 let metrics: Arc<ServiceMetricsAtomic> = registry.for_service("drop_test");
6065 let sink = RpcResponseSink {
6066 inner: tx,
6067 metrics: Some(metrics.clone()),
6068 };
6069 // 5 sends; first 2 buffer, next 3 drop.
6070 for i in 0..5u8 {
6071 sink.send(vec![i]);
6072 }
6073 assert_eq!(
6074 metrics
6075 .streaming_chunks_dropped_total
6076 .load(Ordering::Relaxed),
6077 3,
6078 "expected 3 dropped chunks (capacity=2, sent 5)",
6079 );
6080 }
6081
6082 /// Malformed REQUEST payload on the streaming fold: emits one
6083 /// terminal `UnknownVersion` frame and continues — same
6084 /// keep-the-adapter-alive contract as the unary fold.
6085 #[tokio::test]
6086 async fn streaming_fold_malformed_payload_emits_unknown_version_terminal() {
6087 struct NoopHandler;
6088 #[async_trait::async_trait]
6089 impl RpcStreamingHandler for NoopHandler {
6090 async fn call(
6091 &self,
6092 _ctx: RpcContext,
6093 _sink: RpcResponseSink,
6094 ) -> Result<(), RpcHandlerError> {
6095 Ok(())
6096 }
6097 }
6098 let (emit, captured) = capturing_async_emitter();
6099 let mut fold = RpcServerStreamingFold::new(Arc::new(NoopHandler), emit);
6100 // Garbage tail: valid meta + 0x00 svc_len → Truncated.
6101 let meta = EventMeta::new(DISPATCH_RPC_REQUEST, 0, 1, 1, 0);
6102 let mut buf = Vec::new();
6103 buf.extend_from_slice(&meta.to_bytes());
6104 buf.push(0x00);
6105 let ev = RedexEvent {
6106 entry: RedexEntry::new_heap(0, 0, buf.len() as u32, 0, 0),
6107 payload: bytes::Bytes::from(buf),
6108 };
6109 let result = fold.apply(&ev, &mut ());
6110 assert!(
6111 result.is_ok(),
6112 "malformed payload must NOT kill the adapter",
6113 );
6114 assert!(
6115 wait_until(|| !captured.lock().is_empty(), Duration::from_secs(2)).await,
6116 "synthetic UnknownVersion terminal must arrive",
6117 );
6118 let captured = captured.lock();
6119 assert_eq!(captured[0].2.status, RpcStatus::UnknownVersion);
6120 let hdr = captured[0]
6121 .2
6122 .headers
6123 .iter()
6124 .find(|(n, _)| n == HEADER_NRPC_STREAMING);
6125 assert!(hdr.is_some(), "malformed terminal must include end marker");
6126 }
6127
6128 // ====================================================================
6129 // Phase B — RpcStreamingRequestFold (server-side client-streaming)
6130 // ====================================================================
6131
6132 /// Build a REQUEST_CHUNK event for tests. Mirrors
6133 /// `rpc_request_event` / `rpc_stream_grant_event` shape.
6134 fn rpc_request_chunk_event(
6135 caller_origin: u64,
6136 call_id: u64,
6137 flags: u16,
6138 body: Vec<u8>,
6139 ) -> RedexEvent {
6140 let meta = EventMeta::new(DISPATCH_RPC_REQUEST_CHUNK, 0, caller_origin, call_id, 0);
6141 let payload = RpcRequestChunkPayload {
6142 call_id,
6143 flags,
6144 headers: vec![],
6145 body: body.into(),
6146 };
6147 let mut buf = Vec::new();
6148 buf.extend_from_slice(&meta.to_bytes());
6149 buf.extend_from_slice(&payload.encode());
6150 RedexEvent {
6151 entry: RedexEntry::new_heap(0, 0, buf.len() as u32, 0, 0),
6152 payload: bytes::Bytes::from(buf),
6153 }
6154 }
6155
6156 /// Collecting client-streaming handler: drains the stream into
6157 /// a Vec, returns an Ok response whose body is the count of
6158 /// chunks seen (8-byte LE). Captured chunk bodies are exposed
6159 /// via the `Arc<Mutex<Vec<Bytes>>>` so tests can assert
6160 /// ordering and content.
6161 struct CollectingClientStreamHandler {
6162 seen: Arc<Mutex<Vec<bytes::Bytes>>>,
6163 observed_cancel: Arc<AtomicBool>,
6164 }
6165 #[async_trait::async_trait]
6166 impl RpcClientStreamingHandler for CollectingClientStreamHandler {
6167 async fn call(
6168 &self,
6169 ctx: RpcStreamingContext,
6170 mut requests: RequestStream,
6171 ) -> Result<RpcResponsePayload, RpcHandlerError> {
6172 use futures::StreamExt;
6173 while let Some(chunk) = requests.next().await {
6174 self.seen.lock().push(chunk);
6175 }
6176 // Re-check cancellation after EOF so the test can
6177 // distinguish "clean REQUEST_END" from "CANCEL closed
6178 // the stream early".
6179 if ctx.cancellation.is_cancelled() {
6180 self.observed_cancel
6181 .store(true, std::sync::atomic::Ordering::SeqCst);
6182 }
6183 let count = self.seen.lock().len() as u64;
6184 Ok(RpcResponsePayload {
6185 status: RpcStatus::Ok,
6186 headers: vec![],
6187 body: Bytes::copy_from_slice(&count.to_le_bytes()),
6188 })
6189 }
6190 }
6191
6192 /// 1/6 — happy path: REQUEST + 3 REQUEST_CHUNKs (last has
6193 /// FLAG_END) delivers 4 bodies to the handler in order; the
6194 /// fold emits exactly one terminal RESPONSE carrying the
6195 /// handler's reply.
6196 #[tokio::test]
6197 async fn streaming_request_fold_collects_all_chunks_and_emits_terminal_response() {
6198 let seen = Arc::new(Mutex::new(Vec::new()));
6199 let observed_cancel = Arc::new(AtomicBool::new(false));
6200 let (emit, captured) = capturing_emitter();
6201 let mut fold = RpcStreamingRequestFold::new(
6202 Arc::new(CollectingClientStreamHandler {
6203 seen: seen.clone(),
6204 observed_cancel: observed_cancel.clone(),
6205 }),
6206 emit,
6207 );
6208 // REQUEST with the client-streaming flag, body = "a".
6209 let req = RpcRequestPayload {
6210 service: "agg".to_string(),
6211 deadline_ns: 0,
6212 flags: FLAG_RPC_CLIENT_STREAMING_REQUEST,
6213 headers: vec![],
6214 body: Bytes::from_static(b"a"),
6215 };
6216 fold.apply(&rpc_request_event(0xCAFE, 7, req), &mut ())
6217 .unwrap();
6218 // Wait until the sender is registered (handler task has
6219 // picked up the request and the apply path completed).
6220 assert!(
6221 wait_until(
6222 || fold.sender_keys().contains(&(0xCAFE, 7)),
6223 Duration::from_secs(1)
6224 )
6225 .await
6226 );
6227 // Three more chunks; last sets FLAG_REQUEST_END.
6228 fold.apply(
6229 &rpc_request_chunk_event(0xCAFE, 7, 0, b"b".to_vec()),
6230 &mut (),
6231 )
6232 .unwrap();
6233 fold.apply(
6234 &rpc_request_chunk_event(0xCAFE, 7, 0, b"c".to_vec()),
6235 &mut (),
6236 )
6237 .unwrap();
6238 fold.apply(
6239 &rpc_request_chunk_event(0xCAFE, 7, FLAG_RPC_REQUEST_END, b"d".to_vec()),
6240 &mut (),
6241 )
6242 .unwrap();
6243 // Handler should observe 4 bodies in order and emit one
6244 // terminal RESPONSE whose body encodes the count.
6245 assert!(
6246 wait_until(|| !captured.lock().is_empty(), Duration::from_secs(2)).await,
6247 "expected terminal RESPONSE"
6248 );
6249 let captured = captured.lock();
6250 assert_eq!(captured.len(), 1, "exactly one terminal RESPONSE");
6251 let (origin, call_id, resp) = &captured[0];
6252 assert_eq!(*origin, 0xCAFE);
6253 assert_eq!(*call_id, 7);
6254 assert_eq!(resp.status, RpcStatus::Ok);
6255 assert_eq!(resp.body.as_ref(), 4u64.to_le_bytes());
6256 // And the chunks landed in order.
6257 let seen = seen.lock();
6258 let collected: Vec<&[u8]> = seen.iter().map(|b| b.as_ref()).collect();
6259 assert_eq!(collected, vec![b"a", b"b", b"c", b"d"]);
6260 assert!(
6261 !observed_cancel.load(std::sync::atomic::Ordering::SeqCst),
6262 "clean REQUEST_END must NOT register as a cancellation"
6263 );
6264 }
6265
6266 /// 2/6 — degenerate case: initial REQUEST with both the
6267 /// client-streaming AND request-end flags set. Handler sees
6268 /// exactly one body (the REQUEST's own body) and EOF — the
6269 /// "one-item upload" fast path that saves a trailing CHUNK
6270 /// event.
6271 #[tokio::test]
6272 async fn streaming_request_fold_initial_request_with_end_flag_yields_single_item() {
6273 let seen = Arc::new(Mutex::new(Vec::new()));
6274 let observed_cancel = Arc::new(AtomicBool::new(false));
6275 let (emit, captured) = capturing_emitter();
6276 let mut fold = RpcStreamingRequestFold::new(
6277 Arc::new(CollectingClientStreamHandler {
6278 seen: seen.clone(),
6279 observed_cancel,
6280 }),
6281 emit,
6282 );
6283 let req = RpcRequestPayload {
6284 service: "agg".to_string(),
6285 deadline_ns: 0,
6286 flags: FLAG_RPC_CLIENT_STREAMING_REQUEST | FLAG_RPC_REQUEST_END,
6287 headers: vec![],
6288 body: Bytes::from_static(b"only"),
6289 };
6290 fold.apply(&rpc_request_event(1, 42, req), &mut ()).unwrap();
6291 assert!(
6292 wait_until(|| !captured.lock().is_empty(), Duration::from_secs(2)).await,
6293 "expected terminal RESPONSE"
6294 );
6295 let captured = captured.lock();
6296 assert_eq!(captured.len(), 1);
6297 assert_eq!(captured[0].2.status, RpcStatus::Ok);
6298 assert_eq!(captured[0].2.body.as_ref(), 1u64.to_le_bytes());
6299 assert_eq!(
6300 seen.lock()
6301 .iter()
6302 .map(|b| b.as_ref())
6303 .collect::<Vec<&[u8]>>(),
6304 vec![b"only" as &[u8]]
6305 );
6306 // Sender must NOT have been registered (initial-REQUEST-
6307 // with-END skips the map insert).
6308 assert!(fold.sender_keys().is_empty());
6309 }
6310
6311 /// 3/6 — CANCEL closes the request stream early, flips the
6312 /// cancellation token, and the spawned task overrides the
6313 /// handler's terminal with `RpcStatus::Cancelled`.
6314 #[tokio::test]
6315 async fn streaming_request_fold_cancel_closes_stream_and_overrides_terminal() {
6316 let seen = Arc::new(Mutex::new(Vec::new()));
6317 let observed_cancel = Arc::new(AtomicBool::new(false));
6318 let (emit, captured) = capturing_emitter();
6319 let mut fold = RpcStreamingRequestFold::new(
6320 Arc::new(CollectingClientStreamHandler {
6321 seen: seen.clone(),
6322 observed_cancel: observed_cancel.clone(),
6323 }),
6324 emit,
6325 );
6326 let req = RpcRequestPayload {
6327 service: "agg".to_string(),
6328 deadline_ns: 0,
6329 flags: FLAG_RPC_CLIENT_STREAMING_REQUEST,
6330 headers: vec![],
6331 body: Bytes::from_static(b"first"),
6332 };
6333 fold.apply(&rpc_request_event(2, 17, req), &mut ()).unwrap();
6334 // Wait for the handler to register, then one in-flight
6335 // CHUNK, then CANCEL before the handler ever finishes
6336 // draining.
6337 assert!(
6338 wait_until(
6339 || fold.sender_keys().contains(&(2, 17)),
6340 Duration::from_secs(1)
6341 )
6342 .await
6343 );
6344 fold.apply(
6345 &rpc_request_chunk_event(2, 17, 0, b"second".to_vec()),
6346 &mut (),
6347 )
6348 .unwrap();
6349 fold.apply(&rpc_cancel_event(2, 17), &mut ()).unwrap();
6350 // Terminal must arrive and must be Cancelled (CANCEL-wins
6351 // ordering, same as the response-side fold).
6352 assert!(
6353 wait_until(|| !captured.lock().is_empty(), Duration::from_secs(2)).await,
6354 "expected terminal RESPONSE"
6355 );
6356 let captured = captured.lock();
6357 assert_eq!(captured.len(), 1);
6358 assert_eq!(
6359 captured[0].2.status,
6360 RpcStatus::Cancelled,
6361 "CANCEL must override terminal status"
6362 );
6363 assert!(
6364 observed_cancel.load(std::sync::atomic::Ordering::SeqCst),
6365 "handler must observe cancellation token after stream EOF"
6366 );
6367 // Both maps must be clean post-cancel.
6368 assert!(fold.in_flight_keys().is_empty());
6369 assert!(fold.sender_keys().is_empty());
6370 }
6371
6372 /// 4/6 — handler returns `Err(RpcHandlerError::Application)`
6373 /// → terminal RESPONSE carries the application status code +
6374 /// message body.
6375 #[tokio::test]
6376 async fn streaming_request_fold_application_error_round_trips() {
6377 struct AppErrHandler;
6378 #[async_trait::async_trait]
6379 impl RpcClientStreamingHandler for AppErrHandler {
6380 async fn call(
6381 &self,
6382 _ctx: RpcStreamingContext,
6383 mut requests: RequestStream,
6384 ) -> Result<RpcResponsePayload, RpcHandlerError> {
6385 use futures::StreamExt;
6386 // Drain so the stream's EOF doesn't race the
6387 // error return.
6388 while requests.next().await.is_some() {}
6389 Err(RpcHandlerError::Application {
6390 code: 0xBEEF,
6391 message: "bad input".to_string(),
6392 })
6393 }
6394 }
6395 let (emit, captured) = capturing_emitter();
6396 let mut fold = RpcStreamingRequestFold::new(Arc::new(AppErrHandler), emit);
6397 let req = RpcRequestPayload {
6398 service: "agg".to_string(),
6399 deadline_ns: 0,
6400 flags: FLAG_RPC_CLIENT_STREAMING_REQUEST | FLAG_RPC_REQUEST_END,
6401 headers: vec![],
6402 body: Bytes::new(),
6403 };
6404 fold.apply(&rpc_request_event(3, 100, req), &mut ())
6405 .unwrap();
6406 assert!(
6407 wait_until(|| !captured.lock().is_empty(), Duration::from_secs(2)).await,
6408 "expected terminal RESPONSE"
6409 );
6410 let captured = captured.lock();
6411 assert_eq!(captured.len(), 1);
6412 assert_eq!(captured[0].2.status, RpcStatus::Application(0xBEEF));
6413 assert_eq!(captured[0].2.body.as_ref(), b"bad input");
6414 }
6415
6416 /// 5/6 — handler panic is caught by `catch_unwind`; terminal
6417 /// surfaces as `Internal` carrying the panic message. Same
6418 /// contract as the existing folds — a misbehaving handler
6419 /// can't take down the cortex adapter.
6420 #[tokio::test]
6421 async fn streaming_request_fold_handler_panic_surfaces_as_internal() {
6422 struct PanickyHandler;
6423 #[async_trait::async_trait]
6424 impl RpcClientStreamingHandler for PanickyHandler {
6425 async fn call(
6426 &self,
6427 _ctx: RpcStreamingContext,
6428 _requests: RequestStream,
6429 ) -> Result<RpcResponsePayload, RpcHandlerError> {
6430 panic!("intentional test panic");
6431 }
6432 }
6433 let (emit, captured) = capturing_emitter();
6434 let mut fold = RpcStreamingRequestFold::new(Arc::new(PanickyHandler), emit);
6435 let req = RpcRequestPayload {
6436 service: "agg".to_string(),
6437 deadline_ns: 0,
6438 flags: FLAG_RPC_CLIENT_STREAMING_REQUEST | FLAG_RPC_REQUEST_END,
6439 headers: vec![],
6440 body: Bytes::new(),
6441 };
6442 fold.apply(&rpc_request_event(4, 200, req), &mut ())
6443 .unwrap();
6444 assert!(
6445 wait_until(|| !captured.lock().is_empty(), Duration::from_secs(2)).await,
6446 "expected terminal RESPONSE"
6447 );
6448 let captured = captured.lock();
6449 assert_eq!(captured.len(), 1);
6450 assert_eq!(captured[0].2.status, RpcStatus::Internal);
6451 assert!(
6452 String::from_utf8_lossy(&captured[0].2.body).contains("intentional test panic"),
6453 "panic body should carry the panic message"
6454 );
6455 }
6456
6457 /// 6/6 — duplicate REQUEST with the same `(origin, call_id)`
6458 /// is refused with a synthetic `Internal` terminal frame and
6459 /// does NOT spawn a second handler. Mirror of the regression
6460 /// pinned in the unary + response-streaming folds.
6461 #[tokio::test]
6462 async fn streaming_request_fold_duplicate_request_refuses_without_double_dispatch() {
6463 let invocations = Arc::new(AtomicUsize::new(0));
6464 struct CountingHandler {
6465 invocations: Arc<AtomicUsize>,
6466 }
6467 #[async_trait::async_trait]
6468 impl RpcClientStreamingHandler for CountingHandler {
6469 async fn call(
6470 &self,
6471 _ctx: RpcStreamingContext,
6472 mut requests: RequestStream,
6473 ) -> Result<RpcResponsePayload, RpcHandlerError> {
6474 use futures::StreamExt;
6475 self.invocations.fetch_add(1, Ordering::SeqCst);
6476 // Slow handler to keep the call in-flight while
6477 // the duplicate REQUEST arrives.
6478 tokio::time::sleep(Duration::from_millis(80)).await;
6479 while requests.next().await.is_some() {}
6480 Ok(RpcResponsePayload {
6481 status: RpcStatus::Ok,
6482 headers: vec![],
6483 body: Bytes::new(),
6484 })
6485 }
6486 }
6487 let (emit, captured) = capturing_emitter();
6488 let mut fold = RpcStreamingRequestFold::new(
6489 Arc::new(CountingHandler {
6490 invocations: invocations.clone(),
6491 }),
6492 emit,
6493 );
6494 let req = RpcRequestPayload {
6495 service: "agg".to_string(),
6496 deadline_ns: 0,
6497 flags: FLAG_RPC_CLIENT_STREAMING_REQUEST,
6498 headers: vec![],
6499 body: Bytes::new(),
6500 };
6501 fold.apply(&rpc_request_event(5, 99, req.clone()), &mut ())
6502 .unwrap();
6503 assert!(
6504 wait_until(
6505 || fold.in_flight_keys().contains(&(5, 99)),
6506 Duration::from_secs(1)
6507 )
6508 .await
6509 );
6510 // Duplicate REQUEST: synthetic Internal terminal emitted,
6511 // handler invocation count must stay at 1.
6512 fold.apply(&rpc_request_event(5, 99, req), &mut ()).unwrap();
6513 assert!(
6514 wait_until(|| !captured.lock().is_empty(), Duration::from_secs(1)).await,
6515 "synthetic refusal terminal expected"
6516 );
6517 let refusal = captured.lock()[0].clone();
6518 assert_eq!(refusal.2.status, RpcStatus::Internal);
6519 assert!(String::from_utf8_lossy(&refusal.2.body).contains("duplicate"));
6520 // Finish the first handler so its terminal lands too.
6521 fold.apply(
6522 &rpc_request_chunk_event(5, 99, FLAG_RPC_REQUEST_END, vec![]),
6523 &mut (),
6524 )
6525 .unwrap();
6526 assert!(
6527 wait_until(|| captured.lock().len() >= 2, Duration::from_secs(2)).await,
6528 "first handler should still complete normally"
6529 );
6530 assert_eq!(
6531 invocations.load(Ordering::SeqCst),
6532 1,
6533 "duplicate REQUEST must NOT spawn a second handler",
6534 );
6535 }
6536}