dynamo_mocker/services/
bootstrap.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Framed transport for mocker prefill/decode handoff sessions.
5//!
6//! Session and request ownership live in `dynamo-llm`. This module only
7//! validates framed connections and hands them to that owner.
8
9use std::sync::{Arc, Mutex};
10use std::time::Duration;
11
12use anyhow::{Context, Result, anyhow, bail};
13use bytes::Bytes;
14use futures::{FutureExt, SinkExt, StreamExt};
15use serde::{Deserialize, Serialize};
16use tokio::io::{AsyncReadExt, AsyncWriteExt};
17use tokio::net::{TcpListener, TcpStream};
18use tokio::sync::{Semaphore, mpsc, watch};
19use tokio_util::codec::{Framed, LengthDelimitedCodec};
20use tokio_util::sync::CancellationToken;
21use tokio_util::task::TaskTracker;
22use uuid::Uuid;
23
24use crate::common::handoff::{
25    HandoffActionId, HandoffActionOutcome, HandoffFact, HandoffId, HandoffOrder,
26    IssuedHandoffAction,
27};
28use crate::common::protocols::EngineType;
29
30pub const BOOTSTRAP_PROTOCOL_VERSION: u16 = 1;
31pub const MAX_BOOTSTRAP_FRAME_BYTES: usize = 64 * 1024;
32const MAGIC: [u8; 4] = *b"DMHF";
33const HEADER_BYTES: usize = 8;
34const RENDEZVOUS_TIMEOUT: Duration = Duration::from_secs(30);
35
36#[derive(Clone, Debug)]
37pub struct BootstrapServerConfig {
38    pub max_pending_connections: usize,
39    pub registration_timeout: Duration,
40}
41
42impl Default for BootstrapServerConfig {
43    fn default() -> Self {
44        Self {
45            max_pending_connections: 256,
46            registration_timeout: RENDEZVOUS_TIMEOUT,
47        }
48    }
49}
50
51#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
52pub enum BootstrapParticipantRole {
53    Destination,
54}
55
56#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
57pub struct BootstrapIdentity {
58    pub handoff_id: HandoffId,
59    pub bootstrap_room: u64,
60    pub request_id: Uuid,
61}
62
63#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
64pub struct ParticipantRegistration {
65    pub role: BootstrapParticipantRole,
66    pub dp_rank: u32,
67    pub order: HandoffOrder,
68    pub engine_type: EngineType,
69}
70
71#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
72pub enum BootstrapMessage {
73    Register(ParticipantRegistration),
74    Registered,
75    Action(IssuedHandoffAction),
76    ActionAck {
77        action_id: HandoffActionId,
78        outcome: HandoffActionOutcome,
79    },
80    Fact(HandoffFact),
81    Complete,
82    Abort {
83        message: String,
84    },
85    Overloaded,
86    ProtocolError {
87        message: String,
88    },
89}
90
91#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
92struct BootstrapWireFrame {
93    identity: BootstrapIdentity,
94    message: BootstrapMessage,
95}
96
97pub struct IncomingBootstrapConnection {
98    pub identity: BootstrapIdentity,
99    pub registration: ParticipantRegistration,
100    pub connection: BootstrapConnection,
101}
102
103pub struct BootstrapConnection {
104    identity: BootstrapIdentity,
105    framed: Framed<TcpStream, LengthDelimitedCodec>,
106}
107
108impl BootstrapConnection {
109    pub fn identity(&self) -> &BootstrapIdentity {
110        &self.identity
111    }
112
113    pub async fn send(&mut self, message: BootstrapMessage) -> Result<()> {
114        let payload = serde_json::to_vec(&BootstrapWireFrame {
115            identity: self.identity.clone(),
116            message,
117        })?;
118        if payload.is_empty() || payload.len() > MAX_BOOTSTRAP_FRAME_BYTES {
119            bail!(
120                "bootstrap frame length {} is outside 1..={MAX_BOOTSTRAP_FRAME_BYTES}",
121                payload.len()
122            );
123        }
124        tokio::time::timeout(RENDEZVOUS_TIMEOUT, self.framed.send(Bytes::from(payload)))
125            .await
126            .map_err(|_| anyhow!("bootstrap send timed out"))??;
127        Ok(())
128    }
129
130    pub async fn recv(&mut self) -> Result<Option<BootstrapMessage>> {
131        let Some(frame) = self.framed.next().await.transpose()? else {
132            return Ok(None);
133        };
134        if frame.is_empty() {
135            bail!("bootstrap received an empty frame");
136        }
137        let frame: BootstrapWireFrame =
138            serde_json::from_slice(&frame).context("bootstrap frame contains malformed JSON")?;
139        if frame.identity != self.identity {
140            bail!("bootstrap frame changed session identity");
141        }
142        Ok(Some(frame.message))
143    }
144
145    pub fn peer_closed_now(&self) -> Result<bool> {
146        let mut byte = [0u8; 1];
147        match self.framed.get_ref().peek(&mut byte).now_or_never() {
148            Some(Ok(0)) => Ok(true),
149            Some(Ok(_)) | None => Ok(false),
150            Some(Err(error)) => Err(error.into()),
151        }
152    }
153}
154
155pub struct BootstrapServer {
156    port: u16,
157    incoming_rx: Mutex<Option<mpsc::Receiver<IncomingBootstrapConnection>>>,
158    closed_rx: watch::Receiver<bool>,
159    #[cfg(test)]
160    accepted_with_slot_rx: watch::Receiver<u64>,
161}
162
163impl BootstrapServer {
164    pub async fn start(
165        port: u16,
166        cancel: CancellationToken,
167        config: BootstrapServerConfig,
168    ) -> Result<Arc<Self>> {
169        if config.max_pending_connections == 0 {
170            bail!("bootstrap max_pending_connections must be at least one");
171        }
172        let listener = TcpListener::bind(format!("0.0.0.0:{port}")).await?;
173        let actual_port = listener.local_addr()?.port();
174        let (incoming_tx, incoming_rx) = mpsc::channel(config.max_pending_connections);
175        let permits = Arc::new(Semaphore::new(config.max_pending_connections));
176        let overload_permits = Arc::new(Semaphore::new(1));
177        let (closed_tx, closed_rx) = watch::channel(false);
178        #[cfg(test)]
179        let (accepted_with_slot_tx, accepted_with_slot_rx) = watch::channel(0_u64);
180        let server = Arc::new(Self {
181            port: actual_port,
182            incoming_rx: Mutex::new(Some(incoming_rx)),
183            closed_rx,
184            #[cfg(test)]
185            accepted_with_slot_rx,
186        });
187
188        tokio::spawn(async move {
189            let connections = TaskTracker::new();
190            #[cfg(test)]
191            let mut accepted_with_slot = 0_u64;
192            loop {
193                let accepted = tokio::select! {
194                    biased;
195                    _ = cancel.cancelled() => break,
196                    accepted = listener.accept() => accepted,
197                };
198                let Ok((stream, _)) = accepted else {
199                    tokio::select! {
200                        _ = cancel.cancelled() => break,
201                        _ = tokio::time::sleep(Duration::from_millis(10)) => {}
202                    }
203                    continue;
204                };
205                let Ok(permit) = permits.clone().try_acquire_owned() else {
206                    let Ok(overload_permit) = overload_permits.clone().try_acquire_owned() else {
207                        drop(stream);
208                        continue;
209                    };
210                    let registration_timeout = config.registration_timeout;
211                    let connection_cancel = cancel.clone();
212                    connections.spawn(async move {
213                        let _permit = overload_permit;
214                        let incoming = tokio::select! {
215                            biased;
216                            _ = connection_cancel.cancelled() => return,
217                            incoming = tokio::time::timeout(
218                                registration_timeout,
219                                accept_connection(stream),
220                            ) => incoming,
221                        };
222                        let Ok(Ok(mut incoming)) = incoming else {
223                            return;
224                        };
225                        let _ = incoming.connection.send(BootstrapMessage::Overloaded).await;
226                    });
227                    continue;
228                };
229                #[cfg(test)]
230                {
231                    accepted_with_slot = accepted_with_slot
232                        .checked_add(1)
233                        .expect("bootstrap accepted-connection test counter overflow");
234                    let _ = accepted_with_slot_tx.send(accepted_with_slot);
235                }
236                let incoming_tx = incoming_tx.clone();
237                let registration_timeout = config.registration_timeout;
238                let connection_cancel = cancel.clone();
239                connections.spawn(async move {
240                    let _permit = permit;
241                    let incoming = tokio::select! {
242                        biased;
243                        _ = connection_cancel.cancelled() => return,
244                        incoming = tokio::time::timeout(
245                            registration_timeout,
246                            accept_connection(stream),
247                        ) => incoming,
248                    };
249                    let Ok(incoming) = incoming else {
250                        return;
251                    };
252                    let Ok(mut incoming) = incoming else {
253                        return;
254                    };
255                    if let Err(error) = incoming_tx.try_send(incoming) {
256                        incoming = error.into_inner();
257                        let _ = incoming.connection.send(BootstrapMessage::Overloaded).await;
258                    }
259                });
260            }
261            connections.close();
262            connections.wait().await;
263            let _ = closed_tx.send(true);
264        });
265
266        Ok(server)
267    }
268
269    pub fn take_incoming_receiver(&self) -> Option<mpsc::Receiver<IncomingBootstrapConnection>> {
270        self.incoming_rx
271            .lock()
272            .expect("bootstrap incoming receiver mutex poisoned")
273            .take()
274    }
275
276    pub fn port(&self) -> u16 {
277        self.port
278    }
279
280    pub async fn wait_closed(&self) {
281        let mut closed_rx = self.closed_rx.clone();
282        if *closed_rx.borrow() {
283            return;
284        }
285        let _ = closed_rx.wait_for(|closed| *closed).await;
286    }
287
288    #[cfg(test)]
289    async fn wait_for_accepted_with_slot(&self, expected: u64) {
290        let mut accepted = self.accepted_with_slot_rx.clone();
291        let _ = accepted.wait_for(|count| *count >= expected).await;
292    }
293}
294
295pub async fn connect_to_prefill(
296    host: &str,
297    port: u16,
298    identity: BootstrapIdentity,
299    registration: ParticipantRegistration,
300) -> Result<BootstrapConnection> {
301    let addr = bootstrap_addr(host, port);
302    let mut stream = tokio::time::timeout(RENDEZVOUS_TIMEOUT, TcpStream::connect(&addr))
303        .await
304        .map_err(|_| anyhow!("bootstrap connect timeout to {addr}"))??;
305    tokio::time::timeout(RENDEZVOUS_TIMEOUT, write_header(&mut stream))
306        .await
307        .map_err(|_| anyhow!("bootstrap header send timed out"))??;
308    let mut connection = BootstrapConnection {
309        identity,
310        framed: framed(stream),
311    };
312    connection
313        .send(BootstrapMessage::Register(registration))
314        .await?;
315    Ok(connection)
316}
317
318fn bootstrap_addr(host: &str, port: u16) -> String {
319    let host = host
320        .strip_prefix('[')
321        .and_then(|host| host.strip_suffix(']'))
322        .unwrap_or(host);
323    if host.contains(':') {
324        format!("[{host}]:{port}")
325    } else {
326        format!("{host}:{port}")
327    }
328}
329
330async fn accept_connection(mut stream: TcpStream) -> Result<IncomingBootstrapConnection> {
331    read_header(&mut stream).await?;
332    let mut connection = BootstrapConnection {
333        identity: BootstrapIdentity {
334            handoff_id: HandoffId::default(),
335            bootstrap_room: 0,
336            request_id: Uuid::nil(),
337        },
338        framed: framed(stream),
339    };
340    let Some(frame) = connection.framed.next().await.transpose()? else {
341        bail!("bootstrap connection closed before registration");
342    };
343    let frame: BootstrapWireFrame =
344        serde_json::from_slice(&frame).context("bootstrap registration contains malformed JSON")?;
345    let BootstrapMessage::Register(registration) = frame.message else {
346        bail!("bootstrap first frame must register a participant");
347    };
348    if registration.role != BootstrapParticipantRole::Destination {
349        bail!("bootstrap server accepts only destination participants");
350    }
351    connection.identity = frame.identity.clone();
352    Ok(IncomingBootstrapConnection {
353        identity: frame.identity,
354        registration,
355        connection,
356    })
357}
358
359fn framed(stream: TcpStream) -> Framed<TcpStream, LengthDelimitedCodec> {
360    LengthDelimitedCodec::builder()
361        .little_endian()
362        .length_field_type::<u32>()
363        .max_frame_length(MAX_BOOTSTRAP_FRAME_BYTES)
364        .new_framed(stream)
365}
366
367async fn write_header(stream: &mut TcpStream) -> Result<()> {
368    let mut header = [0u8; HEADER_BYTES];
369    header[..4].copy_from_slice(&MAGIC);
370    header[4..6].copy_from_slice(&BOOTSTRAP_PROTOCOL_VERSION.to_le_bytes());
371    stream.write_all(&header).await?;
372    Ok(())
373}
374
375async fn read_header(stream: &mut TcpStream) -> Result<()> {
376    let mut header = [0u8; HEADER_BYTES];
377    stream.read_exact(&mut header).await?;
378    if header[..4] != MAGIC {
379        bail!("bootstrap protocol magic mismatch");
380    }
381    let version = u16::from_le_bytes([header[4], header[5]]);
382    if version != BOOTSTRAP_PROTOCOL_VERSION {
383        bail!("unsupported bootstrap protocol version {version}");
384    }
385    let flags = u16::from_le_bytes([header[6], header[7]]);
386    if flags != 0 {
387        bail!("bootstrap protocol flags must be zero");
388    }
389    Ok(())
390}
391
392#[cfg(test)]
393mod tests {
394    use super::*;
395
396    fn identity(value: u128) -> BootstrapIdentity {
397        BootstrapIdentity {
398            handoff_id: HandoffId::from(Uuid::from_u128(value)),
399            bootstrap_room: value as u64,
400            request_id: Uuid::from_u128(value + 100),
401        }
402    }
403
404    fn registration() -> ParticipantRegistration {
405        ParticipantRegistration {
406            role: BootstrapParticipantRole::Destination,
407            dp_rank: 0,
408            order: HandoffOrder::SourceFirst,
409            engine_type: EngineType::Vllm,
410        }
411    }
412
413    #[test]
414    fn bootstrap_address_preserves_ipv6_literals() {
415        assert_eq!(bootstrap_addr("[::1]", 1234), "[::1]:1234");
416        assert_eq!(bootstrap_addr("::1", 1234), "[::1]:1234");
417        assert_eq!(bootstrap_addr("127.0.0.1", 1234), "127.0.0.1:1234");
418    }
419
420    #[tokio::test]
421    async fn send_rejects_oversized_frame_before_codec_allocation() {
422        let cancel = CancellationToken::new();
423        let server = BootstrapServer::start(0, cancel.clone(), BootstrapServerConfig::default())
424            .await
425            .unwrap();
426        let mut client =
427            connect_to_prefill("127.0.0.1", server.port(), identity(2), registration())
428                .await
429                .unwrap();
430        let error = client
431            .send(BootstrapMessage::ProtocolError {
432                message: "x".repeat(MAX_BOOTSTRAP_FRAME_BYTES),
433            })
434            .await
435            .unwrap_err();
436        assert!(error.to_string().contains("outside"));
437        cancel.cancel();
438    }
439
440    #[tokio::test]
441    async fn bad_magic_never_enters_the_incoming_queue() {
442        let cancel = CancellationToken::new();
443        let server = BootstrapServer::start(0, cancel.clone(), BootstrapServerConfig::default())
444            .await
445            .unwrap();
446        let mut incoming_rx = server.take_incoming_receiver().unwrap();
447        let mut stream = TcpStream::connect(("127.0.0.1", server.port()))
448            .await
449            .unwrap();
450        stream.write_all(b"NOPE\x01\x00\x00\x00").await.unwrap();
451        let mut byte = [0_u8; 1];
452        let _ = stream.read(&mut byte).await;
453        assert!(matches!(
454            incoming_rx.try_recv(),
455            Err(mpsc::error::TryRecvError::Empty)
456        ));
457        cancel.cancel();
458    }
459
460    #[tokio::test]
461    async fn invalid_headers_and_frames_never_enter_session_ownership() {
462        let cancel = CancellationToken::new();
463        let server = BootstrapServer::start(0, cancel.clone(), BootstrapServerConfig::default())
464            .await
465            .unwrap();
466        let mut incoming_rx = server.take_incoming_receiver().unwrap();
467
468        for header in [
469            [MAGIC[0], MAGIC[1], MAGIC[2], MAGIC[3], 2, 0, 0, 0],
470            [MAGIC[0], MAGIC[1], MAGIC[2], MAGIC[3], 1, 0, 1, 0],
471        ] {
472            let mut stream = TcpStream::connect(("127.0.0.1", server.port()))
473                .await
474                .unwrap();
475            stream.write_all(&header).await.unwrap();
476            let mut byte = [0_u8; 1];
477            let _ = stream.read(&mut byte).await;
478        }
479
480        let mut malformed = TcpStream::connect(("127.0.0.1", server.port()))
481            .await
482            .unwrap();
483        write_header(&mut malformed).await.unwrap();
484        malformed.write_all(&1u32.to_le_bytes()).await.unwrap();
485        malformed.write_all(b"{").await.unwrap();
486        let mut byte = [0_u8; 1];
487        let _ = malformed.read(&mut byte).await;
488
489        let mut oversized = TcpStream::connect(("127.0.0.1", server.port()))
490            .await
491            .unwrap();
492        write_header(&mut oversized).await.unwrap();
493        oversized
494            .write_all(&((MAX_BOOTSTRAP_FRAME_BYTES + 1) as u32).to_le_bytes())
495            .await
496            .unwrap();
497        let _ = oversized.read(&mut byte).await;
498
499        assert!(matches!(
500            incoming_rx.try_recv(),
501            Err(mpsc::error::TryRecvError::Empty)
502        ));
503        cancel.cancel();
504        server.wait_closed().await;
505    }
506
507    #[tokio::test]
508    async fn changed_identity_is_rejected_after_registration() {
509        let cancel = CancellationToken::new();
510        let server = BootstrapServer::start(0, cancel.clone(), BootstrapServerConfig::default())
511            .await
512            .unwrap();
513        let mut incoming_rx = server.take_incoming_receiver().unwrap();
514        let mut client =
515            connect_to_prefill("127.0.0.1", server.port(), identity(3), registration())
516                .await
517                .unwrap();
518        let mut incoming = incoming_rx.recv().await.unwrap();
519
520        client.identity = identity(4);
521        client.send(BootstrapMessage::Complete).await.unwrap();
522        let error = incoming.connection.recv().await.unwrap_err();
523        assert!(error.to_string().contains("changed session identity"));
524        cancel.cancel();
525    }
526
527    #[tokio::test]
528    async fn first_frame_must_be_registration() {
529        let cancel = CancellationToken::new();
530        let server = BootstrapServer::start(0, cancel.clone(), BootstrapServerConfig::default())
531            .await
532            .unwrap();
533        let mut incoming_rx = server.take_incoming_receiver().unwrap();
534        let mut stream = TcpStream::connect(("127.0.0.1", server.port()))
535            .await
536            .unwrap();
537        write_header(&mut stream).await.unwrap();
538        let frame = serde_json::to_vec(&BootstrapWireFrame {
539            identity: identity(5),
540            message: BootstrapMessage::Registered,
541        })
542        .unwrap();
543        let mut framed_stream = framed(stream);
544        framed_stream.send(Bytes::from(frame)).await.unwrap();
545        let _ = framed_stream.next().await;
546
547        assert!(matches!(
548            incoming_rx.try_recv(),
549            Err(mpsc::error::TryRecvError::Empty)
550        ));
551        cancel.cancel();
552    }
553
554    #[tokio::test]
555    async fn full_incoming_queue_returns_overloaded() {
556        let cancel = CancellationToken::new();
557        let server = BootstrapServer::start(
558            0,
559            cancel.clone(),
560            BootstrapServerConfig {
561                max_pending_connections: 1,
562                ..BootstrapServerConfig::default()
563            },
564        )
565        .await
566        .unwrap();
567        let _incoming_rx = server.take_incoming_receiver().unwrap();
568        let _first = connect_to_prefill("127.0.0.1", server.port(), identity(6), registration())
569            .await
570            .unwrap();
571        server.wait_for_accepted_with_slot(1).await;
572        let mut second =
573            connect_to_prefill("127.0.0.1", server.port(), identity(7), registration())
574                .await
575                .unwrap();
576
577        assert_eq!(
578            tokio::time::timeout(Duration::from_secs(1), second.recv())
579                .await
580                .unwrap()
581                .unwrap(),
582            Some(BootstrapMessage::Overloaded)
583        );
584        cancel.cancel();
585    }
586
587    #[tokio::test]
588    async fn half_open_registration_saturation_returns_overloaded() {
589        let cancel = CancellationToken::new();
590        let server = BootstrapServer::start(
591            0,
592            cancel.clone(),
593            BootstrapServerConfig {
594                max_pending_connections: 1,
595                ..BootstrapServerConfig::default()
596            },
597        )
598        .await
599        .unwrap();
600        let _incoming_rx = server.take_incoming_receiver().unwrap();
601        let _half_open = TcpStream::connect(("127.0.0.1", server.port()))
602            .await
603            .unwrap();
604        server.wait_for_accepted_with_slot(1).await;
605
606        let mut rejected =
607            connect_to_prefill("127.0.0.1", server.port(), identity(8), registration())
608                .await
609                .unwrap();
610        assert_eq!(
611            tokio::time::timeout(Duration::from_secs(1), rejected.recv())
612                .await
613                .unwrap()
614                .unwrap(),
615            Some(BootstrapMessage::Overloaded)
616        );
617        cancel.cancel();
618        server.wait_closed().await;
619    }
620
621    #[tokio::test]
622    async fn shutdown_closes_half_open_registration() {
623        let cancel = CancellationToken::new();
624        let server = BootstrapServer::start(0, cancel.clone(), BootstrapServerConfig::default())
625            .await
626            .unwrap();
627        let mut stream = TcpStream::connect(("127.0.0.1", server.port()))
628            .await
629            .unwrap();
630        server.wait_for_accepted_with_slot(1).await;
631        cancel.cancel();
632
633        let mut byte = [0u8; 1];
634        assert_eq!(
635            tokio::time::timeout(Duration::from_secs(1), stream.read(&mut byte))
636                .await
637                .unwrap()
638                .unwrap(),
639            0
640        );
641        server.wait_closed().await;
642    }
643}
dynamo_mocker/services/bootstrap.rs

dynamo_mocker/services/
bootstrap.rs
