ts_pd/
net.rs

1// Copyright (c) 2026 LightSeek Foundation
2//
3// Permission is hereby granted, free of charge, to any person obtaining a copy
4// of this software and associated documentation files (the "Software"), to deal
5// in the Software without restriction, including without limitation the rights
6// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7// copies of the Software, and to permit persons to whom the Software is
8// furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
19// SOFTWARE.
20
21//! The P/D control sockets: one PULL socket a rank's peers push to, and
22//! PUSH sockets to the peers, one per endpoint (`DisaggManagerBase`).
23//!
24//! Socket I/O runs on a small Tokio runtime; the PD threads exchange
25//! messages with it through channels.
26
27use std::{collections::HashMap, sync::mpsc};
28
29use bytes::Bytes;
30use tokio::runtime::Runtime;
31use zeromq::{PullSocket, PushSocket, Socket, SocketRecv, SocketSend, ZmqMessage};
32
33/// A socket failure at startup.
34#[derive(Debug, thiserror::Error)]
35pub enum NetError {
36    /// A socket operation failed.
37    #[error("zmq: {0}")]
38    Zmq(#[from] zeromq::ZmqError),
39    /// The runtime did not start.
40    #[error("starting the PD socket runtime: {0}")]
41    Runtime(#[from] std::io::Error),
42    /// The bound endpoint carries no port.
43    #[error("the bound endpoint {0} has no port")]
44    Port(String),
45}
46
47/// The control sockets of one rank.
48pub struct Sockets {
49    runtime: Runtime,
50    outgoing: tokio::sync::mpsc::UnboundedSender<(String, Vec<Bytes>)>,
51    port: u16,
52}
53
54impl Sockets {
55    /// Bind the PULL socket on `host` at a port the OS picks; every message
56    /// it receives goes to `incoming`, frames in order.
57    pub fn start(host: &str, incoming: mpsc::Sender<Vec<Bytes>>) -> Result<Self, NetError> {
58        let runtime = tokio::runtime::Builder::new_multi_thread()
59            .worker_threads(1)
60            .thread_name("ts-pd-net")
61            .enable_all()
62            .build()?;
63        let address = if host.contains(':') {
64            format!("tcp://[{host}]:0")
65        } else {
66            format!("tcp://{host}:0")
67        };
68        let (pull, port) = runtime.block_on(async {
69            let mut pull = PullSocket::new();
70            let bound = pull.bind(&address).await?.to_string();
71            let port = bound
72                .rsplit(':')
73                .next()
74                .and_then(|port| port.parse().ok())
75                .ok_or(NetError::Port(bound.clone()))?;
76            Ok::<_, NetError>((pull, port))
77        })?;
78        runtime.spawn(async move {
79            let mut pull = pull;
80            loop {
81                match pull.recv().await {
82                    Ok(message) => {
83                        if incoming.send(message.into_vec()).is_err() {
84                            return;
85                        }
86                    }
87                    Err(error) => {
88                        tracing::warn!(%error, "PD control socket closed");
89                        return;
90                    }
91                }
92            }
93        });
94        let (outgoing, mut queue) = tokio::sync::mpsc::unbounded_channel::<(String, Vec<Bytes>)>();
95        runtime.spawn(async move {
96            let mut sockets: HashMap<String, PushSocket> = HashMap::new();
97            while let Some((endpoint, frames)) = queue.recv().await {
98                if !sockets.contains_key(&endpoint) {
99                    let mut push = PushSocket::new();
100                    if let Err(error) = push.connect(&endpoint).await {
101                        tracing::warn!(endpoint, %error, "cannot reach a PD peer");
102                        continue;
103                    }
104                    sockets.insert(endpoint.clone(), push);
105                }
106                let Ok(message) = ZmqMessage::try_from(frames) else {
107                    continue;
108                };
109                let push = sockets.get_mut(&endpoint).expect("connected above");
110                if let Err(error) = push.send(message).await {
111                    tracing::warn!(endpoint, %error, "a PD message did not go out");
112                    sockets.remove(&endpoint);
113                }
114            }
115        });
116        Ok(Self {
117            runtime,
118            outgoing,
119            port,
120        })
121    }
122
123    /// The PULL socket's port.
124    pub fn port(&self) -> u16 {
125        self.port
126    }
127
128    /// Push `frames` to `host:port`; never waits.
129    pub fn send(&self, host: &str, port: u16, frames: Vec<Bytes>) {
130        let endpoint = if host.contains(':') && !host.starts_with('[') {
131            format!("tcp://[{host}]:{port}")
132        } else {
133            format!("tcp://{host}:{port}")
134        };
135        if self.outgoing.send((endpoint, frames)).is_err() {
136            tracing::warn!("the PD socket runtime stopped");
137        }
138    }
139}
140
141impl Drop for Sockets {
142    fn drop(&mut self) {
143        // Let queued messages (a failure notice) leave before the sockets
144        // close.
145        self.runtime
146            .block_on(async { tokio::time::sleep(std::time::Duration::from_millis(50)).await });
147    }
148}
149
150/// An ASCII frame.
151pub fn ascii(value: impl ToString) -> Bytes {
152    Bytes::from(value.to_string())
153}
154
155/// A canonical unsigned ASCII integer frame (`_bounded_ascii_uint`).
156pub fn parse_uint(frame: &[u8], maximum: u64) -> Option<u64> {
157    let digits = maximum.to_string().len();
158    if frame.is_empty()
159        || frame.len() > digits
160        || !frame.iter().all(u8::is_ascii_digit)
161        || (frame.len() > 1 && frame[0] == b'0')
162    {
163        return None;
164    }
165    let value: u64 = std::str::from_utf8(frame).ok()?.parse().ok()?;
166    (value <= maximum).then_some(value)
167}
168
169#[cfg(test)]
170mod tests {
171    use std::time::Duration;
172
173    use super::*;
174
175    #[test]
176    fn a_pushed_message_reaches_the_pull_socket() {
177        let (incoming, received) = mpsc::channel();
178        let receiver = Sockets::start("127.0.0.1", incoming).unwrap();
179        let (unused, _never) = mpsc::channel();
180        let sender = Sockets::start("127.0.0.1", unused).unwrap();
181        sender.send(
182            "127.0.0.1",
183            receiver.port(),
184            vec![ascii(7), Bytes::from_static(b"x")],
185        );
186        let frames = received.recv_timeout(Duration::from_secs(10)).unwrap();
187        assert_eq!(frames, [Bytes::from_static(b"7"), Bytes::from_static(b"x")]);
188    }
189
190    #[test]
191    fn unsigned_frames_are_canonical() {
192        assert_eq!(parse_uint(b"42", 100), Some(42));
193        assert_eq!(parse_uint(b"042", 100), None);
194        assert_eq!(parse_uint(b"101", 100), None);
195        assert_eq!(parse_uint(b"", 100), None);
196        assert_eq!(parse_uint(b"4a", 100), None);
197    }
198}
ts_pd/net.rs

ts_pd/
net.rs
