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(crate) enum RequestPlanePayloadCodec {
60    // TODO(jthomson04): Migrate the default to Msgpack after the 1.3 release.
61    #[default]
62    Json,
63    Msgpack,
64}
65
66impl RequestPlanePayloadCodec {
67    pub(crate) fn configured() -> Self {
68        *REQUEST_PLANE_PAYLOAD_CODEC.get_or_init(Self::from_env)
69    }
70
71    fn from_env() -> Self {
72        match std::env::var(
73            crate::config::environment_names::request_plane::DYN_REQUEST_PLANE_CODEC,
74        )
75        .as_deref()
76        {
77            Err(_) | Ok("") | Ok("json") => Self::Json,
78            Ok("msgpack") => Self::Msgpack,
79            Ok(other) => {
80                tracing::warn!(
81                    env_var =
82                        crate::config::environment_names::request_plane::DYN_REQUEST_PLANE_CODEC,
83                    value = other,
84                    "invalid request plane payload codec, defaulting to json"
85                );
86                Self::Json
87            }
88        }
89    }
90
91    pub(crate) fn name(&self) -> &'static str {
92        match self {
93            Self::Json => "json",
94            Self::Msgpack => "msgpack",
95        }
96    }
97
98    pub(crate) fn encode<T: Serialize>(&self, value: &T) -> Result<Vec<u8>> {
99        match self {
100            Self::Json => Ok(serde_json::to_vec(value)?),
101            Self::Msgpack => Ok(rmp_serde::to_vec_named(value)?),
102        }
103    }
104
105    pub(crate) fn decode<T: DeserializeOwned>(&self, bytes: &[u8]) -> Result<T> {
106        match self {
107            Self::Json => Ok(serde_json::from_slice(bytes)?),
108            Self::Msgpack => Ok(rmp_serde::from_slice(bytes)?),
109        }
110    }
111}
112
113pub trait Codable: PipelineIO + Serialize + for<'de> Deserialize<'de> {}
114impl<T: PipelineIO + Serialize + for<'de> Deserialize<'de>> Codable for T {}
115
116/// `WorkQueueConsumer` is a generic interface for a work queue that can be used to send and receive
117#[async_trait]
118pub trait WorkQueueConsumer {
119    async fn dequeue(&self) -> Result<Bytes, String>;
120}
121
122#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
123#[serde(rename_all = "snake_case")]
124pub enum StreamType {
125    Request,
126    Response,
127}
128
129#[derive(Debug, Clone, Serialize, Deserialize)]
130#[serde(rename_all = "snake_case")]
131pub(crate) enum RequestType {
132    SingleIn,
133    ManyIn,
134}
135
136#[derive(Debug, Clone, Serialize, Deserialize)]
137#[serde(rename_all = "snake_case")]
138pub(crate) enum ResponseType {
139    SingleOut,
140    ManyOut,
141}
142
143#[derive(Debug, Clone, Serialize, Deserialize)]
144pub(crate) struct RequestControlMessage {
145    pub(crate) id: String,
146    pub(crate) request_type: RequestType,
147    pub(crate) response_type: ResponseType,
148    #[serde(default)]
149    pub(crate) payload_codec: RequestPlanePayloadCodec,
150    pub(crate) connection_info: ConnectionInfo,
151    #[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")]
152    pub(crate) metadata: std::collections::BTreeMap<String, String>,
153    /// Wall-clock send timestamp (nanos since UNIX epoch) for transport latency breakdown.
154    /// Uses `SystemTime` so accuracy depends on NTP sync between frontend and backend hosts.
155    /// Reliable for single-machine profiling; treat cross-host values as approximate.
156    #[serde(default, skip_serializing_if = "Option::is_none")]
157    pub(crate) frontend_send_ts_ns: Option<u64>,
158    /// For bidirectional dispatch (`request_type == ManyIn`): connection info the
159    /// worker dials back to in order to receive subsequent request frames. `None`
160    /// for the unary path, which is the wire-compatible default.
161    #[serde(default, skip_serializing_if = "Option::is_none")]
162    pub(crate) request_stream_connection_info: Option<ConnectionInfo>,
163}
164
165#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
166#[serde(rename_all = "snake_case")]
167pub enum ControlMessage {
168    Stop,
169    Kill,
170    Sentinel,
171}
172
173/// This is the first message in a `ResponseStream`. This is not a message that gets process
174/// by the general pipeline, but is a control message that is awaited before the
175/// [`AsyncEngine::generate`] method is allowed to return.
176///
177/// If an error is present, the [`AsyncEngine::generate`] method will return the error instead
178/// of returning the `ResponseStream`.
179#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
180pub struct ResponseStreamPrologue {
181    error: Option<String>,
182}
183
184pub type StreamProvider<T> = tokio::sync::oneshot::Receiver<Result<T, String>>;
185
186/// Owning `Drop` here (rather than on `RegisteredStream`) lets `into_parts()`
187/// move the public fields out by plain destructure.
188struct Cleanup(Option<Box<dyn FnOnce() + Send + 'static>>);
189
190impl Drop for Cleanup {
191    fn drop(&mut self) {
192        if let Some(f) = self.0.take() {
193            f();
194        }
195    }
196}
197
198/// Awaitable handle for a stream sender or receiver. Drop without calling
199/// [`into_parts()`] runs the optional cleanup closure, removing the
200/// registration from the stream server's maps.
201pub struct RegisteredStream<T> {
202    pub connection_info: ConnectionInfo,
203    pub stream_provider: StreamProvider<T>,
204    cleanup: Cleanup,
205}
206
207impl<T> std::fmt::Debug for RegisteredStream<T> {
208    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
209        f.debug_struct("RegisteredStream")
210            .field("connection_info", &self.connection_info)
211            .finish_non_exhaustive()
212    }
213}
214
215impl<T> RegisteredStream<T> {
216    pub(crate) fn new(connection_info: ConnectionInfo, stream_provider: StreamProvider<T>) -> Self {
217        Self {
218            connection_info,
219            stream_provider,
220            cleanup: Cleanup(None),
221        }
222    }
223
224    pub(crate) fn with_cleanup<F>(mut self, cleanup: F) -> Self
225    where
226        F: FnOnce() + Send + 'static,
227    {
228        self.cleanup.0 = Some(Box::new(cleanup));
229        self
230    }
231
232    /// Consume the registration, disarming the RAII cleanup. Caller takes
233    /// responsibility for cleanup if the stream provider is never awaited.
234    pub fn into_parts(self) -> (ConnectionInfo, StreamProvider<T>) {
235        let Self {
236            connection_info,
237            stream_provider,
238            mut cleanup,
239        } = self;
240        cleanup.0.take();
241        (connection_info, stream_provider)
242    }
243}
244
245/// After registering a stream, the [`PendingConnections`] object is returned to the caller. This
246/// object can be used to await the connection to be established.
247pub struct PendingConnections {
248    pub send_stream: Option<RegisteredStream<StreamSender>>,
249    pub recv_stream: Option<RegisteredStream<StreamReceiver>>,
250}
251
252impl PendingConnections {
253    pub fn into_parts(
254        self,
255    ) -> (
256        Option<RegisteredStream<StreamSender>>,
257        Option<RegisteredStream<StreamReceiver>>,
258    ) {
259        (self.send_stream, self.recv_stream)
260    }
261}
262
263/// A [`ResponseService`] implements a services in which a context a specific subject with will
264/// be associated with a stream of responses.
265#[async_trait::async_trait]
266pub trait ResponseService {
267    async fn register(&self, options: StreamOptions) -> PendingConnections;
268}
269
270#[cfg(test)]
271mod registered_stream_tests {
272    use super::*;
273    use std::sync::atomic::{AtomicBool, Ordering};
274
275    fn dummy_conn_info() -> ConnectionInfo {
276        ConnectionInfo {
277            transport: "test".to_string(),
278            info: "{}".to_string(),
279        }
280    }
281
282    /// Drop without `into_parts()` must run the cleanup closure.
283    #[test]
284    fn drop_runs_cleanup() {
285        let flag = Arc::new(AtomicBool::new(false));
286        let flag_clone = flag.clone();
287
288        let (_tx, rx) = tokio::sync::oneshot::channel::<Result<(), String>>();
289        let stream = RegisteredStream::new(dummy_conn_info(), rx).with_cleanup(move || {
290            flag_clone.store(true, Ordering::SeqCst);
291        });
292
293        drop(stream);
294        assert!(
295            flag.load(Ordering::SeqCst),
296            "cleanup must fire when RegisteredStream is dropped"
297        );
298    }
299
300    /// `into_parts()` must disarm the cleanup. After the call, dropping the
301    /// returned halves must NOT trigger the closure -- the caller has taken
302    /// ownership of cleanup responsibility.
303    #[test]
304    fn into_parts_disarms_cleanup() {
305        let flag = Arc::new(AtomicBool::new(false));
306        let flag_clone = flag.clone();
307
308        let (_tx, rx) = tokio::sync::oneshot::channel::<Result<(), String>>();
309        let stream = RegisteredStream::new(dummy_conn_info(), rx).with_cleanup(move || {
310            flag_clone.store(true, Ordering::SeqCst);
311        });
312
313        let (conn, provider) = stream.into_parts();
314        drop(conn);
315        drop(provider);
316
317        assert!(
318            !flag.load(Ordering::SeqCst),
319            "into_parts() must disarm the cleanup closure"
320        );
321    }
322
323    /// `RegisteredStream` with no cleanup configured must drop cleanly.
324    #[test]
325    fn drop_without_cleanup_is_a_noop() {
326        let (_tx, rx) = tokio::sync::oneshot::channel::<Result<(), String>>();
327        let stream: RegisteredStream<()> = RegisteredStream::new(dummy_conn_info(), rx);
328        drop(stream); // must not panic; nothing observable to assert beyond that
329    }
330}
331
332// #[derive(Debug, Clone, Serialize, Deserialize)]
333// struct Handshake {
334//     request_id: String,
335//     worker_id: Option<String>,
336//     error: Option<String>,
337// }
338
339// impl Handshake {
340//     pub fn validate(&self) -> Result<(), String> {
341//         if let Some(e) = &self.error {
342//             return Err(e.clone());
343//         }
344//         Ok(())
345//     }
346// }
347
348// this probably needs to be come a ResponseStreamSender
349// since the prologue in this scenario sender telling the receiver
350// that all is good and it's ready to send
351//
352// in the RequestStreamSender, the prologue would be coming from the
353// receiver, so the sender would have to await the prologue which if
354// was not an error, would indicate the RequestStreamReceiver is read
355// to receive data.
356pub struct StreamSender {
357    tx: tokio::sync::mpsc::Sender<TwoPartMessage>,
358    prologue: Option<ResponseStreamPrologue>,
359}
360
361impl StreamSender {
362    pub async fn send(&self, data: Bytes) -> Result<()> {
363        Ok(self.tx.send(TwoPartMessage::from_data(data)).await?)
364    }
365
366    pub async fn send_control(&self, control: ControlMessage) -> Result<()> {
367        let bytes = serde_json::to_vec(&control)?;
368        Ok(self
369            .tx
370            .send(TwoPartMessage::from_header(bytes.into()))
371            .await?)
372    }
373
374    #[allow(clippy::needless_update)]
375    pub async fn send_prologue(&mut self, error: Option<String>) -> Result<(), String> {
376        // leaving the original logic in place for now
377        // error overrides the dissolved prologue, but the only field on `ResponseStreamPrologue` is `error`
378        // so the second argument can never be used, and the value of error passed by the caller would always be used
379        if let Some(_prologue) = self.prologue.take() {
380            // let prologue = ResponseStreamPrologue { error, ..prologue };
381            let prologue = ResponseStreamPrologue { error };
382            let header_bytes: Bytes = match serde_json::to_vec(&prologue) {
383                Ok(b) => b.into(),
384                Err(err) => {
385                    tracing::error!(%err, "send_prologue: ResponseStreamPrologue did not serialize to a JSON array");
386                    return Err("Invalid prologue".to_string());
387                }
388            };
389            self.tx
390                .send(TwoPartMessage::from_header(header_bytes))
391                .await
392                .map_err(|e| e.to_string())?;
393        } else {
394            panic!("Prologue already sent; or not set; logic error");
395        }
396        Ok(())
397    }
398}
399
400pub struct StreamReceiver {
401    rx: tokio::sync::mpsc::Receiver<Bytes>,
402}
403
404/// Connection Info is encoded as JSON and then again serialized has part of the Transport
405/// Layer. The double serialization is not performance critical as it is only done once per
406/// connection. The primary reason storing the ConnecitonInfo has a JSON string is for type
407/// erasure. The Transport Layer will check the [`ConnectionInfo::transport`] type and then
408/// route it to the appropriate instance of the Transport, which will then deserialize the
409/// [`ConnectionInfo::info`] field to its internal connection info object.
410///
411/// Optionally, this object could become strongly typed for which all possible combinations
412/// of transport and connection info would need to be enumerated.
413#[derive(Debug, Clone, Serialize, Deserialize)]
414pub struct ConnectionInfo {
415    pub transport: String,
416    pub info: String,
417}
418
419/// Default number of frames buffered between the data-plane socket task and the
420/// engine consumer/producer for a single stream. Preserves the historically
421/// hard-coded mpsc channel capacity used by the TCP transport.
422pub const DEFAULT_SEND_BUFFER_COUNT: usize = 64;
423
424/// When registering a new TransportStream on the server, the caller specifies if the
425/// stream is a sender, receiver or both.
426///
427/// Senders and Receivers are with share a Context, but result in separate tcp socket
428/// connections to the server. Internally, we may use bcast channels to coordinate the
429/// internal control messages between the sender and receiver socket connections.
430#[derive(Clone, Builder)]
431pub struct StreamOptions {
432    /// Context
433    pub context: Arc<dyn AsyncEngineContext>,
434
435    /// Register with the server that this connection will have a server-side Sender
436    /// that can be picked up by the Request/Forward pipeline. The downstream side
437    /// dials in via [`crate::pipeline::network::tcp::client::TcpClient::create_request_stream`]
438    /// to receive the frames the server pushes.
439    pub enable_request_stream: bool,
440
441    /// Register with the server that this connection will have a server-side Receiver
442    /// that can be picked up by the Response/Reverse pipeline
443    pub enable_response_stream: bool,
444
445    /// The number of frames buffered between the data-plane socket task and the
446    /// engine consumer/producer before backpressure kicks in. Drives the mpsc
447    /// channel capacity for the per-stream buffer in the TCP transport.
448    #[builder(default = "DEFAULT_SEND_BUFFER_COUNT")]
449    pub send_buffer_count: usize,
450
451    /// The number of messages to buffer before blocking
452    #[builder(default = "8")]
453    pub recv_buffer_count: usize,
454}
455
456impl StreamOptions {
457    pub fn builder() -> StreamOptionsBuilder {
458        StreamOptionsBuilder::default()
459    }
460}
461
462pub struct Egress<Req: PipelineIO, Resp: PipelineIO> {
463    transport_engine: Arc<dyn AsyncTransportEngine<Req, Resp>>,
464}
465
466#[cfg(test)]
467mod tests {
468    use super::{
469        DEFAULT_SEND_BUFFER_COUNT, NetworkStreamWrapper, RequestControlMessage,
470        RequestPlanePayloadCodec, RequestType, ResponseType, StreamOptions,
471    };
472    use crate::engine::AsyncEngineContextProvider;
473    use crate::pipeline::Context;
474    use serde::{Deserialize, Serialize};
475
476    #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
477    struct TestPayload {
478        id: u64,
479        text: String,
480        tokens: Vec<u32>,
481    }
482
483    #[test]
484    fn stream_options_send_buffer_count_defaults_to_64() {
485        let context = Context::new(());
486        let options = StreamOptions::builder()
487            .context(context.context())
488            .enable_request_stream(true)
489            .enable_response_stream(true)
490            .build()
491            .expect("stream options should build");
492
493        assert_eq!(DEFAULT_SEND_BUFFER_COUNT, 64);
494        assert_eq!(options.send_buffer_count, DEFAULT_SEND_BUFFER_COUNT);
495    }
496
497    #[test]
498    fn stream_options_send_buffer_count_overrides_default() {
499        let context = Context::new(());
500        let options = StreamOptions::builder()
501            .context(context.context())
502            .enable_request_stream(true)
503            .enable_response_stream(true)
504            .send_buffer_count(128)
505            .build()
506            .expect("stream options should build");
507
508        assert_eq!(options.send_buffer_count, 128);
509    }
510
511    #[test]
512    fn request_control_message_defaults_missing_metadata() {
513        let json = r#"{
514            "id": "request-123",
515            "request_type": "single_in",
516            "response_type": "many_out",
517            "connection_info": {
518                "transport": "tcp",
519                "info": "{}"
520            }
521        }"#;
522
523        let message: RequestControlMessage =
524            serde_json::from_str(json).expect("control message should deserialize");
525
526        assert_eq!(message.id, "request-123");
527        assert!(matches!(message.request_type, RequestType::SingleIn));
528        assert!(matches!(message.response_type, ResponseType::ManyOut));
529        assert_eq!(message.payload_codec, RequestPlanePayloadCodec::Json);
530        assert_eq!(message.connection_info.transport, "tcp");
531        assert_eq!(message.connection_info.info, "{}");
532        assert!(message.metadata.is_empty());
533        assert!(message.frontend_send_ts_ns.is_none());
534    }
535
536    #[test]
537    fn request_control_message_decodes_msgpack_payload_codec() {
538        let json = r#"{
539            "id": "request-123",
540            "request_type": "single_in",
541            "response_type": "many_out",
542            "payload_codec": "msgpack",
543            "connection_info": {
544                "transport": "tcp",
545                "info": "{}"
546            }
547        }"#;
548
549        let message: RequestControlMessage =
550            serde_json::from_str(json).expect("control message should deserialize");
551
552        assert_eq!(message.payload_codec, RequestPlanePayloadCodec::Msgpack);
553    }
554
555    #[test]
556    fn request_plane_payload_codec_round_trips_response_wrapper_json_and_msgpack() {
557        let wrapper = NetworkStreamWrapper {
558            data: Some(TestPayload {
559                id: 42,
560                text: "line\nquote\"slash\\unicode 中".to_string(),
561                tokens: vec![1, 2, 3, 65535],
562            }),
563            complete_final: false,
564        };
565
566        for codec in [
567            RequestPlanePayloadCodec::Json,
568            RequestPlanePayloadCodec::Msgpack,
569        ] {
570            let encoded = codec.encode(&wrapper).expect("wrapper should encode");
571            let decoded: NetworkStreamWrapper<TestPayload> =
572                codec.decode(&encoded).expect("wrapper should decode");
573            assert_eq!(decoded, wrapper);
574        }
575    }
576}
577
578#[async_trait]
579impl<T: Data, U: Data> AsyncEngine<SingleIn<T>, ManyOut<U>, Error>
580    for Egress<SingleIn<T>, ManyOut<U>>
581where
582    T: Data + Serialize,
583    U: for<'de> Deserialize<'de> + Data,
584{
585    async fn generate(&self, request: SingleIn<T>) -> Result<ManyOut<U>, Error> {
586        self.transport_engine.generate(request).await
587    }
588}
589
590pub struct Ingress<Req: PipelineIO, Resp: PipelineIO> {
591    segment: OnceLock<Arc<SegmentSource<Req, Resp>>>,
592    metrics: OnceLock<Arc<WorkHandlerMetrics>>,
593    /// Endpoint-specific notifier for health check timer resets
594    endpoint_health_check_notifier: OnceLock<Arc<tokio::sync::Notify>>,
595}
596
597impl<Req: PipelineIO + Sync, Resp: PipelineIO> Ingress<Req, Resp> {
598    pub fn new() -> Arc<Self> {
599        Arc::new(Self {
600            segment: OnceLock::new(),
601            metrics: OnceLock::new(),
602            endpoint_health_check_notifier: OnceLock::new(),
603        })
604    }
605
606    pub fn attach(&self, segment: Arc<SegmentSource<Req, Resp>>) -> Result<()> {
607        self.segment
608            .set(segment)
609            .map_err(|_| anyhow::anyhow!("Segment already set"))
610    }
611
612    pub fn add_metrics(
613        &self,
614        endpoint: &crate::component::Endpoint,
615        metrics_labels: Option<&[(&str, &str)]>,
616    ) -> Result<()> {
617        let metrics = WorkHandlerMetrics::from_endpoint(endpoint, metrics_labels)
618            .map_err(|e| anyhow::anyhow!("Failed to create work handler metrics: {}", e))?;
619
620        // Register global transport breakdown metrics (idempotent)
621        crate::metrics::work_handler_perf::ensure_work_handler_perf_metrics_registered(
622            endpoint.get_metrics_registry(),
623        );
624
625        // Register worker-pool saturation metrics (idempotent). These are
626        // process-global and shared across all endpoints attached to the
627        // same shared TCP server.
628        crate::metrics::work_handler_pool::ensure_work_handler_pool_metrics_registered(
629            endpoint.get_metrics_registry(),
630        );
631
632        self.metrics
633            .set(Arc::new(metrics))
634            .map_err(|_| anyhow::anyhow!("Metrics already set"))
635    }
636
637    pub fn link(segment: Arc<SegmentSource<Req, Resp>>) -> Result<Arc<Self>> {
638        let ingress = Ingress::new();
639        ingress.attach(segment)?;
640        Ok(ingress)
641    }
642
643    pub fn for_pipeline(segment: Arc<SegmentSource<Req, Resp>>) -> Result<Arc<Self>> {
644        let ingress = Ingress::new();
645        ingress.attach(segment)?;
646        Ok(ingress)
647    }
648
649    pub fn for_engine(engine: ServiceEngine<Req, Resp>) -> Result<Arc<Self>> {
650        let frontend = SegmentSource::<Req, Resp>::new();
651        let backend = ServiceBackend::from_engine(engine);
652
653        // create the pipeline
654        let pipeline = frontend.link(backend)?.link(frontend)?;
655
656        let ingress = Ingress::new();
657        ingress.attach(pipeline)?;
658
659        Ok(ingress)
660    }
661
662    /// Helper method to access metrics if available
663    fn metrics(&self) -> Option<&Arc<WorkHandlerMetrics>> {
664        self.metrics.get()
665    }
666}
667
668#[async_trait]
669pub trait PushWorkHandler: Send + Sync {
670    async fn handle_payload(
671        &self,
672        payload: Bytes,
673        request_id: Option<String>,
674    ) -> Result<(), PipelineError>;
675
676    /// Add metrics to the handler
677    fn add_metrics(
678        &self,
679        endpoint: &crate::component::Endpoint,
680        metrics_labels: Option<&[(&str, &str)]>,
681    ) -> Result<()>;
682
683    /// Set the endpoint-specific notifier for health check timer resets
684    fn set_endpoint_health_check_notifier(
685        &self,
686        _notifier: Arc<tokio::sync::Notify>,
687    ) -> Result<()> {
688        // Default implementation for backwards compatibility
689        Ok(())
690    }
691}
692
693/*
694/// `NetworkStreamWrapper` is a simple wrapper used to detect proper stream termination
695/// in network communication between ingress and egress components.
696///
697/// **Purpose**: This wrapper solves the problem of detecting whether a stream ended
698/// gracefully or was cut off prematurely (e.g., due to network issues).
699///
700/// **Design Rationale**:
701/// - Cannot use `Annotated` directly because the generic type `U` varies:
702///   - Sometimes `U = Annotated<...>`
703///   - Sometimes `U = LLMEngineOutput<...>`
704/// - Using `Annotated` would require double-wrapping like `Annotated<Annotated<...>>`
705/// - A simple wrapper is cleaner and more straightforward
706///
707/// **Stream Flow**:
708/// ```
709/// At AsyncEngine:
710///   response 1 -> response 2 -> response 3 -> <end>
711///
712/// Between ingress/egress:
713///   response 1 <end=false> -> response 2 <end=false> -> response 3 <end=false> -> (null) <end=true>
714///
715/// At client:
716///   response 1 -> response 2 -> response 3 -> <end>
717/// ```
718///
719/// **Error Handling**:
720/// If the stream is cut off before proper termination, the egress is responsible for
721/// injecting an error response to communicate the incomplete stream to the client:
722/// ```
723/// At AsyncEngine:
724///   response 1 -> ... <without end flag>
725///
726/// At egress:
727///   response 1 <end=false> -> <stream ended without end flag -> convert to error>
728///
729/// At client:
730///   response 1 -> error response
731/// ```
732///
733/// The detection must be done at egress level because premature stream termination
734/// can be due to network issues that only the egress component can detect.
735*/
736/// TODO: Detect end-of-stream using Server-Sent Events (SSE). This will be removed.
737#[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
738pub struct NetworkStreamWrapper<U> {
739    #[serde(skip_serializing_if = "Option::is_none")]
740    pub data: Option<U>,
741    pub complete_final: bool,
742}
dynamo_runtime/pipeline/network.rs

dynamo_runtime/pipeline/
network.rs
