hydro_lang/deploy/maelstrom/

deploy_runtime_maelstrom.rs

//! Runtime support for Maelstrom deployment backend.
//!
//! This module provides the runtime code that runs inside Maelstrom nodes,
//! handling stdin/stdout JSON message passing according to the Maelstrom protocol.

#![allow(
    unused,
    reason = "unused in trybuild but the __staged version is needed"
)]
#![allow(missing_docs, reason = "used internally")]

use std::io::{BufRead, Write};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

use futures::{Stream, StreamExt};
use serde::{Deserialize, Serialize};
use stageleft::{QuotedWithContext, RuntimeData, q};

use crate::forward_handle::ForwardHandle;
use crate::live_collections::boundedness::Unbounded;
use crate::live_collections::keyed_stream::KeyedStream;
use crate::live_collections::stream::{ExactlyOnce, NoOrder, TotalOrder};
use crate::location::dynamic::LocationId;
use crate::location::member_id::TaglessMemberId;
use crate::location::{Cluster, LocationKey, MembershipEvent};
use crate::nondet::nondet;

/// Maelstrom message envelope structure.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MaelstromMessage<T> {
    pub src: String,
    pub dest: String,
    pub body: T,
}

/// Maelstrom init message body.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InitBody {
    #[serde(rename = "type")]
    pub msg_type: String,
    pub msg_id: u64,
    pub node_id: String,
    pub node_ids: Vec<String>,
}

/// Maelstrom `init_ok` response body.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InitOkBody {
    #[serde(rename = "type")]
    pub msg_type: String,
    pub in_reply_to: u64,
}

/// Metadata for a Maelstrom node, populated from the init message.
/// Also manages a shared stdin reader that broadcasts lines to multiple subscribers.
pub struct MaelstromMeta {
    pub node_id: String,
    pub node_ids: Vec<String>,
    stdin_tx: tokio::sync::broadcast::Sender<String>,
    /// Set once this node has received its first post-`init` message from a client or a
    /// peer node. Until then, outbound node-to-node messages are dropped; see
    /// [`MaelstromMeta::peers_ready`] for why this is necessary.
    peers_ready: Arc<AtomicBool>,
}

impl MaelstromMeta {
    /// Subscribe to stdin lines. Each subscriber receives all lines read from stdin.
    /// Multiple subscribers can be created and each will receive a copy of every line.
    pub fn subscribe_stdin(&self) -> tokio_stream::wrappers::BroadcastStream<String> {
        tokio_stream::wrappers::BroadcastStream::new(self.stdin_tx.subscribe())
    }

    /// A flag that becomes `true` once it is safe to send node-to-node messages.
    ///
    /// Maelstrom sets up nodes concurrently: each node's process is launched, registered in
    /// the simulated network, and sent `init` independently, with no barrier in between. A
    /// node that answers `init` quickly may therefore start sending to a peer whose process
    /// Maelstrom has not launched yet. Maelstrom treats a message to an unregistered
    /// destination as a fatal assertion failure (`Invalid dest for message ...`) and
    /// permanently stops reading the *sender's* stdout, so every later response from that
    /// node (e.g. `topology_ok`) is silently lost and the test fails with
    /// `Client read timeout`.
    ///
    /// Jepsen only opens clients after `db/setup!` (launch + `init`) has completed on every
    /// node, so receiving a client message proves all peers are registered. A message from a
    /// peer node proves the same thing transitively: peers also hold their outbound messages
    /// until *they* have been activated, so any peer message we receive was sent after some
    /// node observed a client message. Dropping outbound messages before activation is
    /// permitted because node-to-node channels in the Maelstrom backend are lossy.
    pub fn peers_ready(&self) -> Arc<AtomicBool> {
        self.peers_ready.clone()
    }

    /// Start receiving incoming messages from clients and other nodes, after launching the DFIR.
    pub fn start_receiving(&self) {
        let tx_clone = self.stdin_tx.clone();
        let peers_ready = self.peers_ready.clone();
        let node_ids: std::collections::HashSet<String> = self.node_ids.iter().cloned().collect();

        // Spawn thread to read stdin and broadcast lines
        std::thread::spawn(move || {
            let stdin = std::io::stdin();
            for line in stdin.lock().lines() {
                match line {
                    Ok(l) => {
                        if !peers_ready.load(Ordering::Relaxed)
                            && message_activates_node(&l, &node_ids)
                        {
                            peers_ready.store(true, Ordering::Release);
                        }
                        // Ignore send errors (no receivers)
                        let _ = tx_clone.send(l);
                    }
                    Err(_) => break,
                }
            }
        });
    }
}

