ts_client/
handshake.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 frontend side of the startup handshake (see [`ts_wire::handshake`]).
22
23use std::collections::{BTreeMap, BTreeSet};
24
25use bytes::Bytes;
26use tokio::time::{timeout_at, Instant};
27use tracing::{debug, info};
28use ts_wire::{
29    decode_bytes, encode_frames,
30    handshake::{STATUS_HELLO, STATUS_READY},
31    Dict, EngineCoreReadyResponse, HandshakeAddresses, HandshakeInitMessage, ReadyMessage,
32};
33use zeromq::{PullSocket, RouterSocket, Socket, SocketRecv, SocketSend, ZmqMessage};
34
35use crate::{
36    config::ClientConfig,
37    error::{Error, Result},
38};
39
40/// An engine that completed the handshake.
41#[derive(Debug, Clone)]
42pub struct EngineInfo {
43    /// The engine's index (its routing identity).
44    pub index: u16,
45    /// The configuration it registered with.
46    pub ready: EngineCoreReadyResponse,
47}
48
49/// The frontend sockets, bound, before any engine has connected.
50pub(crate) struct Bound {
51    /// The handshake address as bound, with the port the OS chose for `:0`.
52    pub handshake_address: String,
53    handshake: RouterSocket,
54    input: RouterSocket,
55    output: PullSocket,
56    /// The INIT message sent to every engine.
57    init: Vec<Bytes>,
58    expected: BTreeSet<u16>,
59}
60
61/// Bound sockets with every expected engine registered.
62pub(crate) struct Connected {
63    pub engines: Vec<EngineInfo>,
64    pub input: RouterSocket,
65    pub output: PullSocket,
66}
67
68/// Bind the frontend sockets.
69pub(crate) async fn bind(config: &ClientConfig) -> Result<Bound> {
70    let last_index = u32::from(config.first_engine_index) + u32::from(config.num_engines);
71    if config.num_engines == 0 || last_index > u32::from(u16::MAX) + 1 {
72        return Err(Error::Handshake(format!(
73            "engine indices {}..{last_index} are empty or exceed the two-byte identity",
74            config.first_engine_index
75        )));
76    }
77    let expected: BTreeSet<u16> = (config.first_engine_index..=(last_index - 1) as u16).collect();
78
79    let mut handshake = RouterSocket::new();
80    let handshake_address = handshake.bind(&config.handshake_address).await?.to_string();
81    let mut input = RouterSocket::new();
82    let input_address = input.bind(&config.input_address).await?.to_string();
83    let mut output = PullSocket::new();
84    let output_address = output.bind(&config.output_address).await?.to_string();
85    info!(
86        handshake = %handshake_address,
87        input = %input_address,
88        output = %output_address,
89        engines = ?expected,
90        "frontend sockets bound"
91    );
92
93    let init = encode_frames(&HandshakeInitMessage {
94        addresses: HandshakeAddresses {
95            inputs: vec![input_address],
96            outputs: vec![output_address],
97            coordinator_input: None,
98            coordinator_output: None,
99            frontend_stats_publish_address: None,
100        },
101        parallel_config: Dict::default(),
102    });
103    Ok(Bound {
104        handshake_address,
105        handshake,
106        input,
107        output,
108        init,
109        expected,
110    })
111}
112
113/// Complete the handshake with every expected engine.
114pub(crate) async fn accept(bound: Bound, config: &ClientConfig) -> Result<Connected> {
115    let deadline = Instant::now() + config.handshake_timeout;
116    let Bound {
117        mut handshake,
118        mut input,
119        output,
120        init,
121        expected,
122        ..
123    } = bound;
124
125    // HELLO, then INIT from us, then READY, per engine; engines interleave.
126    let mut awaiting_hello = expected.clone();
127    let mut awaiting_ready = expected.clone();
128    while !awaiting_ready.is_empty() {
129        let message =
130            recv_before(&mut handshake, deadline, config, "READY", &awaiting_ready).await?;
131        let (index, payload) = engine_message(message, &expected)?;
132        let status: ReadyMessage = decode_bytes(&payload)?;
133        match status.status.as_deref() {
134            Some(STATUS_HELLO) => {
135                if !awaiting_hello.remove(&index) {
136                    return Err(Error::Handshake(format!("engine {index} sent HELLO twice")));
137                }
138                let mut reply = ZmqMessage::from(Bytes::copy_from_slice(&index.to_le_bytes()));
139                for frame in &init {
140                    reply.push_back(frame.clone());
141                }
142                handshake.send(reply).await?;
143                debug!(engine = index, "sent INIT");
144            }
145            Some(STATUS_READY) => {
146                if awaiting_hello.contains(&index) || !awaiting_ready.remove(&index) {
147                    return Err(Error::Handshake(format!(
148                        "engine {index} sent READY out of order"
149                    )));
150                }
151                debug!(engine = index, "engine READY");
152            }
153            other => {
154                return Err(Error::Handshake(format!(
155                    "engine {index} sent unknown status {other:?}"
156                )))
157            }
158        }
159    }
160
161    // Each engine then registers its configuration on the input socket.
162    let mut registered: BTreeMap<u16, EngineCoreReadyResponse> = BTreeMap::new();
163    while registered.len() < expected.len() {
164        let pending: BTreeSet<u16> = expected
165            .iter()
166            .filter(|index| !registered.contains_key(index))
167            .copied()
168            .collect();
169        let message = recv_before(&mut input, deadline, config, "registration", &pending).await?;
170        let (index, payload) = engine_message(message, &expected)?;
171        let ready: EngineCoreReadyResponse = decode_bytes(&payload)?;
172        check_rank(config, index, &ready)?;
173        if registered.insert(index, ready).is_some() {
174            return Err(Error::Handshake(format!("engine {index} registered twice")));
175        }
176    }
177    info!(engines = registered.len(), "all engines registered");
178
179    Ok(Connected {
180        engines: registered
181            .into_iter()
182            .map(|(index, ready)| EngineInfo { index, ready })
183            .collect(),
184        input,
185        output,
186    })
187}
188
189async fn recv_before(
190    socket: &mut RouterSocket,
191    deadline: Instant,
192    config: &ClientConfig,
193    stage: &'static str,
194    pending: &BTreeSet<u16>,
195) -> Result<ZmqMessage> {
196    match timeout_at(deadline, socket.recv()).await {
197        Ok(message) => Ok(message?),
198        Err(_) => Err(Error::HandshakeTimeout {
199            stage,
200            pending: pending.iter().copied().collect(),
201            timeout: config.handshake_timeout,
202        }),
203    }
204}
205
206/// Split a ROUTER message into the sending engine's index and its payload.
207fn engine_message(message: ZmqMessage, expected: &BTreeSet<u16>) -> Result<(u16, Bytes)> {
208    let frames = message.into_vec();
209    let [identity, payload] = frames.as_slice() else {
210        return Err(Error::Handshake(format!(
211            "expected an identity and a payload, got {} frames",
212            frames.len()
213        )));
214    };
215    let index = match identity.as_ref() {
216        [low, high] => u16::from_le_bytes([*low, *high]),
217        other => {
218            return Err(Error::Handshake(format!(
219                "peer identity {other:02x?} is not a two-byte engine index"
220            )))
221        }
222    };
223    if !expected.contains(&index) {
224        return Err(Error::Handshake(format!(
225            "unexpected engine {index}; expecting {expected:?}"
226        )));
227    }
228    Ok((index, payload.clone()))
229}
230
231/// Every engine must report the DP layout the client was configured for.
232fn check_rank(config: &ClientConfig, index: u16, ready: &EngineCoreReadyResponse) -> Result<()> {
233    let rank = i64::from(index - config.first_engine_index);
234    if ready.data_parallel_size != i64::from(config.num_engines) || ready.data_parallel_rank != rank
235    {
236        return Err(Error::Handshake(format!(
237            "engine {index} reports DP rank {} of {}, expected rank {rank} of {}",
238            ready.data_parallel_rank, ready.data_parallel_size, config.num_engines
239        )));
240    }
241    Ok(())
242}
ts_client/handshake.rs

ts_client/
handshake.rs
