hydro_lang/deploy/maelstrom/
deploy_runtime_maelstrom.rs

1//! Runtime support for Maelstrom deployment backend.
2//!
3//! This module provides the runtime code that runs inside Maelstrom nodes,
4//! handling stdin/stdout JSON message passing according to the Maelstrom protocol.
5
6#![allow(
7    unused,
8    reason = "unused in trybuild but the __staged version is needed"
9)]
10#![allow(missing_docs, reason = "used internally")]
11
12use std::io::{BufRead, Write};
13
14use futures::{Stream, StreamExt};
15use serde::{Deserialize, Serialize};
16use stageleft::{QuotedWithContext, RuntimeData, q};
17
18use crate::forward_handle::ForwardHandle;
19use crate::live_collections::boundedness::Unbounded;
20use crate::live_collections::keyed_stream::KeyedStream;
21use crate::live_collections::stream::{ExactlyOnce, NoOrder, TotalOrder};
22use crate::location::dynamic::LocationId;
23use crate::location::member_id::TaglessMemberId;
24use crate::location::{Cluster, LocationKey, MembershipEvent};
25use crate::nondet::nondet;
26
27/// Maelstrom message envelope structure.
28#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct MaelstromMessage<T> {
30    pub src: String,
31    pub dest: String,
32    pub body: T,
33}
34
35/// Maelstrom init message body.
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct InitBody {
38    #[serde(rename = "type")]
39    pub msg_type: String,
40    pub msg_id: u64,
41    pub node_id: String,
42    pub node_ids: Vec<String>,
43}
44
45/// Maelstrom init_ok response body.
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct InitOkBody {
48    #[serde(rename = "type")]
49    pub msg_type: String,
50    pub in_reply_to: u64,
51}
52
53/// Metadata for a Maelstrom node, populated from the init message.
54/// Also manages a shared stdin reader that broadcasts lines to multiple subscribers.
55pub struct MaelstromMeta {
56    pub node_id: String,
57    pub node_ids: Vec<String>,
58    stdin_tx: tokio::sync::broadcast::Sender<String>,
59}
60
61impl MaelstromMeta {
62    /// Subscribe to stdin lines. Each subscriber receives all lines read from stdin.
63    /// Multiple subscribers can be created and each will receive a copy of every line.
64    pub fn subscribe_stdin(&self) -> tokio_stream::wrappers::BroadcastStream<String> {
65        tokio_stream::wrappers::BroadcastStream::new(self.stdin_tx.subscribe())
66    }
67
68    /// Start receiving incoming messages from clients and other nodes, after launching the DFIR.
69    pub fn start_receiving(&self) {
70        let tx_clone = self.stdin_tx.clone();
71
72        // Spawn thread to read stdin and broadcast lines
73        std::thread::spawn(move || {
74            let stdin = std::io::stdin();
75            for line in stdin.lock().lines() {
76                match line {
77                    Ok(l) => {
78                        // Ignore send errors (no receivers)
79                        let _ = tx_clone.send(l);
80                    }
81                    Err(_) => break,
82                }
83            }
84        });
85    }
86}
87
88/// Initialize a Maelstrom node by reading the init message from stdin.
89/// Returns the node metadata and sends init_ok response.
90/// Also spawns a background thread to read stdin and broadcast lines to subscribers.
91pub fn maelstrom_init() -> MaelstromMeta {
92    let stdin = std::io::stdin();
93    let mut stdout = std::io::stdout();
94
95    // Read until the init message arrives. Maelstrom delivers node-to-node
96    // messages to a node's stdin as soon as the process starts, so a
97    // fast-starting peer may already be broadcasting `hydro_data` messages
98    // before the init RPC for this node is sent — the first line on stdin is
99    // not necessarily `init`. Such pre-init messages are dropped: inter-node
100    // channels in the Maelstrom backend are lossy, so dropped messages are
101    // handled the same as network loss.
102    let msg: MaelstromMessage<InitBody> = loop {
103        let mut line = String::new();
104        let bytes_read = stdin
105            .lock()
106            .read_line(&mut line)
107            .expect("Failed to read init message");
108
109        if bytes_read == 0 {
110            // stdin reached EOF before an init message arrived. This happens when
111            // Maelstrom tears the node down (e.g. after an init timeout) or during
112            // a warmup invocation with no input. Exit cleanly instead of panicking
113            // so the node is not reported as "crashed", which would mask the real
114            // error (such as the init timeout).
115            eprintln!("stdin closed before an init message was received; exiting");
116            std::process::exit(0);
117        }
118
119        let parsed: MaelstromMessage<serde_json::Value> =
120            serde_json::from_str(&line).expect("Failed to parse message while waiting for init");
121
122        if parsed.body.get("type").and_then(|t| t.as_str()) == Some("init") {
123            let body: InitBody =
124                serde_json::from_value(parsed.body).expect("Failed to parse init message");
125            break MaelstromMessage {
126                src: parsed.src,
127                dest: parsed.dest,
128                body,
129            };
130        } else {
131            eprintln!(
132                "dropping message received before init (lossy channel): {}",
133                line.trim_end()
134            );
135        }
136    };
137
138    // Set up broadcast channel for stdin lines
139    let (stdin_tx, _) = tokio::sync::broadcast::channel::<String>(1024);
140
141    let meta = MaelstromMeta {
142        node_id: msg.body.node_id.clone(),
143        node_ids: msg.body.node_ids.clone(),
144        stdin_tx,
145    };
146
147    // Send init_ok response
148    let response = MaelstromMessage {
149        src: msg.body.node_id,
150        dest: msg.src,
151        body: InitOkBody {
152            msg_type: "init_ok".to_owned(),
153            in_reply_to: msg.body.msg_id,
154        },
155    };
156
157    let response_json = serde_json::to_string(&response).expect("Failed to serialize init_ok");
158    writeln!(stdout, "{}", response_json).expect("Failed to write init_ok");
159    stdout.flush().expect("Failed to flush stdout");
160
161    meta
162}
163
164/// Get the cluster member IDs from the Maelstrom metadata.
165/// The `meta` parameter is a RuntimeData reference to the MaelstromMeta that will be
166/// available at runtime as `__hydro_lang_maelstrom_meta`.
167pub(super) fn cluster_members<'a>(
168    meta: RuntimeData<&'a MaelstromMeta>,
169    _of_cluster: LocationKey,
170) -> impl QuotedWithContext<'a, &'a [TaglessMemberId], ()> + Clone + 'a {
171    q!({
172        let members: &'static [TaglessMemberId] = Box::leak(
173            meta.node_ids
174                .iter()
175                .map(|id| TaglessMemberId::from_maelstrom_node_id(id.clone()))
176                .collect::<Vec<TaglessMemberId>>()
177                .into_boxed_slice(),
178        );
179        members
180    })
181}
182
183/// Get the self ID for this cluster member.
184pub(super) fn cluster_self_id<'a>(
185    meta: RuntimeData<&'a MaelstromMeta>,
186) -> impl QuotedWithContext<'a, TaglessMemberId, ()> + Clone + 'a {
187    q!(TaglessMemberId::from_maelstrom_node_id(
188        meta.node_id.clone()
189    ))
190}
191
192/// Get the cluster membership stream (static for Maelstrom - all members join at start).
193/// This references `__hydro_lang_maelstrom_meta` which will be in scope at runtime.
194pub(super) fn cluster_membership_stream<'a>(
195    _location_id: &LocationId,
196) -> impl QuotedWithContext<'a, Box<dyn Stream<Item = (TaglessMemberId, MembershipEvent)> + Unpin>, ()>
197{
198    let meta: RuntimeData<&MaelstromMeta> = RuntimeData::new("__hydro_lang_maelstrom_meta");
199    q!(Box::new(futures::stream::iter(
200        meta.node_ids
201            .iter()
202            .map(|id| (
203                TaglessMemberId::from_maelstrom_node_id(id.clone()),
204                MembershipEvent::Joined
205            ))
206            .collect::<Vec<_>>()
207    ))
208        as Box<
209            dyn futures::Stream<Item = (TaglessMemberId, MembershipEvent)> + Unpin,
210        >)
211}
212
213/// Create sink and source for m2m (cluster member to cluster member) communication.
214/// Messages are routed through Maelstrom's network via stdin/stdout.
215pub(super) fn deploy_maelstrom_m2m(meta: RuntimeData<&MaelstromMeta>) -> (syn::Expr, syn::Expr) {
216    // Sink: serialize and write to stdout with Maelstrom message envelope
217    let sink_expr = q!({
218        let node_id = meta.node_id.clone();
219        sinktools::map(
220            move |(dest_id, payload): (TaglessMemberId, bytes::Bytes)| {
221                let msg = serde_json::json!({
222                    "src": node_id,
223                    "dest": dest_id.get_maelstrom_node_id(),
224                    "body": {
225                        "type": "hydro_data",
226                        "data": payload.to_vec()
227                    }
228                });
229                serde_json::to_string(&msg).unwrap() + "\n"
230            },
231            futures::sink::unfold((), |(), line: String| {
232                Box::pin(async move {
233                    print!("{}", line);
234                    std::io::stdout().flush().ok();
235                    Ok::<_, std::io::Error>(())
236                })
237            }),
238        )
239    })
240    .splice_untyped_ctx(&());
241
242    // Source: subscribe to the shared stdin broadcast stream
243    let source_expr = q!({
244        let node_ids: std::collections::HashSet<String> = meta.node_ids.iter().cloned().collect();
245        let lines = meta.subscribe_stdin();
246        futures::StreamExt::filter_map(lines, move |line_result| {
247            let node_ids = node_ids.clone();
248            Box::pin(async move {
249                let line = line_result.ok()?;
250                let mut msg =
251                    serde_json::from_str::<MaelstromMessage<serde_json::Value>>(&line).ok()?;
252                // Only process messages from other nodes (not clients)
253                if msg
254                    .body
255                    .get("type")
256                    .is_some_and(|t| t.as_str() == Some("hydro_data"))
257                {
258                    let deser: Vec<u8> =
259                        serde_json::from_value(msg.body.get_mut("data").unwrap().take()).unwrap();
260                    Some(Ok::<_, std::io::Error>((
261                        TaglessMemberId::from_maelstrom_node_id(msg.src),
262                        bytes::BytesMut::from(&deser[..]),
263                    )))
264                } else {
265                    None
266                }
267            })
268        })
269    })
270    .splice_untyped_ctx(&());
271
272    (sink_expr, source_expr)
273}
274
275/// Creates a stream of client messages from Maelstrom stdin.
276/// Returns tuples of (client_id, message_body) where client_id is the source client
277/// and message_body is the JSON value of the message body.
278///
279/// This function is meant to be used with `source_stream` on a Cluster location.
280pub fn maelstrom_client_source(
281    meta: &MaelstromMeta,
282) -> impl Stream<Item = (String, serde_json::Value)> + Unpin {
283    use std::collections::HashSet;
284
285    let node_ids: HashSet<String> = meta.node_ids.iter().cloned().collect();
286    let lines = meta.subscribe_stdin();
287
288    Box::pin(lines.filter_map(move |line_result| {
289        let node_ids = node_ids.clone();
290        async move {
291            let line = line_result.ok()?;
292            let msg: MaelstromMessage<serde_json::Value> = serde_json::from_str(&line).ok()?;
293            // Only process messages from clients (not other nodes)
294            if !node_ids.contains(&msg.src) {
295                Some((msg.src, msg.body))
296            } else {
297                None
298            }
299        }
300    }))
301}
302
303/// Sends a response to a Maelstrom client via stdout.
304///
305/// This function is meant to be used with `for_each` on a stream of responses.
306pub fn maelstrom_send_response(node_id: &str, client_id: &str, body: serde_json::Value) {
307    use std::io::Write;
308
309    let msg = MaelstromMessage {
310        src: node_id.to_owned(),
311        dest: client_id.to_owned(),
312        body,
313    };
314
315    let json = serde_json::to_string(&msg).expect("Failed to serialize response");
316    println!("{}", json);
317    std::io::stdout().flush().ok();
318}
hydro_lang/deploy/maelstrom/deploy_runtime_maelstrom.rs

hydro_lang/deploy/maelstrom/
deploy_runtime_maelstrom.rs