/// Whether receiving `line` proves that every node in the cluster has been registered in the
/// Maelstrom network (see [`MaelstromMeta::peers_ready`]). This is the case for messages from
/// peer nodes and from Jepsen clients (`c1`, `c2`, ...). Replies from Maelstrom's built-in
/// services (`lin-kv`, `seq-kv`, ...) do not count: services are registered before any node is
/// launched, so a service reply can arrive while peers are still being set up.
fn message_activates_node(line: &str, node_ids: &std::collections::HashSet<String>) -> bool {
    let Ok(msg) = serde_json::from_str::<MaelstromMessage<serde_json::Value>>(line) else {
        return false;
    };
    node_ids.contains(&msg.src) || msg.src.starts_with('c')
}

/// Initialize a Maelstrom node by reading the init message from stdin.
/// Returns the node metadata and sends `init_ok` response.
/// Also spawns a background thread to read stdin and broadcast lines to subscribers.
pub fn maelstrom_init() -> MaelstromMeta {
    let stdin = std::io::stdin();
    let mut stdout = std::io::stdout();

    // Read until the init message arrives. Maelstrom delivers node-to-node
    // messages to a node's stdin as soon as the process starts, so a peer
    // may already be sending `hydro_data` messages before the init RPC for
    // this node is sent (e.g. after this node was killed and restarted by a
    // nemesis) — the first line on stdin is not necessarily `init`. Such
    // pre-init messages are dropped: inter-node channels in the Maelstrom
    // backend are lossy, so dropped messages are handled the same as network
    // loss.
    let msg: MaelstromMessage<InitBody> = loop {
        let mut line = String::new();
        let bytes_read = stdin
            .lock()
            .read_line(&mut line)
            .expect("Failed to read init message");

        if bytes_read == 0 {
            // stdin reached EOF before an init message arrived. This happens when
            // Maelstrom tears the node down (e.g. after an init timeout) or during
            // a warmup invocation with no input. Exit cleanly instead of panicking
            // so the node is not reported as "crashed", which would mask the real
            // error (such as the init timeout).
            eprintln!("stdin closed before an init message was received; exiting");
            std::process::exit(0);
        }

        let parsed: MaelstromMessage<serde_json::Value> =
            serde_json::from_str(&line).expect("Failed to parse message while waiting for init");

        if parsed.body.get("type").and_then(|t| t.as_str()) == Some("init") {
            let body: InitBody =
                serde_json::from_value(parsed.body).expect("Failed to parse init message");
            break MaelstromMessage {
                src: parsed.src,
                dest: parsed.dest,
                body,
            };
        }
        eprintln!(
            "dropping message received before init (lossy channel): {}",
            line.trim_end()
        );
    };

    // Set up broadcast channel for stdin lines
    let (stdin_tx, _) = tokio::sync::broadcast::channel::<String>(1024);

    let meta = MaelstromMeta {
        node_id: msg.body.node_id.clone(),
        node_ids: msg.body.node_ids.clone(),
        stdin_tx,
        peers_ready: Arc::new(AtomicBool::new(false)),
    };

    // Send init_ok response
    let response = MaelstromMessage {
        src: msg.body.node_id,
        dest: msg.src,
        body: InitOkBody {
            msg_type: "init_ok".to_owned(),
            in_reply_to: msg.body.msg_id,
        },
    };

    let response_json = serde_json::to_string(&response).expect("Failed to serialize init_ok");
    writeln!(stdout, "{}", response_json).expect("Failed to write init_ok");
    stdout.flush().expect("Failed to flush stdout");

    meta
}

/// Get the cluster member IDs from the Maelstrom metadata.
/// The `meta` parameter is a `RuntimeData` reference to the `MaelstromMeta` that will be
/// available at runtime as `__hydro_lang_maelstrom_meta`.
pub(super) fn cluster_members<'a>(
    meta: RuntimeData<&'a MaelstromMeta>,
    _of_cluster: LocationKey,
) -> impl QuotedWithContext<'a, &'a [TaglessMemberId], ()> + Clone + 'a {
    q!({
        let members: &'static [TaglessMemberId] = Box::leak(
            meta.node_ids
                .iter()
                .map(|id| TaglessMemberId::from_maelstrom_node_id(id.clone()))
                .collect::<Vec<TaglessMemberId>>()
                .into_boxed_slice(),
        );
        members
    })
}

/// Get the self ID for this cluster member.
pub(super) fn cluster_self_id<'a>(
    meta: RuntimeData<&'a MaelstromMeta>,
) -> impl QuotedWithContext<'a, TaglessMemberId, ()> + Clone + 'a {
    q!(TaglessMemberId::from_maelstrom_node_id(
        meta.node_id.clone()
    ))
}

