dynamo_runtime/pipeline/
network.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Network layer for distributed communication
5//!
6//! Provides request distribution across multiple transport protocols:
7//! - HTTP/2 for standard deployments
8//! - TCP with length-prefixed protocol for high-performance scenarios
9//! - NATS for legacy/messaging-based deployments
10
11pub mod codec;
12pub mod egress;
13pub mod ingress;
14pub mod manager;
15pub mod tcp;
16
17use crate::SystemHealth;
18use std::sync::{Arc, OnceLock};
19
20use anyhow::Result;
21use async_trait::async_trait;
22use bytes::Bytes;
23use codec::{TwoPartCodec, TwoPartMessage, TwoPartMessageType};
24use derive_builder::Builder;
25use futures::StreamExt;
26// io::Cursor, TryStreamExt
27use super::{AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, ResponseStream};
28use serde::{Deserialize, Serialize, de::DeserializeOwned};
29
30use super::{
31    AsyncTransportEngine, Context, Data, Error, ManyIn, ManyOut, PipelineError, PipelineIO,
32    SegmentSource, ServiceBackend, ServiceEngine, SingleIn, Source, context,
33};
34use crate::metrics::MetricsHierarchy;
35use crate::metrics::prometheus_names::work_handler;
36use crate::protocols::maybe_error::MaybeError;
37use ingress::push_handler::WorkHandlerMetrics;
38use prometheus::{CounterVec, Histogram, IntCounter, IntCounterVec, IntGauge};
39
40/// Shared default maximum TCP message size across request-plane components.
41pub(crate) const DEFAULT_TCP_MAX_MESSAGE_SIZE: usize = 32 * 1024 * 1024;
42
43static TCP_MAX_MESSAGE_SIZE: OnceLock<usize> = OnceLock::new();
44static REQUEST_PLANE_PAYLOAD_CODEC: OnceLock<RequestPlanePayloadCodec> = OnceLock::new();
45
46/// Read the configured TCP max message size once and share it across client,
47/// server, and zero-copy decoder code paths.
48pub(crate) fn get_tcp_max_message_size() -> usize {
49    *TCP_MAX_MESSAGE_SIZE.get_or_init(|| {
50        std::env::var("DYN_TCP_MAX_MESSAGE_SIZE")
51            .ok()
52            .and_then(|s| s.parse::<usize>().ok())
53            .unwrap_or(DEFAULT_TCP_MAX_MESSAGE_SIZE)
54    })
55}
56
57#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
58#[serde(rename_all = "snake_case")]
59pub enum RequestPlanePayloadCodec {
60    /// The serde default deliberately remains JSON for wire compatibility with
61    /// control messages produced before the payload codec field existed.
62    #[default]
63    Json,
64    Msgpack,
65}
66
67impl RequestPlanePayloadCodec {
68    pub fn configured() -> Self {
69        *REQUEST_PLANE_PAYLOAD_CODEC.get_or_init(Self::from_env)
70    }
71
72    fn from_env() -> Self {
73        let value =
74            std::env::var(crate::config::environment_names::request_plane::DYN_REQUEST_PLANE_CODEC)
75                .ok();
76        Self::from_config_value(value.as_deref())
77    }
78
79    fn from_config_value(value: Option<&str>) -> Self {
80        match value {
81            None | Some("") | Some("msgpack") => Self::Msgpack,
82            Some("json") => Self::Json,
83            Some(other) => {
84                tracing::warn!(
85                    env_var =
86                        crate::config::environment_names::request_plane::DYN_REQUEST_PLANE_CODEC,
87                    value = other,
88                    "invalid request plane payload codec, defaulting to msgpack"
89                );
90                Self::Msgpack
91            }
92        }
93    }
94
95    pub fn name(&self) -> &'static str {
96        match self {
97            Self::Json => "json",
98            Self::Msgpack => "msgpack",
99        }
100    }
101
102    pub fn encode<T: Serialize>(&self, value: &T) -> Result<Vec<u8>> {
103        match self {
104            Self::Json => Ok(serde_json::to_vec(value)?),
105            Self::Msgpack => Ok(rmp_serde::to_vec_named(value)?),
106        }
107    }
108
109    pub fn decode<T: DeserializeOwned>(&self, bytes: &[u8]) -> Result<T> {
110        match self {
111            Self::Json => Ok(serde_json::from_slice(bytes)?),
112            Self::Msgpack => Ok(rmp_serde::from_slice(bytes)?),
113        }
114    }
115}
116
117pub trait Codable: PipelineIO + Serialize + for<'de> Deserialize<'de> {}
118impl<T: PipelineIO + Serialize + for<'de> Deserialize<'de>> Codable for T {}
119
120/// `WorkQueueConsumer` is a generic interface for a work queue that can be used to send and receive
121#[async_trait]
122pub trait WorkQueueConsumer {
123    async fn dequeue(&self) -> Result<Bytes, String>;
124}
125
126#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
127#[serde(rename_all = "snake_case")]
128pub enum StreamType {
129    Request,
130    Response,
131}
132
133#[derive(Debug, Clone, Serialize, Deserialize)]
134#[serde(rename_all = "snake_case")]
135pub(crate) enum RequestType {
136    SingleIn,
137    ManyIn,
138}
139
140#[derive(Debug, Clone, Serialize, Deserialize)]
141#[serde(rename_all = "snake_case")]
142pub(crate) enum ResponseType {
143    SingleOut,
144    ManyOut,
145}
146
147#[derive(Debug, Clone, Serialize, Deserialize)]
148pub(crate) struct RequestControlMessage {
149    pub(crate) id: String,
150    pub(crate) request_type: RequestType,
151    pub(crate) response_type: ResponseType,
152    #[serde(default)]
153    pub(crate) payload_codec: RequestPlanePayloadCodec,
154    pub(crate) connection_info: ConnectionInfo,
155    #[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")]
156    pub(crate) metadata: std::collections::BTreeMap<String, String>,
157    /// Wall-clock send timestamp (nanos since UNIX epoch) for transport latency breakdown.
158    /// Uses `SystemTime` so accuracy depends on NTP sync between frontend and backend hosts.
159    /// Reliable for single-machine profiling; treat cross-host values as approximate.
160    #[serde(default, skip_serializing_if = "Option::is_none")]
161    pub(crate) frontend_send_ts_ns: Option<u64>,
162    /// For bidirectional dispatch (`request_type == ManyIn`): connection info the
163    /// worker dials back to in order to receive subsequent request frames. `None`
164    /// for the unary path, which is the wire-compatible default.
165    #[serde(default, skip_serializing_if = "Option::is_none")]
166    pub(crate) request_stream_connection_info: Option<ConnectionInfo>,
167}
168
169#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
170#[serde(rename_all = "snake_case")]
171pub enum ControlMessage {
172    Stop,
173    Kill,
174    Sentinel,
175}
176
177/// This is the first message in a `ResponseStream`. This is not a message that gets process
178/// by the general pipeline, but is a control message that is awaited before the
179/// [`AsyncEngine::generate`] method is allowed to return.
180///
181/// If an error is present, the [`AsyncEngine::generate`] method will return the error instead
182/// of returning the `ResponseStream`.
183#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
184pub struct ResponseStreamPrologue {
185    error: Option<String>,
186}
187
188pub type StreamProvider<T> = tokio::sync::oneshot::Receiver<Result<T, String>>;
189
190/// Owning `Drop` here (rather than on `RegisteredStream`) lets `into_parts()`
191/// move the public fields out by plain destructure.
192struct Cleanup(Option<Box<dyn FnOnce() + Send + 'static>>);
193
194impl Drop for Cleanup {
195    fn drop(&mut self) {
196        if let Some(f) = self.0.take() {
197            f();
198        }
199    }
200}
201
202/// Awaitable handle for a stream sender or receiver. Drop without calling
203/// [`into_parts()`] runs the optional cleanup closure, removing the
204/// registration from the stream server's maps.
205pub struct RegisteredStream<T> {
206    pub connection_info: ConnectionInfo,
207    pub stream_provider: StreamProvider<T>,
208    cleanup: Cleanup,
209}
210
211impl<T> std::fmt::Debug for RegisteredStream<T> {
212    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
213        f.debug_struct("RegisteredStream")
214            .field("connection_info", &self.connection_info)
215            .finish_non_exhaustive()
216    }
217}
218
219impl<T> RegisteredStream<T> {
220    pub(crate) fn new(connection_info: ConnectionInfo, stream_provider: StreamProvider<T>) -> Self {
221        Self {
222            connection_info,
223            stream_provider,
224            cleanup: Cleanup(None),
225        }
226    }
227
228    pub(crate) fn with_cleanup<F>(mut self, cleanup: F) -> Self
229    where
230        F: FnOnce() + Send + 'static,
231    {
232        self.cleanup.0 = Some(Box::new(cleanup));
233        self
234    }
235
236    /// Consume the registration, disarming the RAII cleanup. Caller takes
237    /// responsibility for cleanup if the stream provider is never awaited.
238    pub fn into_parts(self) -> (ConnectionInfo, StreamProvider<T>) {
239        let Self {
240            connection_info,
241            stream_provider,
242            mut cleanup,
243        } = self;
244        cleanup.0.take();
245        (connection_info, stream_provider)
246    }
247}
248
249/// After registering a stream, the [`PendingConnections`] object is returned to the caller. This
250/// object can be used to await the connection to be established.
251pub struct PendingConnections {
252    pub send_stream: Option<RegisteredStream<StreamSender>>,
253    pub recv_stream: Option<RegisteredStream<StreamReceiver>>,
254}
255
256impl PendingConnections {
257    pub fn into_parts(
258        self,
259    ) -> (
260        Option<RegisteredStream<StreamSender>>,
261        Option<RegisteredStream<StreamReceiver>>,
262    ) {
263        (self.send_stream, self.recv_stream)
264    }
265}
266
267/// A [`ResponseService`] implements a services in which a context a specific subject with will
268/// be associated with a stream of responses.
269#[async_trait::async_trait]
270pub trait ResponseService {
271    async fn register(&self, options: StreamOptions) -> PendingConnections;
272}
273
274#[cfg(test)]
275mod registered_stream_tests {
276    use super::*;
277    use std::sync::atomic::{AtomicBool, Ordering};
278
279    fn dummy_conn_info() -> ConnectionInfo {
280        ConnectionInfo {
281            transport: "test".to_string(),
282            info: "{}".to_string(),
283        }
284    }
285
286    /// Drop without `into_parts()` must run the cleanup closure.
287    #[test]
288    fn drop_runs_cleanup() {
289        let flag = Arc::new(AtomicBool::new(false));
290        let flag_clone = flag.clone();
291
292        let (_tx, rx) = tokio::sync::oneshot::channel::<Result<(), String>>();
293        let stream = RegisteredStream::new(dummy_conn_info(), rx).with_cleanup(move || {
294            flag_clone.store(true, Ordering::SeqCst);
295        });
296
297        drop(stream);
298        assert!(
299            flag.load(Ordering::SeqCst),
300            "cleanup must fire when RegisteredStream is dropped"
301        );
302    }
303
304    /// `into_parts()` must disarm the cleanup. After the call, dropping the
305    /// returned halves must NOT trigger the closure -- the caller has taken
306    /// ownership of cleanup responsibility.
307    #[test]
308    fn into_parts_disarms_cleanup() {
309        let flag = Arc::new(AtomicBool::new(false));
310        let flag_clone = flag.clone();
311
312        let (_tx, rx) = tokio::sync::oneshot::channel::<Result<(), String>>();
313        let stream = RegisteredStream::new(dummy_conn_info(), rx).with_cleanup(move || {
314            flag_clone.store(true, Ordering::SeqCst);
315        });
316
317        let (conn, provider) = stream.into_parts();
318        drop(conn);
319        drop(provider);
320
321        assert!(
322            !flag.load(Ordering::SeqCst),
323            "into_parts() must disarm the cleanup closure"
324        );
325    }
326
327    /// `RegisteredStream` with no cleanup configured must drop cleanly.
328    #[test]
329    fn drop_without_cleanup_is_a_noop() {
330        let (_tx, rx) = tokio::sync::oneshot::channel::<Result<(), String>>();
331        let stream: RegisteredStream<()> = RegisteredStream::new(dummy_conn_info(), rx);
332        drop(stream); // must not panic; nothing observable to assert beyond that
333    }
334}
335
336// #[derive(Debug, Clone, Serialize, Deserialize)]
337// struct Handshake {
338//     request_id: String,
339//     worker_id: Option<String>,
340//     error: Option<String>,
341// }
342
343// impl Handshake {
344//     pub fn validate(&self) -> Result<(), String> {
345//         if let Some(e) = &self.error {
346//             return Err(e.clone());
347//         }
348//         Ok(())
349//     }
350// }
351
352// this probably needs to be come a ResponseStreamSender
353// since the prologue in this scenario sender telling the receiver
354// that all is good and it's ready to send
355//
356// in the RequestStreamSender, the prologue would be coming from the
357// receiver, so the sender would have to await the prologue which if
358// was not an error, would indicate the RequestStreamReceiver is read
359// to receive data.
360pub struct StreamSender {
361    tx: tokio::sync::mpsc::Sender<TwoPartMessage>,
362    prologue: Option<ResponseStreamPrologue>,
363}
364
365impl StreamSender {
366    pub async fn send(&self, data: Bytes) -> Result<()> {
367        Ok(self.tx.send(TwoPartMessage::from_data(data)).await?)
368    }
369
370    pub async fn send_control(&self, control: ControlMessage) -> Result<()> {
371        let bytes = serde_json::to_vec(&control)?;
372        Ok(self
373            .tx
374            .send(TwoPartMessage::from_header(bytes.into()))
375            .await?)
376    }
377
378    #[allow(clippy::needless_update)]
379    pub async fn send_prologue(&mut self, error: Option<String>) -> Result<(), String> {
380        // leaving the original logic in place for now
381        // error overrides the dissolved prologue, but the only field on `ResponseStreamPrologue` is `error`
382        // so the second argument can never be used, and the value of error passed by the caller would always be used
383        if let Some(_prologue) = self.prologue.take() {
384            // let prologue = ResponseStreamPrologue { error, ..prologue };
385            let prologue = ResponseStreamPrologue { error };
386            let header_bytes: Bytes = match serde_json::to_vec(&prologue) {
387                Ok(b) => b.into(),
388                Err(err) => {
389                    tracing::error!(%err, "send_prologue: ResponseStreamPrologue did not serialize to a JSON array");
390                    return Err("Invalid prologue".to_string());
391                }
392            };
393            self.tx
394                .send(TwoPartMessage::from_header(header_bytes))
395                .await
396                .map_err(|e| e.to_string())?;
397        } else {
398            panic!("Prologue already sent; or not set; logic error");
399        }
400        Ok(())
401    }
402}
403
404pub struct StreamReceiver {
405    rx: tokio::sync::mpsc::Receiver<Bytes>,
406}
407
408/// Connection Info is encoded as JSON and then again serialized has part of the Transport
409/// Layer. The double serialization is not performance critical as it is only done once per
410/// connection. The primary reason storing the ConnecitonInfo has a JSON string is for type
411/// erasure. The Transport Layer will check the [`ConnectionInfo::transport`] type and then
412/// route it to the appropriate instance of the Transport, which will then deserialize the
413/// [`ConnectionInfo::info`] field to its internal connection info object.
414///
415/// Optionally, this object could become strongly typed for which all possible combinations
416/// of transport and connection info would need to be enumerated.
417#[derive(Debug, Clone, Serialize, Deserialize)]
418pub struct ConnectionInfo {
419    pub transport: String,
420    pub info: String,
421}
422
423/// Default number of frames buffered between the data-plane socket task and the
424/// engine consumer/producer for a single stream. Preserves the historically
425/// hard-coded mpsc channel capacity used by the TCP transport.
426pub const DEFAULT_SEND_BUFFER_COUNT: usize = 64;
427
428/// When registering a new TransportStream on the server, the caller specifies if the
429/// stream is a sender, receiver or both.
430///
431/// Senders and Receivers are with share a Context, but result in separate tcp socket
432/// connections to the server. Internally, we may use bcast channels to coordinate the
433/// internal control messages between the sender and receiver socket connections.
434#[derive(Clone, Builder)]
435pub struct StreamOptions {
436    /// Context
437    pub context: Arc<dyn AsyncEngineContext>,
438
439    /// Register with the server that this connection will have a server-side Sender
440    /// that can be picked up by the Request/Forward pipeline. The downstream side
441    /// dials in via [`crate::pipeline::network::tcp::client::TcpClient::create_request_stream`]
442    /// to receive the frames the server pushes.
443    pub enable_request_stream: bool,
444
445    /// Register with the server that this connection will have a server-side Receiver
446    /// that can be picked up by the Response/Reverse pipeline
447    pub enable_response_stream: bool,
448
449    /// The number of frames buffered between the data-plane socket task and the
450    /// engine consumer/producer before backpressure kicks in. Drives the mpsc
451    /// channel capacity for the per-stream buffer in the TCP transport.
452    #[builder(default = "DEFAULT_SEND_BUFFER_COUNT")]
453    pub send_buffer_count: usize,
454
455    /// The number of messages to buffer before blocking
456    #[builder(default = "8")]
457    pub recv_buffer_count: usize,
458}
459
460impl StreamOptions {
461    pub fn builder() -> StreamOptionsBuilder {
462        StreamOptionsBuilder::default()
463    }
464}
465
466pub struct Egress<Req: PipelineIO, Resp: PipelineIO> {
467    transport_engine: Arc<dyn AsyncTransportEngine<Req, Resp>>,
468}
469
470#[cfg(test)]
471mod tests {
472    use super::{
473        DEFAULT_SEND_BUFFER_COUNT, NetworkStreamWrapper, RequestControlMessage,
474        RequestPlanePayloadCodec, RequestType, ResponseType, StreamOptions,
475    };
476    use crate::engine::AsyncEngineContextProvider;
477    use crate::pipeline::Context;
478    use serde::{Deserialize, Serialize};
479
480    #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
481    struct TestPayload {
482        id: u64,
483        text: String,
484        tokens: Vec<u32>,
485    }
486
487    #[test]
488    fn stream_options_send_buffer_count_defaults_to_64() {
489        let context = Context::new(());
490        let options = StreamOptions::builder()
491            .context(context.context())
492            .enable_request_stream(true)
493            .enable_response_stream(true)
494            .build()
495            .expect("stream options should build");
496
497        assert_eq!(DEFAULT_SEND_BUFFER_COUNT, 64);
498        assert_eq!(options.send_buffer_count, DEFAULT_SEND_BUFFER_COUNT);
499    }
500
501    #[test]
502    fn stream_options_send_buffer_count_overrides_default() {
503        let context = Context::new(());
504        let options = StreamOptions::builder()
505            .context(context.context())
506            .enable_request_stream(true)
507            .enable_response_stream(true)
508            .send_buffer_count(128)
509            .build()
510            .expect("stream options should build");
511
512        assert_eq!(options.send_buffer_count, 128);
513    }
514
515    #[test]
516    fn legacy_frontend_control_message_defaults_payload_codec_to_json() {
517        let json = r#"{
518            "id": "request-123",
519            "request_type": "single_in",
520            "response_type": "many_out",
521            "connection_info": {
522                "transport": "tcp",
523                "info": "{}"
524            }
525        }"#;
526
527        let message: RequestControlMessage =
528            serde_json::from_str(json).expect("control message should deserialize");
529
530        assert_eq!(message.id, "request-123");
531        assert!(matches!(message.request_type, RequestType::SingleIn));
532        assert!(matches!(message.response_type, ResponseType::ManyOut));
533        assert_eq!(message.payload_codec, RequestPlanePayloadCodec::Json);
534        assert_eq!(message.connection_info.transport, "tcp");
535        assert_eq!(message.connection_info.info, "{}");
536        assert!(message.metadata.is_empty());
537        assert!(message.frontend_send_ts_ns.is_none());
538
539        let payload = br#"{"id":7,"text":"legacy","tokens":[1,2]}"#;
540        let decoded: TestPayload = message
541            .payload_codec
542            .decode(payload)
543            .expect("worker should decode the legacy frontend's JSON payload");
544        assert_eq!(
545            decoded,
546            TestPayload {
547                id: 7,
548                text: "legacy".to_string(),
549                tokens: vec![1, 2],
550            }
551        );
552    }
553
554    #[test]
555    fn request_control_message_decodes_msgpack_payload_codec() {
556        let json = r#"{
557            "id": "request-123",
558            "request_type": "single_in",
559            "response_type": "many_out",
560            "payload_codec": "msgpack",
561            "connection_info": {
562                "transport": "tcp",
563                "info": "{}"
564            }
565        }"#;
566
567        let message: RequestControlMessage =
568            serde_json::from_str(json).expect("control message should deserialize");
569
570        assert_eq!(message.payload_codec, RequestPlanePayloadCodec::Msgpack);
571    }
572
573    #[test]
574    fn request_plane_payload_codec_configuration_defaults_to_msgpack() {
575        assert_eq!(
576            RequestPlanePayloadCodec::from_config_value(None),
577            RequestPlanePayloadCodec::Msgpack
578        );
579        assert_eq!(
580            RequestPlanePayloadCodec::from_config_value(Some("")),
581            RequestPlanePayloadCodec::Msgpack
582        );
583        assert_eq!(
584            RequestPlanePayloadCodec::from_config_value(Some("invalid")),
585            RequestPlanePayloadCodec::Msgpack
586        );
587    }
588
589    #[test]
590    fn request_plane_payload_codec_configuration_honors_explicit_overrides() {
591        assert_eq!(
592            RequestPlanePayloadCodec::from_config_value(Some("json")),
593            RequestPlanePayloadCodec::Json
594        );
595        assert_eq!(
596            RequestPlanePayloadCodec::from_config_value(Some("msgpack")),
597            RequestPlanePayloadCodec::Msgpack
598        );
599    }
600
601    #[test]
602    fn request_plane_payload_codec_round_trips_response_wrapper_json_and_msgpack() {
603        let wrapper = NetworkStreamWrapper {
604            data: Some(TestPayload {
605                id: 42,
606                text: "line\nquote\"slash\\unicode 中".to_string(),
607                tokens: vec![1, 2, 3, 65535],
608            }),
609            complete_final: false,
610        };
611
612        for codec in [
613            RequestPlanePayloadCodec::Json,
614            RequestPlanePayloadCodec::Msgpack,
615        ] {
616            let encoded = codec.encode(&wrapper).expect("wrapper should encode");
617            let decoded: NetworkStreamWrapper<TestPayload> =
618                codec.decode(&encoded).expect("wrapper should decode");
619            assert_eq!(decoded, wrapper);
620        }
621    }
622}
623
624#[async_trait]
625impl<T: Data, U: Data> AsyncEngine<SingleIn<T>, ManyOut<U>, Error>
626    for Egress<SingleIn<T>, ManyOut<U>>
627where
628    T: Data + Serialize,
629    U: for<'de> Deserialize<'de> + Data,
630{
631    async fn generate(&self, request: SingleIn<T>) -> Result<ManyOut<U>, Error> {
632        self.transport_engine.generate(request).await
633    }
634}
635
636/// Result of encoding one response item for the request plane.
637pub struct EncodedResponseFrame {
638    pub bytes: Bytes,
639    pub is_error: bool,
640    /// Stop consuming the engine stream after publishing this frame. The
641    /// normal complete-final frame is still sent.
642    pub stop_stream: bool,
643}
644
645/// Converts request-plane bytes into the item consumed by an ingress engine.
646pub trait IngressRequestDecoder<T>: Send + Sync + 'static
647where
648    T: Data,
649{
650    fn decode_request(
651        &self,
652        payload_codec: RequestPlanePayloadCodec,
653        bytes: Bytes,
654    ) -> impl std::future::Future<Output = std::result::Result<T, PipelineError>> + Send;
655}
656
657/// Converts an ingress engine response into its complete on-wire frame.
658pub trait IngressResponseEncoder<U>: Send + Sync + 'static
659where
660    U: Data,
661{
662    fn encode_response(
663        &self,
664        payload_codec: RequestPlanePayloadCodec,
665        response: Option<U>,
666        complete_final: bool,
667    ) -> impl std::future::Future<Output = std::result::Result<EncodedResponseFrame, PipelineError>> + Send;
668}
669
670/// Complete request/response payload adapter for an ingress engine.
671pub trait IngressPayloadAdapter<T, U>:
672    IngressRequestDecoder<T> + IngressResponseEncoder<U>
673where
674    T: Data,
675    U: Data,
676{
677}
678
679impl<T, U, Adapter> IngressPayloadAdapter<T, U> for Adapter
680where
681    T: Data,
682    U: Data,
683    Adapter: IngressRequestDecoder<T> + IngressResponseEncoder<U>,
684{
685}
686
687/// Default adapter for ordinary Rust request and response types.
688#[derive(Debug, Default)]
689pub struct SerdeIngressPayloadAdapter;
690
691impl<T> IngressRequestDecoder<T> for SerdeIngressPayloadAdapter
692where
693    T: Data + DeserializeOwned,
694{
695    #[inline]
696    fn decode_request(
697        &self,
698        payload_codec: RequestPlanePayloadCodec,
699        bytes: Bytes,
700    ) -> impl std::future::Future<Output = std::result::Result<T, PipelineError>> + Send {
701        let decoded = payload_codec.decode(&bytes).map_err(|err| {
702            PipelineError::DeserializationError(format!(
703                "Failed deserializing {} request payload: {}",
704                payload_codec.name(),
705                err
706            ))
707        });
708        std::future::ready(decoded)
709    }
710}
711
712impl<U> IngressResponseEncoder<U> for SerdeIngressPayloadAdapter
713where
714    U: Data + Serialize + MaybeError,
715{
716    #[inline]
717    fn encode_response(
718        &self,
719        payload_codec: RequestPlanePayloadCodec,
720        response: Option<U>,
721        complete_final: bool,
722    ) -> impl std::future::Future<Output = std::result::Result<EncodedResponseFrame, PipelineError>> + Send
723    {
724        let is_error = response
725            .as_ref()
726            .is_some_and(|response| response.err().is_some());
727        let wrapper = NetworkStreamWrapper {
728            data: response,
729            complete_final,
730        };
731        let encoded = payload_codec.encode(&wrapper).map_err(|err| {
732            PipelineError::SerializationError(format!(
733                "Failed serializing {} request-plane response: {}",
734                payload_codec.name(),
735                err
736            ))
737        });
738        std::future::ready(encoded.map(|bytes| EncodedResponseFrame {
739            bytes: bytes.into(),
740            is_error,
741            stop_stream: false,
742        }))
743    }
744}
745
746pub struct Ingress<Req: PipelineIO, Resp: PipelineIO, Adapter = SerdeIngressPayloadAdapter> {
747    segment: OnceLock<Arc<SegmentSource<Req, Resp>>>,
748    metrics: OnceLock<Arc<WorkHandlerMetrics>>,
749    /// Endpoint-specific notifier for health check timer resets
750    endpoint_health_check_notifier: OnceLock<Arc<tokio::sync::Notify>>,
751    payload_adapter: Arc<Adapter>,
752}
753
754impl<Req: PipelineIO + Sync, Resp: PipelineIO> Ingress<Req, Resp> {
755    pub fn new() -> Arc<Self> {
756        Ingress::new_with_adapter(SerdeIngressPayloadAdapter)
757    }
758
759    pub fn link(segment: Arc<SegmentSource<Req, Resp>>) -> Result<Arc<Self>> {
760        let ingress = Ingress::new();
761        ingress.attach(segment)?;
762        Ok(ingress)
763    }
764
765    pub fn for_pipeline(segment: Arc<SegmentSource<Req, Resp>>) -> Result<Arc<Self>> {
766        let ingress = Ingress::new();
767        ingress.attach(segment)?;
768        Ok(ingress)
769    }
770
771    pub fn for_engine(engine: ServiceEngine<Req, Resp>) -> Result<Arc<Self>> {
772        Self::for_engine_with_adapter(engine, SerdeIngressPayloadAdapter)
773    }
774}
775
776impl<Req, Resp, Adapter> Ingress<Req, Resp, Adapter>
777where
778    Req: PipelineIO + Sync,
779    Resp: PipelineIO,
780    Adapter: Send + Sync + 'static,
781{
782    pub fn new_with_adapter(payload_adapter: Adapter) -> Arc<Self> {
783        Arc::new(Self {
784            segment: OnceLock::new(),
785            metrics: OnceLock::new(),
786            endpoint_health_check_notifier: OnceLock::new(),
787            payload_adapter: Arc::new(payload_adapter),
788        })
789    }
790
791    pub fn attach(&self, segment: Arc<SegmentSource<Req, Resp>>) -> Result<()> {
792        self.segment
793            .set(segment)
794            .map_err(|_| anyhow::anyhow!("Segment already set"))
795    }
796
797    pub fn add_metrics(
798        &self,
799        endpoint: &crate::component::Endpoint,
800        metrics_labels: Option<&[(&str, &str)]>,
801    ) -> Result<()> {
802        let metrics = WorkHandlerMetrics::from_endpoint(endpoint, metrics_labels)
803            .map_err(|e| anyhow::anyhow!("Failed to create work handler metrics: {}", e))?;
804
805        // Register global transport breakdown metrics (idempotent)
806        crate::metrics::work_handler_perf::ensure_work_handler_perf_metrics_registered(
807            endpoint.get_metrics_registry(),
808        );
809
810        // Register worker-pool saturation metrics (idempotent). These are
811        // process-global and shared across all endpoints attached to the
812        // same shared TCP server.
813        crate::metrics::work_handler_pool::ensure_work_handler_pool_metrics_registered(
814            endpoint.get_metrics_registry(),
815        );
816
817        self.metrics
818            .set(Arc::new(metrics))
819            .map_err(|_| anyhow::anyhow!("Metrics already set"))
820    }
821
822    pub fn for_engine_with_adapter(
823        engine: ServiceEngine<Req, Resp>,
824        payload_adapter: Adapter,
825    ) -> Result<Arc<Self>> {
826        let frontend = SegmentSource::<Req, Resp>::new();
827        let backend = ServiceBackend::from_engine(engine);
828
829        // create the pipeline
830        let pipeline = frontend.link(backend)?.link_terminal(frontend)?;
831
832        let ingress = Ingress::new_with_adapter(payload_adapter);
833        ingress.attach(pipeline)?;
834
835        Ok(ingress)
836    }
837
838    /// Helper method to access metrics if available
839    fn metrics(&self) -> Option<&Arc<WorkHandlerMetrics>> {
840        self.metrics.get()
841    }
842}
843
844#[async_trait]
845pub trait PushWorkHandler: Send + Sync {
846    async fn handle_payload(
847        &self,
848        payload: Bytes,
849        request_id: Option<String>,
850    ) -> Result<(), PipelineError>;
851
852    /// Add metrics to the handler
853    fn add_metrics(
854        &self,
855        endpoint: &crate::component::Endpoint,
856        metrics_labels: Option<&[(&str, &str)]>,
857    ) -> Result<()>;
858
859    /// Set the endpoint-specific notifier for health check timer resets
860    fn set_endpoint_health_check_notifier(
861        &self,
862        _notifier: Arc<tokio::sync::Notify>,
863    ) -> Result<()> {
864        // Default implementation for backwards compatibility
865        Ok(())
866    }
867}
868
869/*
870/// `NetworkStreamWrapper` is a simple wrapper used to detect proper stream termination
871/// in network communication between ingress and egress components.
872///
873/// **Purpose**: This wrapper solves the problem of detecting whether a stream ended
874/// gracefully or was cut off prematurely (e.g., due to network issues).
875///
876/// **Design Rationale**:
877/// - Cannot use `Annotated` directly because the generic type `U` varies:
878///   - Sometimes `U = Annotated<...>`
879///   - Sometimes `U = LLMEngineOutput<...>`
880/// - Using `Annotated` would require double-wrapping like `Annotated<Annotated<...>>`
881/// - A simple wrapper is cleaner and more straightforward
882///
883/// **Stream Flow**:
884/// ```
885/// At AsyncEngine:
886///   response 1 -> response 2 -> response 3 -> <end>
887///
888/// Between ingress/egress:
889///   response 1 <end=false> -> response 2 <end=false> -> response 3 <end=false> -> (null) <end=true>
890///
891/// At client:
892///   response 1 -> response 2 -> response 3 -> <end>
893/// ```
894///
895/// **Error Handling**:
896/// If the stream is cut off before proper termination, the egress is responsible for
897/// injecting an error response to communicate the incomplete stream to the client:
898/// ```
899/// At AsyncEngine:
900///   response 1 -> ... <without end flag>
901///
902/// At egress:
903///   response 1 <end=false> -> <stream ended without end flag -> convert to error>
904///
905/// At client:
906///   response 1 -> error response
907/// ```
908///
909/// The detection must be done at egress level because premature stream termination
910/// can be due to network issues that only the egress component can detect.
911*/
912/// TODO: Detect end-of-stream using Server-Sent Events (SSE). This will be removed.
913#[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
914pub struct NetworkStreamWrapper<U> {
915    #[serde(skip_serializing_if = "Option::is_none")]
916    pub data: Option<U>,
917    pub complete_final: bool,
918}
dynamo_runtime/pipeline/network.rs

dynamo_runtime/pipeline/
network.rs
