dynamo_runtime/pipeline/network/
tcp.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! TCP Transport Module
5//!
6//! The TCP Transport module consists of two main components: Client and Server. The Client is
7//! the downstream node that is responsible for connecting back to the upstream node (Server).
8//!
9//! Both Client and Server are given a Stream object that they can specialize for their specific
10//! needs, i.e. if they are SingleIn/ManyIn or SingleOut/ManyOut.
11//!
12//! The Request object will carry the Transport Type and Connection details, i.e. how the receiver
13//! of a Request is able to communicate back to the source of the Request.
14//!
15//! There are two types of TcpStream:
16//! - CallHome stream - the address for the listening socket is forward via some mechanism which then
17//!   connects back to the source of the CallHome stream. To match the socket with an awaiting data
18//!   stream, the CallHomeHandshake is used.
19
20pub mod client;
21pub mod server;
22
23use super::ControlMessage;
24use serde::{Deserialize, Serialize};
25
26#[allow(unused_imports)]
27use super::{
28    ConnectionInfo, PendingConnections, RegisteredStream, ResponseService, StreamOptions,
29    StreamReceiver, StreamSender, StreamType, codec::TwoPartCodec,
30};
31
32const TCP_TRANSPORT: &str = "tcp_server";
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct TcpStreamConnectionInfo {
36    pub address: String,
37    pub subject: String,
38    pub context: String,
39    pub stream_type: StreamType,
40}
41
42impl From<TcpStreamConnectionInfo> for ConnectionInfo {
43    fn from(info: TcpStreamConnectionInfo) -> Self {
44        // Need to consider the below. If failure should be fatal, keep the below with .expect()
45        // But if there is a default value, we can use:
46        // unwrap_or_else(|e| {
47        //     eprintln!("Failed to serialize TcpStreamConnectionInfo: {:?}", e);
48        //     "{}".to_string() // Provide a fallback empty JSON string or default value
49        ConnectionInfo {
50            transport: TCP_TRANSPORT.to_string(),
51            info: serde_json::to_string(&info)
52                .expect("Failed to serialize TcpStreamConnectionInfo"),
53        }
54    }
55}
56
57impl TryFrom<ConnectionInfo> for TcpStreamConnectionInfo {
58    type Error = anyhow::Error;
59
60    fn try_from(info: ConnectionInfo) -> Result<Self, Self::Error> {
61        if info.transport != TCP_TRANSPORT {
62            return Err(anyhow::anyhow!(
63                "Invalid transport; TcpClient requires the transport to be `tcp_server`; however {} was passed",
64                info.transport
65            ));
66        }
67
68        serde_json::from_str(&info.info)
69            .map_err(|e| anyhow::anyhow!("Failed parse ConnectionInfo: {:?}", e))
70    }
71}
72
73/// First message sent over a CallHome stream which will map the newly created socket to a specific
74/// response data stream which was registered with the same subject.
75///
76/// This is a transport specific message as part of forming/completing a CallHome TcpStream.
77#[derive(Debug, Clone, Serialize, Deserialize)]
78struct CallHomeHandshake {
79    subject: String,
80    stream_type: StreamType,
81}
82
83#[cfg(test)]
84mod tests {
85    use crate::engine::AsyncEngineContextProvider;
86
87    use super::*;
88    use crate::pipeline::Context;
89
90    #[derive(Debug, Clone, Serialize, Deserialize)]
91    struct TestMessage {
92        foo: String,
93    }
94
95    #[tokio::test]
96    async fn test_tcp_stream_client_server() {
97        println!("Test Started");
98        let options = server::ServerOptions::builder().port(9124).build().unwrap();
99        println!("Test Started");
100        let server = server::TcpStreamServer::new(options).await.unwrap();
101        println!("Server created");
102
103        let context_rank0 = Context::new(());
104
105        let options = StreamOptions::builder()
106            .context(context_rank0.context())
107            .enable_request_stream(false)
108            .enable_response_stream(true)
109            .build()
110            .unwrap();
111
112        let pending_connection = server.register(options).await;
113
114        let connection_info = pending_connection
115            .recv_stream
116            .as_ref()
117            .unwrap()
118            .connection_info
119            .clone();
120
121        // set up the other rank
122        let context_rank1 = Context::with_id((), context_rank0.id().to_string());
123
124        // connect to the server socket
125        let mut send_stream =
126            client::TcpClient::create_response_stream(context_rank1.context(), connection_info)
127                .await
128                .unwrap();
129        println!("Client connected");
130
131        // the client can now setup it's end of the stream and if it errors, it can send a message
132        // to the server to stop the stream
133        //
134        // this step must be done before the next step on the server can complete, i.e.
135        // the server's stream is now blocked on receiving the prologue message
136        //
137        // let's improve this and use an enum like Ok/Err; currently, None means good-to-go, and
138        // Some(String) means an error happened on this downstream node and we need to alert the
139        // upstream node that an error occurred
140        send_stream.send_prologue(None).await.unwrap();
141
142        // [server] next - now pending connections should be connected
143        let recv_stream = pending_connection
144            .recv_stream
145            .unwrap()
146            .stream_provider
147            .await
148            .unwrap();
149
150        println!("Server paired");
151
152        let msg = TestMessage {
153            foo: "bar".to_string(),
154        };
155
156        let payload = serde_json::to_vec(&msg).unwrap();
157
158        send_stream.send(payload.into()).await.unwrap();
159
160        println!("Client sent message");
161
162        let data = recv_stream.unwrap().rx.recv().await.unwrap();
163
164        println!("Server received message");
165
166        let recv_msg = serde_json::from_slice::<TestMessage>(&data).unwrap();
167
168        assert_eq!(msg.foo, recv_msg.foo);
169        println!("message match");
170
171        drop(send_stream);
172
173        // let data = recv_stream.rx.recv().await;
174
175        // assert!(data.is_none());
176    }
177}
dynamo_runtime/pipeline/network/tcp.rs

dynamo_runtime/pipeline/network/
tcp.rs