/// Get the cluster membership stream (static for Maelstrom - all members join at start).
/// This references `__hydro_lang_maelstrom_meta` which will be in scope at runtime.
pub(super) fn cluster_membership_stream<'a>(
    _location_id: &LocationId,
) -> impl QuotedWithContext<'a, Box<dyn Stream<Item = (TaglessMemberId, MembershipEvent)> + Unpin>, ()>
{
    let meta: RuntimeData<&MaelstromMeta> = RuntimeData::new("__hydro_lang_maelstrom_meta");
    q!(Box::new(futures::stream::iter(
        meta.node_ids
            .iter()
            .map(|id| (
                TaglessMemberId::from_maelstrom_node_id(id.clone()),
                MembershipEvent::Joined
            ))
            .collect::<Vec<_>>()
    ))
        as Box<
            dyn futures::Stream<Item = (TaglessMemberId, MembershipEvent)> + Unpin,
        >)
}

/// Create sink and source for m2m (cluster member to cluster member) communication.
/// Messages are routed through Maelstrom's network via stdin/stdout.
pub(super) fn deploy_maelstrom_m2m(meta: RuntimeData<&MaelstromMeta>) -> (syn::Expr, syn::Expr) {
    // Sink: serialize and write to stdout with Maelstrom message envelope.
    // Messages are dropped until the node has been activated by an incoming
    // client or peer message; see `MaelstromMeta::peers_ready`.
    let sink_expr = q!({
        let node_id = meta.node_id.clone();
        let peers_ready = meta.peers_ready();
        sinktools::map(
            move |(dest_id, payload): (TaglessMemberId, bytes::Bytes)| {
                let msg = serde_json::json!({
                    "src": node_id,
                    "dest": dest_id.get_maelstrom_node_id(),
                    "body": {
                        "type": "hydro_data",
                        "data": payload.to_vec()
                    }
                });
                let mut line = serde_json::to_string(&msg).unwrap();
                line.push('\n');
                line
            },
            futures::sink::unfold(
                peers_ready,
                |peers_ready: std::sync::Arc<std::sync::atomic::AtomicBool>, line: String| {
                    Box::pin(async move {
                        if peers_ready.load(std::sync::atomic::Ordering::Acquire) {
                            print!("{}", line);
                            std::io::stdout().flush().ok();
                        } else {
                            eprintln!(
                                "dropping message sent before peers are known to be up (lossy channel): {}",
                                line.trim_end()
                            );
                        }
                        Ok::<_, std::io::Error>(peers_ready)
                    })
                },
            ),
        )
    })
    .splice_untyped_ctx(&());

    // Source: subscribe to the shared stdin broadcast stream
    let source_expr = q!({
        let node_ids: std::collections::HashSet<String> = meta.node_ids.iter().cloned().collect();
        let lines = meta.subscribe_stdin();
        futures::StreamExt::filter_map(lines, move |line_result| {
            let node_ids = node_ids.clone();
            Box::pin(async move {
                let line = line_result.ok()?;
                let mut msg =
                    serde_json::from_str::<MaelstromMessage<serde_json::Value>>(&line).ok()?;
                // Only process messages from other nodes (not clients)
                if msg
                    .body
                    .get("type")
                    .is_some_and(|t| t.as_str() == Some("hydro_data"))
                {
                    let deser: Vec<u8> =
                        serde_json::from_value(msg.body.get_mut("data").unwrap().take()).unwrap();
                    Some(Ok::<_, std::io::Error>((
                        TaglessMemberId::from_maelstrom_node_id(msg.src),
                        bytes::BytesMut::from(&deser[..]),
                    )))
                } else {
                    None
                }
            })
        })
    })
    .splice_untyped_ctx(&());

    (sink_expr, source_expr)
}

/// Creates a stream of client messages from Maelstrom stdin.
/// Returns tuples of `(client_id, message_body)` where `client_id` is the source client
/// and `message_body` is the JSON value of the message body.
///
/// This function is meant to be used with `source_stream` on a Cluster location.
pub fn maelstrom_client_source(
    meta: &MaelstromMeta,
) -> impl Stream<Item = (String, serde_json::Value)> + Unpin {
    use std::collections::HashSet;

    let node_ids: HashSet<String> = meta.node_ids.iter().cloned().collect();
    let lines = meta.subscribe_stdin();

    Box::pin(lines.filter_map(move |line_result| {
        let node_ids = node_ids.clone();
        async move {
            let line = line_result.ok()?;
            let msg: MaelstromMessage<serde_json::Value> = serde_json::from_str(&line).ok()?;
            // Only process messages from clients (not other nodes)
            if !node_ids.contains(&msg.src) {
                Some((msg.src, msg.body))
            } else {
                None
            }
        }
    }))
}

/// Sends a response to a Maelstrom client via stdout.
///
/// This function is meant to be used with `for_each` on a stream of responses.
pub fn maelstrom_send_response(node_id: &str, client_id: &str, body: serde_json::Value) {
    use std::io::Write;

    let msg = MaelstromMessage {
        src: node_id.to_owned(),
        dest: client_id.to_owned(),
        body,
    };

    let json = serde_json::to_string(&msg).expect("Failed to serialize response");
    println!("{}", json);
    std::io::stdout().flush().ok();
}